package engine

import (
	"reflect"
	"testing"

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

	"github.com/google/gopacket/layers"
)

func newTestStream() *tcpStream {
	return &tcpStream{
		info:         ruleset.StreamInfo{Props: analyzer.CombinedPropMap{}},
		synAckWindow: -1,
		synAckWscale: -1,
	}
}

func TestTCPWindowScale(t *testing.T) {
	mss := layers.TCPOption{OptionType: layers.TCPOptionKindMSS, OptionLength: 4, OptionData: []byte{0x05, 0xb4}}
	tests := []struct {
		name string
		opts []layers.TCPOption
		want int
	}{
		{
			name: "present",
			opts: []layers.TCPOption{mss, {OptionType: layers.TCPOptionKindWindowScale, OptionLength: 3, OptionData: []byte{7}}},
			want: 7,
		},
		{name: "absent", opts: []layers.TCPOption{mss}, want: -1},
		{
			name: "malformed is ignored",
			opts: []layers.TCPOption{{OptionType: layers.TCPOptionKindWindowScale, OptionLength: 2, OptionData: nil}},
			want: -1,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := tcpWindowScale(&layers.TCP{Options: tt.opts}); got != tt.want {
				t.Errorf("wscale = %d, want %d", got, tt.want)
			}
		})
	}
}

// Whichever signal materializes the prop map, the handshake fields must come
// along with it. Otherwise a rule like `tcpmeta != nil && tcpmeta.synack_window
// <= 16` sees a map without that key and fails at runtime with a nil comparison.
func TestTCPMetaAlwaysCarriesHandshakeFields(t *testing.T) {
	tests := []struct {
		name   string
		window int // -1 for a stream where no SYN-ACK was observed
		wscale int
		want   analyzer.PropMap
	}{
		{
			name:   "normal handshake, short first segment",
			window: 64240,
			wscale: 7,
			want:   analyzer.PropMap{"synack_window": 64240, "synack_wscale": 7, "first_seg_len": 21},
		},
		{
			name:   "rewritten window",
			window: 1,
			wscale: 7,
			want:   analyzer.PropMap{"synack_window": 1, "synack_wscale": 7, "first_seg_len": 21},
		},
		{
			name:   "no window scale negotiated",
			window: 64240,
			wscale: -1,
			want:   analyzer.PropMap{"synack_window": 64240, "first_seg_len": 21},
		},
		{
			name:   "return path not visible",
			window: -1,
			wscale: -1,
			want:   analyzer.PropMap{"first_seg_len": 21},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := newTestStream()
			s.synAckWindow, s.synAckWscale = tt.window, tt.wscale
			if !s.recordFirstSegLen(21) {
				t.Fatal("first segment not recorded")
			}
			if got := s.info.Props[tcpMetaPropKey]; !reflect.DeepEqual(got, tt.want) {
				t.Errorf("props = %#v, want %#v", got, tt.want)
			}
		})
	}
}

func TestTCPRecordFirstSegLen(t *testing.T) {
	tests := []struct {
		name     string
		segLen   int
		recorded bool
	}{
		{name: "one byte segment", segLen: 1, recorded: true},
		{name: "at the threshold", segLen: tcpFirstSegLenMax, recorded: true},
		{name: "just over the threshold", segLen: tcpFirstSegLenMax + 1, recorded: false},
		{name: "normal first segment", segLen: 517, recorded: false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := newTestStream()
			if got := s.recordFirstSegLen(tt.segLen); got != tt.recorded {
				t.Errorf("recorded = %v, want %v", got, tt.recorded)
			}
			m := s.info.Props[tcpMetaPropKey]
			if !tt.recorded {
				// A normal connection must not allocate a prop map at all.
				if m != nil {
					t.Errorf("props = %#v, want none", m)
				}
				return
			}
			if got := m["first_seg_len"]; got != tt.segLen {
				t.Errorf("first_seg_len = %v, want %d", got, tt.segLen)
			}
		})
	}
}

func TestTCPRecordFirstSegLenIsOneShot(t *testing.T) {
	s := newTestStream()
	s.recordFirstSegLen(1)
	if s.recordFirstSegLen(5) {
		t.Error("a later segment was recorded as the first one")
	}
	if got := s.info.Props[tcpMetaPropKey]["first_seg_len"]; got != 1 {
		t.Errorf("first_seg_len = %v, want 1", got)
	}
}

// A stream that opens normally must not be able to look like a split one just
// because a short segment shows up later.
func TestTCPNormalFirstSegmentClosesTheWindow(t *testing.T) {
	s := newTestStream()
	if s.recordFirstSegLen(517) {
		t.Fatal("a normal first segment should not be recorded")
	}
	if s.recordFirstSegLen(1) {
		t.Error("a later short segment was recorded as the first one")
	}
	if m := s.info.Props[tcpMetaPropKey]; m != nil {
		t.Errorf("props = %#v, want none", m)
	}
}

func TestTCPSynAckWindowThreshold(t *testing.T) {
	tests := []struct {
		window uint16
		record bool
	}{
		{window: 1, record: true},
		{window: tcpSynAckWindowMax, record: true},
		{window: tcpSynAckWindowMax + 1, record: false},
		{window: 64240, record: false},
	}
	for _, tt := range tests {
		if got := int(tt.window) <= tcpSynAckWindowMax; got != tt.record {
			t.Errorf("window %d: recorded = %v, want %v", tt.window, got, tt.record)
		}
	}
}