package tcp

import (
	"math/rand"
	"testing"

	"github.com/apernet/OpenGFW/analyzer"
)

// feedFET runs a sequence of client-to-server segments through a fresh stream
// and returns the props it settled on, or nil if it never reached a verdict.
func feedFET(t *testing.T, segments ...[]byte) analyzer.PropMap {
	t.Helper()
	s := newFETStream(nil)
	for _, seg := range segments {
		u, done := s.Feed(false, false, false, 0, seg)
		if done {
			if u == nil {
				return nil
			}
			return u.M
		}
	}
	if u := s.Close(false); u != nil {
		return u.M
	}
	return nil
}

// encryptedSample returns deterministic high-entropy bytes, standing in for the
// first segment of a fully encrypted proxy protocol such as Shadowsocks.
func encryptedSample(n int) []byte {
	r := rand.New(rand.NewSource(1))
	b := make([]byte, n)
	for i := range b {
		b[i] = byte(r.Intn(256))
	}
	return b
}

func tlsClientHelloSample() []byte {
	return append([]byte{0x16, 0x03, 0x01, 0x02, 0x00, 0x01, 0x00, 0x01, 0xfc}, encryptedSample(64)...)
}

// A peer can force a one-byte first segment by advertising a tiny TCP receive
// window in its SYN-ACK. The verdict must not depend on where the boundary
// falls: classifying the first segment alone used to exempt these streams.
func TestFETVerdictIsSegmentationIndependent(t *testing.T) {
	samples := map[string][]byte{
		"encrypted": encryptedSample(64),
		"tls":       tlsClientHelloSample(),
		"http":      []byte("GET /index.html HTTP/1.1\r\nHost: example.com\r\n\r\n"),
	}
	for name, data := range samples {
		t.Run(name, func(t *testing.T) {
			whole := feedFET(t, data)
			if whole == nil {
				t.Fatal("no verdict on unsplit data")
			}
			split := feedFET(t, data[:1], data[1:])
			if got, want := split.Get("yes"), whole.Get("yes"); got != want {
				t.Errorf("split verdict yes = %v, unsplit = %v", got, want)
			}
		})
	}
}

func TestFETDetectsEncryptedFirstSegment(t *testing.T) {
	data := encryptedSample(64)
	if got := feedFET(t, data).Get("yes"); got != true {
		t.Fatalf("unsplit: yes = %v, want true", got)
	}
	if got := feedFET(t, data[:1], data[1:]).Get("yes"); got != true {
		t.Errorf("one-byte first segment: yes = %v, want true", got)
	}
}

func TestFETExemptsTLS(t *testing.T) {
	data := tlsClientHelloSample()
	if got := feedFET(t, data).Get("yes"); got != false {
		t.Fatalf("unsplit: yes = %v, want false", got)
	}
	if got := feedFET(t, data[:1], data[1:]).Get("yes"); got != false {
		t.Errorf("one-byte first segment: yes = %v, want false", got)
	}
}

// Waiting for more data must be bounded, or a peer that sends one byte and then
// stops would hold the stream open until the TCP timeout, keeping every packet
// of the connection in userspace.
func TestFETGivesUpAfterMaxSegments(t *testing.T) {
	s := newFETStream(nil)
	for i := 0; i < fetMaxSegments-1; i++ {
		if _, done := s.Feed(false, false, false, 0, []byte{0x00}); done {
			t.Fatalf("verdict reached after %d segments, want to keep waiting", i+1)
		}
	}
	u, done := s.Feed(false, false, false, 0, []byte{0x00})
	if !done || u == nil {
		t.Fatalf("no verdict after %d segments", fetMaxSegments)
	}
}

// The metrics only describe a client's opening bytes, so a server that speaks
// first should end the stream rather than be classified as if it were a client.
func TestFETServerSpeaksFirst(t *testing.T) {
	s := newFETStream(nil)
	u, done := s.Feed(true, false, false, 0, encryptedSample(64))
	if !done {
		t.Error("stream not done after server-first data")
	}
	if u != nil {
		t.Errorf("props = %#v, want none", u.M)
	}
}