From 78fd54db66702498a661d27e3666bcb2a499fc21 Mon Sep 17 00:00:00 2001
From: pseusys <shveitsar215@gmail.com>
Date: Sun, 11 Feb 2024 12:35:52 +0100
Subject: [PATCH] test added, tag size fixed

---
 ecies_test.go | 13 ++++++++++---
 symm.go       | 38 +++++++++++++++++---------------------
 2 files changed, 27 insertions(+), 24 deletions(-)

diff --git a/ecies_test.go b/ecies_test.go
index c0842b9..950f0e6 100644
--- a/ecies_test.go
+++ b/ecies_test.go
@@ -56,15 +56,15 @@ func BenchmarkDecrypt(b *testing.B) {
 	}
 }
 
-func TestEncryptAndDecrypt(t *testing.T) {
+func testEncryptAndDecryptParameters(conf Config, t *testing.T) {
 	privkey := NewPrivateKeyFromBytes(testingReceiverPrivkey)
 
-	ciphertext, err := Encrypt(privkey.PublicKey, []byte(testingMessage))
+	ciphertext, err := EncryptConf(privkey.PublicKey, []byte(testingMessage), conf)
 	if !assert.NoError(t, err) {
 		return
 	}
 
-	plaintext, err := Decrypt(privkey, ciphertext)
+	plaintext, err := DecryptConf(privkey, ciphertext, conf)
 	if !assert.NoError(t, err) {
 		return
 	}
@@ -72,6 +72,13 @@ func TestEncryptAndDecrypt(t *testing.T) {
 	assert.Equal(t, testingMessage, string(plaintext))
 }
 
+func TestEncryptAndDecrypt(t *testing.T) {
+	testEncryptAndDecryptParameters(DEFAULT_CONFIG, t)
+	testEncryptAndDecryptParameters(Config{symmetricAlgorithm: "aes-256-gcm", symmetricNonceLength: 12}, t)
+	testEncryptAndDecryptParameters(Config{symmetricAlgorithm: "aes-256-gcm", symmetricNonceLength: 16}, t)
+	testEncryptAndDecryptParameters(Config{symmetricAlgorithm: "xchacha20"}, t)
+}
+
 func TestPublicKeyDecompression(t *testing.T) {
 	// Generate public key
 	privkey, err := GenerateKey()
diff --git a/symm.go b/symm.go
index df89996..e1f1dd5 100644
--- a/symm.go
+++ b/symm.go
@@ -10,47 +10,42 @@ import (
 	"golang.org/x/crypto/chacha20poly1305"
 )
 
-func generateSymmCipher(key []byte, conf Config) (cipher.AEAD, int, error) {
+func generateSymmCipher(key []byte, conf Config) (cipher.AEAD, error) {
 	var err error
 	var aead cipher.AEAD
-	var nonceLength int
 
 	switch conf.symmetricAlgorithm {
 	case "aes-256-gcm":
-		nonceLength = conf.symmetricNonceLength
-
 		block, err := aes.NewCipher(key)
 		if err != nil {
-			return nil, -1, fmt.Errorf("cannot create new AES block: %w", err)
+			return nil, fmt.Errorf("cannot create new AES block: %w", err)
 		}
 
-		aead, err = cipher.NewGCMWithNonceSize(block, nonceLength)
+		aead, err = cipher.NewGCMWithNonceSize(block, conf.symmetricNonceLength)
 		if err != nil {
-			return nil, -1, fmt.Errorf("cannot create AES GCM: %w", err)
+			return nil, fmt.Errorf("cannot create AES GCM: %w", err)
 		}
 	case "xchacha20":
 		aead, err = chacha20poly1305.NewX(key)
 		if err != nil {
-			return nil, -1, fmt.Errorf("cannot create XChaCha20: %w", err)
+			return nil, fmt.Errorf("cannot create XChaCha20: %w", err)
 		}
-
-		nonceLength = aead.NonceSize()
 	default:
-		return nil, -1, fmt.Errorf("unknown cipher: %s", conf.symmetricAlgorithm)
+		return nil, fmt.Errorf("unknown cipher: %s", conf.symmetricAlgorithm)
 	}
 
-	return aead, nonceLength, nil
+	return aead, nil
 }
 
 func EncryptSymm(key []byte, msg []byte, conf Config) ([]byte, error) {
 	var ct bytes.Buffer
 
-	aead, nonceLength, err := generateSymmCipher(key, conf)
+	aead, err := generateSymmCipher(key, conf)
 	if err != nil {
 		return nil, err
 	}
 
-	nonce := make([]byte, nonceLength)
+	nonce := make([]byte, aead.NonceSize())
 	if _, err := rand.Read(nonce); err != nil {
 		return nil, fmt.Errorf("cannot read random bytes for nonce: %w", err)
 	}
@@ -59,7 +54,7 @@ func EncryptSymm(key []byte, msg []byte, conf Config) ([]byte, error) {
 
 	ciphertext := aead.Seal(nil, nonce, msg, nil)
 
-	tag := ciphertext[len(ciphertext)-aead.NonceSize():]
+	tag := ciphertext[len(ciphertext)-aead.Overhead():]
 	ct.Write(tag)
 	ciphertext = ciphertext[:len(ciphertext)-len(tag)]
 	ct.Write(ciphertext)
@@ -68,22 +63,23 @@ func EncryptSymm(key []byte, msg []byte, conf Config) ([]byte, error) {
 }
 
 func DecryptSymm(key []byte, msg []byte, conf Config) ([]byte, error) {
-	aead, nonceLength, err := generateSymmCipher(key, conf)
+	aead, err := generateSymmCipher(key, conf)
 	if err != nil {
 		return nil, err
 	}
 
 	// Message cannot be less than length of public key (65) + nonce + tag (16)
-	if len(msg) <= (nonceLength + 16) {
+	if len(msg) <= (aead.NonceSize() + aead.Overhead()) {
 		return nil, fmt.Errorf("invalid length of message")
 	}
 
-	// AES decryption part
-	nonce := msg[:nonceLength]
-	tag := msg[nonceLength : nonceLength+16]
+	// Symmetrical decryption part
+	nonce := msg[:aead.NonceSize()]
+	tag := msg[aead.NonceSize() : aead.NonceSize()+aead.Overhead()]
+	msg = msg[aead.NonceSize()+aead.Overhead():]
 
 	// Create Golang-accepted ciphertext
-	ciphertext := bytes.Join([][]byte{msg[nonceLength+16:], tag}, nil)
+	ciphertext := bytes.Join([][]byte{msg, tag}, nil)
 
 	plaintext, err := aead.Open(nil, nonce, ciphertext, nil)
 	if err != nil {