Update

internal/crypto.go

@@ -6,6 +6,7 @@ import (
 	"crypto/ed25519"
 	"crypto/rand"
 	"crypto/sha256"
+	"errors"
 	"io"
 
 	"golang.org/x/crypto/argon2"
@@ -18,20 +19,12 @@ const (
 	nonceSize = 12
 )
 
-// IdentityKeyPair represents a user's long-term identity key
 type IdentityKeyPair struct {
 	PublicKey  [32]byte
 	PrivateKey [32]byte
 }
 
-// PrekeyBundle contains public keys for establishing a session
-type PrekeyBundle struct {
-	IdentityKey     [32]byte
-	Prekey          [32]byte
-	PrekeySignature []byte
-}
-
-// generateIdentityKeyPair creates a new Curve25519 key pair for identity
+// Creates a new Curve25519 key pair for identity
 func generateIdentityKeyPair() (*IdentityKeyPair, error) {
 	var privateKey [32]byte
 	if _, err := io.ReadFull(rand.Reader, privateKey[:]); err != nil {
@@ -47,7 +40,7 @@ func generateIdentityKeyPair() (*IdentityKeyPair, error) {
 	}, nil
 }
 
-// generateSignedPrekey creates a prekey signed with an Ed25519 signing key
+// Creates a prekey signed with an Ed25519 signing key
 func generateSignedPrekey() ([32]byte, []byte, error) {
 	var prekey [32]byte
 	var prekeyPriv [32]byte
@@ -64,23 +57,20 @@ func generateSignedPrekey() ([32]byte, []byte, error) {
 
 	// Sign the prekey
 	signature := ed25519.Sign(signingPriv, prekey[:])
-
-	// In practice, we'd store signingPub to verify later. For now, we include it in signature
 	combinedSig := append(signingPub, signature...)
 
 	return prekey, combinedSig, nil
 }
 
-// performDH does an ECDH key exchange
+// Does an ECDH key exchange
 func performDH(privateKey, publicKey [32]byte) ([32]byte, error) {
 	var sharedSecret [32]byte
 	curve25519.ScalarMult(&sharedSecret, &privateKey, &publicKey)
 	return sharedSecret, nil
 }
 
-// deriveSharedSecret combines multiple DH outputs to create a shared secret (simplified X3DH)
+// Combines multiple DH outputs to create a shared secret
 func deriveSharedSecret(dh1, dh2, dh3 [32]byte) []byte {
-	// KDF: hash all DH outputs together
 	h := sha256.New()
 	h.Write(dh1[:])
 	h.Write(dh2[:])
@@ -88,7 +78,7 @@ func deriveSharedSecret(dh1, dh2, dh3 [32]byte) []byte {
 	return h.Sum(nil)
 }
 
-// encryptUserKey encrypts a user's private key with their passphrase
+// Encrypts a user's private key with their passphrase
 func encryptUserKey(privateKey [32]byte, passphrase string) ([]byte, []byte, []byte, error) {
 	salt := make([]byte, saltSize)
 	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
@@ -104,7 +94,7 @@ func encryptUserKey(privateKey [32]byte, passphrase string) ([]byte, []byte, []b
 	return ciphertext, nonce, salt, nil
 }
 
-// decryptUserKey decrypts a user's private key with their passphrase
+// Decrypts a user's private key with their passphrase
 func decryptUserKey(ciphertext, nonce, salt []byte, passphrase string) ([32]byte, error) {
 	var privateKey [32]byte
 	key := deriveKey(passphrase, salt)
@@ -112,7 +102,7 @@ func decryptUserKey(ciphertext, nonce, salt []byte, passphrase string) ([32]byte
 	if err != nil {
 		return privateKey, err
 	}
-	copy(privateKey[:], plain)
+	copy(privateKey[:], []byte(plain))
 	return privateKey, nil
 }
 
@@ -152,3 +142,68 @@ func decryptMsg(ciphertext, nonce, key []byte) (string, error) {
 	}
 	return string(plain), nil
 }
+
+// Format: EphemeralPub (32) || Nonce (12) || Ciphertext (...)
+func EncryptKeyForUser(recipientPubKey [32]byte, payload []byte) ([]byte, error) {
+	var ephPriv, ephPub [32]byte
+	if _, err := io.ReadFull(rand.Reader, ephPriv[:]); err != nil {
+		return nil, err
+	}
+	curve25519.ScalarBaseMult(&ephPub, &ephPriv)
+
+	sharedSecret, _ := performDH(ephPriv, recipientPubKey)
+
+	kdf := sha256.Sum256(sharedSecret[:])
+	encryptionKey := kdf[:]
+
+	block, err := aes.NewCipher(encryptionKey)
+	if err != nil {
+		return nil, err
+	}
+	gcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, err
+	}
+
+	nonce := make([]byte, gcm.NonceSize())
+	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
+		return nil, err
+	}
+
+	ciphertext := gcm.Seal(nil, nonce, payload, nil)
+
+	out := make([]byte, 0, 32+len(nonce)+len(ciphertext))
+	out = append(out, ephPub[:]...)
+	out = append(out, nonce...)
+	out = append(out, ciphertext...)
+
+	return out, nil
+}
+
+// Decrypts a blob using the user's private identity key.
+func DecryptKeyForUser(myPrivKey [32]byte, blob []byte) ([]byte, error) {
+	if len(blob) < 32+12 {
+		return nil, errors.New("invalid key blob size")
+	}
+
+	var ephPub [32]byte
+	copy(ephPub[:], blob[:32])
+	nonce := blob[32 : 32+12]
+	ciphertext := blob[32+12:]
+
+	sharedSecret, _ := performDH(myPrivKey, ephPub)
+
+	kdf := sha256.Sum256(sharedSecret[:])
+	decryptionKey := kdf[:]
+
+	block, err := aes.NewCipher(decryptionKey)
+	if err != nil {
+		return nil, err
+	}
+	gcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, err
+	}
+
+	return gcm.Open(nil, nonce, ciphertext, nil)
+}