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

p2p/security/noise/handshake.go

@@ -78,11 +78,7 @@ func (s *secureSession) runHandshake(ctx context.Context) (err error) {
 	if s.initiator {
 		// stage 0 //
 		// Handshake Msg Len = len(DH ephemeral key)
-		var ed []byte
-		if s.initiatorEarlyDataHandler != nil {
-			ed = s.initiatorEarlyDataHandler.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)
 		}
 
@@ -103,7 +99,11 @@ func (s *secureSession) runHandshake(ctx context.Context) (err error) {
 
 		// stage 2 //
 		// Handshake Msg Len = len(DHT static key) +  MAC(static key is encrypted) + len(Payload) + MAC(payload is encrypted)
-		payload, err := s.generateHandshakePayload(kp, nil)
+		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
 		}
@@ -113,15 +113,9 @@ func (s *secureSession) runHandshake(ctx context.Context) (err error) {
 		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.responderEarlyDataHandler != nil {
-			if err := s.responderEarlyDataHandler.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) +
@@ -143,9 +137,16 @@ func (s *secureSession) runHandshake(ctx context.Context) (err error) {
 		if err != nil {
 			return fmt.Errorf("error reading handshake message: %w", err)
 		}
-		// we don't expect any early data on this message
-		_, err = s.handleRemoteHandshakePayload(plaintext, hs.PeerStatic())
-		return err
+		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
 	}
 }
 

p2p/security/noise/session_transport.go

@@ -22,27 +22,24 @@ func Prologue(prologue []byte) SessionOption {
 	}
 }
 
-// EarlyDataHandler defines what the application payload is for either the first
-// (if initiator) or second (if responder) handshake message. And defines the
+// EarlyDataHandler defines what the application payload is for either the second
+// (if initiator) or third (if responder) handshake message, and defines the
 // logic for handling the other side's early data. Note the early data in the
-// first handshake message is **unencrypted**, but will be retroactively
-// authenticated if the handshake completes.
+// second handshake message is encrypted, but the peer is not authenticated at that point.
 type EarlyDataHandler interface {
-	// Send for the initiator is called for the client before sending the first
-	// handshake message. Defines the application payload for the first message.
-	// This payload is sent **unencrypted**.
-	// Send for the responder is called before sending the second handshake message. This is encrypted.
+	// 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 for the initiator is called when the second handshake message
 	// from the responder is received.
-	// Received for the responder is called when the first handshake message
+	// Received for the responder is called when the third handshake message
 	// from the initiator is received.
 	Received(context.Context, net.Conn, []byte) error
 }
 
 // EarlyData sets the `EarlyDataHandler` for the initiator and responder roles.
-// See `EarlyDataHandler` for more details. Note: an initiator's early data will
-// be sent **unencrypted** in the first handshake message.
+// See `EarlyDataHandler` for more details.
 func EarlyData(initiator, responder EarlyDataHandler) SessionOption {
 	return func(s *SessionTransport) error {
 		s.initiatorEarlyDataHandler = initiator

p2p/security/noise/transport_test.go

@@ -465,6 +465,12 @@ func TestEarlyDataAccepted(t *testing.T) {
 
 		conn, err := initTransport.SecureOutbound(context.Background(), respConn, tpt.localID)
 		require.NoError(t, err)
+		select {
+		case <-time.After(500 * time.Millisecond):
+			t.Fatal("timeout")
+		case err := <-errChan:
+			require.NoError(t, err)
+		}
 		defer conn.Close()
 	}
 
@@ -494,7 +500,6 @@ func TestEarlyDataAccepted(t *testing.T) {
 
 func TestEarlyDataRejected(t *testing.T) {
 	handshake := func(t *testing.T, client, server EarlyDataHandler) (clientErr, serverErr error) {
-		t.Helper()
 		initTransport, err := newTestTransport(t, crypto.Ed25519, 2048).WithSessionOptions(EarlyData(client, nil))
 		require.NoError(t, err)
 		tpt := newTestTransport(t, crypto.Ed25519, 2048)
@@ -509,7 +514,13 @@ func TestEarlyDataRejected(t *testing.T) {
 			errChan <- err
 		}()
 
-		_, clientErr = initTransport.SecureOutbound(context.Background(), respConn, tpt.localID)
+		// As early data is sent with the last handshake message, the handshake will appear
+		// to succeed for the client.
+		var conn sec.SecureConn
+		conn, clientErr = initTransport.SecureOutbound(context.Background(), respConn, tpt.localID)
+		if clientErr == nil {
+			_, clientErr = conn.Read([]byte{0})
+		}
 
 		select {
 		case <-time.After(500 * time.Millisecond):
@@ -541,7 +552,7 @@ func TestEarlyDataRejected(t *testing.T) {
 	})
 }
 
-func TestEarlyDataAcceptedWithNoHandler(t *testing.T) {
+func TestEarlyfffDataAcceptedWithNoHandler(t *testing.T) {
 	clientEDH := &earlyDataHandler{
 		send: func(ctx context.Context, conn net.Conn, id peer.ID) []byte { return []byte("foobar") },
 	}