package main

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"io"
	"testing"

	"golang.org/x/crypto/hkdf"
)

// Verifies the decryptMessage path still produces the correct plaintext after
// wrapping the ECDH/HKDF secrets in memguard.LockedBuffer, using the exact
// wire format the browser client (static/crypto.js) and relay produce.
func TestDecryptMessageAfterMemguard(t *testing.T) {
	readerPriv, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	ephPriv, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}

	shared, err := ephPriv.ECDH(readerPriv.PublicKey())
	if err != nil {
		t.Fatal(err)
	}
	aesKey := make([]byte, 32)
	hkdfR := hkdf.New(sha256.New, shared, make([]byte, 32), []byte("nymdrop-v1"))
	if _, err := io.ReadFull(hkdfR, aesKey); err != nil {
		t.Fatal(err)
	}
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		t.Fatal(err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		t.Fatal(err)
	}
	iv := make([]byte, 12)
	if _, err := rand.Read(iv); err != nil {
		t.Fatal(err)
	}
	want := "E2E memguard verification message"
	ciphertext := gcm.Seal(nil, iv, []byte(want), nil)

	packet := append(append(append([]byte{}, ephPriv.PublicKey().Bytes()...), iv...), ciphertext...)
	noncePrefix := hex.EncodeToString(make([]byte, 16)) + ":"
	raw := append([]byte(noncePrefix), packet...)

	got, err := decryptMessage(raw, readerPriv)
	if err != nil {
		t.Fatalf("decryptMessage: %v", err)
	}
	if got != want {
		t.Fatalf("plaintext mismatch: got %q want %q", got, want)
	}
}