nip44: some more refactors and fix max length case.

fixes https://github.com/nbd-wtf/go-nostr/issues/134

nip44/nip44.go

@@ -25,24 +25,24 @@ const (
 )
 
 type encryptOptions struct {
-	err  error
+	err   error
-	salt []byte
+	nonce []byte
 }
 
-func WithCustomSalt(salt []byte) func(opts *encryptOptions) {
+// Deprecated: use WithCustomNonce instead of WithCustomSalt, so the naming is less confusing
+var WithCustomSalt = WithCustomNonce
+
+func WithCustomNonce(salt []byte) func(opts *encryptOptions) {
 	return func(opts *encryptOptions) {
 		if len(salt) != 32 {
 			opts.err = errors.New("salt must be 32 bytes")
 		}
-		opts.salt = salt
+		opts.nonce = salt
 	}
 }
 
 func Encrypt(plaintext string, conversationKey []byte, applyOptions ...func(opts *encryptOptions)) (string, error) {
-	opts := encryptOptions{
+	opts := encryptOptions{}
-		salt: nil,
-	}
-
 	for _, apply := range applyOptions {
 		apply(&opts)
 	}
@@ -51,53 +51,59 @@ func Encrypt(plaintext string, conversationKey []byte, applyOptions ...func(opts
 		return "", opts.err
 	}
 
-	salt := opts.salt
+	nonce := opts.nonce
-	if salt == nil {
+	if nonce == nil {
-		salt := make([]byte, 32)
+		nonce := make([]byte, 32)
-		if _, err := rand.Read(salt); err != nil {
+		if _, err := rand.Read(nonce); err != nil {
 			return "", err
 		}
 	}
 
-	enc, nonce, auth, err := messageKeys(conversationKey, salt)
+	enc, cc20nonce, auth, err := messageKeys(conversationKey, nonce)
 	if err != nil {
 		return "", err
 	}
 
-	padded, err := pad(plaintext)
+	plain := []byte(plaintext)
+	size := len(plain)
+	if size < MinPlaintextSize || size > MaxPlaintextSize {
+		return "", errors.New("plaintext should be between 1b and 64kB")
+	}
+
+	padding := calcPadding(size)
+	padded := make([]byte, 2+padding)
+	binary.BigEndian.PutUint16(padded, uint16(size))
+	copy(padded[2:], plain)
+
+	ciphertext, err := chacha(enc, cc20nonce, []byte(padded))
 	if err != nil {
 		return "", err
 	}
 
-	ciphertext, err := chacha20_(enc, nonce, []byte(padded))
+	mac, err := sha256Hmac(auth, ciphertext, nonce)
 	if err != nil {
 		return "", err
 	}
 
-	hmac_, err := sha256Hmac(auth, ciphertext, salt)
+	concat := make([]byte, 1+32+len(ciphertext)+32)
-	if err != nil {
-		return "", err
-	}
-
-	concat := make([]byte, 1+len(salt)+len(ciphertext)+len(hmac_))
 	concat[0] = version
-	copy(concat[1:], salt)
+	copy(concat[1:], nonce)
-	copy(concat[1+len(salt):], ciphertext)
+	copy(concat[1+32:], ciphertext)
-	copy(concat[1+len(salt)+len(ciphertext):], hmac_)
+	copy(concat[1+32+len(ciphertext):], mac)
 
 	return base64.StdEncoding.EncodeToString(concat), nil
 }
 
-func Decrypt(ciphertext string, conversationKey []byte) (string, error) {
+func Decrypt(b64ciphertextWrapped string, conversationKey []byte) (string, error) {
-	cLen := len(ciphertext)
+	cLen := len(b64ciphertextWrapped)
 	if cLen < 132 || cLen > 87472 {
 		return "", errors.New(fmt.Sprintf("invalid payload length: %d", cLen))
 	}
-	if ciphertext[0:1] == "#" {
+	if b64ciphertextWrapped[0:1] == "#" {
 		return "", errors.New("unknown version")
 	}
 
-	decoded, err := base64.StdEncoding.DecodeString(ciphertext)
+	decoded, err := base64.StdEncoding.DecodeString(b64ciphertextWrapped)
 	if err != nil {
 		return "", errors.New("invalid base64")
 	}
@@ -111,33 +117,33 @@ func Decrypt(ciphertext string, conversationKey []byte) (string, error) {
 		return "", errors.New(fmt.Sprintf("invalid data length: %d", dLen))
 	}
 
-	salt, ciphertext_, hmac_ := decoded[1:33], decoded[33:dLen-32], decoded[dLen-32:]
+	nonce, ciphertext, givenMac := decoded[1:33], decoded[33:dLen-32], decoded[dLen-32:]
-	enc, nonce, auth, err := messageKeys(conversationKey, salt)
+	enc, cc20nonce, auth, err := messageKeys(conversationKey, nonce)
 	if err != nil {
 		return "", err
 	}
 
-	hmac, err := sha256Hmac(auth, ciphertext_, salt)
+	expectedMac, err := sha256Hmac(auth, ciphertext, nonce)
 	if err != nil {
 		return "", err
 	}
 
-	if !bytes.Equal(hmac_, hmac) {
+	if !bytes.Equal(givenMac, expectedMac) {
 		return "", errors.New("invalid hmac")
 	}
 
-	padded, err := chacha20_(enc, nonce, ciphertext_)
+	padded, err := chacha(enc, cc20nonce, ciphertext)
 	if err != nil {
 		return "", err
 	}
 
 	unpaddedLen := binary.BigEndian.Uint16(padded[0:2])
-	if unpaddedLen < uint16(MinPlaintextSize) ||
+	if unpaddedLen < uint16(MinPlaintextSize) || unpaddedLen > uint16(MaxPlaintextSize) ||
-		unpaddedLen > uint16(MaxPlaintextSize) || len(padded) != 2+calcPadding(int(unpaddedLen)) {
+		len(padded) != 2+calcPadding(int(unpaddedLen)) {
 		return "", errors.New("invalid padding")
 	}
 
-	unpadded := padded[2 : unpaddedLen+2]
+	unpadded := padded[2:][:unpaddedLen]
 	if len(unpadded) == 0 || len(unpadded) != int(unpaddedLen) {
 		return "", errors.New("invalid padding")
 	}
@@ -157,7 +163,7 @@ func GenerateConversationKey(pub string, sk string) ([]byte, error) {
 	return hkdf.Extract(sha256.New, shared, []byte("nip44-v2")), nil
 }
 
-func chacha20_(key []byte, nonce []byte, message []byte) ([]byte, error) {
+func chacha(key []byte, nonce []byte, message []byte) ([]byte, error) {
 	cipher, err := chacha20.NewUnauthenticatedCipher(key, nonce)
 	if err != nil {
 		return nil, err
@@ -168,32 +174,32 @@ func chacha20_(key []byte, nonce []byte, message []byte) ([]byte, error) {
 	return dst, nil
 }
 
-func sha256Hmac(key []byte, ciphertext []byte, aad []byte) ([]byte, error) {
+func sha256Hmac(key []byte, ciphertext []byte, nonce []byte) ([]byte, error) {
-	if len(aad) != 32 {
+	if len(nonce) != 32 {
-		return nil, errors.New("aad data must be 32 bytes")
+		return nil, errors.New("nonce aad must be 32 bytes")
 	}
 	h := hmac.New(sha256.New, key)
-	h.Write(aad)
+	h.Write(nonce)
 	h.Write(ciphertext)
 	return h.Sum(nil), nil
 }
 
-func messageKeys(conversationKey []byte, salt []byte) ([]byte, []byte, []byte, error) {
+func messageKeys(conversationKey []byte, nonce []byte) ([]byte, []byte, []byte, error) {
 	if len(conversationKey) != 32 {
 		return nil, nil, nil, errors.New("conversation key must be 32 bytes")
 	}
-	if len(salt) != 32 {
+	if len(nonce) != 32 {
-		return nil, nil, nil, errors.New("salt must be 32 bytes")
+		return nil, nil, nil, errors.New("nonce must be 32 bytes")
 	}
 
-	r := hkdf.Expand(sha256.New, conversationKey, salt)
+	r := hkdf.Expand(sha256.New, conversationKey, nonce)
 	enc := make([]byte, 32)
 	if _, err := io.ReadFull(r, enc); err != nil {
 		return nil, nil, nil, err
 	}
 
-	nonce := make([]byte, 12)
+	cc20nonce := make([]byte, 12)
-	if _, err := io.ReadFull(r, nonce); err != nil {
+	if _, err := io.ReadFull(r, cc20nonce); err != nil {
 		return nil, nil, nil, err
 	}
 
@@ -202,21 +208,7 @@ func messageKeys(conversationKey []byte, salt []byte) ([]byte, []byte, []byte, e
 		return nil, nil, nil, err
 	}
 
-	return enc, nonce, auth, nil
+	return enc, cc20nonce, auth, nil
-}
-
-func pad(s string) ([]byte, error) {
-	sb := []byte(s)
-	sbLen := len(sb)
-	if sbLen < 1 || sbLen > MaxPlaintextSize {
-		return nil, errors.New("plaintext should be between 1b and 64kB")
-	}
-	padding := calcPadding(sbLen)
-	result := make([]byte, 2)
-	binary.BigEndian.PutUint16(result, uint16(sbLen))
-	result = append(result, sb...)
-	result = append(result, make([]byte, padding-sbLen)...)
-	return result, nil
 }
 
 func calcPadding(sLen int) int {

nip44/nip44_test.go

@@ -1,9 +1,12 @@
 package nip44
 
 import (
+	"crypto/rand"
 	"crypto/sha256"
 	"encoding/hex"
+	"fmt"
 	"hash"
+	"strings"
 	"testing"
 
 	"github.com/nbd-wtf/go-nostr"
@@ -30,41 +33,7 @@ func assertCryptPriv(t *testing.T, sk1 string, sk2 string, conversationKey strin
 	if ok = assert.NoErrorf(t, err, "hex decode failed for salt: %v", err); !ok {
 		return
 	}
-	actual, err = Encrypt(plaintext, k1, WithCustomSalt(s))
+	actual, err = Encrypt(plaintext, k1, WithCustomNonce(s))
-	if ok = assert.NoError(t, err, "encryption failed: %v", err); !ok {
-		return
-	}
-	if ok = assert.Equalf(t, expected, actual, "wrong encryption"); !ok {
-		return
-	}
-	decrypted, err = Decrypt(expected, k1)
-	if ok = assert.NoErrorf(t, err, "decryption failed: %v", err); !ok {
-		return
-	}
-	assert.Equal(t, decrypted, plaintext, "wrong decryption")
-}
-
-func assertCryptPub(t *testing.T, sk1 string, pub2 string, conversationKey string, salt string, plaintext string, expected string) {
-	var (
-		k1        []byte
-		s         []byte
-		actual    string
-		decrypted string
-		ok        bool
-		err       error
-	)
-	k1, err = hex.DecodeString(conversationKey)
-	if ok = assert.NoErrorf(t, err, "hex decode failed for conversation key: %v", err); !ok {
-		return
-	}
-	if ok = assertConversationKeyGenerationPub(t, sk1, pub2, conversationKey); !ok {
-		return
-	}
-	s, err = hex.DecodeString(salt)
-	if ok = assert.NoErrorf(t, err, "hex decode failed for salt: %v", err); !ok {
-		return
-	}
-	actual, err = Encrypt(plaintext, k1, WithCustomSalt(s))
 	if ok = assert.NoError(t, err, "encryption failed: %v", err); !ok {
 		return
 	}
@@ -209,7 +178,7 @@ func assertCryptLong(t *testing.T, conversationKey string, salt string, pattern
 	if ok = assert.Equalf(t, plaintextSha256, actualPlaintextSha256, "invalid plaintext sha256 hash: %v", err); !ok {
 		return
 	}
-	actualPayload, err = Encrypt(plaintext, convKey, WithCustomSalt(convSalt))
+	actualPayload, err = Encrypt(plaintext, convKey, WithCustomNonce(convSalt))
 	if ok = assert.NoErrorf(t, err, "encryption failed: %v", err); !ok {
 		return
 	}
@@ -1169,3 +1138,60 @@ func TestMessageKeyGeneration033(t *testing.T) {
 		"47b89da97f68d389867b5d8a2d7ba55715a30e3d88a3cc11f3646bc2af5580ef",
 	)
 }
+
+func TestMaxLength(t *testing.T) {
+	sk1 := nostr.GeneratePrivateKey()
+	sk2 := nostr.GeneratePrivateKey()
+	pub2, _ := nostr.GetPublicKey(sk2)
+	salt := make([]byte, 32)
+	rand.Read(salt)
+	conversationKey, _ := GenerateConversationKey(pub2, sk1)
+	plaintext := strings.Repeat("a", MaxPlaintextSize)
+	encrypted, err := Encrypt(plaintext, conversationKey, WithCustomNonce(salt))
+	if err != nil {
+		t.Error(err)
+	}
+
+	assertCryptPub(t,
+		sk1,
+		pub2,
+		fmt.Sprintf("%x", conversationKey),
+		fmt.Sprintf("%x", salt),
+		plaintext,
+		encrypted,
+	)
+}
+
+func assertCryptPub(t *testing.T, sk1 string, pub2 string, conversationKey string, salt string, plaintext string, expected string) {
+	var (
+		k1        []byte
+		s         []byte
+		actual    string
+		decrypted string
+		ok        bool
+		err       error
+	)
+	k1, err = hex.DecodeString(conversationKey)
+	if ok = assert.NoErrorf(t, err, "hex decode failed for conversation key: %v", err); !ok {
+		return
+	}
+	if ok = assertConversationKeyGenerationPub(t, sk1, pub2, conversationKey); !ok {
+		return
+	}
+	s, err = hex.DecodeString(salt)
+	if ok = assert.NoErrorf(t, err, "hex decode failed for salt: %v", err); !ok {
+		return
+	}
+	actual, err = Encrypt(plaintext, k1, WithCustomNonce(s))
+	if ok = assert.NoError(t, err, "encryption failed: %v", err); !ok {
+		return
+	}
+	if ok = assert.Equalf(t, expected, actual, "wrong encryption"); !ok {
+		return
+	}
+	decrypted, err = Decrypt(expected, k1)
+	if ok = assert.NoErrorf(t, err, "decryption failed: %v", err); !ok {
+		return
+	}
+	assert.Equal(t, decrypted, plaintext, "wrong decryption")
+}