package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"testing"
	"time"

	"dragontcp/internal/cover"
	"dragontcp/internal/protocol"
	"dragontcp/internal/wire"
)

func TestParseOpenAllowsEmptyToken(t *testing.T) {
	host := "example.com"
	p := make([]byte, 6+len(host))
	binary.BigEndian.PutUint16(p[0:2], 0)
	binary.BigEndian.PutUint16(p[2:4], uint16(len(host)))
	binary.BigEndian.PutUint16(p[4:6], 443)
	copy(p[6:], host)
	token, gotHost, port, err := parseOpen(p)
	if err != nil {
		t.Fatal(err)
	}
	if token != "" || gotHost != host || port != 443 {
		t.Fatalf("got token=%q host=%q port=%d", token, gotHost, port)
	}
}

func TestBinaryProfileProbeEndToEnd(t *testing.T) {
	for n := 0; n < 256; n += 8 {
		mask := byte(n)
		server, client := net.Pipe()
		clientResult := make(chan error, 1)
		go func() {
			defer client.Close()
			var sid wire.SessionID
			payload := make([]byte, 11)
			copy(payload[:4], wire.ProbeMagic[:])
			payload[4] = wire.ProbeKeepalive
			if err := wire.WriteRequestProfile(client, wire.ModeProbe, sid, 1, payload, mask); err != nil {
				clientResult <- err
				return
			}
			status, _, err := wire.ReadResponseProfile(client, mask)
			if err == nil && status != wire.StatusOK {
				err = fmt.Errorf("status=%d", status)
			}
			clientResult <- err
		}()

		profiled, isXOR, gotMask, err := sniffWire(server)
		if err != nil || isXOR || gotMask != mask {
			t.Fatalf("mask %02x sniff: xor=%t gotMask=%02x err=%v", mask, isXOR, gotMask, err)
		}
		req, err := wire.ReadRequestProfile(profiled, gotMask)
		if err == nil {
			err = processWireRequest(profiled, req, "", false, nil, 0, nil, 1024, 0, 0, nil)
		}
		if err != nil {
			t.Fatalf("mask %02x server: %v", mask, err)
		}
		if err := <-clientResult; err != nil {
			t.Fatalf("mask %02x client: %v", mask, err)
		}
		_ = server.Close()
	}
}

func TestXORProfileProbeEndToEnd(t *testing.T) {
	for n := 0; n < 256; n++ {
		mask := byte(n)
		if ('U'^mask)&7 < 5 {
			continue
		}
		server, client := net.Pipe()
		clientResult := make(chan error, 1)
		go func() {
			defer client.Close()
			if err := protocol.WriteRequestFrameProfile(client, 7, []byte("CPROBE -"), mask); err != nil {
				clientResult <- err
				return
			}
			id, payload, err := protocol.ReadResponseFrameProfile(client, mask)
			if err == nil && (id != 7 || string(payload) != "PROBEOK") {
				err = fmt.Errorf("id=%d payload=%q", id, payload)
			}
			clientResult <- err
		}()

		profiled, isXOR, gotMask, err := sniffWire(server)
		if err != nil || !isXOR || gotMask != mask {
			t.Fatalf("mask %02x sniff: xor=%t gotMask=%02x err=%v", mask, isXOR, gotMask, err)
		}
		handleXOR(profiled, gotMask, "", false, nil, 0, nil, 1024, 8, time.Millisecond, nil)
		if err := <-clientResult; err != nil {
			t.Fatalf("mask %02x client: %v", mask, err)
		}
		_ = server.Close()
	}
}

func TestCoveredProfilesProbeEndToEnd(t *testing.T) {
	for _, padding := range []uint16{0, 64, cover.MaxPadding} {
		for _, xor := range []bool{false, true} {
			profile := cover.Profile{Enabled: true, ID: 0x91e7, Padding: padding, HeaderMask: 0x6b, XOR: xor}
			server, client := net.Pipe()
			clientResult := make(chan error, 1)
			go func() {
				defer client.Close()
				if err := cover.WritePreface(client, profile); err != nil {
					clientResult <- err
					return
				}
				if xor {
					if err := protocol.WriteRequestFrameProfile(client, 11, []byte("CPROBE -"), profile.HeaderMask); err != nil {
						clientResult <- err
						return
					}
					id, payload, err := protocol.ReadResponseFrameProfile(client, profile.HeaderMask)
					if err == nil && (id != 11 || string(payload) != "PROBEOK") {
						err = fmt.Errorf("id=%d payload=%q", id, payload)
					}
					clientResult <- err
					return
				}

				var sid wire.SessionID
				payload := make([]byte, 11)
				copy(payload[:4], wire.ProbeMagic[:])
				payload[4] = wire.ProbeKeepalive
				if err := wire.WriteRequestProfile(client, wire.ModeProbe, sid, 3, payload, profile.HeaderMask); err != nil {
					clientResult <- err
					return
				}
				status, _, err := wire.ReadResponseProfile(client, profile.HeaderMask)
				if err == nil && status != wire.StatusOK {
					err = fmt.Errorf("status=%d", status)
				}
				clientResult <- err
			}()

			profiled, gotXOR, gotMask, err := sniffWire(server)
			if err != nil || gotXOR != xor || gotMask != profile.HeaderMask {
				t.Fatalf("padding=%d xor=%t sniff got xor=%t mask=%02x err=%v", padding, xor, gotXOR, gotMask, err)
			}
			if xor {
				handleXOR(profiled, gotMask, "", false, nil, 0, nil, 1024, 8, time.Millisecond, nil)
			} else {
				req, readErr := wire.ReadRequestProfile(profiled, gotMask)
				if readErr == nil {
					readErr = processWireRequest(profiled, req, "", false, nil, 0, nil, 1024, 0, 0, nil)
				}
				if readErr != nil {
					t.Fatalf("padding=%d binary server: %v", padding, readErr)
				}
			}
			if err := <-clientResult; err != nil {
				t.Fatalf("padding=%d xor=%t client: %v", padding, xor, err)
			}
			_ = server.Close()
		}
	}
}

func TestSniffWireRecognizesAllHeaderProfiles(t *testing.T) {
	test := func(firstTwo []byte, wantXOR bool, wantMask byte) {
		server, client := net.Pipe()
		defer server.Close()
		go func() {
			initial := make([]byte, 12)
			copy(initial, firstTwo)
			_, _ = client.Write(initial)
			_ = client.Close()
		}()

		profiled, gotXOR, gotMask, err := sniffWire(server)
		if err != nil {
			t.Fatalf("header=%x: %v", firstTwo, err)
		}
		if gotXOR != wantXOR || gotMask != wantMask {
			t.Fatalf("header=%x got xor=%t mask=%02x, want xor=%t mask=%02x", firstTwo, gotXOR, gotMask, wantXOR, wantMask)
		}
		replayed := make([]byte, 2)
		if _, err := io.ReadFull(profiled, replayed); err != nil || !bytes.Equal(replayed, firstTwo) {
			t.Fatalf("header=%x replay=%x err=%v", firstTwo, replayed, err)
		}
	}

	for n := 0; n < 256; n += 8 {
		mask := byte(n)
		test([]byte{mask, 0xa7}, false, mask)
	}
	for n := 0; n < 256; n++ {
		mask := byte(n)
		if ('U'^mask)&7 >= 5 {
			test([]byte{'U' ^ mask, 'P' ^ mask}, true, mask)
		}
	}
}