diff --git a/p2p/dial.go b/p2p/dial.go index 936887a1a1..b8855136f3 100644 --- a/p2p/dial.go +++ b/p2p/dial.go @@ -17,7 +17,6 @@ package p2p import ( - "container/heap" "errors" "fmt" "net" @@ -29,9 +28,10 @@ import ( ) const ( - // This is the amount of time spent waiting in between - // redialing a certain node. - dialHistoryExpiration = 30 * time.Second + // This is the amount of time spent waiting in between redialing a certain node. The + // limit is a bit higher than inboundThrottleTime to prevent failing dials in small + // private networks. + dialHistoryExpiration = inboundThrottleTime + 5*time.Second // Discovery lookups are throttled and can only run // once every few seconds. @@ -78,7 +78,7 @@ type dialstate struct { lookupBuf []*enode.Node // current discovery lookup results randomNodes []*enode.Node // filled from Table static map[enode.ID]*dialTask - hist *dialHistory + hist expHeap start time.Time // time when the dialer was first used bootnodes []*enode.Node // default dials when there are no peers @@ -91,15 +91,6 @@ type discoverTable interface { ReadRandomNodes([]*enode.Node) int } -// the dial history remembers recent dials. -type dialHistory []pastDial - -// pastDial is an entry in the dial history. -type pastDial struct { - id enode.ID - exp time.Time -} - type task interface { Do(*Server) } @@ -136,7 +127,6 @@ func newDialState(self enode.ID, static []*enode.Node, bootnodes []*enode.Node, dialing: make(map[enode.ID]connFlag), bootnodes: make([]*enode.Node, len(bootnodes)), randomNodes: make([]*enode.Node, maxdyn/2), - hist: new(dialHistory), } copy(s.bootnodes, bootnodes) for _, n := range static { @@ -154,9 +144,6 @@ func (s *dialstate) addStatic(n *enode.Node) { func (s *dialstate) removeStatic(n *enode.Node) { // This removes a task so future attempts to connect will not be made. delete(s.static, n.ID()) - // This removes a previous dial timestamp so that application - // can force a server to reconnect with chosen peer immediately. - s.hist.remove(n.ID()) } func (s *dialstate) newTasks(nRunning int, peers map[enode.ID]*Peer, now time.Time) []task { @@ -246,7 +233,7 @@ func (s *dialstate) newTasks(nRunning int, peers map[enode.ID]*Peer, now time.Ti // This should prevent cases where the dialer logic is not ticked // because there are no pending events. if nRunning == 0 && len(newtasks) == 0 && s.hist.Len() > 0 { - t := &waitExpireTask{s.hist.min().exp.Sub(now)} + t := &waitExpireTask{s.hist.nextExpiry().Sub(now)} newtasks = append(newtasks, t) } return newtasks @@ -271,7 +258,7 @@ func (s *dialstate) checkDial(n *enode.Node, peers map[enode.ID]*Peer) error { return errSelf case s.netrestrict != nil && !s.netrestrict.Contains(n.IP()): return errNotWhitelisted - case s.hist.contains(n.ID()): + case s.hist.contains(string(n.ID().Bytes())): return errRecentlyDialed } return nil @@ -280,7 +267,7 @@ func (s *dialstate) checkDial(n *enode.Node, peers map[enode.ID]*Peer) error { func (s *dialstate) taskDone(t task, now time.Time) { switch t := t.(type) { case *dialTask: - s.hist.add(t.dest.ID(), now.Add(dialHistoryExpiration)) + s.hist.add(string(t.dest.ID().Bytes()), now.Add(dialHistoryExpiration)) delete(s.dialing, t.dest.ID()) case *discoverTask: s.lookupRunning = false @@ -385,49 +372,3 @@ func (t waitExpireTask) Do(*Server) { func (t waitExpireTask) String() string { return fmt.Sprintf("wait for dial hist expire (%v)", t.Duration) } - -// Use only these methods to access or modify dialHistory. -func (h dialHistory) min() pastDial { - return h[0] -} -func (h *dialHistory) add(id enode.ID, exp time.Time) { - heap.Push(h, pastDial{id, exp}) - -} -func (h *dialHistory) remove(id enode.ID) bool { - for i, v := range *h { - if v.id == id { - heap.Remove(h, i) - return true - } - } - return false -} -func (h dialHistory) contains(id enode.ID) bool { - for _, v := range h { - if v.id == id { - return true - } - } - return false -} -func (h *dialHistory) expire(now time.Time) { - for h.Len() > 0 && h.min().exp.Before(now) { - heap.Pop(h) - } -} - -// heap.Interface boilerplate -func (h dialHistory) Len() int { return len(h) } -func (h dialHistory) Less(i, j int) bool { return h[i].exp.Before(h[j].exp) } -func (h dialHistory) Swap(i, j int) { h[i], h[j] = h[j], h[i] } -func (h *dialHistory) Push(x interface{}) { - *h = append(*h, x.(pastDial)) -} -func (h *dialHistory) Pop() interface{} { - old := *h - n := len(old) - x := old[n-1] - *h = old[0 : n-1] - return x -} diff --git a/p2p/dial_test.go b/p2p/dial_test.go index f41ab77524..48d356144b 100644 --- a/p2p/dial_test.go +++ b/p2p/dial_test.go @@ -525,50 +525,6 @@ func TestDialStateStaticDial(t *testing.T) { }) } -// This test checks that static peers will be redialed immediately if they were re-added to a static list. -func TestDialStaticAfterReset(t *testing.T) { - wantStatic := []*enode.Node{ - newNode(uintID(1), nil), - newNode(uintID(2), nil), - } - - rounds := []round{ - // Static dials are launched for the nodes that aren't yet connected. - { - peers: nil, - new: []task{ - &dialTask{flags: staticDialedConn, dest: newNode(uintID(1), nil)}, - &dialTask{flags: staticDialedConn, dest: newNode(uintID(2), nil)}, - }, - }, - // No new dial tasks, all peers are connected. - { - peers: []*Peer{ - {rw: &conn{flags: staticDialedConn, node: newNode(uintID(1), nil)}}, - {rw: &conn{flags: staticDialedConn, node: newNode(uintID(2), nil)}}, - }, - done: []task{ - &dialTask{flags: staticDialedConn, dest: newNode(uintID(1), nil)}, - &dialTask{flags: staticDialedConn, dest: newNode(uintID(2), nil)}, - }, - new: []task{ - &waitExpireTask{Duration: 30 * time.Second}, - }, - }, - } - dTest := dialtest{ - init: newDialState(enode.ID{}, wantStatic, nil, fakeTable{}, 0, nil), - rounds: rounds, - } - runDialTest(t, dTest) - for _, n := range wantStatic { - dTest.init.removeStatic(n) - dTest.init.addStatic(n) - } - // without removing peers they will be considered recently dialed - runDialTest(t, dTest) -} - // This test checks that past dials are not retried for some time. func TestDialStateCache(t *testing.T) { wantStatic := []*enode.Node{ diff --git a/p2p/server.go b/p2p/server.go index f5d4eda4a1..b103084f03 100644 --- a/p2p/server.go +++ b/p2p/server.go @@ -22,6 +22,7 @@ import ( "crypto/ecdsa" "encoding/hex" "errors" + "fmt" "net" "sort" "sync" @@ -50,6 +51,9 @@ const ( defaultMaxPendingPeers = 50 defaultDialRatio = 3 + // This time limits inbound connection attempts per source IP. + inboundThrottleTime = 30 * time.Second + // Maximum time allowed for reading a complete message. // This is effectively the amount of time a connection can be idle. frameReadTimeout = 30 * time.Second @@ -159,6 +163,7 @@ type Server struct { // the whole protocol stack. newTransport func(net.Conn) transport newPeerHook func(*Peer) + listenFunc func(network, addr string) (net.Listener, error) lock sync.Mutex // protects running running bool @@ -168,24 +173,26 @@ type Server struct { ntab discoverTable listener net.Listener ourHandshake *protoHandshake - lastLookup time.Time DiscV5 *discv5.Network + loopWG sync.WaitGroup // loop, listenLoop + peerFeed event.Feed + log log.Logger - // These are for Peers, PeerCount (and nothing else). - peerOp chan peerOpFunc - peerOpDone chan struct{} + // Channels into the run loop. + quit chan struct{} + addstatic chan *enode.Node + removestatic chan *enode.Node + addtrusted chan *enode.Node + removetrusted chan *enode.Node + peerOp chan peerOpFunc + peerOpDone chan struct{} + delpeer chan peerDrop + checkpointPostHandshake chan *conn + checkpointAddPeer chan *conn - quit chan struct{} - addstatic chan *enode.Node - removestatic chan *enode.Node - addtrusted chan *enode.Node - removetrusted chan *enode.Node - posthandshake chan *conn - addpeer chan *conn - delpeer chan peerDrop - loopWG sync.WaitGroup // loop, listenLoop - peerFeed event.Feed - log log.Logger + // State of run loop and listenLoop. + lastLookup time.Time + inboundHistory expHeap } type peerOpFunc func(map[enode.ID]*Peer) @@ -429,13 +436,16 @@ func (srv *Server) Start() (err error) { if srv.newTransport == nil { srv.newTransport = newRLPX } + if srv.listenFunc == nil { + srv.listenFunc = net.Listen + } if srv.Dialer == nil { srv.Dialer = TCPDialer{&net.Dialer{Timeout: defaultDialTimeout}} } srv.quit = make(chan struct{}) - srv.addpeer = make(chan *conn) srv.delpeer = make(chan peerDrop) - srv.posthandshake = make(chan *conn) + srv.checkpointPostHandshake = make(chan *conn) + srv.checkpointAddPeer = make(chan *conn) srv.addstatic = make(chan *enode.Node) srv.removestatic = make(chan *enode.Node) srv.addtrusted = make(chan *enode.Node) @@ -571,27 +581,28 @@ func (srv *Server) setupDiscovery() error { } func (srv *Server) setupListening() error { - // Launch the TCP listener. - listener, err := net.Listen("tcp", srv.ListenAddr) + // Launch the listener. + listener, err := srv.listenFunc("tcp", srv.ListenAddr) if err != nil { return err } - laddr := listener.Addr().(*net.TCPAddr) - srv.ListenAddr = laddr.String() srv.listener = listener - srv.localnode.Set(enr.TCP(laddr.Port)) + srv.ListenAddr = listener.Addr().String() + + // Update the local node record and map the TCP listening port if NAT is configured. + if tcp, ok := listener.Addr().(*net.TCPAddr); ok { + srv.localnode.Set(enr.TCP(tcp.Port)) + if !tcp.IP.IsLoopback() && srv.NAT != nil { + srv.loopWG.Add(1) + go func() { + nat.Map(srv.NAT, srv.quit, "tcp", tcp.Port, tcp.Port, "ethereum p2p") + srv.loopWG.Done() + }() + } + } srv.loopWG.Add(1) go srv.listenLoop() - - // Map the TCP listening port if NAT is configured. - if !laddr.IP.IsLoopback() && srv.NAT != nil { - srv.loopWG.Add(1) - go func() { - nat.Map(srv.NAT, srv.quit, "tcp", laddr.Port, laddr.Port, "ethereum p2p") - srv.loopWG.Done() - }() - } return nil } @@ -659,12 +670,14 @@ running: case <-srv.quit: // The server was stopped. Run the cleanup logic. break running + case n := <-srv.addstatic: // This channel is used by AddPeer to add to the // ephemeral static peer list. Add it to the dialer, // it will keep the node connected. srv.log.Trace("Adding static node", "node", n) dialstate.addStatic(n) + case n := <-srv.removestatic: // This channel is used by RemovePeer to send a // disconnect request to a peer and begin the @@ -674,6 +687,7 @@ running: if p, ok := peers[n.ID()]; ok { p.Disconnect(DiscRequested) } + case n := <-srv.addtrusted: // This channel is used by AddTrustedPeer to add an enode // to the trusted node set. @@ -683,6 +697,7 @@ running: if p, ok := peers[n.ID()]; ok { p.rw.set(trustedConn, true) } + case n := <-srv.removetrusted: // This channel is used by RemoveTrustedPeer to remove an enode // from the trusted node set. @@ -694,10 +709,12 @@ running: if p, ok := peers[n.ID()]; ok { p.rw.set(trustedConn, false) } + case op := <-srv.peerOp: // This channel is used by Peers and PeerCount. op(peers) srv.peerOpDone <- struct{}{} + case t := <-taskdone: // A task got done. Tell dialstate about it so it // can update its state and remove it from the active @@ -705,7 +722,8 @@ running: srv.log.Trace("Dial task done", "task", t) dialstate.taskDone(t, time.Now()) delTask(t) - case c := <-srv.posthandshake: + + case c := <-srv.checkpointPostHandshake: // A connection has passed the encryption handshake so // the remote identity is known (but hasn't been verified yet). if trusted[c.node.ID()] { @@ -714,14 +732,14 @@ running: } // TODO: track in-progress inbound node IDs (pre-Peer) to avoid dialing them. select { - case c.cont <- srv.encHandshakeChecks(peers, inboundCount, c): + case c.cont <- srv.postHandshakeChecks(peers, inboundCount, c): case <-srv.quit: break running } - case c := <-srv.addpeer: + case c := <-srv.checkpointAddPeer: // At this point the connection is past the protocol handshake. // Its capabilities are known and the remote identity is verified. - err := srv.protoHandshakeChecks(peers, inboundCount, c) + err := srv.addPeerChecks(peers, inboundCount, c) if err == nil { // The handshakes are done and it passed all checks. p := newPeer(srv.log, c, srv.Protocols) @@ -746,6 +764,7 @@ running: case <-srv.quit: break running } + case pd := <-srv.delpeer: // A peer disconnected. d := common.PrettyDuration(mclock.Now() - pd.created) @@ -780,17 +799,7 @@ running: } } -func (srv *Server) protoHandshakeChecks(peers map[enode.ID]*Peer, inboundCount int, c *conn) error { - // Drop connections with no matching protocols. - if len(srv.Protocols) > 0 && countMatchingProtocols(srv.Protocols, c.caps) == 0 { - return DiscUselessPeer - } - // Repeat the encryption handshake checks because the - // peer set might have changed between the handshakes. - return srv.encHandshakeChecks(peers, inboundCount, c) -} - -func (srv *Server) encHandshakeChecks(peers map[enode.ID]*Peer, inboundCount int, c *conn) error { +func (srv *Server) postHandshakeChecks(peers map[enode.ID]*Peer, inboundCount int, c *conn) error { switch { case !c.is(trustedConn|staticDialedConn) && len(peers) >= srv.MaxPeers: return DiscTooManyPeers @@ -805,9 +814,20 @@ func (srv *Server) encHandshakeChecks(peers map[enode.ID]*Peer, inboundCount int } } +func (srv *Server) addPeerChecks(peers map[enode.ID]*Peer, inboundCount int, c *conn) error { + // Drop connections with no matching protocols. + if len(srv.Protocols) > 0 && countMatchingProtocols(srv.Protocols, c.caps) == 0 { + return DiscUselessPeer + } + // Repeat the post-handshake checks because the + // peer set might have changed since those checks were performed. + return srv.postHandshakeChecks(peers, inboundCount, c) +} + func (srv *Server) maxInboundConns() int { return srv.MaxPeers - srv.maxDialedConns() } + func (srv *Server) maxDialedConns() int { if srv.NoDiscovery || srv.NoDial { return 0 @@ -835,7 +855,7 @@ func (srv *Server) listenLoop() { } for { - // Wait for a handshake slot before accepting. + // Wait for a free slot before accepting. <-slots var ( @@ -854,21 +874,16 @@ func (srv *Server) listenLoop() { break } - // Reject connections that do not match NetRestrict. - if srv.NetRestrict != nil { - if tcp, ok := fd.RemoteAddr().(*net.TCPAddr); ok && !srv.NetRestrict.Contains(tcp.IP) { - srv.log.Debug("Rejected conn (not whitelisted in NetRestrict)", "addr", fd.RemoteAddr()) - fd.Close() - slots <- struct{}{} - continue - } + remoteIP := netutil.AddrIP(fd.RemoteAddr()) + if err := srv.checkInboundConn(fd, remoteIP); err != nil { + srv.log.Debug("Rejected inbound connnection", "addr", fd.RemoteAddr(), "err", err) + fd.Close() + slots <- struct{}{} + continue } - - var ip net.IP - if tcp, ok := fd.RemoteAddr().(*net.TCPAddr); ok { - ip = tcp.IP + if remoteIP != nil { + fd = newMeteredConn(fd, true, remoteIP) } - fd = newMeteredConn(fd, true, ip) srv.log.Trace("Accepted connection", "addr", fd.RemoteAddr()) go func() { srv.SetupConn(fd, inboundConn, nil) @@ -877,6 +892,22 @@ func (srv *Server) listenLoop() { } } +func (srv *Server) checkInboundConn(fd net.Conn, remoteIP net.IP) error { + if remoteIP != nil { + // Reject connections that do not match NetRestrict. + if srv.NetRestrict != nil && !srv.NetRestrict.Contains(remoteIP) { + return fmt.Errorf("not whitelisted in NetRestrict") + } + // Reject Internet peers that try too often. + srv.inboundHistory.expire(time.Now()) + if !netutil.IsLAN(remoteIP) && srv.inboundHistory.contains(remoteIP.String()) { + return fmt.Errorf("too many attempts") + } + srv.inboundHistory.add(remoteIP.String(), time.Now().Add(inboundThrottleTime)) + } + return nil +} + // SetupConn runs the handshakes and attempts to add the connection // as a peer. It returns when the connection has been added as a peer // or the handshakes have failed. @@ -898,6 +929,7 @@ func (srv *Server) setupConn(c *conn, flags connFlag, dialDest *enode.Node) erro if !running { return errServerStopped } + // If dialing, figure out the remote public key. var dialPubkey *ecdsa.PublicKey if dialDest != nil { @@ -906,7 +938,8 @@ func (srv *Server) setupConn(c *conn, flags connFlag, dialDest *enode.Node) erro return errors.New("dial destination doesn't have a secp256k1 public key") } } - // Run the encryption handshake. + + // Run the RLPx handshake. remotePubkey, err := c.doEncHandshake(srv.PrivateKey, dialPubkey) if err != nil { srv.log.Trace("Failed RLPx handshake", "addr", c.fd.RemoteAddr(), "conn", c.flags, "err", err) @@ -925,12 +958,13 @@ func (srv *Server) setupConn(c *conn, flags connFlag, dialDest *enode.Node) erro conn.handshakeDone(c.node.ID()) } clog := srv.log.New("id", c.node.ID(), "addr", c.fd.RemoteAddr(), "conn", c.flags) - err = srv.checkpoint(c, srv.posthandshake) + err = srv.checkpoint(c, srv.checkpointPostHandshake) if err != nil { - clog.Trace("Rejected peer before protocol handshake", "err", err) + clog.Trace("Rejected peer", "err", err) return err } - // Run the protocol handshake + + // Run the capability negotiation handshake. phs, err := c.doProtoHandshake(srv.ourHandshake) if err != nil { clog.Trace("Failed proto handshake", "err", err) @@ -941,14 +975,15 @@ func (srv *Server) setupConn(c *conn, flags connFlag, dialDest *enode.Node) erro return DiscUnexpectedIdentity } c.caps, c.name = phs.Caps, phs.Name - err = srv.checkpoint(c, srv.addpeer) + err = srv.checkpoint(c, srv.checkpointAddPeer) if err != nil { clog.Trace("Rejected peer", "err", err) return err } - // If the checks completed successfully, runPeer has now been - // launched by run. - clog.Trace("connection set up", "inbound", dialDest == nil) + + // If the checks completed successfully, the connection has been added as a peer and + // runPeer has been launched. + clog.Trace("Connection set up", "inbound", dialDest == nil) return nil } diff --git a/p2p/server_test.go b/p2p/server_test.go index 68b616e03e..e8bc627e1d 100644 --- a/p2p/server_test.go +++ b/p2p/server_test.go @@ -19,6 +19,7 @@ package p2p import ( "crypto/ecdsa" "errors" + "io" "math/rand" "net" "reflect" @@ -380,19 +381,19 @@ func TestServerAtCap(t *testing.T) { // Inject a few connections to fill up the peer set. for i := 0; i < 10; i++ { c := newconn(randomID()) - if err := srv.checkpoint(c, srv.addpeer); err != nil { + if err := srv.checkpoint(c, srv.checkpointAddPeer); err != nil { t.Fatalf("could not add conn %d: %v", i, err) } } // Try inserting a non-trusted connection. anotherID := randomID() c := newconn(anotherID) - if err := srv.checkpoint(c, srv.posthandshake); err != DiscTooManyPeers { + if err := srv.checkpoint(c, srv.checkpointPostHandshake); err != DiscTooManyPeers { t.Error("wrong error for insert:", err) } // Try inserting a trusted connection. c = newconn(trustedID) - if err := srv.checkpoint(c, srv.posthandshake); err != nil { + if err := srv.checkpoint(c, srv.checkpointPostHandshake); err != nil { t.Error("unexpected error for trusted conn @posthandshake:", err) } if !c.is(trustedConn) { @@ -402,14 +403,14 @@ func TestServerAtCap(t *testing.T) { // Remove from trusted set and try again srv.RemoveTrustedPeer(newNode(trustedID, nil)) c = newconn(trustedID) - if err := srv.checkpoint(c, srv.posthandshake); err != DiscTooManyPeers { + if err := srv.checkpoint(c, srv.checkpointPostHandshake); err != DiscTooManyPeers { t.Error("wrong error for insert:", err) } // Add anotherID to trusted set and try again srv.AddTrustedPeer(newNode(anotherID, nil)) c = newconn(anotherID) - if err := srv.checkpoint(c, srv.posthandshake); err != nil { + if err := srv.checkpoint(c, srv.checkpointPostHandshake); err != nil { t.Error("unexpected error for trusted conn @posthandshake:", err) } if !c.is(trustedConn) { @@ -626,3 +627,100 @@ func randomID() (id enode.ID) { } return id } + +// This test checks that inbound connections are throttled by IP. +func TestServerInboundThrottle(t *testing.T) { + const timeout = 5 * time.Second + newTransportCalled := make(chan struct{}) + srv := &Server{ + Config: Config{ + PrivateKey: newkey(), + ListenAddr: "127.0.0.1:0", + MaxPeers: 10, + NoDial: true, + NoDiscovery: true, + Protocols: []Protocol{discard}, + Logger: testlog.Logger(t, log.LvlTrace), + }, + newTransport: func(fd net.Conn) transport { + newTransportCalled <- struct{}{} + return newRLPX(fd) + }, + listenFunc: func(network, laddr string) (net.Listener, error) { + fakeAddr := &net.TCPAddr{IP: net.IP{95, 33, 21, 2}, Port: 4444} + return listenFakeAddr(network, laddr, fakeAddr) + }, + } + if err := srv.Start(); err != nil { + t.Fatal("can't start: ", err) + } + defer srv.Stop() + + // Dial the test server. + conn, err := net.DialTimeout("tcp", srv.ListenAddr, timeout) + if err != nil { + t.Fatalf("could not dial: %v", err) + } + select { + case <-newTransportCalled: + // OK + case <-time.After(timeout): + t.Error("newTransport not called") + } + conn.Close() + + // Dial again. This time the server should close the connection immediately. + connClosed := make(chan struct{}) + conn, err = net.DialTimeout("tcp", srv.ListenAddr, timeout) + if err != nil { + t.Fatalf("could not dial: %v", err) + } + defer conn.Close() + go func() { + conn.SetDeadline(time.Now().Add(timeout)) + buf := make([]byte, 10) + if n, err := conn.Read(buf); err != io.EOF || n != 0 { + t.Errorf("expected io.EOF and n == 0, got error %q and n == %d", err, n) + } + connClosed <- struct{}{} + }() + select { + case <-connClosed: + // OK + case <-newTransportCalled: + t.Error("newTransport called for second attempt") + case <-time.After(timeout): + t.Error("connection not closed within timeout") + } +} + +func listenFakeAddr(network, laddr string, remoteAddr net.Addr) (net.Listener, error) { + l, err := net.Listen(network, laddr) + if err == nil { + l = &fakeAddrListener{l, remoteAddr} + } + return l, err +} + +// fakeAddrListener is a listener that creates connections with a mocked remote address. +type fakeAddrListener struct { + net.Listener + remoteAddr net.Addr +} + +type fakeAddrConn struct { + net.Conn + remoteAddr net.Addr +} + +func (l *fakeAddrListener) Accept() (net.Conn, error) { + c, err := l.Listener.Accept() + if err != nil { + return nil, err + } + return &fakeAddrConn{c, l.remoteAddr}, nil +} + +func (c *fakeAddrConn) RemoteAddr() net.Addr { + return c.remoteAddr +} diff --git a/p2p/util.go b/p2p/util.go new file mode 100644 index 0000000000..2a6edf5cee --- /dev/null +++ b/p2p/util.go @@ -0,0 +1,82 @@ +// Copyright 2019 The go-ethereum Authors +// This file is part of the go-ethereum library. +// +// The go-ethereum library is free software: you can redistribute it and/or modify +// it under the terms of the GNU Lesser General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// The go-ethereum library is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Lesser General Public License for more details. +// +// You should have received a copy of the GNU Lesser General Public License +// along with the go-ethereum library. If not, see . + +package p2p + +import ( + "container/heap" + "time" +) + +// expHeap tracks strings and their expiry time. +type expHeap []expItem + +// expItem is an entry in addrHistory. +type expItem struct { + item string + exp time.Time +} + +// nextExpiry returns the next expiry time. +func (h *expHeap) nextExpiry() time.Time { + return (*h)[0].exp +} + +// add adds an item and sets its expiry time. +func (h *expHeap) add(item string, exp time.Time) { + heap.Push(h, expItem{item, exp}) +} + +// remove removes an item. +func (h *expHeap) remove(item string) bool { + for i, v := range *h { + if v.item == item { + heap.Remove(h, i) + return true + } + } + return false +} + +// contains checks whether an item is present. +func (h expHeap) contains(item string) bool { + for _, v := range h { + if v.item == item { + return true + } + } + return false +} + +// expire removes items with expiry time before 'now'. +func (h *expHeap) expire(now time.Time) { + for h.Len() > 0 && h.nextExpiry().Before(now) { + heap.Pop(h) + } +} + +// heap.Interface boilerplate +func (h expHeap) Len() int { return len(h) } +func (h expHeap) Less(i, j int) bool { return h[i].exp.Before(h[j].exp) } +func (h expHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } +func (h *expHeap) Push(x interface{}) { *h = append(*h, x.(expItem)) } +func (h *expHeap) Pop() interface{} { + old := *h + n := len(old) + x := old[n-1] + *h = old[0 : n-1] + return x +}