brontide/listener: allow parallel handshakes

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()
 }