From 782a8088ebd08e6abbf26a47a103dc207e5b94c3 Mon Sep 17 00:00:00 2001 From: Conner Fromknecht Date: Sun, 1 Apr 2018 07:00:33 -0700 Subject: [PATCH] brontide/listener: allow parallel handshakes --- brontide/listener.go | 142 ++++++++++++++++++++++++++++++++++++------- 1 file changed, 121 insertions(+), 21 deletions(-) diff --git a/brontide/listener.go b/brontide/listener.go index 1cc171f66..603a6b1a8 100644 --- a/brontide/listener.go +++ b/brontide/listener.go @@ -1,6 +1,7 @@ package brontide import ( + "errors" "io" "net" "time" @@ -8,6 +9,10 @@ import ( "github.com/roasbeef/btcd/btcec" ) +// defaultHandshakes is the maximum number of handshakes that can be done in +// parallel. +const defaultHandshakes = 1000 + // Listener is an implementation of a net.Conn which executes an authenticated // key exchange and message encryption protocol dubbed "Machine" after // initial connection acceptance. See the Machine struct for additional @@ -17,6 +22,10 @@ type Listener struct { localStatic *btcec.PrivateKey tcp *net.TCPListener + + handshakeSema chan struct{} + conns chan maybeConn + quit chan struct{} } // A compile-time assertion to ensure that Conn meets the net.Listener interface. @@ -36,23 +45,57 @@ func NewListener(localStatic *btcec.PrivateKey, listenAddr string) (*Listener, return nil, err } - return &Listener{ - localStatic: localStatic, - tcp: l, - }, nil + brontideListener := &Listener{ + localStatic: localStatic, + tcp: l, + handshakeSema: make(chan struct{}, defaultHandshakes), + conns: make(chan maybeConn), + quit: make(chan struct{}), + } + + for i := 0; i < defaultHandshakes; i++ { + brontideListener.handshakeSema <- struct{}{} + } + + go brontideListener.listen() + + return brontideListener, nil } -// Accept waits for and returns the next connection to the listener. All -// incoming connections are authenticated via the three act Brontide -// key-exchange scheme. This function will fail with a non-nil error in the -// case that either the handshake breaks down, or the remote peer doesn't know -// our static public key. +// listen accepts connection from the underlying tcp conn, then performs +// the brontinde handshake procedure asynchronously. A maximum of +// defaultHandshakes will be active at any given time. // -// Part of the net.Listener interface. -func (l *Listener) Accept() (net.Conn, error) { - conn, err := l.tcp.Accept() - if err != nil { - return nil, err +// NOTE: This method must be run as a goroutine. +func (l *Listener) listen() { + for { + select { + case <-l.handshakeSema: + case <-l.quit: + return + } + + conn, err := l.tcp.Accept() + if err != nil { + l.rejectConn(err) + l.handshakeSema <- struct{}{} + continue + } + + go l.doHandshake(conn) + } +} + +// doHandshake asynchronously performs the brontide handshake, so that it does +// not block the main accept loop. This prevents peers that delay writing to the +// connection from block other connection attempts. +func (l *Listener) doHandshake(conn net.Conn) { + defer func() { l.handshakeSema <- struct{}{} }() + + select { + case <-l.quit: + return + default: } brontideConn := &Conn{ @@ -71,11 +114,13 @@ func (l *Listener) Accept() (net.Conn, error) { var actOne [ActOneSize]byte if _, err := io.ReadFull(conn, actOne[:]); err != nil { brontideConn.conn.Close() - return nil, err + l.rejectConn(err) + return } if err := brontideConn.noise.RecvActOne(actOne); err != nil { brontideConn.conn.Close() - return nil, err + l.rejectConn(err) + return } // Next, progress the handshake processes by sending over our ephemeral @@ -83,11 +128,19 @@ func (l *Listener) Accept() (net.Conn, error) { actTwo, err := brontideConn.noise.GenActTwo() if err != nil { brontideConn.conn.Close() - return nil, err + l.rejectConn(err) + return } if _, err := conn.Write(actTwo[:]); err != nil { brontideConn.conn.Close() - return nil, err + l.rejectConn(err) + return + } + + select { + case <-l.quit: + return + default: } // We'll ensure that we get ActTwo from the remote peer in a timely @@ -101,18 +154,59 @@ func (l *Listener) Accept() (net.Conn, error) { var actThree [ActThreeSize]byte if _, err := io.ReadFull(conn, actThree[:]); err != nil { brontideConn.conn.Close() - return nil, err + l.rejectConn(err) + return } if err := brontideConn.noise.RecvActThree(actThree); err != nil { brontideConn.conn.Close() - return nil, err + l.rejectConn(err) + return } // We'll reset the deadline as it's no longer critical beyond the // initial handshake. conn.SetReadDeadline(time.Time{}) - return brontideConn, nil + l.acceptConn(brontideConn) +} + +// maybeConn holds either a brontide connection or an error returned from the +// handshake. +type maybeConn struct { + conn *Conn + err error +} + +// acceptConn returns a connection that successfully performed a handshake. +func (l *Listener) acceptConn(conn *Conn) { + select { + case l.conns <- maybeConn{conn: conn}: + case <-l.quit: + } +} + +// rejectConn returns any errors encountered during connection or handshake. +func (l *Listener) rejectConn(err error) { + select { + case l.conns <- maybeConn{err: err}: + case <-l.quit: + } +} + +// Accept waits for and returns the next connection to the listener. All +// incoming connections are authenticated via the three act Brontide +// key-exchange scheme. This function will fail with a non-nil error in the +// case that either the handshake breaks down, or the remote peer doesn't know +// our static public key. +// +// Part of the net.Listener interface. +func (l *Listener) Accept() (net.Conn, error) { + select { + case result := <-l.conns: + return result.conn, result.err + case <-l.quit: + return nil, errors.New("brontide connection closed") + } } // Close closes the listener. Any blocked Accept operations will be unblocked @@ -120,6 +214,12 @@ func (l *Listener) Accept() (net.Conn, error) { // // Part of the net.Listener interface. func (l *Listener) Close() error { + select { + case <-l.quit: + default: + close(l.quit) + } + return l.tcp.Close() }