1 files changed, 51 insertions, 8 deletions

diff --git a/crypto.go b/crypto.go
index 402d06b..2d8605c 100644
--- a/crypto.go
+++ b/crypto.go
@@ -3,13 +3,16 @@ package axon
 import (
         "crypto/ed25519"
         "crypto/rand"
+        "crypto/sha256"
         "crypto/sha512"
         "errors"
         "fmt"
+        "io"
         "math/big"
 
         "golang.org/x/crypto/chacha20poly1305"
         "golang.org/x/crypto/curve25519"
+        "golang.org/x/crypto/hkdf"
 )
 
 // p is the prime for Curve25519: 2^255 - 19.
@@ -59,7 +62,9 @@ func Ed25519PubKeyToX25519(pubKey ed25519.PublicKey) ([32]byte, error) {
 
         den := new(big.Int).Sub(one, y)
         den.Mod(den, p)
-        den.ModInverse(den, p)
+        if den.ModInverse(den, p) == nil {
+                return [32]byte{}, errors.New("axon: degenerate public key (y=1)")
+        }
 
         u := new(big.Int).Mul(num, den)
         u.Mod(u, p)
@@ -91,16 +96,43 @@ func DHSharedSecret(localPriv ed25519.PrivateKey, remotePub ed25519.PublicKey) (
         return out, nil
 }
 
+// dmKey derives a symmetric key from the raw DH shared secret using
+// HKDF-SHA256. The salt is nil (HKDF uses a zero-filled hash-length salt)
+// and the info string binds the key to the axon DM context.
+func dmKey(shared [32]byte) ([32]byte, error) {
+        r := hkdf.New(sha256.New, shared[:], nil, []byte("axon-dm-v1"))
+        var key [32]byte
+        if _, err := io.ReadFull(r, key[:]); err != nil {
+                return [32]byte{}, fmt.Errorf("axon: hkdf: %w", err)
+        }
+        return key, nil
+}
+
+// dmAD builds the associated data for DM encryption by concatenating the
+// sender and recipient Ed25519 public keys, preventing key-confusion attacks.
+func dmAD(senderPub, recipientPub ed25519.PublicKey) []byte {
+        ad := make([]byte, 0, len(senderPub)+len(recipientPub))
+        ad = append(ad, senderPub...)
+        ad = append(ad, recipientPub...)
+        return ad
+}
+
 // EncryptDM encrypts plaintext for a DM using ChaCha20-Poly1305. The
-// encryption key is the X25519 shared secret derived from senderPriv and
-// recipientPub. The returned blob is: nonce (12 bytes) || ciphertext.
+// encryption key is derived via HKDF-SHA256 from the X25519 shared secret
+// of senderPriv and recipientPub. The sender and recipient public keys are
+// bound as associated data. The returned blob is: nonce (12 bytes) || ciphertext.
 func EncryptDM(senderPriv ed25519.PrivateKey, recipientPub ed25519.PublicKey, plaintext []byte) ([]byte, error) {
         shared, err := DHSharedSecret(senderPriv, recipientPub)
         if err != nil {
                 return nil, fmt.Errorf("axon: dh: %w", err)
         }
 
-        aead, err := chacha20poly1305.New(shared[:])
+        key, err := dmKey(shared)
+        if err != nil {
+                return nil, err
+        }
+
+        aead, err := chacha20poly1305.New(key[:])
         if err != nil {
                 return nil, fmt.Errorf("axon: create aead: %w", err)
         }
@@ -110,7 +142,10 @@ func EncryptDM(senderPriv ed25519.PrivateKey, recipientPub ed25519.PublicKey, pl
                 return nil, fmt.Errorf("axon: generate nonce: %w", err)
         }
 
-        ct := aead.Seal(nil, nonce, plaintext, nil)
+        senderPub := senderPriv.Public().(ed25519.PublicKey)
+        ad := dmAD(senderPub, recipientPub)
+
+        ct := aead.Seal(nil, nonce, plaintext, ad)
         return append(nonce, ct...), nil
 }
 
@@ -123,7 +158,12 @@ func DecryptDM(recipientPriv ed25519.PrivateKey, senderPub ed25519.PublicKey, bl
                 return nil, fmt.Errorf("axon: dh: %w", err)
         }
 
-        aead, err := chacha20poly1305.New(shared[:])
+        key, err := dmKey(shared)
+        if err != nil {
+                return nil, err
+        }
+
+        aead, err := chacha20poly1305.New(key[:])
         if err != nil {
                 return nil, fmt.Errorf("axon: create aead: %w", err)
         }
@@ -136,9 +176,12 @@ func DecryptDM(recipientPriv ed25519.PrivateKey, senderPub ed25519.PublicKey, bl
         nonce := blob[:nonceSize]
         ct := blob[nonceSize:]
 
-        pt, err := aead.Open(nil, nonce, ct, nil)
+        recipientPub := recipientPriv.Public().(ed25519.PublicKey)
+        ad := dmAD(senderPub, recipientPub)
+
+        pt, err := aead.Open(nil, nonce, ct, ad)
         if err != nil {
-                return nil, fmt.Errorf("axon: decrypt: authentication failed")
+                return nil, fmt.Errorf("axon: decrypt: %w", err)
         }
         return pt, nil
 }