More duplicated code reduction, interface improvement

lib/transport/ntcp/aead.go

@@ -7,6 +7,23 @@ import (
 	"golang.org/x/crypto/chacha20poly1305"
 )
 
+// AEADOperator defines the interface for AEAD operations in the NTCP2 protocol
+type AEADOperator interface {
+	// EncryptWithAssociatedData encrypts data using the provided key and associated data
+	EncryptWithAssociatedData(key, data, associatedData []byte, nonceCounter uint64) ([]byte, error)
+
+	// DecryptWithAssociatedData decrypts data using the provided key and associated data
+	DecryptWithAssociatedData(key, data, associatedData []byte, nonceCounter uint64) ([]byte, error)
+
+	// EncryptWithDerivedKey encrypts data, deriving the key from raw key material first
+	EncryptWithDerivedKey(keyMaterial, data, associatedData []byte, nonceCounter uint64) ([]byte, error)
+
+	// DecryptWithDerivedKey decrypts data, deriving the key from raw key material first
+	DecryptWithDerivedKey(keyMaterial, data, associatedData []byte, nonceCounter uint64) ([]byte, error)
+}
+
+var _ AEADOperator = (*NTCP2Session)(nil)
+
 // PerformAEADOperation handles both encryption and decryption using ChaCha20-Poly1305
 func (c *NTCP2Session) PerformAEADOperation(
 	keyMaterial []byte, // Raw key material to derive key from
@@ -48,3 +65,61 @@ func (c *NTCP2Session) PerformAEADOperation(
 
 	return result, nil
 }
+
+// EncryptWithAssociatedData encrypts data using ChaCha20-Poly1305 with the provided key and associated data
+func (c *NTCP2Session) EncryptWithAssociatedData(
+	key []byte,
+	data []byte,
+	associatedData []byte,
+	nonceCounter uint64,
+) ([]byte, error) {
+	return c.PerformAEADOperation(key, data, associatedData, nonceCounter, true)
+}
+
+// DecryptWithAssociatedData decrypts data using ChaCha20-Poly1305 with the provided key and associated data
+func (c *NTCP2Session) DecryptWithAssociatedData(
+	key []byte,
+	data []byte,
+	associatedData []byte,
+	nonceCounter uint64,
+) ([]byte, error) {
+	return c.PerformAEADOperation(key, data, associatedData, nonceCounter, false)
+}
+
+// PerformAEADWithDerivedKey performs AEAD operation, deriving the key from raw key material first
+func (c *NTCP2Session) PerformAEADWithDerivedKey(
+	keyMaterial []byte, // Raw key material to derive key from
+	data []byte,
+	associatedData []byte,
+	nonceCounter uint64,
+	encrypt bool,
+) ([]byte, error) {
+	// 1. Derive key using KDF
+	key, err := c.deriveChacha20Key(keyMaterial)
+	if err != nil {
+		return nil, oops.Errorf("failed to derive ChaCha20 key: %w", err)
+	}
+
+	// 2. Perform the AEAD operation
+	return c.PerformAEADOperation(key, data, associatedData, nonceCounter, encrypt)
+}
+
+// EncryptWithDerivedKey encrypts data, deriving the key from raw key material first
+func (c *NTCP2Session) EncryptWithDerivedKey(
+	keyMaterial []byte,
+	data []byte,
+	associatedData []byte,
+	nonceCounter uint64,
+) ([]byte, error) {
+	return c.PerformAEADWithDerivedKey(keyMaterial, data, associatedData, nonceCounter, true)
+}
+
+// DecryptWithDerivedKey decrypts data, deriving the key from raw key material first
+func (c *NTCP2Session) DecryptWithDerivedKey(
+	keyMaterial []byte,
+	data []byte,
+	associatedData []byte,
+	nonceCounter uint64,
+) ([]byte, error) {
+	return c.PerformAEADWithDerivedKey(keyMaterial, data, associatedData, nonceCounter, false)
+}

lib/transport/ntcp/handshake/handshake.go

@@ -30,6 +30,8 @@ type HandshakeState struct {
 	RemotePaddingLen int
 	// chachaKey is the derived ChaCha20 symmetric key for the session
 	ChachaKey []byte
+	// HandshakeHash is the cumulative hash of the handshake
+	HandshakeHash []byte
 	// sharedSecret is the Diffie-Hellman shared secret computed during handshake
 	SharedSecret []byte
 	// timestamp is the Unix timestamp when handshake was initiated

lib/transport/ntcp/session.go

@@ -10,6 +10,7 @@ import (
 	"github.com/go-i2p/go-i2p/lib/crypto/aes"
 	"github.com/go-i2p/go-i2p/lib/transport/noise"
 	"github.com/go-i2p/go-i2p/lib/transport/ntcp/handshake"
+	"github.com/go-i2p/go-i2p/lib/transport/ntcp/kdf"
 	"github.com/go-i2p/go-i2p/lib/transport/obfs"
 	"github.com/go-i2p/go-i2p/lib/transport/padding"
 	"github.com/go-i2p/go-i2p/lib/util/time/sntp"
@@ -37,6 +38,14 @@ import (
 type NTCP2Session struct {
 	*noise.NoiseSession
 	*NTCP2Transport
+	// Session keys for encrypted communication
+	inboundKey  []byte // Key for decrypting incoming messages
+	outboundKey []byte // Key for encrypting outgoing messages
+	// Keys for length obfuscation using SipHash
+	lengthEncryptKey1 []byte // First SipHash key (k1)
+	lengthEncryptKey2 []byte // Second SipHash key (k2)
+	// Key for frame obfuscation in data phase
+	framingKey      []byte
 	paddingStrategy padding.PaddingStrategy
 }
 
@@ -173,7 +182,72 @@ func (c *NTCP2Session) computeSharedSecret(ephemeralKey, param []byte) ([]byte,
 
 // deriveSessionKeys computes the session keys from the completed handshake
 func (c *NTCP2Session) deriveSessionKeys(hs *handshake.HandshakeState) error {
-	// Use shared secrets to derive session keys
-	// TODO: Implement key derivation according to NTCP2 spec
+	// Create KDF context if not already present
+	kdfContext := kdf.NewNTCP2KDF()
+
+	// If we have a handshake hash from the handshake state, use it
+	if len(hs.HandshakeHash) > 0 {
+		kdfContext.HandshakeHash = hs.HandshakeHash
+	}
+
+	// If we have a chaining key from the handshake state, use it
+	if len(hs.ChachaKey) > 0 {
+		kdfContext.ChainingKey = hs.ChachaKey
+	}
+
+	// Derive the final session keys for bidirectional communication
+	keyAB, keyBA, err := kdfContext.DeriveKeys()
+	if err != nil {
+		return oops.Errorf("failed to derive session keys: %w", err)
+	}
+
+	// Set the session keys based on whether we're the initiator or responder
+	if hs.IsInitiator {
+		// For initiator (Alice), outbound = Alice->Bob, inbound = Bob->Alice
+		c.outboundKey = keyAB
+		c.inboundKey = keyBA
+	} else {
+		// For responder (Bob), outbound = Bob->Alice, inbound = Alice->Bob
+		c.outboundKey = keyBA
+		c.inboundKey = keyAB
+	}
+
+	// Derive SipHash keys for length obfuscation
+	sipHashKey, err := kdfContext.DeriveSipHashKey()
+	if err != nil {
+		return oops.Errorf("failed to derive SipHash keys: %w", err)
+	}
+
+	// SipHash requires two 8-byte keys (k1, k2) and an 8-byte IV
+	// The sipHashKey is 16 bytes - first 8 bytes are k1, next 8 bytes are k2
+	if len(sipHashKey) < 16 {
+		return oops.Errorf("derived SipHash key too short: %d bytes", len(sipHashKey))
+	}
+
+	// Set up length obfuscation
+	c.lengthEncryptKey1 = sipHashKey[:8]
+	c.lengthEncryptKey2 = sipHashKey[8:16]
+
+	// Derive framing key for data phase
+	framingKey, err := kdfContext.DeriveFramingKey()
+	if err != nil {
+		return oops.Errorf("failed to derive framing key: %w", err)
+	}
+	c.framingKey = framingKey
+
+	// Clear sensitive key material from the KDF context
+	// to prevent leaking it in memory
+	for i := range kdfContext.ChainingKey {
+		kdfContext.ChainingKey[i] = 0
+	}
+
+	// For additional security, also clear the handshake state keys
+	// that are no longer needed
+	for i := range hs.ChachaKey {
+		hs.ChachaKey[i] = 0
+	}
+	hs.ChachaKey = nil
+
+	log.Debugf("NTCP2: Session keys derived successfully")
 	return nil
 }

lib/transport/ntcp/session_request.go

@@ -14,7 +14,6 @@ import (
 	"github.com/go-i2p/go-i2p/lib/transport/ntcp/kdf"
 	"github.com/go-i2p/go-i2p/lib/transport/ntcp/messages"
 	"github.com/samber/oops"
-	"golang.org/x/crypto/chacha20poly1305"
 )
 
 func (s *NTCP2Session) CreateSessionRequest() (*messages.SessionRequest, error) {
@@ -271,17 +270,8 @@ func (c *NTCP2Session) encryptSessionRequestOptions(
 	kdfContext.MixHash(obfuscatedX)
 
 	// Create AEAD cipher
-	aead, err := chacha20poly1305.New(chacha20Key)
-	if err != nil {
-		return nil, oops.Errorf("failed to create ChaCha20-Poly1305 cipher: %v", err)
-	}
-
-	// Prepare the nonce (all zeros for first message)
-	nonce := make([]byte, chacha20poly1305.NonceSize)
-
-	// Encrypt options block using associated data
 	optionsData := sessionRequestMessage.Options.Data()
-	ciphertext := aead.Seal(nil, nonce, optionsData, obfuscatedX)
+	ciphertext, err := c.EncryptWithAssociatedData(chacha20Key, optionsData, obfuscatedX, 0)
 
 	return ciphertext, nil
 }

lib/transport/ntcp/session_request_new.go

@@ -43,10 +43,17 @@ func (p *SessionRequestProcessor) EncryptPayload(
 ) ([]byte, error) {
 	req, ok := message.(*messages.SessionRequest)
 	if !ok {
-		return nil, oops.Errorf("expected SessionRequest message")
+		return nil, oops.Errorf("expected SessionRequest message, got %T", message)
 	}
 
-	return p.NTCP2Session.encryptSessionRequestOptions(req, obfuscatedKey)
+	// Use the central AEAD operation instead of custom encryption
+	// The key material would be derived in this case
+	return p.NTCP2Session.EncryptWithDerivedKey(
+		hs.LocalEphemeral.Bytes(),
+		req.Options.Data(),
+		obfuscatedKey, // Using obfuscated key as associated data
+		0,             // First message uses nonce counter 0
+	)
 }
 
 // MessageType implements handshake.HandshakeMessageProcessor.