package session

import (
	"bytes"
	"errors"
	"testing"

	"github.com/aura/singbox-aura/aura/frame"
)

func TestReplayWindowBasicMonotonic(t *testing.T) {
	w := NewReplay(2)
	for _, s := range []uint64{2, 3, 4} {
		if err := w.CheckAndSet(s); err != nil {
			t.Fatalf("seq %d: unexpected %v", s, err)
		}
	}
	for _, s := range []uint64{2, 3, 4} {
		var e *ErrReplay
		if err := w.CheckAndSet(s); !errors.As(err, &e) {
			t.Fatalf("seq %d: want replay, got %v", s, err)
		}
	}
}

func TestReplayWindowOutOfOrderWithinWindow(t *testing.T) {
	w := NewReplay(0)
	if err := w.CheckAndSet(0); err != nil {
		t.Fatal(err)
	}
	if err := w.CheckAndSet(10); err != nil {
		t.Fatal(err)
	}
	if err := w.CheckAndSet(5); err != nil {
		t.Fatalf("5 inside window: %v", err)
	}
	if err := w.CheckAndSet(5); err == nil {
		t.Fatal("replay of 5 must be rejected")
	}
	if err := w.CheckAndSet(10); err == nil {
		t.Fatal("replay of 10 must be rejected")
	}
	if err := w.CheckAndSet(11); err != nil {
		t.Fatalf("new high 11: %v", err)
	}
}

func TestReplayWindowRejectsTooOld(t *testing.T) {
	w := NewReplay(0)
	if err := w.CheckAndSet(0); err != nil {
		t.Fatal(err)
	}
	if err := w.CheckAndSet(200); err != nil {
		t.Fatal(err)
	}
	if err := w.CheckAndSet(1); err == nil {
		t.Fatal("far below window must be rejected")
	}
	if err := w.CheckAndSet(200 - ReplayWindow); err == nil {
		t.Fatal("at the floor of the window must be rejected")
	}
	if err := w.CheckAndSet(200 - ReplayWindow + 1); err != nil {
		t.Fatalf("just inside window: %v", err)
	}
}

func TestDatagramRoundtripReorderAndReplay(t *testing.T) {
	key := make([]byte, 32)
	for i := range key {
		key[i] = 11
	}
	tx, err := NewDatagramSender(key, 2)
	if err != nil {
		t.Fatal(err)
	}
	rx, err := NewDatagramReceiver(key, 2)
	if err != nil {
		t.Fatal(err)
	}

	d0 := tx.Seal(&frame.Frame{Kind: frame.FrameData, StreamID: 0, Payload: []byte("pkt-a")})
	d1 := tx.Seal(&frame.Frame{Kind: frame.FrameData, StreamID: 0, Payload: []byte("pkt-b")})

	// Out-of-order delivery within the window.
	gotB, err := rx.Open(d1)
	if err != nil {
		t.Fatalf("open d1: %v", err)
	}
	if gotB.Kind != frame.FrameData || string(gotB.Payload) != "pkt-b" {
		t.Fatalf("d1: %+v", gotB)
	}
	gotA, err := rx.Open(d0)
	if err != nil {
		t.Fatalf("open d0: %v", err)
	}
	if string(gotA.Payload) != "pkt-a" {
		t.Fatalf("d0: %+v", gotA)
	}

	if _, err := rx.Open(d1); err == nil {
		t.Fatal("replay of d1 must be rejected")
	}

	bad := tx.Seal(&frame.Frame{Kind: frame.FramePing, Seq: 7})
	bad[len(bad)-1] ^= 1
	if _, err := rx.Open(bad); err == nil {
		t.Fatal("tampered ciphertext must fail")
	}
}

func TestSenderNextSeqAdvances(t *testing.T) {
	tx, err := NewDatagramSender(bytes.Repeat([]byte{1}, 32), 2)
	if err != nil {
		t.Fatal(err)
	}
	if tx.NextSeq() != 2 {
		t.Fatalf("initial next seq %d", tx.NextSeq())
	}
	_ = tx.Seal(&frame.Frame{Kind: frame.FramePing, Seq: 1})
	if tx.NextSeq() != 3 {
		t.Fatalf("after seal: %d", tx.NextSeq())
	}
}