package tcp

import (
	"reflect"
	"testing"

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

// tlsExampleClientHello is taken from <https://tls12.xargs.org/#client-hello/annotated>.
var tlsExampleClientHello = []byte{
	0x16, 0x03, 0x01, 0x00, 0xa5, 0x01, 0x00, 0x00, 0xa1, 0x03, 0x03, 0x00,
	0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c,
	0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18,
	0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, 0x00, 0x00, 0x20, 0xcc, 0xa8,
	0xcc, 0xa9, 0xc0, 0x2f, 0xc0, 0x30, 0xc0, 0x2b, 0xc0, 0x2c, 0xc0, 0x13,
	0xc0, 0x09, 0xc0, 0x14, 0xc0, 0x0a, 0x00, 0x9c, 0x00, 0x9d, 0x00, 0x2f,
	0x00, 0x35, 0xc0, 0x12, 0x00, 0x0a, 0x01, 0x00, 0x00, 0x58, 0x00, 0x00,
	0x00, 0x18, 0x00, 0x16, 0x00, 0x00, 0x13, 0x65, 0x78, 0x61, 0x6d, 0x70,
	0x6c, 0x65, 0x2e, 0x75, 0x6c, 0x66, 0x68, 0x65, 0x69, 0x6d, 0x2e, 0x6e,
	0x65, 0x74, 0x00, 0x05, 0x00, 0x05, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x0a, 0x00, 0x0a, 0x00, 0x08, 0x00, 0x1d, 0x00, 0x17, 0x00, 0x18, 0x00,
	0x19, 0x00, 0x0b, 0x00, 0x02, 0x01, 0x00, 0x00, 0x0d, 0x00, 0x12, 0x00,
	0x10, 0x04, 0x01, 0x04, 0x03, 0x05, 0x01, 0x05, 0x03, 0x06, 0x01, 0x06,
	0x03, 0x02, 0x01, 0x02, 0x03, 0xff, 0x01, 0x00, 0x01, 0x00, 0x00, 0x12,
	0x00, 0x00,
}

var tlsExampleClientHelloProps = analyzer.PropMap{
	"ciphers":     []uint16{52392, 52393, 49199, 49200, 49195, 49196, 49171, 49161, 49172, 49162, 156, 157, 47, 53, 49170, 10},
	"compression": []uint8{0},
	"random":      []uint8{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31},
	"session":     []uint8{},
	"sni":         "example.ulfheim.net",
	"version":     uint16(771),
}

func TestTlsStreamParsing_ClientHello(t *testing.T) {
	s := newTLSStream(nil)
	u, _ := s.Feed(false, false, false, 0, tlsExampleClientHello)
	if got := u.M.Get("req"); !reflect.DeepEqual(got, tlsExampleClientHelloProps) {
		t.Errorf("parsed = %v, want %v", got, tlsExampleClientHelloProps)
	}
	if got := u.M.Get("req_records"); got != 1 {
		t.Errorf("req_records = %v, want 1", got)
	}
}

func TestTlsStreamParsing_ServerHello(t *testing.T) {
	// example packet taken from <https://tls12.xargs.org/#server-hello/annotated>
	serverHello := []byte{
		0x16, 0x03, 0x03, 0x00, 0x31, 0x02, 0x00, 0x00, 0x2d, 0x03, 0x03, 0x70,
		0x71, 0x72, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7a, 0x7b, 0x7c,
		0x7d, 0x7e, 0x7f, 0x80, 0x81, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88,
		0x89, 0x8a, 0x8b, 0x8c, 0x8d, 0x8e, 0x8f, 0x00, 0xc0, 0x13, 0x00, 0x00,
		0x05, 0xff, 0x01, 0x00, 0x01, 0x00,
	}
	want := analyzer.PropMap{
		"cipher":      uint16(49171),
		"compression": uint8(0),
		"random":      []uint8{112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143},
		"session":     []uint8{},
		"version":     uint16(771),
	}

	s := newTLSStream(nil)
	u, _ := s.Feed(true, false, false, 0, serverHello)
	got := u.M.Get("resp")
	if !reflect.DeepEqual(got, want) {
		t.Errorf("%d B parsed = %v, want %v", len(serverHello), got, want)
	}
}

// tlsFragment re-encodes the handshake message in a single-record hello as
// records carrying at most size bytes of it each.
func tlsFragment(hello []byte, size int) []byte {
	var out []byte
	for msg := hello[5:]; len(msg) > 0; {
		n := min(size, len(msg))
		out = append(out, hello[0], hello[1], hello[2], byte(n>>8), byte(n))
		out = append(out, msg[:n]...)
		msg = msg[n:]
	}
	return out
}

// feedTLS feeds each segment in turn and returns the properties as the engine
// would have merged them, along with whether the stream said it was done.
func feedTLS(s *tlsStream, rev bool, segs ...[]byte) (analyzer.PropMap, bool) {
	props := analyzer.PropMap{}
	done := false
	for _, seg := range segs {
		var u *analyzer.PropUpdate
		u, done = s.Feed(rev, false, false, 0, seg)
		if u != nil {
			for k, v := range u.M {
				props[k] = v
			}
		}
	}
	return props, done
}

// Splitting the ClientHello across records must not hide it: the parse result
// has to be the same wherever the record boundaries fall.
func TestTlsClientHelloRecordFragmentation(t *testing.T) {
	msgLen := len(tlsExampleClientHello) - 5
	for _, size := range []int{1, 2, 3, 4, 5, 9, 50, msgLen - 1} {
		data := tlsFragment(tlsExampleClientHello, size)
		wantRecords := (msgLen + size - 1) / size
		for name, segs := range map[string][][]byte{
			"one segment":  {data},
			"byte by byte": splitEvery(data, 1),
			"7 byte segs":  splitEvery(data, 7),
		} {
			props, _ := feedTLS(newTLSStream(nil), false, segs...)
			if got := props.Get("req"); !reflect.DeepEqual(got, tlsExampleClientHelloProps) {
				t.Errorf("size %d, %s: parsed = %v, want %v", size, name, got, tlsExampleClientHelloProps)
			}
			if got := props.Get("req_records"); got != wantRecords {
				t.Errorf("size %d, %s: req_records = %v, want %d", size, name, got, wantRecords)
			}
		}
	}
}

// Fragmentation is published before the hello is complete, so that a rule can
// act on it even if the hello never completes.
func TestTlsClientHelloFragmentationReportedEarly(t *testing.T) {
	data := tlsFragment(tlsExampleClientHello, 1)
	s := newTLSStream(nil)
	props, done := feedTLS(s, false, data[:12])
	if done {
		t.Fatal("stream done after two records")
	}
	if props.Get("req") != nil {
		t.Errorf("req published early: %v", props.Get("req"))
	}
	if got := props.Get("req_records"); got != 2 {
		t.Errorf("req_records = %v, want 2", got)
	}

	// Further records are not published one by one...
	if u, _ := s.Feed(false, false, false, 0, data[12:60]); u != nil {
		t.Errorf("unexpected update %v", u.M)
	}
	// ...but the count reached is, if the stream gets cut off.
	u := s.Close(true)
	if u == nil {
		t.Fatal("no update on close")
	}
	if got := u.M.Get("req_records"); got != 10 {
		t.Errorf("req_records on close = %v, want 10", got)
	}
}

func TestTlsClientHelloRejectsInterleavedRecords(t *testing.T) {
	data := tlsFragment(tlsExampleClientHello, 50)
	// Turn the second record into an alert.
	data[5+50] = 0x15
	props, done := feedTLS(newTLSStream(nil), false, data)
	if !done {
		t.Error("stream not done")
	}
	if props.Get("req") != nil {
		t.Errorf("req = %v, want nil", props.Get("req"))
	}
}

func TestTlsStreamIgnoresUnfragmentedPartialHelloOnClose(t *testing.T) {
	s := newTLSStream(nil)
	if props, _ := feedTLS(s, false, tlsExampleClientHello[:100]); len(props) != 0 {
		t.Errorf("props = %v, want none", props)
	}
	if u := s.Close(true); u != nil {
		t.Errorf("close update = %v, want nil", u.M)
	}
}

func splitEvery(data []byte, n int) [][]byte {
	var segs [][]byte
	for len(data) > 0 {
		k := min(n, len(data))
		segs = append(segs, data[:k])
		data = data[k:]
	}
	return segs
}