package server

import (
	"bytes"
	"crypto/tls"
	"net"
	"testing"
	"time"
)

// TestReadClientHello captures a real ClientHello produced by crypto/tls and
// checks the parser and the byte-exact replay capture.
func TestReadClientHello(t *testing.T) {
	client, serverConn := net.Pipe()
	defer client.Close()

	var written bytes.Buffer
	go func() {
		// The client writes its ClientHello and then blocks waiting for the
		// server; we only need the write side.
		tlsClient := tls.Client(captureConn{client, &written}, &tls.Config{
			ServerName:         "api.demo.opentunnel.xyz",
			NextProtos:         []string{"h2", "http/1.1"},
			InsecureSkipVerify: true,
		})
		tlsClient.Handshake() // fails: nobody answers
	}()

	serverConn.SetReadDeadline(time.Now().Add(5 * time.Second))
	raw, sni, alpn, err := readClientHello(serverConn)
	if err != nil {
		t.Fatalf("readClientHello: %v", err)
	}
	if sni != "api.demo.opentunnel.xyz" {
		t.Errorf("sni = %q", sni)
	}
	if len(alpn) != 2 || alpn[0] != "h2" || alpn[1] != "http/1.1" {
		t.Errorf("alpn = %v", alpn)
	}
	if !bytes.HasPrefix(written.Bytes(), raw) {
		t.Errorf("captured %d bytes are not a prefix of the %d written", len(raw), written.Len())
	}
}

func TestReadClientHelloLengthBit2(t *testing.T) {
	// Regression: a handshake length with bit 2 set used to compute `need`
	// four bytes short, truncating the extensions block and losing SNI/ALPN.
	sni := "demo.com" // body length 76: bit 2 is set
	record := syntheticClientHello(sni, []string{"h2", "http/1.1"})
	if bodyLen := len(record) - 9; bodyLen&4 == 0 {
		t.Fatalf("test hello body length %d does not exercise the bug", bodyLen)
	}

	client, serverConn := net.Pipe()
	go func() {
		client.Write(append(record, 0xAA)) // sentinel after the hello
		client.Close()
	}()
	serverConn.SetReadDeadline(time.Now().Add(2 * time.Second))
	raw, gotSNI, alpn, err := readClientHello(serverConn)
	if err != nil {
		t.Fatalf("readClientHello: %v", err)
	}
	if gotSNI != sni {
		t.Errorf("sni = %q, want %q", gotSNI, sni)
	}
	if len(alpn) != 2 || alpn[0] != "h2" || alpn[1] != "http/1.1" {
		t.Errorf("alpn = %v", alpn)
	}
	if len(raw) != len(record) {
		t.Errorf("raw = %d bytes, want %d", len(raw), len(record))
	}
	rest := make([]byte, 1)
	if _, err := serverConn.Read(rest); err != nil || rest[0] != 0xAA {
		t.Errorf("read past ClientHello: %v %x", err, rest)
	}
}

// syntheticClientHello builds one TLS record carrying a minimal ClientHello
// with the given SNI and ALPN protocols.
func syntheticClientHello(sni string, protocols []string) []byte {
	body := []byte{0x03, 0x03}               // legacy version
	body = append(body, make([]byte, 32)...) // random
	body = append(body, 0)                   // session id length
	body = append(body, 0, 2, 0x13, 0x01)    // cipher suites
	body = append(body, 1, 0)                // compression methods

	extension := func(typ uint16, data []byte) []byte {
		out := []byte{byte(typ >> 8), byte(typ), byte(len(data) >> 8), byte(len(data))}
		return append(out, data...)
	}

	name := []byte(sni)
	sniList := []byte{0, 0, byte(len(name))} // name type 0, name length
	sniList = append(sniList, name...)
	sniData := append([]byte{byte(len(sniList) >> 8), byte(len(sniList))}, sniList...)

	var alpnList []byte
	for _, protocol := range protocols {
		alpnList = append(alpnList, byte(len(protocol)))
		alpnList = append(alpnList, protocol...)
	}
	alpnData := append([]byte{byte(len(alpnList) >> 8), byte(len(alpnList))}, alpnList...)

	exts := extension(0, sniData)
	exts = append(exts, extension(16, alpnData)...)
	body = append(body, byte(len(exts)>>8), byte(len(exts)))
	body = append(body, exts...)

	handshake := []byte{1, byte(len(body) >> 16), byte(len(body) >> 8), byte(len(body))}
	handshake = append(handshake, body...)
	record := []byte{22, 0x03, 0x01, byte(len(handshake) >> 8), byte(len(handshake))}
	return append(record, handshake...)
}

// captureConn records everything written to it.
type captureConn struct {
	net.Conn
	w *bytes.Buffer
}

func (c captureConn) Write(p []byte) (int, error) {
	c.w.Write(p)
	return c.Conn.Write(p)
}

func TestReadClientHelloGarbage(t *testing.T) {
	client, serverConn := net.Pipe()
	go func() {
		client.Write([]byte("not tls at all"))
		client.Close()
	}()
	serverConn.SetReadDeadline(time.Now().Add(2 * time.Second))
	if _, _, _, err := readClientHello(serverConn); err == nil {
		t.Error("expected an error for non-TLS bytes")
	}
}

func TestReadClientHelloTrickle(t *testing.T) {
	// A ClientHello trickling in byte by byte still parses.
	client, serverConn := net.Pipe()
	var written bytes.Buffer
	go func() {
		tlsClient := tls.Client(captureConn{client, &written}, &tls.Config{
			ServerName:         "demo.opentunnel.xyz",
			InsecureSkipVerify: true,
		})
		tlsClient.Handshake()
	}()
	serverConn.SetReadDeadline(time.Now().Add(5 * time.Second))
	slow := &slowConn{Conn: serverConn} // one-byte reads, a trickling sender
	raw, sni, _, err := readClientHello(slow)
	if err != nil {
		t.Fatalf("readClientHello: %v", err)
	}
	if sni != "demo.opentunnel.xyz" {
		t.Errorf("sni = %q", sni)
	}
	if !bytes.HasPrefix(written.Bytes(), raw) {
		t.Error("capture mismatch")
	}
}

type slowConn struct{ net.Conn }

func (c *slowConn) Read(p []byte) (int, error) {
	return c.Conn.Read(p[:1])
}