noise: make it possible for the server to send early data (#1750)

* noise: make it possible for the server to send early data

* noise: use separate early data handlers for initiator and responder

* noise: use a 2 kb buffer for all handshake operations

* Add comments around EarlyData(Handler)

* noise: send early data with 2nd and 3rd handshake message

* Update p2p/security/noise/session_transport.go

Co-authored-by: Marco Munizaga <git@marcopolo.io>

p2p/security/noise/crypto_test.go

@@ -93,7 +93,7 @@ func TestCryptoFailsIfHandshakeIncomplete(t *testing.T) {
 	init, resp := net.Pipe()
 	_ = resp.Close()
 
-	session, _ := newSecureSession(initTransport, context.TODO(), init, "remote-peer", nil, nil, true)
+	session, _ := newSecureSession(initTransport, context.TODO(), init, "remote-peer", nil, nil, nil, true)
 	_, err := session.encrypt(nil, []byte("hi"))
 	if err == nil {
 		t.Error("expected encryption error when handshake incomplete")

p2p/security/noise/handshake.go

@@ -10,8 +10,6 @@ import (
 	"runtime/debug"
 	"time"
 
-	"golang.org/x/crypto/chacha20poly1305"
-
 	"github.com/libp2p/go-libp2p/core/crypto"
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/libp2p/go-libp2p/p2p/security/noise/pb"
@@ -64,11 +62,6 @@ func (s *secureSession) runHandshake(ctx context.Context) (err error) {
 		return fmt.Errorf("error initializing handshake state: %w", err)
 	}
 
-	payload, err := s.generateHandshakePayload(kp)
-	if err != nil {
-		return err
-	}
-
 	// set a deadline to complete the handshake, if one has been supplied.
 	// clear it after we're done.
 	if deadline, ok := ctx.Deadline(); ok {
@@ -78,22 +71,14 @@ func (s *secureSession) runHandshake(ctx context.Context) (err error) {
 		}
 	}
 
-	// We can re-use this buffer for all handshake messages as its size
-	// will be the size of the maximum handshake message for the Noise XX pattern.
-	// Also, since we prefix every noise handshake message with its length, we need to account for
-	// it when we fetch the buffer from the pool
-	maxMsgSize := 2*noise.DH25519.DHLen() + len(payload) + 2*chacha20poly1305.Overhead
-	hbuf := pool.Get(maxMsgSize + LengthPrefixLength)
+	// We can re-use this buffer for all handshake messages.
+	hbuf := pool.Get(2 << 10)
 	defer pool.Put(hbuf)
 
 	if s.initiator {
 		// stage 0 //
 		// Handshake Msg Len = len(DH ephemeral key)
-		var ed []byte
-		if s.earlyDataHandler != nil {
-			ed = s.earlyDataHandler.Send(ctx, s.insecureConn, s.remoteID)
-		}
-		if err := s.sendHandshakeMessage(hs, ed, hbuf); err != nil {
+		if err := s.sendHandshakeMessage(hs, nil, hbuf); err != nil {
 			return fmt.Errorf("error sending handshake message: %w", err)
 		}
 
@@ -102,31 +87,47 @@ func (s *secureSession) runHandshake(ctx context.Context) (err error) {
 		if err != nil {
 			return fmt.Errorf("error reading handshake message: %w", err)
 		}
-		if err := s.handleRemoteHandshakePayload(plaintext, hs.PeerStatic()); err != nil {
+		rcvdEd, err := s.handleRemoteHandshakePayload(plaintext, hs.PeerStatic())
+		if err != nil {
 			return err
 		}
+		if s.initiatorEarlyDataHandler != nil {
+			if err := s.initiatorEarlyDataHandler.Received(ctx, s.insecureConn, rcvdEd); err != nil {
+				return err
+			}
+		}
 
 		// stage 2 //
 		// Handshake Msg Len = len(DHT static key) +  MAC(static key is encrypted) + len(Payload) + MAC(payload is encrypted)
+		var ed []byte
+		if s.initiatorEarlyDataHandler != nil {
+			ed = s.initiatorEarlyDataHandler.Send(ctx, s.insecureConn, s.remoteID)
+		}
+		payload, err := s.generateHandshakePayload(kp, ed)
+		if err != nil {
+			return err
+		}
 		if err := s.sendHandshakeMessage(hs, payload, hbuf); err != nil {
 			return fmt.Errorf("error sending handshake message: %w", err)
 		}
 		return nil
 	} else {
 		// stage 0 //
-		initialPayload, err := s.readHandshakeMessage(hs)
-		if err != nil {
+		if _, err := s.readHandshakeMessage(hs); err != nil {
 			return fmt.Errorf("error reading handshake message: %w", err)
 		}
-		if s.earlyDataHandler != nil {
-			if err := s.earlyDataHandler.Received(ctx, s.insecureConn, initialPayload); err != nil {
-				return err
-			}
-		}
 
 		// stage 1 //
 		// Handshake Msg Len = len(DH ephemeral key) + len(DHT static key) +  MAC(static key is encrypted) + len(Payload) +
 		// MAC(payload is encrypted)
+		var ed []byte
+		if s.responderEarlyDataHandler != nil {
+			ed = s.responderEarlyDataHandler.Send(ctx, s.insecureConn, s.remoteID)
+		}
+		payload, err := s.generateHandshakePayload(kp, ed)
+		if err != nil {
+			return err
+		}
 		if err := s.sendHandshakeMessage(hs, payload, hbuf); err != nil {
 			return fmt.Errorf("error sending handshake message: %w", err)
 		}
@@ -136,7 +137,16 @@ func (s *secureSession) runHandshake(ctx context.Context) (err error) {
 		if err != nil {
 			return fmt.Errorf("error reading handshake message: %w", err)
 		}
-		return s.handleRemoteHandshakePayload(plaintext, hs.PeerStatic())
+		rcvdEd, err := s.handleRemoteHandshakePayload(plaintext, hs.PeerStatic())
+		if err != nil {
+			return err
+		}
+		if s.responderEarlyDataHandler != nil {
+			if err := s.responderEarlyDataHandler.Received(ctx, s.insecureConn, rcvdEd); err != nil {
+				return err
+			}
+		}
+		return nil
 	}
 }
 
@@ -214,8 +224,8 @@ func (s *secureSession) readHandshakeMessage(hs *noise.HandshakeState) ([]byte,
 
 // generateHandshakePayload creates a libp2p handshake payload with a
 // signature of our static noise key.
-func (s *secureSession) generateHandshakePayload(localStatic noise.DHKey) ([]byte, error) {
-	// obtain the public key from the handshake session so we can sign it with
+func (s *secureSession) generateHandshakePayload(localStatic noise.DHKey, data []byte) ([]byte, error) {
+	// obtain the public key from the handshake session, so we can sign it with
 	// our libp2p secret key.
 	localKeyRaw, err := crypto.MarshalPublicKey(s.LocalPublicKey())
 	if err != nil {
@@ -230,10 +240,11 @@ func (s *secureSession) generateHandshakePayload(localStatic noise.DHKey) ([]byt
 	}
 
 	// create payload
-	payload := new(pb.NoiseHandshakePayload)
-	payload.IdentityKey = localKeyRaw
-	payload.IdentitySig = signedPayload
-	payloadEnc, err := proto.Marshal(payload)
+	payloadEnc, err := proto.Marshal(&pb.NoiseHandshakePayload{
+		IdentityKey: localKeyRaw,
+		IdentitySig: signedPayload,
+		Data:        data,
+	})
 	if err != nil {
 		return nil, fmt.Errorf("error marshaling handshake payload: %w", err)
 	}
@@ -242,44 +253,45 @@ func (s *secureSession) generateHandshakePayload(localStatic noise.DHKey) ([]byt
 
 // handleRemoteHandshakePayload unmarshals the handshake payload object sent
 // by the remote peer and validates the signature against the peer's static Noise key.
-func (s *secureSession) handleRemoteHandshakePayload(payload []byte, remoteStatic []byte) error {
+// It returns the data attached to the payload.
+func (s *secureSession) handleRemoteHandshakePayload(payload []byte, remoteStatic []byte) ([]byte, error) {
 	// unmarshal payload
 	nhp := new(pb.NoiseHandshakePayload)
 	err := proto.Unmarshal(payload, nhp)
 	if err != nil {
-		return fmt.Errorf("error unmarshaling remote handshake payload: %w", err)
+		return nil, fmt.Errorf("error unmarshaling remote handshake payload: %w", err)
 	}
 
 	// unpack remote peer's public libp2p key
 	remotePubKey, err := crypto.UnmarshalPublicKey(nhp.GetIdentityKey())
 	if err != nil {
-		return err
+		return nil, err
 	}
 	id, err := peer.IDFromPublicKey(remotePubKey)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	// check the peer ID for:
 	// * all outbound connection
 	// * inbound connections, if we know which peer we want to connect to (SecureInbound called with a peer ID)
 	if (s.initiator && s.remoteID != id) || (!s.initiator && s.remoteID != "" && s.remoteID != id) {
 		// use Pretty() as it produces the full b58-encoded string, rather than abbreviated forms.
-		return fmt.Errorf("peer id mismatch: expected %s, but remote key matches %s", s.remoteID.Pretty(), id.Pretty())
+		return nil, fmt.Errorf("peer id mismatch: expected %s, but remote key matches %s", s.remoteID.Pretty(), id.Pretty())
 	}
 
 	// verify payload is signed by asserted remote libp2p key.
 	sig := nhp.GetIdentitySig()
 	msg := append([]byte(payloadSigPrefix), remoteStatic...)
 	ok, err := remotePubKey.Verify(msg, sig)
 	if err != nil {
-		return fmt.Errorf("error verifying signature: %w", err)
+		return nil, fmt.Errorf("error verifying signature: %w", err)
 	} else if !ok {
-		return fmt.Errorf("handshake signature invalid")
+		return nil, fmt.Errorf("handshake signature invalid")
 	}
 
 	// set remote peer key and id
 	s.remoteID = id
 	s.remoteKey = remotePubKey
-	return nil
+	return nhp.Data, nil
 }

p2p/security/noise/session.go

@@ -36,22 +36,24 @@ type secureSession struct {
 	dec *noise.CipherState
 
 	// noise prologue
-	prologue         []byte
-	earlyDataHandler EarlyDataHandler
+	prologue []byte
+
+	initiatorEarlyDataHandler, responderEarlyDataHandler EarlyDataHandler
 }
 
 // newSecureSession creates a Noise session over the given insecureConn Conn, using
 // the libp2p identity keypair from the given Transport.
-func newSecureSession(tpt *Transport, ctx context.Context, insecure net.Conn, remote peer.ID, prologue []byte, edh EarlyDataHandler, initiator bool) (*secureSession, error) {
+func newSecureSession(tpt *Transport, ctx context.Context, insecure net.Conn, remote peer.ID, prologue []byte, initiatorEDH, responderEDH EarlyDataHandler, initiator bool) (*secureSession, error) {
 	s := &secureSession{
-		insecureConn:     insecure,
-		insecureReader:   bufio.NewReader(insecure),
-		initiator:        initiator,
-		localID:          tpt.localID,
-		localKey:         tpt.privateKey,
-		remoteID:         remote,
-		prologue:         prologue,
-		earlyDataHandler: edh,
+		insecureConn:              insecure,
+		insecureReader:            bufio.NewReader(insecure),
+		initiator:                 initiator,
+		localID:                   tpt.localID,
+		localKey:                  tpt.privateKey,
+		remoteID:                  remote,
+		prologue:                  prologue,
+		initiatorEarlyDataHandler: initiatorEDH,
+		responderEarlyDataHandler: responderEDH,
 	}
 
 	// the go-routine we create to run the handshake will

p2p/security/noise/session_transport.go

@@ -22,18 +22,28 @@ func Prologue(prologue []byte) SessionOption {
 	}
 }
 
-// EarlyDataHandler allows attaching an (unencrypted) application payload to the first handshake message.
-// While unencrypted, the integrity of this early data is retroactively authenticated on completion of the handshake.
+// EarlyDataHandler defines what the application payload is for either the second
+// (if responder) or third (if initiator) handshake message, and defines the
+// logic for handling the other side's early data. Note the early data in the
+// second handshake message is encrypted, but the peer is not authenticated at that point.
 type EarlyDataHandler interface {
-	// Send is called for the client before sending the first handshake message.
+	// Send for the initiator is called for the client before sending the third
+	// handshake message. Defines the application payload for the third message.
+	// Send for the responder is called before sending the second handshake message.
 	Send(context.Context, net.Conn, peer.ID) []byte
-	// Received is called for the server when the first handshake message from the client is received.
+	// Received for the initiator is called when the second handshake message
+	// from the responder is received.
+	// Received for the responder is called when the third handshake message
+	// from the initiator is received.
 	Received(context.Context, net.Conn, []byte) error
 }
 
-func EarlyData(h EarlyDataHandler) SessionOption {
+// EarlyData sets the `EarlyDataHandler` for the initiator and responder roles.
+// See `EarlyDataHandler` for more details.
+func EarlyData(initiator, responder EarlyDataHandler) SessionOption {
 	return func(s *SessionTransport) error {
-		s.earlyDataHandler = h
+		s.initiatorEarlyDataHandler = initiator
+		s.responderEarlyDataHandler = responder
 		return nil
 	}
 }
@@ -45,14 +55,15 @@ var _ sec.SecureTransport = &SessionTransport{}
 type SessionTransport struct {
 	t *Transport
 	// options
-	prologue         []byte
-	earlyDataHandler EarlyDataHandler
+	prologue []byte
+
+	initiatorEarlyDataHandler, responderEarlyDataHandler EarlyDataHandler
 }
 
 // SecureInbound runs the Noise handshake as the responder.
 // If p is empty, connections from any peer are accepted.
 func (i *SessionTransport) SecureInbound(ctx context.Context, insecure net.Conn, p peer.ID) (sec.SecureConn, error) {
-	c, err := newSecureSession(i.t, ctx, insecure, p, i.prologue, i.earlyDataHandler, false)
+	c, err := newSecureSession(i.t, ctx, insecure, p, i.prologue, i.initiatorEarlyDataHandler, i.responderEarlyDataHandler, false)
 	if err != nil {
 		addr, maErr := manet.FromNetAddr(insecure.RemoteAddr())
 		if maErr == nil {
@@ -64,5 +75,5 @@ func (i *SessionTransport) SecureInbound(ctx context.Context, insecure net.Conn,
 
 // SecureOutbound runs the Noise handshake as the initiator.
 func (i *SessionTransport) SecureOutbound(ctx context.Context, insecure net.Conn, p peer.ID) (sec.SecureConn, error) {
-	return newSecureSession(i.t, ctx, insecure, p, i.prologue, i.earlyDataHandler, true)
+	return newSecureSession(i.t, ctx, insecure, p, i.prologue, i.initiatorEarlyDataHandler, i.responderEarlyDataHandler, true)
 }

p2p/security/noise/transport.go

@@ -41,7 +41,7 @@ func New(privkey crypto.PrivKey) (*Transport, error) {
 // SecureInbound runs the Noise handshake as the responder.
 // If p is empty, connections from any peer are accepted.
 func (t *Transport) SecureInbound(ctx context.Context, insecure net.Conn, p peer.ID) (sec.SecureConn, error) {
-	c, err := newSecureSession(t, ctx, insecure, p, nil, nil, false)
+	c, err := newSecureSession(t, ctx, insecure, p, nil, nil, nil, false)
 	if err != nil {
 		addr, maErr := manet.FromNetAddr(insecure.RemoteAddr())
 		if maErr == nil {
@@ -53,7 +53,7 @@ func (t *Transport) SecureInbound(ctx context.Context, insecure net.Conn, p peer
 
 // SecureOutbound runs the Noise handshake as the initiator.
 func (t *Transport) SecureOutbound(ctx context.Context, insecure net.Conn, p peer.ID) (sec.SecureConn, error) {
-	return newSecureSession(t, ctx, insecure, p, nil, nil, true)
+	return newSecureSession(t, ctx, insecure, p, nil, nil, nil, true)
 }
 
 func (t *Transport) WithSessionOptions(opts ...SessionOption) (sec.SecureTransport, error) {