1
1
package io .iohk .ethereum .network .rlpx
2
2
3
- import akka .actor .{ActorSystem , Props }
3
+ import java .net .{InetSocketAddress , URI }
4
+
5
+ import akka .actor .{ActorRef , ActorSystem , Props }
4
6
import akka .io .Tcp
5
7
import akka .testkit .{TestActorRef , TestProbe }
6
8
import akka .util .ByteString
@@ -20,7 +22,7 @@ class RLPxConnectionHandlerSpec extends FlatSpec with Matchers with MockFactory
20
22
21
23
it should " write messages send to TCP connection" in new TestSetup {
22
24
23
- setupRLPxConnection ()
25
+ setupIncomingRLPxConnection ()
24
26
25
27
(mockMessageCodec.encodeMessage _).expects(Ping (): MessageSerializable ).returning(ByteString (" ping encoded" ))
26
28
rlpxConnection ! RLPxConnectionHandler .SendMessage (Ping ())
@@ -32,7 +34,7 @@ class RLPxConnectionHandlerSpec extends FlatSpec with Matchers with MockFactory
32
34
33
35
(mockMessageCodec.encodeMessage _).expects(Ping (): MessageSerializable ).returning(ByteString (" ping encoded" )).anyNumberOfTimes()
34
36
35
- setupRLPxConnection ()
37
+ setupIncomingRLPxConnection ()
36
38
37
39
// Send first message
38
40
rlpxConnection ! RLPxConnectionHandler .SendMessage (Ping ())
@@ -51,7 +53,7 @@ class RLPxConnectionHandlerSpec extends FlatSpec with Matchers with MockFactory
51
53
52
54
(mockMessageCodec.encodeMessage _).expects(Ping (): MessageSerializable ).returning(ByteString (" ping encoded" )).anyNumberOfTimes()
53
55
54
- setupRLPxConnection ()
56
+ setupIncomingRLPxConnection ()
55
57
56
58
// Send several messages
57
59
rlpxConnection ! RLPxConnectionHandler .SendMessage (Ping ())
@@ -76,9 +78,7 @@ class RLPxConnectionHandlerSpec extends FlatSpec with Matchers with MockFactory
76
78
it should " close the connection when Ack timeout happens" in new TestSetup {
77
79
(mockMessageCodec.encodeMessage _).expects(Ping (): MessageSerializable ).returning(ByteString (" ping encoded" )).anyNumberOfTimes()
78
80
79
- setupRLPxConnection()
80
-
81
- rlpxConnectionParent watch rlpxConnection
81
+ setupIncomingRLPxConnection()
82
82
83
83
rlpxConnection ! RLPxConnectionHandler .SendMessage (Ping ())
84
84
connection.expectMsg(Tcp .Write (ByteString (" ping encoded" ), RLPxConnectionHandler .Ack ))
@@ -90,7 +90,7 @@ class RLPxConnectionHandlerSpec extends FlatSpec with Matchers with MockFactory
90
90
it should " ignore timeout of old messages" in new TestSetup {
91
91
(mockMessageCodec.encodeMessage _).expects(Ping (): MessageSerializable ).returning(ByteString (" ping encoded" )).anyNumberOfTimes()
92
92
93
- setupRLPxConnection ()
93
+ setupIncomingRLPxConnection ()
94
94
95
95
rlpxConnection ! RLPxConnectionHandler .SendMessage (Ping ()) // With SEQ number 0
96
96
rlpxConnection ! RLPxConnectionHandler .SendMessage (Ping ()) // With SEQ number 1
@@ -111,6 +111,44 @@ class RLPxConnectionHandlerSpec extends FlatSpec with Matchers with MockFactory
111
111
connection.expectMsg(Tcp .Write (ByteString (" ping encoded" ), RLPxConnectionHandler .Ack ))
112
112
}
113
113
114
+ it should " close the connection if the AuthHandshake init message's MAC is invalid" in new TestSetup {
115
+ // Incomming connection arrives
116
+ rlpxConnection ! RLPxConnectionHandler .HandleConnection (connection.ref)
117
+ connection.expectMsgClass(classOf [Tcp .Register ])
118
+
119
+ // AuthHandshaker throws exception on initial message
120
+ (mockHandshaker.handleInitialMessage _).expects(* ).onCall{_ : ByteString => throw new Exception (" MAC invalid" )}
121
+ (mockHandshaker.handleInitialMessageV4 _).expects(* ).onCall{_ : ByteString => throw new Exception (" MAC invalid" )}
122
+
123
+ val data = ByteString ((0 until AuthHandshaker .InitiatePacketLength ).map(_.toByte).toArray)
124
+ rlpxConnection ! Tcp .Received (data)
125
+ rlpxConnectionParent.expectMsg(RLPxConnectionHandler .ConnectionFailed )
126
+ rlpxConnectionParent.expectTerminated(rlpxConnection)
127
+ }
128
+
129
+ it should " close the connection if the AuthHandshake response message's MAC is invalid" in new TestSetup {
130
+ // Outgoing connection request arrives
131
+ rlpxConnection ! RLPxConnectionHandler .ConnectTo (uri)
132
+ tcpActorProbe.expectMsg(Tcp .Connect (inetAddress))
133
+
134
+ // The TCP connection results are handled
135
+ val initPacket = ByteString (" Init packet" )
136
+ (mockHandshaker.initiate _).expects(uri).returning(initPacket -> mockHandshaker)
137
+
138
+ tcpActorProbe.reply(Tcp .Connected (inetAddress, inetAddress))
139
+ tcpActorProbe.expectMsg(Tcp .Register (rlpxConnection))
140
+ tcpActorProbe.expectMsg(Tcp .Write (initPacket))
141
+
142
+ // AuthHandshaker handles the response message (that throws an invalid MAC)
143
+ (mockHandshaker.handleResponseMessage _).expects(* ).onCall{_ : ByteString => throw new Exception (" MAC invalid" )}
144
+ (mockHandshaker.handleResponseMessageV4 _).expects(* ).onCall{_ : ByteString => throw new Exception (" MAC invalid" )}
145
+
146
+ val data = ByteString ((0 until AuthHandshaker .ResponsePacketLength ).map(_.toByte).toArray)
147
+ rlpxConnection ! Tcp .Received (data)
148
+ rlpxConnectionParent.expectMsg(RLPxConnectionHandler .ConnectionFailed )
149
+ rlpxConnectionParent.expectTerminated(rlpxConnection)
150
+ }
151
+
114
152
trait TestSetup extends MockFactory with SecureRandomBuilder {
115
153
implicit val system = ActorSystem (" RLPxHandlerSpec_System" )
116
154
@@ -124,19 +162,27 @@ class RLPxConnectionHandlerSpec extends FlatSpec with Matchers with MockFactory
124
162
val connection = TestProbe ()
125
163
val mockMessageCodec = mock[MessageCodec ]
126
164
165
+ val uri = new URI (" enode://18a551bee469c2e02de660ab01dede06503c986f6b8520cb5a65ad122df88b17b285e3fef09a40a0d44f99e014f8616cf1ebc2e094f96c6e09e2f390f5d34857@47.90.36.129:30303" )
166
+ val inetAddress = new InetSocketAddress (uri.getHost, uri.getPort)
167
+
127
168
val rlpxConfiguration = new RLPxConfiguration {
128
169
override val waitForTcpAckTimeout : FiniteDuration = Timeouts .normalTimeout
129
- override val waitForHandshakeTimeout : FiniteDuration = Timeouts .normalTimeout
170
+
171
+ // unused
172
+ override val waitForHandshakeTimeout : FiniteDuration = Timeouts .veryLongTimeout
130
173
}
131
174
175
+ val tcpActorProbe = TestProbe ()
132
176
val rlpxConnectionParent = TestProbe ()
133
177
val rlpxConnection = TestActorRef (
134
- Props (new RLPxConnectionHandler (
135
- mockMessageDecoder, protocolVersion, mockHandshaker, (_, _, _) => mockMessageCodec, rlpxConfiguration)),
178
+ Props (new RLPxConnectionHandler (mockMessageDecoder, protocolVersion, mockHandshaker, (_, _, _) => mockMessageCodec, rlpxConfiguration) {
179
+ override def tcpActor : ActorRef = tcpActorProbe.ref
180
+ }),
136
181
rlpxConnectionParent.ref)
182
+ rlpxConnectionParent watch rlpxConnection
137
183
138
184
// Setup for RLPxConnection, after it the RLPxConnectionHandler is in a handshaked state
139
- def setupRLPxConnection (): Unit = {
185
+ def setupIncomingRLPxConnection (): Unit = {
140
186
// Start setting up connection
141
187
rlpxConnection ! RLPxConnectionHandler .HandleConnection (connection.ref)
142
188
connection.expectMsgClass(classOf [Tcp .Register ])
0 commit comments