package wire

import (
	"bytes"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"testing"
)

func fixedEnvelope(t *testing.T) Envelope {
	t.Helper()
	ek, err := hex.DecodeString("8520f0098930a754748b7ddcb43ef75a0dbf3a0d26381af4eba4a98eaa9b4e6a")
	if err != nil {
		t.Fatal(err)
	}
	return Envelope{
		V:    Version,
		Node: "nd_7f3a9c2e",
		Req:  "3f2504e0-4f89-41d3-9a0c-0305e82c3301",
		App:  "ap_5b1d",
		TS:   1727289600,
		EK:   ek,
		Alg:  Suite,
		P:    []byte{0x00, 0x01, 0x02, 0xfe, 0xff},
	}
}

// TestCanonicalRequestVector pins the exact bytes. This is the file every implementation must agree
// with: if the Node relay or the Swift app produces anything else here, signatures will fail in
// production with an error that says nothing about why. Keeping the vector in the repo turns that
// into a failing test instead.
func TestCanonicalRequestVector(t *testing.T) {
	canonical, err := CanonicalRequest(fixedEnvelope(t))
	if err != nil {
		t.Fatal(err)
	}
	got := hex.EncodeToString(canonical)
	const want = "" +
		"6d73612d72656c61792d76312d72657100" + // "msa-relay-v1-req" NUL
		"000000080000000000000001" + // v = 1 as uint64be, length-prefixed
		"0000000b6e645f3766336139633265" + // node
		"0000002433663235303465302d346638392d343164332d396130632d303330356538326333333031" + // req
		"0000000661705f35623164" + // app
		"000000080000000000000000" + // ts placeholder, replaced below
		"000000208520f0098930a754748b7ddcb43ef75a0dbf3a0d26381af4eba4a98eaa9b4e6a" + // ek
		"00000005000102feff" // payload
	_ = want

	// The literal above is easier to get wrong than to derive, so assert the structural properties
	// that actually matter, then pin the whole thing against a recomputed expectation.
	if !bytes.HasPrefix(canonical, []byte("msa-relay-v1-req\x00")) {
		t.Fatalf("canonical does not start with the domain: %x", canonical[:20])
	}
	rawEK := fixedEnvelope(t).EK
	if !bytes.Contains(canonical, rawEK) {
		t.Fatal("canonical is missing the raw ek bytes")
	}
	if bytes.Contains(canonical, []byte(base64.StdEncoding.EncodeToString(rawEK))) {
		t.Fatal("canonical contains base64 ek text; the signature must cover raw bytes")
	}

	// Field order is part of the contract, so check each marker's offset is increasing.
	positions := []int{
		bytes.Index(canonical, []byte("msa-relay-v1-req")),
		bytes.Index(canonical, []byte("nd_7f3a9c2e")),
		bytes.Index(canonical, []byte("3f2504e0-4f89-41d3-9a0c-0305e82c3301")),
		bytes.Index(canonical, []byte("ap_5b1d")),
	}
	for index := 1; index < len(positions); index++ {
		if positions[index] <= positions[index-1] {
			t.Fatalf("field order is wrong: %v", positions)
		}
	}
	if got != expectCanonicalRequest() {
		t.Fatalf("canonical request changed:\n got %s\nwant %s", got, expectCanonicalRequest())
	}
}

// expectCanonicalRequest is the frozen vector. It is spelled out rather than computed so that a
// change to the encoding shows up as a diff here instead of silently invalidating every paired
// device in the field.
func expectCanonicalRequest() string {
	parts := []string{
		hex.EncodeToString([]byte("msa-relay-v1-req")),
		"00",
		"00000008", "0000000000000001",
		"0000000b", hex.EncodeToString([]byte("nd_7f3a9c2e")),
		"00000024", hex.EncodeToString([]byte("3f2504e0-4f89-41d3-9a0c-0305e82c3301")),
		"00000007", hex.EncodeToString([]byte("ap_5b1d")),
		"00000008", "0000000066f45900",
		"00000020", "8520f0098930a754748b7ddcb43ef75a0dbf3a0d26381af4eba4a98eaa9b4e6a",
		"00000005", "000102feff",
	}
	out := ""
	for _, part := range parts {
		out += part
	}
	return out
}

func TestCanonicalResponseOmitsAppID(t *testing.T) {
	envelope := fixedEnvelope(t)
	request, err := CanonicalRequest(envelope)
	if err != nil {
		t.Fatal(err)
	}
	response, err := CanonicalResponse(envelope)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Equal(request, response) {
		t.Fatal("request and response canonical forms are identical; a signature could be replayed across directions")
	}
	if bytes.Contains(response, []byte("ap_5b1d")) {
		t.Fatal("response canonical form includes the app id")
	}
	if !bytes.HasPrefix(response, []byte("msa-relay-v1-resp\x00")) {
		t.Fatalf("response canonical form has the wrong domain: %x", response[:20])
	}
}

func TestCanonicalRejectsMalformedEnvelopes(t *testing.T) {
	cases := map[string]func(*Envelope){
		"missing node":  func(e *Envelope) { e.Node = "" },
		"missing req":   func(e *Envelope) { e.Req = "  " },
		"missing suite": func(e *Envelope) { e.Alg = "" },
		"unknown suite": func(e *Envelope) { e.Alg = "AES128" },
		"short ek":      func(e *Envelope) { e.EK = e.EK[:16] },
		"empty ek":      func(e *Envelope) { e.EK = nil },
	}
	for name, mutate := range cases {
		envelope := fixedEnvelope(t)
		mutate(&envelope)
		if _, err := CanonicalRequest(envelope); err == nil {
			t.Errorf("%s: canonical form accepted it", name)
		}
	}
}

func TestSignAndVerifyRoundTrip(t *testing.T) {
	public, private, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	envelope := fixedEnvelope(t)
	if err := SignRequest(&envelope, private); err != nil {
		t.Fatal(err)
	}
	if err := VerifyRequest(envelope, public); err != nil {
		t.Fatalf("VerifyRequest: %v", err)
	}

	// Every routing field must be covered. If any of these mutations still verified, the relay
	// could change it in flight.
	mutations := map[string]func(*Envelope){
		"node": func(e *Envelope) { e.Node = "nd_attacker" },
		"req":  func(e *Envelope) { e.Req = "00000000-0000-4000-8000-000000000000" },
		"app":  func(e *Envelope) { e.App = "ap_other" },
		"ts":   func(e *Envelope) { e.TS = e.TS + 1 },
		"ek":   func(e *Envelope) { e.EK = bytes.Repeat([]byte{7}, PublicBytes) },
		"p":    func(e *Envelope) { e.P = append(append([]byte{}, e.P...), 0) },
		"alg":  func(e *Envelope) { e.Alg = "other" },
		"v":    func(e *Envelope) { e.V = 2 },
	}
	for name, mutate := range mutations {
		changed := envelope
		changed.EK = append([]byte{}, envelope.EK...)
		changed.P = append([]byte{}, envelope.P...)
		changed.S = append([]byte{}, envelope.S...)
		mutate(&changed)
		if err := VerifyRequest(changed, public); err == nil {
			t.Errorf("signature still verified after mutating %s", name)
		}
	}

	// A response signature must not satisfy a request verification even over identical bytes.
	response := fixedEnvelope(t)
	if err := SignResponse(&response, private); err != nil {
		t.Fatal(err)
	}
	if err := VerifyRequest(response, public); err == nil {
		t.Fatal("a response signature verified as a request signature")
	}
}

func TestVerifyRejectsWrongKey(t *testing.T) {
	_, private, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	otherPublic, _, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	envelope := fixedEnvelope(t)
	if err := SignRequest(&envelope, private); err != nil {
		t.Fatal(err)
	}
	if err := VerifyRequest(envelope, otherPublic); !errors.Is(err, ErrSignature) {
		t.Fatalf("VerifyRequest with the wrong key = %v, want ErrSignature", err)
	}
	if err := VerifyRequest(envelope, otherPublic[:16]); err == nil {
		t.Fatal("accepted a truncated public key")
	}
}

// TestSharedSecretMatchesRFC7748 is the X25519 known-answer test. Without it a silent change to the
// scalar clamping or the endianness would produce a working system that interoperates with nobody.
func TestSharedSecretMatchesRFC7748(t *testing.T) {
	// RFC 7748 section 6.1: alice private, bob public, expected shared secret.
	alicePrivate := mustHex(t, "77076d0a7318a57d3c16c17251b26645df4c2f87ebc0992ab177fba51db92c2a")
	bobPublic := mustHex(t, "de9edb7d7b7dc1b4d35b61c2ece435373f8343c85b78674dadfc7e146f882b4f")
	want := mustHex(t, "4a5d9d5ba4ce2de1728e3bf480350f25e07e21c947d19e3376f09b3c1e161742")
	got, err := SharedSecret(alicePrivate, bobPublic)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(got, want) {
		t.Fatalf("shared secret = %x, want %x", got, want)
	}
}

func TestSharedSecretRejectsBadKeys(t *testing.T) {
	valid := bytes.Repeat([]byte{7}, PublicBytes)
	if _, err := SharedSecret(nil, valid); !errors.Is(err, ErrBadKey) {
		t.Errorf("nil private key err = %v, want ErrBadKey", err)
	}
	if _, err := SharedSecret(bytes.Repeat([]byte{1}, KeyBytes), nil); !errors.Is(err, ErrBadPublicKey) {
		t.Errorf("nil peer key err = %v, want ErrBadPublicKey", err)
	}
	// The all-zero public key is the classic low-order point; it must not yield a shared secret.
	if _, err := SharedSecret(bytes.Repeat([]byte{9}, KeyBytes), make([]byte, PublicBytes)); err == nil {
		t.Error("accepted the all-zero peer public key")
	}
}

func TestSealOpenRoundTrip(t *testing.T) {
	key := bytes.Repeat([]byte{0x2a}, KeyBytes)
	iv := bytes.Repeat([]byte{0x0b}, IVBytes)
	plaintext := []byte(`{"m":"GET","p":"/api/health"}`)
	sealed, err := Seal(key, iv, plaintext)
	if err != nil {
		t.Fatal(err)
	}
	if len(sealed) != len(plaintext)+TagBytes {
		t.Fatalf("sealed length = %d, want %d", len(sealed), len(plaintext)+TagBytes)
	}
	opened, err := Open(key, iv, sealed)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(opened, plaintext) {
		t.Fatalf("opened = %q, want %q", opened, plaintext)
	}

	// Tampering anywhere, including in the tag, must fail closed.
	for index := range sealed {
		damaged := append([]byte{}, sealed...)
		damaged[index] ^= 0x01
		if _, err := Open(key, iv, damaged); !errors.Is(err, ErrOpen) {
			t.Fatalf("flipping byte %d still authenticated", index)
		}
	}
	if _, err := Open(bytes.Repeat([]byte{0x2b}, KeyBytes), iv, sealed); !errors.Is(err, ErrOpen) {
		t.Fatal("a different key decrypted the payload")
	}
	if _, err := Open(key, bytes.Repeat([]byte{0x0c}, IVBytes), sealed); !errors.Is(err, ErrOpen) {
		t.Fatal("a different iv decrypted the payload")
	}
	if _, err := Open(key, iv, sealed[:TagBytes-1]); !errors.Is(err, ErrOpen) {
		t.Fatal("a truncated payload was accepted")
	}
}

// TestDerivedKeysAreDirectionalAndRequestBound checks the properties that make key reuse safe to
// reason about: the two directions differ, and two requests never share an IV.
func TestDerivedKeysAreDirectionalAndRequestBound(t *testing.T) {
	shared := bytes.Repeat([]byte{0x5c}, KeyBytes)
	requestKey, requestIV, err := DeriveKey(shared, "nd_a", "req-1", false)
	if err != nil {
		t.Fatal(err)
	}
	responseKey, responseIV, err := DeriveKey(shared, "nd_a", "req-1", true)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Equal(requestKey, responseKey) {
		t.Fatal("both directions derived the same key")
	}
	if bytes.Equal(requestIV, responseIV) {
		t.Fatal("both directions derived the same iv")
	}
	if len(requestKey) != KeyBytes || len(requestIV) != IVBytes {
		t.Fatalf("derived sizes are wrong: key %d iv %d", len(requestKey), len(requestIV))
	}

	otherKey, otherIV, err := DeriveKey(shared, "nd_a", "req-2", false)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Equal(requestIV, otherIV) {
		t.Fatal("two request ids derived the same iv")
	}
	if bytes.Equal(requestKey, otherKey) {
		t.Fatal("two request ids derived the same key")
	}
	_ = otherIV

	otherNodeKey, _, err := DeriveKey(shared, "nd_b", "req-1", false)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Equal(requestKey, otherNodeKey) {
		t.Fatal("two nodes derived the same key from the same secret")
	}
}

// TestHKDFMatchesRFC5869 pins the KDF itself against the published test case, so an implementation
// slip cannot hide behind the round-trip tests above.
func TestHKDFMatchesRFC5869(t *testing.T) {
	ikm := mustHex(t, "0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b")
	salt := mustHex(t, "000102030405060708090a0b0c")
	info := mustHex(t, "f0f1f2f3f4f5f6f7f8f9")
	want := mustHex(t, "3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865")
	got, err := HKDF(ikm, salt, info, 42)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(got, want) {
		t.Fatalf("hkdf = %x, want %x", got, want)
	}
	if _, err := HKDF(ikm, salt, info, 0); err == nil {
		t.Fatal("accepted a zero-length output")
	}
	if _, err := HKDF(ikm, salt, info, 255*32+1); err == nil {
		t.Fatal("accepted an output longer than HKDF allows")
	}
}

func TestEnvelopeJSONRoundTrip(t *testing.T) {
	original := fixedEnvelope(t)
	original.S = bytes.Repeat([]byte{0x11}, ed25519.SignatureSize)
	encoded, err := json.Marshal(original)
	if err != nil {
		t.Fatal(err)
	}
	var decoded Envelope
	if err := json.Unmarshal(encoded, &decoded); err != nil {
		t.Fatal(err)
	}
	if decoded.Node != original.Node || decoded.Req != original.Req || decoded.TS != original.TS {
		t.Fatalf("string fields round-tripped wrong: %+v", decoded)
	}
	if !bytes.Equal(decoded.EK, original.EK) || !bytes.Equal(decoded.P, original.P) || !bytes.Equal(decoded.S, original.S) {
		t.Fatal("byte fields did not round-trip")
	}
	// The wire form must be base64 text, not an array of numbers.
	if !bytes.Contains(encoded, []byte(`"ek":"`)) {
		t.Fatalf("ek is not encoded as base64: %s", encoded)
	}
	if bytes.Contains(encoded, []byte(`"ek":[`)) {
		t.Fatalf("ek leaked as a byte array: %s", encoded)
	}
}

func TestEnvelopeJSONRejectsMalformedBase64(t *testing.T) {
	raw := []byte(`{"v":1,"node":"nd_x","req":"r","ts":1,"ek":"!!!!","alg":"` + Suite + `","p":"AA==","s":""}`)
	var envelope Envelope
	if err := json.Unmarshal(raw, &envelope); err == nil {
		t.Fatal("accepted invalid base64 in ek")
	}
}

// TestIntentDigestIsKeyOrderInsensitive is the property that lets the phone and the agent agree on
// a digest without agreeing on JSON key order.
func TestIntentDigestIsKeyOrderInsensitive(t *testing.T) {
	first, err := IntentDigest("post", "/api/sms/send", []byte(`{"b":2,"a":1,"nested":{"z":true,"y":[1,2,3]}}`))
	if err != nil {
		t.Fatal(err)
	}
	second, err := IntentDigest("POST", "/api/sms/send", []byte("{\n  \"nested\": {\"y\": [1, 2, 3], \"z\": true},\n  \"a\": 1,\n  \"b\": 2\n}"))
	if err != nil {
		t.Fatal(err)
	}
	if !VerifyIntent(first, second) {
		t.Fatalf("intent digest depends on key order or whitespace: %x vs %x", first, second)
	}

	// Changing any part of the binding must change the digest.
	variants := [][]byte{}
	for _, mutate := range []func() ([]byte, error){
		func() ([]byte, error) {
			return IntentDigest("POST", "/api/sms/send", []byte(`{"a":1,"b":3,"nested":{"z":true,"y":[1,2,3]}}`))
		},
		func() ([]byte, error) { return IntentDigest("POST", "/api/sms/batch-delete", []byte(`{"a":1,"b":2}`)) },
		func() ([]byte, error) { return IntentDigest("GET", "/api/sms/send", []byte(`{"a":1,"b":2}`)) },
		func() ([]byte, error) { return IntentDigest("POST", "/api/sms/send", nil) },
	} {
		digest, err := mutate()
		if err != nil {
			t.Fatal(err)
		}
		variants = append(variants, digest)
	}
	for index, variant := range variants {
		if VerifyIntent(first, variant) {
			t.Errorf("variant %d produced the same intent digest", index)
		}
	}

	// No body and an explicitly empty body are the same absence of a body.
	empty, err := IntentDigest("POST", "/api/sms/send", []byte("   "))
	if err != nil {
		t.Fatal(err)
	}
	if !VerifyIntent(empty, variants[3]) {
		t.Fatal("blank and absent bodies canonicalised differently")
	}
}

func TestCanonicalJSONRules(t *testing.T) {
	cases := []struct{ in, want string }{
		{`{"b":1,"a":2}`, `{"a":2,"b":1}`},
		{`[3,1,2]`, `[3,1,2]`},
		{`  null `, `null`},
		{`true`, `true`},
		{`"<&>"`, `"<&>"`},
		{`{"x":100}`, `{"x":100}`},
		{``, ``},
		{`   `, ``},
		{`{"a":{"c":[1,{"b":2,"a":1}],"b":0},"A":9}`, `{"A":9,"a":{"b":0,"c":[1,{"a":1,"b":2}]}}`},
	}
	for _, c := range cases {
		got, err := CanonicalJSON([]byte(c.in))
		if err != nil {
			t.Errorf("CanonicalJSON(%q): %v", c.in, err)
			continue
		}
		if string(got) != c.want {
			t.Errorf("CanonicalJSON(%q) = %q, want %q", c.in, got, c.want)
		}
	}

	// HTML escaping must be off, otherwise a body containing & or < hashes differently in Go than
	// in Node or Swift.
	got, err := CanonicalJSON([]byte(`{"text":"a<b&c>d"}`))
	if err != nil {
		t.Fatal(err)
	}
	if string(got) != `{"text":"a<b&c>d"}` {
		t.Fatalf("HTML characters were escaped: %s", got)
	}
}

func TestCanonicalJSONRejectsGarbage(t *testing.T) {
	for _, raw := range []string{`{`, `{"a":1} trailing`, `{"a":+1}`, `{"a":.5}`, `{"a":NaN}`, `{"a":1.0}`, `{"a":1e3}`, `{"a":01}`, `{"a":-01}`, `{"a":-0}`} {
		if _, err := CanonicalJSON([]byte(raw)); err == nil {
			t.Errorf("CanonicalJSON(%q) accepted it", raw)
		}
	}
}

func mustHex(t *testing.T, value string) []byte {
	t.Helper()
	raw, err := hex.DecodeString(value)
	if err != nil {
		t.Fatalf("bad hex in test: %v", err)
	}
	return raw
}