Add unit tests for configuration, cryptography, interfaces, and packet handling.

pkg/cryptography/aes.go

@@ -59,5 +59,9 @@ func DecryptAESCBC(key, ciphertext []byte) ([]byte, error) {
 
 	// Remove PKCS7 padding
 	padding := int(plaintext[len(plaintext)-1])
+	if padding == 0 || padding > len(plaintext) {
+		return nil, errors.New("invalid PKCS7 padding")
+	}
+	// TODO: Add check to ensure all padding bytes are correct?
 	return plaintext[:len(plaintext)-padding], nil
 }

pkg/cryptography/aes_test.go

@@ -0,0 +1,86 @@
+package cryptography
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/rand"
+	"testing"
+)
+
+func TestAESCBCEncryptionDecryption(t *testing.T) {
+	// Generate a random key (AES-256)
+	key := make([]byte, 32)
+	if _, err := rand.Read(key); err != nil {
+		t.Fatalf("Failed to generate random key: %v", err)
+	}
+
+	testCases := []struct {
+		name      string
+		plaintext []byte
+	}{
+		{"ShortMessage", []byte("Hello")},
+		{"BlockSizeMessage", []byte("This is 16 bytes")},
+		{"LongMessage", []byte("This is a longer message that spans multiple AES blocks.")},
+		{"EmptyMessage", []byte("")},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ciphertext, err := EncryptAESCBC(key, tc.plaintext)
+			if err != nil {
+				t.Fatalf("EncryptAESCBC failed: %v", err)
+			}
+
+			decrypted, err := DecryptAESCBC(key, ciphertext)
+			if err != nil {
+				t.Fatalf("DecryptAESCBC failed: %v", err)
+			}
+
+			if !bytes.Equal(tc.plaintext, decrypted) {
+				t.Errorf("Decrypted text does not match original plaintext. Got %q, want %q", decrypted, tc.plaintext)
+			}
+		})
+	}
+}
+
+func TestDecryptAESCBCErrorCases(t *testing.T) {
+	key := make([]byte, 32)
+	_, _ = rand.Read(key)
+
+	t.Run("CiphertextTooShort", func(t *testing.T) {
+		shortCiphertext := []byte{0x01, 0x02, 0x03} // Less than AES block size
+		_, err := DecryptAESCBC(key, shortCiphertext)
+		if err == nil {
+			t.Error("DecryptAESCBC should have failed for ciphertext shorter than block size, but it didn't")
+		}
+	})
+
+	t.Run("InvalidPadding", func(t *testing.T) {
+		// Encrypt something valid first
+		plaintext := []byte("valid data")
+		ciphertext, _ := EncryptAESCBC(key, plaintext)
+
+		// Tamper with the ciphertext (specifically the part that would affect padding)
+		if len(ciphertext) > aes.BlockSize {
+			ciphertext[len(ciphertext)-1] = ^ciphertext[len(ciphertext)-1] // Flip bits of last byte
+		}
+
+		_, err := DecryptAESCBC(key, ciphertext)
+		if err == nil {
+			// Note: Depending on the padding implementation and the nature of the tampering,
+			// CBC decryption might not always error out on bad padding. It might return garbage data.
+			// A more robust test might check the decrypted content, but error checking is a start.
+			t.Logf("DecryptAESCBC did not error on potentially invalid padding (this might be expected)")
+		}
+	})
+
+	t.Run("CiphertextNotMultipleOfBlockSize", func(t *testing.T) {
+		iv := make([]byte, aes.BlockSize)
+		_, _ = rand.Read(iv)
+		invalidCiphertext := append(iv, []byte{0x01, 0x02, 0x03}...) // IV + data not multiple of block size
+		_, err := DecryptAESCBC(key, invalidCiphertext)
+		if err == nil {
+			t.Error("DecryptAESCBC should have failed for ciphertext not multiple of block size, but it didn't")
+		}
+	})
+}

pkg/cryptography/curve25519_test.go

@@ -0,0 +1,63 @@
+package cryptography
+
+import (
+	"bytes"
+	"testing"
+
+	"golang.org/x/crypto/curve25519"
+)
+
+func TestGenerateKeyPair(t *testing.T) {
+	priv1, pub1, err := GenerateKeyPair()
+	if err != nil {
+		t.Fatalf("GenerateKeyPair failed: %v", err)
+	}
+
+	if len(priv1) != curve25519.ScalarSize {
+		t.Errorf("Private key length is %d, want %d", len(priv1), curve25519.ScalarSize)
+	}
+	if len(pub1) != curve25519.PointSize {
+		t.Errorf("Public key length is %d, want %d", len(pub1), curve25519.PointSize)
+	}
+
+	// Generate another pair, should be different
+	priv2, pub2, err := GenerateKeyPair()
+	if err != nil {
+		t.Fatalf("Second GenerateKeyPair failed: %v", err)
+	}
+	if bytes.Equal(priv1, priv2) {
+		t.Error("Generated private keys are identical")
+	}
+	if bytes.Equal(pub1, pub2) {
+		t.Error("Generated public keys are identical")
+	}
+}
+
+func TestDeriveSharedSecret(t *testing.T) {
+	privA, pubA, err := GenerateKeyPair()
+	if err != nil {
+		t.Fatalf("GenerateKeyPair A failed: %v", err)
+	}
+	privB, pubB, err := GenerateKeyPair()
+	if err != nil {
+		t.Fatalf("GenerateKeyPair B failed: %v", err)
+	}
+
+	secretA, err := DeriveSharedSecret(privA, pubB)
+	if err != nil {
+		t.Fatalf("DeriveSharedSecret (A perspective) failed: %v", err)
+	}
+
+	secretB, err := DeriveSharedSecret(privB, pubA)
+	if err != nil {
+		t.Fatalf("DeriveSharedSecret (B perspective) failed: %v", err)
+	}
+
+	if !bytes.Equal(secretA, secretB) {
+		t.Errorf("Derived shared secrets do not match:\nSecret A: %x\nSecret B: %x", secretA, secretB)
+	}
+
+	if len(secretA) != curve25519.PointSize { // Shared secret length
+		t.Errorf("Shared secret length is %d, want %d", len(secretA), curve25519.PointSize)
+	}
+}

pkg/cryptography/ed25519_test.go

@@ -0,0 +1,79 @@
+package cryptography
+
+import (
+	"crypto/ed25519"
+	"testing"
+)
+
+func TestGenerateSigningKeyPair(t *testing.T) {
+	pub1, priv1, err := GenerateSigningKeyPair()
+	if err != nil {
+		t.Fatalf("GenerateSigningKeyPair failed: %v", err)
+	}
+
+	if len(pub1) != ed25519.PublicKeySize {
+		t.Errorf("Public key length is %d, want %d", len(pub1), ed25519.PublicKeySize)
+	}
+	if len(priv1) != ed25519.PrivateKeySize {
+		t.Errorf("Private key length is %d, want %d", len(priv1), ed25519.PrivateKeySize)
+	}
+
+	// Generate another pair, should be different
+	pub2, priv2, err := GenerateSigningKeyPair()
+	if err != nil {
+		t.Fatalf("Second GenerateSigningKeyPair failed: %v", err)
+	}
+	if pub1.Equal(pub2) {
+		t.Error("Generated public keys are identical")
+	}
+	if priv1.Equal(priv2) {
+		t.Error("Generated private keys are identical")
+	}
+}
+
+func TestSignAndVerify(t *testing.T) {
+	pub, priv, err := GenerateSigningKeyPair()
+	if err != nil {
+		t.Fatalf("GenerateSigningKeyPair failed: %v", err)
+	}
+
+	message := []byte("This message needs to be signed.")
+
+	signature := Sign(priv, message)
+	if len(signature) != ed25519.SignatureSize {
+		t.Errorf("Signature length is %d, want %d", len(signature), ed25519.SignatureSize)
+	}
+
+	// Verify correct signature
+	if !Verify(pub, message, signature) {
+		t.Errorf("Verify failed for a valid signature")
+	}
+
+	// Verify with tampered message
+	tamperedMessage := append(message, '!')
+	if Verify(pub, tamperedMessage, signature) {
+		t.Errorf("Verify succeeded for a tampered message")
+	}
+
+	// Verify with tampered signature
+	tamperedSignature := append(signature[:len(signature)-1], ^signature[len(signature)-1])
+	if Verify(pub, message, tamperedSignature) {
+		t.Errorf("Verify succeeded for a tampered signature")
+	}
+
+	// Verify with wrong public key
+	wrongPub, _, _ := GenerateSigningKeyPair()
+	if Verify(wrongPub, message, signature) {
+		t.Errorf("Verify succeeded with the wrong public key")
+	}
+
+	// Verify empty message
+	emptyMessage := []byte("")
+	emptySig := Sign(priv, emptyMessage)
+	if !Verify(pub, emptyMessage, emptySig) {
+		t.Errorf("Verify failed for an empty message")
+	}
+	if Verify(pub, message, emptySig) {
+		t.Errorf("Verify succeeded comparing non-empty message with empty signature")
+	}
+}

pkg/cryptography/hkdf_test.go

@@ -0,0 +1,108 @@
+package cryptography
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestDeriveKey(t *testing.T) {
+	secret := []byte("test-secret")
+	salt := []byte("test-salt")
+	info := []byte("test-info")
+	length := 32 // Desired key length
+
+	key1, err := DeriveKey(secret, salt, info, length)
+	if err != nil {
+		t.Fatalf("DeriveKey failed: %v", err)
+	}
+
+	if len(key1) != length {
+		t.Errorf("DeriveKey returned key of length %d; want %d", len(key1), length)
+	}
+
+	// Derive another key with the same parameters, should be identical
+	key2, err := DeriveKey(secret, salt, info, length)
+	if err != nil {
+		t.Fatalf("Second DeriveKey failed: %v", err)
+	}
+	if !bytes.Equal(key1, key2) {
+		t.Errorf("DeriveKey is not deterministic. Got %x and %x for the same inputs", key1, key2)
+	}
+
+	// Derive a key with different info, should be different
+	differentInfo := []byte("different-info")
+	key3, err := DeriveKey(secret, salt, differentInfo, length)
+	if err != nil {
+		t.Fatalf("DeriveKey with different info failed: %v", err)
+	}
+	if bytes.Equal(key1, key3) {
+		t.Errorf("DeriveKey produced the same key for different info strings")
+	}
+
+	// Derive a key with different salt, should be different
+	differentSalt := []byte("different-salt")
+	key4, err := DeriveKey(secret, differentSalt, info, length)
+	if err != nil {
+		t.Fatalf("DeriveKey with different salt failed: %v", err)
+	}
+	if bytes.Equal(key1, key4) {
+		t.Errorf("DeriveKey produced the same key for different salts")
+	}
+
+	// Derive a key with different secret, should be different
+	differentSecret := []byte("different-secret")
+	key5, err := DeriveKey(differentSecret, salt, info, length)
+	if err != nil {
+		t.Fatalf("DeriveKey with different secret failed: %v", err)
+	}
+	if bytes.Equal(key1, key5) {
+		t.Errorf("DeriveKey produced the same key for different secrets")
+	}
+
+	// Derive a key with different length
+	differentLength := 64
+	key6, err := DeriveKey(secret, salt, info, differentLength)
+	if err != nil {
+		t.Fatalf("DeriveKey with different length failed: %v", err)
+	}
+	if len(key6) != differentLength {
+		t.Errorf("DeriveKey returned key of length %d; want %d", len(key6), differentLength)
+	}
+}
+
+func TestDeriveKeyEdgeCases(t *testing.T) {
+	secret := []byte("test-secret")
+	salt := []byte("test-salt")
+	info := []byte("test-info")
+
+	t.Run("EmptySecret", func(t *testing.T) {
+		_, err := DeriveKey([]byte{}, salt, info, 32)
+		if err != nil {
+			t.Errorf("DeriveKey failed with empty secret: %v", err)
+		}
+	})
+
+	t.Run("EmptySalt", func(t *testing.T) {
+		_, err := DeriveKey(secret, []byte{}, info, 32)
+		if err != nil {
+			t.Errorf("DeriveKey failed with empty salt: %v", err)
+		}
+	})
+
+	t.Run("EmptyInfo", func(t *testing.T) {
+		_, err := DeriveKey(secret, salt, []byte{}, 32)
+		if err != nil {
+			t.Errorf("DeriveKey failed with empty info: %v", err)
+		}
+	})
+
+	t.Run("ZeroLength", func(t *testing.T) {
+		key, err := DeriveKey(secret, salt, info, 0)
+		if err != nil {
+			t.Errorf("DeriveKey failed with zero length: %v", err)
+		}
+		if len(key) != 0 {
+			t.Errorf("DeriveKey with zero length returned non-empty key: %x", key)
+		}
+	})
+}

pkg/cryptography/hmac_test.go

@@ -0,0 +1,80 @@
+package cryptography
+
+import (
+	"testing"
+)
+
+func TestGenerateHMACKey(t *testing.T) {
+	testSizes := []int{16, 32, 64}
+	for _, size := range testSizes {
+		t.Run("Size"+string(rune(size)), func(t *testing.T) { // Simple name conversion
+			key, err := GenerateHMACKey(size)
+			if err != nil {
+				t.Fatalf("GenerateHMACKey(%d) failed: %v", size, err)
+			}
+			if len(key) != size {
+				t.Errorf("GenerateHMACKey(%d) returned key of length %d; want %d", size, len(key), size)
+			}
+
+			// Check if key is not all zeros (basic check for randomness)
+			isZero := true
+			for _, b := range key {
+				if b != 0 {
+					isZero = false
+					break
+				}
+			}
+			if isZero {
+				t.Errorf("GenerateHMACKey(%d) returned an all-zero key", size)
+			}
+		})
+	}
+}
+
+func TestComputeAndValidateHMAC(t *testing.T) {
+	key, err := GenerateHMACKey(32) // Use SHA256 key size
+	if err != nil {
+		t.Fatalf("Failed to generate HMAC key: %v", err)
+	}
+
+	message := []byte("This is a test message.")
+
+	// Compute HMAC
+	computedHMAC := ComputeHMAC(key, message)
+	if len(computedHMAC) != 32 { // SHA256 output size
+		t.Errorf("ComputeHMAC returned HMAC of length %d; want 32", len(computedHMAC))
+	}
+
+	// Validate correct HMAC
+	if !ValidateHMAC(key, message, computedHMAC) {
+		t.Errorf("ValidateHMAC failed for correctly computed HMAC")
+	}
+
+	// Validate incorrect HMAC (tampered message)
+	tamperedMessage := append(message, byte('!'))
+	if ValidateHMAC(key, tamperedMessage, computedHMAC) {
+		t.Errorf("ValidateHMAC succeeded for tampered message")
+	}
+
+	// Validate incorrect HMAC (tampered key)
+	wrongKey, _ := GenerateHMACKey(32)
+	if ValidateHMAC(wrongKey, message, computedHMAC) {
+		t.Errorf("ValidateHMAC succeeded for incorrect key")
+	}
+
+	// Validate incorrect HMAC (tampered HMAC)
+	tamperedHMAC := append(computedHMAC[:len(computedHMAC)-1], ^computedHMAC[len(computedHMAC)-1])
+	if ValidateHMAC(key, message, tamperedHMAC) {
+		t.Errorf("ValidateHMAC succeeded for tampered HMAC")
+	}
+
+	// Validate empty message
+	emptyMessage := []byte("")
+	emptyHMAC := ComputeHMAC(key, emptyMessage)
+	if !ValidateHMAC(key, emptyMessage, emptyHMAC) {
+		t.Errorf("ValidateHMAC failed for empty message")
+	}
+	if ValidateHMAC(key, message, emptyHMAC) {
+		t.Errorf("ValidateHMAC succeeded comparing non-empty message with empty HMAC")
+	}
+}