server_test.go 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694
  1. // Copyright 2014 The go-ethereum Authors
  2. // This file is part of the go-ethereum library.
  3. //
  4. // The go-ethereum library is free software: you can redistribute it and/or modify
  5. // it under the terms of the GNU Lesser General Public License as published by
  6. // the Free Software Foundation, either version 3 of the License, or
  7. // (at your option) any later version.
  8. //
  9. // The go-ethereum library is distributed in the hope that it will be useful,
  10. // but WITHOUT ANY WARRANTY; without even the implied warranty of
  11. // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  12. // GNU Lesser General Public License for more details.
  13. //
  14. // You should have received a copy of the GNU Lesser General Public License
  15. // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
  16. package p2p
  17. import (
  18. "crypto/ecdsa"
  19. "crypto/sha256"
  20. "errors"
  21. "io"
  22. "io/ioutil"
  23. "math/rand"
  24. "net"
  25. "os"
  26. "path"
  27. "reflect"
  28. "testing"
  29. "time"
  30. "github.com/ethereum/go-ethereum/crypto"
  31. "github.com/ethereum/go-ethereum/internal/testlog"
  32. "github.com/ethereum/go-ethereum/log"
  33. "github.com/ethereum/go-ethereum/p2p/enode"
  34. "github.com/ethereum/go-ethereum/p2p/enr"
  35. "github.com/ethereum/go-ethereum/p2p/rlpx"
  36. "github.com/ethereum/go-ethereum/params"
  37. "github.com/stretchr/testify/assert"
  38. )
  39. type testTransport struct {
  40. *rlpxTransport
  41. rpub *ecdsa.PublicKey
  42. closeErr error
  43. }
  44. func newTestTransport(rpub *ecdsa.PublicKey, fd net.Conn, dialDest *ecdsa.PublicKey) transport {
  45. wrapped := newRLPX(fd, dialDest).(*rlpxTransport)
  46. wrapped.conn.InitWithSecrets(rlpx.Secrets{
  47. AES: make([]byte, 16),
  48. MAC: make([]byte, 16),
  49. EgressMAC: sha256.New(),
  50. IngressMAC: sha256.New(),
  51. })
  52. return &testTransport{rpub: rpub, rlpxTransport: wrapped}
  53. }
  54. func (c *testTransport) doEncHandshake(prv *ecdsa.PrivateKey) (*ecdsa.PublicKey, error) {
  55. return c.rpub, nil
  56. }
  57. func (c *testTransport) doProtoHandshake(our *protoHandshake) (*protoHandshake, error) {
  58. pubkey := crypto.FromECDSAPub(c.rpub)[1:]
  59. return &protoHandshake{ID: pubkey, Name: "test"}, nil
  60. }
  61. func (c *testTransport) close(err error) {
  62. c.conn.Close()
  63. c.closeErr = err
  64. }
  65. func startTestServer(t *testing.T, remoteKey *ecdsa.PublicKey, pf func(*Peer)) *Server {
  66. config := Config{
  67. Name: "test",
  68. MaxPeers: 10,
  69. ListenAddr: "127.0.0.1:0",
  70. NoDiscovery: true,
  71. PrivateKey: newkey(),
  72. Logger: testlog.Logger(t, log.LvlTrace),
  73. }
  74. server := &Server{
  75. Config: config,
  76. newPeerHook: pf,
  77. newTransport: func(fd net.Conn, dialDest *ecdsa.PublicKey) transport {
  78. return newTestTransport(remoteKey, fd, dialDest)
  79. },
  80. }
  81. if err := server.Start(); err != nil {
  82. t.Fatalf("Could not start server: %v", err)
  83. }
  84. return server
  85. }
  86. func TestServerListen(t *testing.T) {
  87. // start the test server
  88. connected := make(chan *Peer)
  89. remid := &newkey().PublicKey
  90. srv := startTestServer(t, remid, func(p *Peer) {
  91. if p.ID() != enode.PubkeyToIDV4(remid) {
  92. t.Error("peer func called with wrong node id")
  93. }
  94. connected <- p
  95. })
  96. defer close(connected)
  97. defer srv.Stop()
  98. // dial the test server
  99. conn, err := net.DialTimeout("tcp", srv.ListenAddr, 5*time.Second)
  100. if err != nil {
  101. t.Fatalf("could not dial: %v", err)
  102. }
  103. defer conn.Close()
  104. select {
  105. case peer := <-connected:
  106. if peer.LocalAddr().String() != conn.RemoteAddr().String() {
  107. t.Errorf("peer started with wrong conn: got %v, want %v",
  108. peer.LocalAddr(), conn.RemoteAddr())
  109. }
  110. peers := srv.Peers()
  111. if !reflect.DeepEqual(peers, []*Peer{peer}) {
  112. t.Errorf("Peers mismatch: got %v, want %v", peers, []*Peer{peer})
  113. }
  114. case <-time.After(1 * time.Second):
  115. t.Error("server did not accept within one second")
  116. }
  117. }
  118. func TestServerDial(t *testing.T) {
  119. // run a one-shot TCP server to handle the connection.
  120. listener, err := net.Listen("tcp", "127.0.0.1:0")
  121. if err != nil {
  122. t.Fatalf("could not setup listener: %v", err)
  123. }
  124. defer listener.Close()
  125. accepted := make(chan net.Conn, 1)
  126. go func() {
  127. conn, err := listener.Accept()
  128. if err != nil {
  129. return
  130. }
  131. accepted <- conn
  132. }()
  133. // start the server
  134. connected := make(chan *Peer)
  135. remid := &newkey().PublicKey
  136. srv := startTestServer(t, remid, func(p *Peer) { connected <- p })
  137. defer close(connected)
  138. defer srv.Stop()
  139. // tell the server to connect
  140. tcpAddr := listener.Addr().(*net.TCPAddr)
  141. node := enode.NewV4(remid, tcpAddr.IP, tcpAddr.Port, 0)
  142. srv.AddPeer(node)
  143. select {
  144. case conn := <-accepted:
  145. defer conn.Close()
  146. select {
  147. case peer := <-connected:
  148. if peer.ID() != enode.PubkeyToIDV4(remid) {
  149. t.Errorf("peer has wrong id")
  150. }
  151. if peer.Name() != "test" {
  152. t.Errorf("peer has wrong name")
  153. }
  154. if peer.RemoteAddr().String() != conn.LocalAddr().String() {
  155. t.Errorf("peer started with wrong conn: got %v, want %v",
  156. peer.RemoteAddr(), conn.LocalAddr())
  157. }
  158. peers := srv.Peers()
  159. if !reflect.DeepEqual(peers, []*Peer{peer}) {
  160. t.Errorf("Peers mismatch: got %v, want %v", peers, []*Peer{peer})
  161. }
  162. // Test AddTrustedPeer/RemoveTrustedPeer and changing Trusted flags
  163. // Particularly for race conditions on changing the flag state.
  164. if peer := srv.Peers()[0]; peer.Info().Network.Trusted {
  165. t.Errorf("peer is trusted prematurely: %v", peer)
  166. }
  167. done := make(chan bool)
  168. go func() {
  169. srv.AddTrustedPeer(node)
  170. if peer := srv.Peers()[0]; !peer.Info().Network.Trusted {
  171. t.Errorf("peer is not trusted after AddTrustedPeer: %v", peer)
  172. }
  173. srv.RemoveTrustedPeer(node)
  174. if peer := srv.Peers()[0]; peer.Info().Network.Trusted {
  175. t.Errorf("peer is trusted after RemoveTrustedPeer: %v", peer)
  176. }
  177. done <- true
  178. }()
  179. // Trigger potential race conditions
  180. peer = srv.Peers()[0]
  181. _ = peer.Inbound()
  182. _ = peer.Info()
  183. <-done
  184. case <-time.After(1 * time.Second):
  185. t.Error("server did not launch peer within one second")
  186. }
  187. case <-time.After(1 * time.Second):
  188. t.Error("server did not connect within one second")
  189. }
  190. }
  191. // This test checks that RemovePeer disconnects the peer if it is connected.
  192. func TestServerRemovePeerDisconnect(t *testing.T) {
  193. srv1 := &Server{Config: Config{
  194. PrivateKey: newkey(),
  195. MaxPeers: 1,
  196. NoDiscovery: true,
  197. Logger: testlog.Logger(t, log.LvlTrace).New("server", "1"),
  198. }}
  199. srv2 := &Server{Config: Config{
  200. PrivateKey: newkey(),
  201. MaxPeers: 1,
  202. NoDiscovery: true,
  203. NoDial: true,
  204. ListenAddr: "127.0.0.1:0",
  205. Logger: testlog.Logger(t, log.LvlTrace).New("server", "2"),
  206. }}
  207. srv1.Start()
  208. defer srv1.Stop()
  209. srv2.Start()
  210. defer srv2.Stop()
  211. if !syncAddPeer(srv1, srv2.Self()) {
  212. t.Fatal("peer not connected")
  213. }
  214. srv1.RemovePeer(srv2.Self())
  215. if srv1.PeerCount() > 0 {
  216. t.Fatal("removed peer still connected")
  217. }
  218. }
  219. // This test checks that connections are disconnected just after the encryption handshake
  220. // when the server is at capacity. Trusted connections should still be accepted.
  221. func TestServerAtCap(t *testing.T) {
  222. trustedNode := newkey()
  223. trustedID := enode.PubkeyToIDV4(&trustedNode.PublicKey)
  224. srv := &Server{
  225. Config: Config{
  226. PrivateKey: newkey(),
  227. MaxPeers: 10,
  228. NoDial: true,
  229. NoDiscovery: true,
  230. TrustedNodes: []*enode.Node{newNode(trustedID, "")},
  231. Logger: testlog.Logger(t, log.LvlTrace),
  232. },
  233. }
  234. if err := srv.Start(); err != nil {
  235. t.Fatalf("could not start: %v", err)
  236. }
  237. defer srv.Stop()
  238. newconn := func(id enode.ID) *conn {
  239. fd, _ := net.Pipe()
  240. tx := newTestTransport(&trustedNode.PublicKey, fd, nil)
  241. node := enode.SignNull(new(enr.Record), id)
  242. return &conn{fd: fd, transport: tx, flags: inboundConn, node: node, cont: make(chan error)}
  243. }
  244. // Inject a few connections to fill up the peer set.
  245. for i := 0; i < 10; i++ {
  246. c := newconn(randomID())
  247. if err := srv.checkpoint(c, srv.checkpointAddPeer); err != nil {
  248. t.Fatalf("could not add conn %d: %v", i, err)
  249. }
  250. }
  251. // Try inserting a non-trusted connection.
  252. anotherID := randomID()
  253. c := newconn(anotherID)
  254. if err := srv.checkpoint(c, srv.checkpointPostHandshake); err != DiscTooManyPeers {
  255. t.Error("wrong error for insert:", err)
  256. }
  257. // Try inserting a trusted connection.
  258. c = newconn(trustedID)
  259. if err := srv.checkpoint(c, srv.checkpointPostHandshake); err != nil {
  260. t.Error("unexpected error for trusted conn @posthandshake:", err)
  261. }
  262. if !c.is(trustedConn) {
  263. t.Error("Server did not set trusted flag")
  264. }
  265. // Remove from trusted set and try again
  266. srv.RemoveTrustedPeer(newNode(trustedID, ""))
  267. c = newconn(trustedID)
  268. if err := srv.checkpoint(c, srv.checkpointPostHandshake); err != DiscTooManyPeers {
  269. t.Error("wrong error for insert:", err)
  270. }
  271. // Add anotherID to trusted set and try again
  272. srv.AddTrustedPeer(newNode(anotherID, ""))
  273. c = newconn(anotherID)
  274. if err := srv.checkpoint(c, srv.checkpointPostHandshake); err != nil {
  275. t.Error("unexpected error for trusted conn @posthandshake:", err)
  276. }
  277. if !c.is(trustedConn) {
  278. t.Error("Server did not set trusted flag")
  279. }
  280. }
  281. func TestServerPeerLimits(t *testing.T) {
  282. srvkey := newkey()
  283. clientkey := newkey()
  284. clientnode := enode.NewV4(&clientkey.PublicKey, nil, 0, 0)
  285. var tp = &setupTransport{
  286. pubkey: &clientkey.PublicKey,
  287. phs: protoHandshake{
  288. ID: crypto.FromECDSAPub(&clientkey.PublicKey)[1:],
  289. // Force "DiscUselessPeer" due to unmatching caps
  290. // Caps: []Cap{discard.cap()},
  291. },
  292. }
  293. srv := &Server{
  294. Config: Config{
  295. PrivateKey: srvkey,
  296. MaxPeers: 0,
  297. NoDial: true,
  298. NoDiscovery: true,
  299. Protocols: []Protocol{discard},
  300. Logger: testlog.Logger(t, log.LvlTrace),
  301. },
  302. newTransport: func(fd net.Conn, dialDest *ecdsa.PublicKey) transport { return tp },
  303. }
  304. if err := srv.Start(); err != nil {
  305. t.Fatalf("couldn't start server: %v", err)
  306. }
  307. defer srv.Stop()
  308. // Check that server is full (MaxPeers=0)
  309. flags := dynDialedConn
  310. dialDest := clientnode
  311. conn, _ := net.Pipe()
  312. srv.SetupConn(conn, flags, dialDest)
  313. if tp.closeErr != DiscTooManyPeers {
  314. t.Errorf("unexpected close error: %q", tp.closeErr)
  315. }
  316. conn.Close()
  317. srv.AddTrustedPeer(clientnode)
  318. // Check that server allows a trusted peer despite being full.
  319. conn, _ = net.Pipe()
  320. srv.SetupConn(conn, flags, dialDest)
  321. if tp.closeErr == DiscTooManyPeers {
  322. t.Errorf("failed to bypass MaxPeers with trusted node: %q", tp.closeErr)
  323. }
  324. if tp.closeErr != DiscUselessPeer {
  325. t.Errorf("unexpected close error: %q", tp.closeErr)
  326. }
  327. conn.Close()
  328. srv.RemoveTrustedPeer(clientnode)
  329. // Check that server is full again.
  330. conn, _ = net.Pipe()
  331. srv.SetupConn(conn, flags, dialDest)
  332. if tp.closeErr != DiscTooManyPeers {
  333. t.Errorf("unexpected close error: %q", tp.closeErr)
  334. }
  335. conn.Close()
  336. }
  337. func TestServerSetupConn(t *testing.T) {
  338. var (
  339. clientkey, srvkey = newkey(), newkey()
  340. clientpub = &clientkey.PublicKey
  341. srvpub = &srvkey.PublicKey
  342. )
  343. tests := []struct {
  344. dontstart bool
  345. tt *setupTransport
  346. flags connFlag
  347. dialDest *enode.Node
  348. wantCloseErr error
  349. wantCalls string
  350. }{
  351. {
  352. dontstart: true,
  353. tt: &setupTransport{pubkey: clientpub},
  354. wantCalls: "close,",
  355. wantCloseErr: errServerStopped,
  356. },
  357. {
  358. tt: &setupTransport{pubkey: clientpub, encHandshakeErr: errors.New("read error")},
  359. flags: inboundConn,
  360. wantCalls: "doEncHandshake,close,",
  361. wantCloseErr: errors.New("read error"),
  362. },
  363. {
  364. tt: &setupTransport{pubkey: clientpub, phs: protoHandshake{ID: randomID().Bytes()}},
  365. dialDest: enode.NewV4(clientpub, nil, 0, 0),
  366. flags: dynDialedConn,
  367. wantCalls: "doEncHandshake,doProtoHandshake,close,",
  368. wantCloseErr: DiscUnexpectedIdentity,
  369. },
  370. {
  371. tt: &setupTransport{pubkey: clientpub, protoHandshakeErr: errors.New("foo")},
  372. dialDest: enode.NewV4(clientpub, nil, 0, 0),
  373. flags: dynDialedConn,
  374. wantCalls: "doEncHandshake,doProtoHandshake,close,",
  375. wantCloseErr: errors.New("foo"),
  376. },
  377. {
  378. tt: &setupTransport{pubkey: srvpub, phs: protoHandshake{ID: crypto.FromECDSAPub(srvpub)[1:]}},
  379. flags: inboundConn,
  380. wantCalls: "doEncHandshake,close,",
  381. wantCloseErr: DiscSelf,
  382. },
  383. {
  384. tt: &setupTransport{pubkey: clientpub, phs: protoHandshake{ID: crypto.FromECDSAPub(clientpub)[1:]}},
  385. flags: inboundConn,
  386. wantCalls: "doEncHandshake,doProtoHandshake,close,",
  387. wantCloseErr: DiscUselessPeer,
  388. },
  389. }
  390. for i, test := range tests {
  391. t.Run(test.wantCalls, func(t *testing.T) {
  392. cfg := Config{
  393. PrivateKey: srvkey,
  394. MaxPeers: 10,
  395. NoDial: true,
  396. NoDiscovery: true,
  397. Protocols: []Protocol{discard},
  398. Logger: testlog.Logger(t, log.LvlTrace),
  399. }
  400. srv := &Server{
  401. Config: cfg,
  402. newTransport: func(fd net.Conn, dialDest *ecdsa.PublicKey) transport { return test.tt },
  403. log: cfg.Logger,
  404. }
  405. if !test.dontstart {
  406. if err := srv.Start(); err != nil {
  407. t.Fatalf("couldn't start server: %v", err)
  408. }
  409. defer srv.Stop()
  410. }
  411. p1, _ := net.Pipe()
  412. srv.SetupConn(p1, test.flags, test.dialDest)
  413. if !reflect.DeepEqual(test.tt.closeErr, test.wantCloseErr) {
  414. t.Errorf("test %d: close error mismatch: got %q, want %q", i, test.tt.closeErr, test.wantCloseErr)
  415. }
  416. if test.tt.calls != test.wantCalls {
  417. t.Errorf("test %d: calls mismatch: got %q, want %q", i, test.tt.calls, test.wantCalls)
  418. }
  419. })
  420. }
  421. }
  422. func TestServerSetupConn_whenNotInRaftCluster(t *testing.T) {
  423. var (
  424. clientkey, srvkey = newkey(), newkey()
  425. clientpub = &clientkey.PublicKey
  426. )
  427. clientNode := enode.NewV4(clientpub, nil, 0, 0)
  428. srv := &Server{
  429. Config: Config{
  430. PrivateKey: srvkey,
  431. NoDiscovery: true,
  432. },
  433. newTransport: func(fd net.Conn, key *ecdsa.PublicKey) transport { return newTestTransport(clientpub, fd, key) },
  434. log: log.New(),
  435. checkPeerInRaft: func(node *enode.Node) bool {
  436. return false
  437. },
  438. }
  439. if err := srv.Start(); err != nil {
  440. t.Fatalf("couldn't start server: %v", err)
  441. }
  442. defer srv.Stop()
  443. p1, _ := net.Pipe()
  444. err := srv.SetupConn(p1, inboundConn, clientNode)
  445. assert.IsType(t, &peerError{}, err)
  446. perr := err.(*peerError)
  447. t.Log(perr.Error())
  448. assert.Equal(t, errNotInRaftCluster, perr.code)
  449. }
  450. func TestServerSetupConn_whenNotPermissioned(t *testing.T) {
  451. tmpDir, err := ioutil.TempDir("", "")
  452. if err != nil {
  453. t.Fatal(err)
  454. }
  455. defer func() { _ = os.RemoveAll(tmpDir) }()
  456. if err := ioutil.WriteFile(path.Join(tmpDir, params.PERMISSIONED_CONFIG), []byte("[]"), 0644); err != nil {
  457. t.Fatal(err)
  458. }
  459. var (
  460. clientkey, srvkey = newkey(), newkey()
  461. clientpub = &clientkey.PublicKey
  462. )
  463. clientNode := enode.NewV4(clientpub, nil, 0, 0)
  464. srv := &Server{
  465. Config: Config{
  466. PrivateKey: srvkey,
  467. NoDiscovery: true,
  468. DataDir: tmpDir,
  469. EnableNodePermission: true,
  470. },
  471. newTransport: func(fd net.Conn, key *ecdsa.PublicKey) transport { return newTestTransport(clientpub, fd, key) },
  472. log: log.New(),
  473. }
  474. if err := srv.Start(); err != nil {
  475. t.Fatalf("couldn't start server: %v", err)
  476. }
  477. defer srv.Stop()
  478. p1, _ := net.Pipe()
  479. err = srv.SetupConn(p1, inboundConn, clientNode)
  480. assert.IsType(t, &peerError{}, err)
  481. perr := err.(*peerError)
  482. t.Log(perr.Error())
  483. assert.Equal(t, errPermissionDenied, perr.code)
  484. }
  485. type setupTransport struct {
  486. pubkey *ecdsa.PublicKey
  487. encHandshakeErr error
  488. phs protoHandshake
  489. protoHandshakeErr error
  490. calls string
  491. closeErr error
  492. }
  493. func (c *setupTransport) doEncHandshake(prv *ecdsa.PrivateKey) (*ecdsa.PublicKey, error) {
  494. c.calls += "doEncHandshake,"
  495. return c.pubkey, c.encHandshakeErr
  496. }
  497. func (c *setupTransport) doProtoHandshake(our *protoHandshake) (*protoHandshake, error) {
  498. c.calls += "doProtoHandshake,"
  499. if c.protoHandshakeErr != nil {
  500. return nil, c.protoHandshakeErr
  501. }
  502. return &c.phs, nil
  503. }
  504. func (c *setupTransport) close(err error) {
  505. c.calls += "close,"
  506. c.closeErr = err
  507. }
  508. // setupConn shouldn't write to/read from the connection.
  509. func (c *setupTransport) WriteMsg(Msg) error {
  510. panic("WriteMsg called on setupTransport")
  511. }
  512. func (c *setupTransport) ReadMsg() (Msg, error) {
  513. panic("ReadMsg called on setupTransport")
  514. }
  515. func newkey() *ecdsa.PrivateKey {
  516. key, err := crypto.GenerateKey()
  517. if err != nil {
  518. panic("couldn't generate key: " + err.Error())
  519. }
  520. return key
  521. }
  522. func randomID() (id enode.ID) {
  523. for i := range id {
  524. id[i] = byte(rand.Intn(255))
  525. }
  526. return id
  527. }
  528. // This test checks that inbound connections are throttled by IP.
  529. func TestServerInboundThrottle(t *testing.T) {
  530. const timeout = 5 * time.Second
  531. newTransportCalled := make(chan struct{})
  532. srv := &Server{
  533. Config: Config{
  534. PrivateKey: newkey(),
  535. ListenAddr: "127.0.0.1:0",
  536. MaxPeers: 10,
  537. NoDial: true,
  538. NoDiscovery: true,
  539. Protocols: []Protocol{discard},
  540. Logger: testlog.Logger(t, log.LvlTrace),
  541. },
  542. newTransport: func(fd net.Conn, dialDest *ecdsa.PublicKey) transport {
  543. newTransportCalled <- struct{}{}
  544. return newRLPX(fd, dialDest)
  545. },
  546. listenFunc: func(network, laddr string) (net.Listener, error) {
  547. fakeAddr := &net.TCPAddr{IP: net.IP{95, 33, 21, 2}, Port: 4444}
  548. return listenFakeAddr(network, laddr, fakeAddr)
  549. },
  550. }
  551. if err := srv.Start(); err != nil {
  552. t.Fatal("can't start: ", err)
  553. }
  554. defer srv.Stop()
  555. // Dial the test server.
  556. conn, err := net.DialTimeout("tcp", srv.ListenAddr, timeout)
  557. if err != nil {
  558. t.Fatalf("could not dial: %v", err)
  559. }
  560. select {
  561. case <-newTransportCalled:
  562. // OK
  563. case <-time.After(timeout):
  564. t.Error("newTransport not called")
  565. }
  566. conn.Close()
  567. // Dial again. This time the server should close the connection immediately.
  568. connClosed := make(chan struct{}, 1)
  569. conn, err = net.DialTimeout("tcp", srv.ListenAddr, timeout)
  570. if err != nil {
  571. t.Fatalf("could not dial: %v", err)
  572. }
  573. defer conn.Close()
  574. go func() {
  575. conn.SetDeadline(time.Now().Add(timeout))
  576. buf := make([]byte, 10)
  577. if n, err := conn.Read(buf); err != io.EOF || n != 0 {
  578. t.Errorf("expected io.EOF and n == 0, got error %q and n == %d", err, n)
  579. }
  580. connClosed <- struct{}{}
  581. }()
  582. select {
  583. case <-connClosed:
  584. // OK
  585. case <-newTransportCalled:
  586. t.Error("newTransport called for second attempt")
  587. case <-time.After(timeout):
  588. t.Error("connection not closed within timeout")
  589. }
  590. }
  591. func listenFakeAddr(network, laddr string, remoteAddr net.Addr) (net.Listener, error) {
  592. l, err := net.Listen(network, laddr)
  593. if err == nil {
  594. l = &fakeAddrListener{l, remoteAddr}
  595. }
  596. return l, err
  597. }
  598. // fakeAddrListener is a listener that creates connections with a mocked remote address.
  599. type fakeAddrListener struct {
  600. net.Listener
  601. remoteAddr net.Addr
  602. }
  603. type fakeAddrConn struct {
  604. net.Conn
  605. remoteAddr net.Addr
  606. }
  607. func (l *fakeAddrListener) Accept() (net.Conn, error) {
  608. c, err := l.Listener.Accept()
  609. if err != nil {
  610. return nil, err
  611. }
  612. return &fakeAddrConn{c, l.remoteAddr}, nil
  613. }
  614. func (c *fakeAddrConn) RemoteAddr() net.Addr {
  615. return c.remoteAddr
  616. }
  617. func syncAddPeer(srv *Server, node *enode.Node) bool {
  618. var (
  619. ch = make(chan *PeerEvent)
  620. sub = srv.SubscribeEvents(ch)
  621. timeout = time.After(2 * time.Second)
  622. )
  623. defer sub.Unsubscribe()
  624. srv.AddPeer(node)
  625. for {
  626. select {
  627. case ev := <-ch:
  628. if ev.Type == PeerEventTypeAdd && ev.Peer == node.ID() {
  629. return true
  630. }
  631. case <-timeout:
  632. return false
  633. }
  634. }
  635. }