package server

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"time"
)

// clientHelloLimit bounds the bytes read while looking for a ClientHello
// (the reference uses 64 KiB and 10 seconds).
const (
	clientHelloLimit   = 64 * 1024
	clientHelloTimeout = 10 * time.Second
)

// readClientHello reads TLS records from conn until the ClientHello is
// complete, without consuming anything beyond it. It returns the exact bytes
// read (for replay or forwarding), the SNI server name, and the offered ALPN
// protocols. Reads are exact (no read-ahead), so conn is handed on intact.
func readClientHello(conn net.Conn) (raw []byte, sni string, alpn []string, err error) {
	conn.SetReadDeadline(time.Now().Add(clientHelloTimeout))
	defer conn.SetReadDeadline(time.Time{})

	var hello []byte // handshake message fragments, starting with the 4-byte header
	need := -1       // total ClientHello bytes once the header is known

	for len(raw) < clientHelloLimit {
		header := make([]byte, 5)
		if _, err = io.ReadFull(conn, header); err != nil {
			return nil, "", nil, fmt.Errorf("reading record header: %w", err)
		}
		if header[0] != 22 { // handshake
			return nil, "", nil, fmt.Errorf("expected a handshake record, got %d", header[0])
		}
		bodyLen := int(binary.BigEndian.Uint16(header[3:]))
		if len(raw)+5+bodyLen > clientHelloLimit {
			return nil, "", nil, errors.New("ClientHello exceeds 64 KiB")
		}
		body := make([]byte, bodyLen)
		if _, err = io.ReadFull(conn, body); err != nil {
			return nil, "", nil, fmt.Errorf("reading record body: %w", err)
		}
		raw = append(raw, header...)
		raw = append(raw, body...)
		hello = append(hello, body...)

		if need < 0 {
			if len(hello) < 4 {
				continue
			}
			if hello[0] != 1 { // client_hello
				return nil, "", nil, fmt.Errorf("expected ClientHello, got handshake type %d", hello[0])
			}
			need = 4 + int(hello[1])<<16 | int(hello[2])<<8 | int(hello[3])
		}
		if len(hello) >= need {
			sni, alpn = parseClientHello(hello[4:need])
			return raw, sni, alpn, nil
		}
	}
	return nil, "", nil, errors.New("ClientHello exceeds 64 KiB")
}

// parseClientHello extracts SNI and ALPN from a ClientHello body (everything
// after the 4-byte handshake header). Truncation is an error; a missing
// extension simply yields an empty result, as the relay closes such
// connections anyway.
func parseClientHello(body []byte) (sni string, alpn []string) {
	p := parser{data: body}
	p.skip(2)              // version
	p.skip(32)             // random
	p.skip(int(p.u8()))    // session id
	p.skip(int(p.u16()))   // cipher suites
	p.skip(int(p.u8()))    // compression methods
	extLen := int(p.u16()) // extensions block
	if p.err {
		return "", nil
	}
	exts := parser{data: p.take(extLen)}
	for !exts.empty() {
		typ := exts.u16()
		data := exts.take(int(exts.u16()))
		if exts.err {
			break
		}
		switch typ {
		case 0: // server_name
			d := parser{data: data}
			list := parser{data: d.take(int(d.u16()))}
			for !list.empty() {
				nameType := list.u8()
				name := list.take(int(list.u16()))
				if list.err {
					break
				}
				if nameType == 0 {
					sni = string(name)
				}
			}
		case 16: // application_layer_protocol_negotiation
			d := parser{data: data}
			list := parser{data: d.take(int(d.u16()))}
			for !list.empty() {
				proto := list.take(int(list.u8()))
				if list.err {
					break
				}
				alpn = append(alpn, string(proto))
			}
		}
	}
	return sni, alpn
}

// parser is a cursor over a byte slice; any out-of-bounds read sets err and
// yields zero values from then on.
type parser struct {
	data []byte
	err  bool
}

func (p *parser) empty() bool { return p.err || len(p.data) == 0 }

func (p *parser) take(n int) []byte {
	if n < 0 || len(p.data) < n {
		p.err = true
		return nil
	}
	out := p.data[:n]
	p.data = p.data[n:]
	return out
}

func (p *parser) skip(n int) { p.take(n) }

func (p *parser) u8() byte {
	b := p.take(1)
	if b == nil {
		return 0
	}
	return b[0]
}

func (p *parser) u16() uint16 {
	b := p.take(2)
	if b == nil {
		return 0
	}
	return binary.BigEndian.Uint16(b)
}