package main

// The stateless-token contract (stateless.go): ceremony proofs and sealed
// WebAuthn sessions must round-trip under the right key and fail closed under
// everything else — expiry, tampering, truncation, a different server key, and
// (for proofs) any attempt to move a token between users.

import (
	"encoding/base64"
	"encoding/json"
	"io"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/go-webauthn/webauthn/webauthn"
)

func testKeys(t *testing.T, secret string) *stateKeys {
	t.Helper()
	k, err := newStateKeys([]byte(secret))
	if err != nil {
		t.Fatalf("newStateKeys: %v", err)
	}
	return k
}

func TestCeremonyProofRoundTrip(t *testing.T) {
	k := testKeys(t, "secret-a")
	proof := k.mintCeremonyProof("user-one")
	if proof == "" {
		t.Fatalf("empty proof")
	}
	if uid, ok := k.ceremonyUser(proof); !ok || uid != "user-one" {
		t.Fatalf("ceremonyUser: want (user-one, true), got (%q, %v)", uid, ok)
	}
	// The uid rides inside, MAC'd: a proof for A never validates as B.
	if uid, _ := k.ceremonyUser(k.mintCeremonyProof("user-two")); uid != "user-two" {
		t.Fatalf("proof carries the wrong uid: %q", uid)
	}
}

func TestCeremonyProofExpiry(t *testing.T) {
	k := testKeys(t, "secret-a")
	// Expiry is enforced server-side, from the MAC'd timestamp in the token.
	if _, ok := k.ceremonyUser(k.mintCeremonyProofAt("u", time.Now().Add(-time.Second))); ok {
		t.Fatalf("expired proof accepted")
	}
	if _, ok := k.ceremonyUser(k.mintCeremonyProofAt("u", time.Now().Add(time.Minute))); !ok {
		t.Fatalf("unexpired proof refused")
	}
}

func TestCeremonyProofTamperRejected(t *testing.T) {
	k := testKeys(t, "secret-a")
	proof := k.mintCeremonyProof("user-one")
	raw, err := base64.RawURLEncoding.DecodeString(proof)
	if err != nil {
		t.Fatalf("decode: %v", err)
	}
	// Flip one byte at every position — expiry, uid, and MAC are all covered;
	// no single-byte change may survive.
	for i := range raw {
		mut := append([]byte(nil), raw...)
		mut[i] ^= 0x01
		if _, ok := k.ceremonyUser(base64.RawURLEncoding.EncodeToString(mut)); ok {
			t.Fatalf("tampered proof accepted (byte %d)", i)
		}
	}
	// Truncation and garbage fail closed too.
	for _, bad := range []string{"", "x", "not!base64url!", proof[:len(proof)/2],
		base64.RawURLEncoding.EncodeToString(raw[:8+proofMACLen])} {
		if _, ok := k.ceremonyUser(bad); ok {
			t.Fatalf("malformed proof %q accepted", bad)
		}
	}
}

func TestCeremonyProofWrongKeyRejected(t *testing.T) {
	a, b := testKeys(t, "secret-a"), testKeys(t, "secret-b")
	if _, ok := b.ceremonyUser(a.mintCeremonyProof("u")); ok {
		t.Fatalf("proof minted under one key verified under another")
	}
	// Two ephemeral keys (empty secret) must not honour each other either.
	e1, e2 := testKeys(t, ""), testKeys(t, "")
	if _, ok := e2.ceremonyUser(e1.mintCeremonyProof("u")); ok {
		t.Fatalf("two ephemeral keys verified each other's proofs")
	}
}

func TestCeremonyProofCrossUserSpliceRejected(t *testing.T) {
	k := testKeys(t, "secret-a")
	// Same-length uids, so a payload splice is structurally plausible: take
	// alice's MAC and staple it onto a payload claiming bob. The MAC covers the
	// uid, so it must fail.
	alice := k.mintCeremonyProof("alice")
	bob := k.mintCeremonyProof("bobby")
	ra, _ := base64.RawURLEncoding.DecodeString(alice)
	rb, _ := base64.RawURLEncoding.DecodeString(bob)
	spliced := append(append([]byte(nil), rb[:len(rb)-proofMACLen]...), ra[len(ra)-proofMACLen:]...)
	if uid, ok := k.ceremonyUser(base64.RawURLEncoding.EncodeToString(spliced)); ok {
		t.Fatalf("spliced proof accepted as %q", uid)
	}
}

func TestSealOpenRoundTrip(t *testing.T) {
	k := testKeys(t, "secret-a")
	msg := []byte(`{"hello":"world"}`)
	tok1, err := k.seal(msg)
	if err != nil {
		t.Fatalf("seal: %v", err)
	}
	if got, ok := k.open(tok1); !ok || string(got) != string(msg) {
		t.Fatalf("open: want (%q, true), got (%q, %v)", msg, got, ok)
	}
	// Fresh random nonce per seal: two seals of one plaintext never collide.
	if tok2, _ := k.seal(msg); tok2 == tok1 {
		t.Fatalf("two seals produced identical tokens — nonce reuse")
	}
}

func TestSealOpenRejectsTamperAndWrongKey(t *testing.T) {
	a, b := testKeys(t, "secret-a"), testKeys(t, "secret-b")
	tok, err := a.seal([]byte("payload"))
	if err != nil {
		t.Fatalf("seal: %v", err)
	}
	if _, ok := b.open(tok); ok {
		t.Fatalf("sealed under one key, opened under another")
	}
	raw, _ := base64.RawURLEncoding.DecodeString(tok)
	for i := range raw {
		mut := append([]byte(nil), raw...)
		mut[i] ^= 0x01
		if _, ok := a.open(base64.RawURLEncoding.EncodeToString(mut)); ok {
			t.Fatalf("tampered sealed token opened (byte %d)", i)
		}
	}
	for _, bad := range []string{"", "!", tok[:10]} {
		if _, ok := a.open(bad); ok {
			t.Fatalf("malformed sealed token %q opened", bad)
		}
	}
}

// The full HTTP surface: a begin's sealed session must be refused on finish
// when tampered with, presented to the other ceremony's finish, or expired —
// and the whole flow must work with NO server-side session state at all.
func TestFinishRejectsBadSessions(t *testing.T) {
	ts, s := newTestServer(t)

	begin := func(path string) string {
		r := do(t, "POST", ts.URL+path, "", nil, nil)
		if r.StatusCode != http.StatusOK {
			t.Fatalf("%s: want 200, got %d", path, r.StatusCode)
		}
		var out struct {
			Session string `json:"session"`
		}
		body, _ := io.ReadAll(r.Body)
		if err := json.Unmarshal(body, &out); err != nil || out.Session == "" {
			t.Fatalf("%s: no session in %s", path, body)
		}
		return out.Session
	}
	finish := func(path, session string) int {
		r := do(t, "POST", ts.URL+path+"?session="+session, "", nil, []byte("{}"))
		return r.StatusCode
	}

	reg := begin("/api/register/begin")
	login := begin("/api/login/begin")

	// Kind is checked: a register session can't drive login/finish (nor the
	// reverse) — 400 before any WebAuthn verification runs.
	if got := finish("/api/login/finish", reg); got != http.StatusBadRequest {
		t.Fatalf("register session on login/finish: want 400, got %d", got)
	}
	if got := finish("/api/register/finish", login); got != http.StatusBadRequest {
		t.Fatalf("login session on register/finish: want 400, got %d", got)
	}

	// Tampering with the sealed param is a 400, never a crash or an accept.
	raw, err := base64.RawURLEncoding.DecodeString(reg)
	if err != nil {
		t.Fatalf("session should be base64url: %v", err)
	}
	raw[len(raw)-1] ^= 0x01
	if got := finish("/api/register/finish", base64.RawURLEncoding.EncodeToString(raw)); got != http.StatusBadRequest {
		t.Fatalf("tampered session: want 400, got %d", got)
	}

	// An expired session is refused even though the seal verifies. (Sealed by
	// hand: sealSession always stamps a fresh expiry, which is the point.)
	expired := sessionEntry{Data: &webauthn.SessionData{}, Kind: "register",
		Expires: time.Now().Add(-time.Minute).Unix()}
	bts, _ := json.Marshal(expired)
	sealed, err := s.keys.seal(bts)
	if err != nil {
		t.Fatalf("seal: %v", err)
	}
	if got := finish("/api/register/finish", sealed); got != http.StatusBadRequest {
		t.Fatalf("expired session: want 400, got %d", got)
	}

	// The session param must never leak ceremony internals in the clear: the
	// challenge from the begin options must not be findable in the token.
	if strings.Contains(reg, "challenge") {
		t.Fatalf("sealed session leaks plaintext")
	}
}