// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package krb5

import (
	"bytes"
	"cmp"
	"encoding/binary"
	"slices"
	"sync"
	"testing"
	"time"
)

func testKeys() []byte {
	serviceKey := make([]byte, 32)
	for i := range serviceKey {
		serviceKey[i] = byte(i)
	}
	return serviceKey
}

func TestGSSExchange(t *testing.T) {
	serviceKey := testKeys()
	initiator, token, err := ClientInit(EtypeAES256, serviceKey,
		"EXAMPLE.ORG", "nfs", "petr@EXAMPLE.ORG")
	if err != nil {
		t.Fatalf("client init: %v", err)
	}
	if len(token) == 0 || token[0] != 0x60 {
		t.Fatalf("context token % x", token[:8])
	}
	acceptor, reply, err := AcceptInit(token, serviceKey)
	if err != nil {
		t.Fatalf("accept init: %v", err)
	}
	if !bytes.Equal(acceptor.Key, initiator.Key) {
		t.Fatal("the session keys differ")
	}
	if len(reply) == 0 {
		t.Fatal("no AP-REP")
	}

	// krb5i: a MIC over the message binds it to the header and the
	// sequence counters move in lock step.
	message := []byte("the rpc call header bytes")
	mic, err := initiator.GetMIC(message)
	if err != nil {
		t.Fatal(err)
	}
	if mic[0] != 0x04 || mic[1] != 0x04 {
		t.Fatalf("mic token id % x", mic[:2])
	}
	if err := acceptor.VerifyMIC(message, mic); err != nil {
		t.Fatalf("verify: %v", err)
	}
	// The mirrored direction: the acceptor's own MIC verifies on the
	// initiator, and the sequence counters advanced.
	mic2, err := acceptor.GetMIC(message)
	if err != nil {
		t.Fatal(err)
	}
	if err := initiator.VerifyMIC(message, mic2); err != nil {
		t.Fatalf("verify reverse: %v", err)
	}
	// A replay of the first token is refused.
	if err := acceptor.VerifyMIC(message, mic); err != ErrContext {
		t.Fatalf("replay accepted: %v", err)
	}
	// Tampering with a fresh token is refused on the checksum.
	tampered := append([]byte{}, message...)
	tampered[0] ^= 0x80
	tamperedMIC, err := initiator.GetMIC(tampered)
	if err != nil {
		t.Fatal(err)
	}
	tamperedMIC[20] ^= 0xff
	if err := acceptor.VerifyMIC(tampered, tamperedMIC); err != ErrIntegrity {
		t.Fatalf("tampering accepted: %v", err)
	}

	// krb5p: the Wrap token hides and binds the payload.
	secret := []byte("the compound arguments")
	wrapped, err := initiator.Wrap(secret)
	if err != nil {
		t.Fatal(err)
	}
	if wrapped[0] != 0x05 || wrapped[1] != 0x04 {
		t.Fatalf("wrap token id % x", wrapped[:2])
	}
	if bytes.Contains(wrapped, secret) {
		t.Fatal("the plaintext leaked into the wrap token")
	}
	opened, err := acceptor.Unwrap(wrapped)
	if err != nil {
		t.Fatalf("unwrap: %v", err)
	}
	if !bytes.Equal(opened, secret) {
		t.Fatalf("unwrapped %q", opened)
	}
	wrapped[20] ^= 0xff
	if _, err := acceptor.Unwrap(wrapped); err == nil {
		t.Fatal("tampering accepted")
	}
}

// craftAPREQ builds an establishment token the way ClientInit does, with
// the ticket endtime and the authenticator time under the caller's
// control, so the acceptor's checks can be driven from the tests.
func craftAPREQ(t *testing.T, endtime, ctime time.Time) []byte {
	t.Helper()
	serviceKey := testKeys()
	session := make([]byte, 32)
	for i := range session {
		session[i] = byte(255 - i)
	}
	ticket, err := buildTicket("EXAMPLE.ORG", "nfs", session, serviceKey, EtypeAES256, endtime)
	if err != nil {
		t.Fatal(err)
	}
	auth, err := buildAuthenticator("EXAMPLE.ORG", "petr@EXAMPLE.ORG", session, EtypeAES256, 0, ctime)
	if err != nil {
		t.Fatal(err)
	}
	apreq := appTag(14, derTLV(tagSequence, append(append(
		derInt(5), derInt(14)...),
		append(bitFlags(0), append(ticket, auth...)...)...)))
	return appTag(0, append(append([]byte{}, derOID...), apreq...))
}

func TestAcceptInitReplay(t *testing.T) {
	serviceKey := testKeys()
	token := craftAPREQ(t, time.Now().Add(time.Hour), time.Now())
	if _, _, err := AcceptInit(token, serviceKey); err != nil {
		t.Fatalf("first accept: %v", err)
	}
	if _, _, err := AcceptInit(token, serviceKey); err != ErrReplay {
		t.Fatalf("replayed accept: %v", err)
	}
}

func TestAcceptInitSkew(t *testing.T) {
	serviceKey := testKeys()
	old := craftAPREQ(t, time.Now().Add(time.Hour), time.Now().Add(-2*clockSkew))
	if _, _, err := AcceptInit(old, serviceKey); err != ErrSkew {
		t.Fatalf("an authenticator outside the window: %v", err)
	}
	future := craftAPREQ(t, time.Now().Add(time.Hour), time.Now().Add(2*clockSkew))
	if _, _, err := AcceptInit(future, serviceKey); err != ErrSkew {
		t.Fatalf("an authenticator from the future: %v", err)
	}
}

func TestAcceptInitExpiredTicket(t *testing.T) {
	serviceKey := testKeys()
	token := craftAPREQ(t, time.Now().Add(-time.Hour), time.Now())
	if _, _, err := AcceptInit(token, serviceKey); err != ErrExpired {
		t.Fatalf("an expired ticket: %v", err)
	}
}

// The per-message token operations of one context serve concurrent
// callers without losing sequence numbers, which the race detector
// watches over.
func TestContextConcurrentTokens(t *testing.T) {
	serviceKey := testKeys()
	initiator, token, err := ClientInit(EtypeAES256, serviceKey,
		"EXAMPLE.ORG", "nfs", "petr@EXAMPLE.ORG")
	if err != nil {
		t.Fatal(err)
	}
	acceptor, _, err := AcceptInit(token, serviceKey)
	if err != nil {
		t.Fatal(err)
	}

	const workers, rounds = 4, 25
	mics := make([][]byte, workers*rounds)
	wrapped := make([][]byte, workers*rounds)
	var wg sync.WaitGroup
	for w := range workers {
		wg.Go(func() {
			for i := range rounds {
				mic, err := initiator.GetMIC([]byte("concurrent"))
				if err != nil {
					t.Error(err)
					return
				}
				mics[w*rounds+i] = mic
				box, err := initiator.Wrap([]byte("secret"))
				if err != nil {
					t.Error(err)
					return
				}
				wrapped[w*rounds+i] = box
			}
		})
	}
	wg.Wait()

	// Every sequence number was handed out exactly once, and the tokens
	// verify in order on the other side. The MIC and Wrap tokens share the
	// send counter, so the two families interleave in one space.
	type sent struct {
		seq  uint64
		mic  []byte
		wrap []byte
	}
	total := make([]sent, 0, 2*workers*rounds)
	for _, mic := range mics {
		total = append(total, sent{seq: binary.BigEndian.Uint64(mic[8:]), mic: mic})
	}
	for _, box := range wrapped {
		total = append(total, sent{seq: binary.BigEndian.Uint64(box[8:]), wrap: box})
	}
	slices.SortFunc(total, func(a, b sent) int { return cmp.Compare(a.seq, b.seq) })
	if len(total) != 2*workers*rounds {
		t.Fatalf("%d tokens, want %d", len(total), 2*workers*rounds)
	}
	for i, tk := range total {
		if tk.seq != uint64(i) {
			t.Fatalf("token %d carries sequence %d", i, tk.seq)
		}
		if tk.mic != nil {
			if err := acceptor.VerifyMIC([]byte("concurrent"), tk.mic); err != nil {
				t.Fatalf("verify: %v", err)
			}
			continue
		}
		opened, err := acceptor.Unwrap(tk.wrap)
		if err != nil {
			t.Fatalf("unwrap: %v", err)
		}
		if string(opened) != "secret" {
			t.Fatalf("unwrapped %q", opened)
		}
	}
}