// SPDX-FileCopyrightText: 2026 Echolot contributors
// SPDX-License-Identifier: GPL-3.0-or-later

// Package tcpecho implements spec §4 TCP echo and its TLS variant on the same
// port. Plain connections get a JSON greeting (observed source, negotiated
// MSS and TCP options from TCP_INFO — the mtu.mss_observed evidence) then a
// byte echo. A connection that opens with a TLS handshake (first byte 0x16)
// and ALPN "elt-echo" gets, additionally, the ClientHello it sent back raw +
// as a JA4 fingerprint (sec.clienthello_echo) before the echo.
package tcpecho

import (
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"io"
	"net"
	"sync"
	"time"
)

// ConnRecord is what the observations API reports per connection (spec §6).
type ConnRecord struct {
	ConnectedAt time.Time `json:"connected_at"`
	Src         string    `json:"src"`
	MSS         int       `json:"mss"`
	Options     []string  `json:"options"`
	TLS         bool      `json:"tls"`
	JA4         string    `json:"ja4,omitempty"`
	ALPN        string    `json:"alpn,omitempty"`
}

type Server struct {
	// TLSConfig enables the elt-echo TLS variant; nil disables it (plain echo
	// only). "elt-echo" is appended to NextProtos at Serve time.
	TLSConfig *tls.Config

	mu     sync.Mutex
	recent []ConnRecord // ring, newest last
}

const recentCap = 1024

func (s *Server) record(r ConnRecord) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if len(s.recent) >= recentCap {
		s.recent = s.recent[1:]
	}
	s.recent = append(s.recent, r)
}

// RecentFor returns records whose source IP matches ip.
func (s *Server) RecentFor(ip string) []ConnRecord {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out []ConnRecord
	for _, r := range s.recent {
		if h, _, err := net.SplitHostPort(r.Src); err == nil && h == ip {
			out = append(out, r)
		}
	}
	return out
}

func (s *Server) Serve(ln net.Listener) error {
	for {
		conn, err := ln.Accept()
		if err != nil {
			return err
		}
		go s.handle(conn)
	}
}

// prefixConn replays already-read bytes before continuing with the underlying
// connection — used to hand the peeked ClientHello record to tls.Server.
type prefixConn struct {
	net.Conn
	prefix []byte
}

func (p *prefixConn) Read(b []byte) (int, error) {
	if len(p.prefix) > 0 {
		n := copy(b, p.prefix)
		p.prefix = p.prefix[n:]
		return n, nil
	}
	return p.Conn.Read(b)
}

func (s *Server) handle(conn net.Conn) {
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(5 * time.Minute))

	// TCP_INFO must be read from the raw *net.TCPConn, before any wrapping.
	info := tcpInfo(conn)
	rec := ConnRecord{
		ConnectedAt: time.Now().UTC(),
		Src:         conn.RemoteAddr().String(),
		MSS:         info.MSS,
		Options:     info.Options,
	}

	// Multiplex TLS vs plain on one port. Plain echo is server-speaks-first
	// (the client waits for the greeting), while a TLS client sends its
	// ClientHello immediately — so peek the first byte with a short deadline:
	// a byte that arrives fast and is 0x16 means TLS; a timeout means a plain
	// client waiting to be greeted.
	// 500ms tolerates ~1s RTT (incl. satellite) before a TLS ClientHello would
	// be misread as a silent plain client; plain clients simply wait this long
	// for the greeting they're already waiting for.
	first := make([]byte, 1)
	_ = conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
	n, err := io.ReadFull(conn, first)
	_ = conn.SetDeadline(time.Now().Add(5 * time.Minute)) // reset for the session
	switch {
	case err == nil && first[0] == 0x16 && s.TLSConfig != nil:
		s.handleTLS(conn, first, rec)
		return
	case err == nil:
		s.plainEcho(conn, first, rec) // client spoke first (rare) — replay it
		return
	case n == 0 && isTimeout(err):
		s.plainEcho(conn, nil, rec) // client waiting for greeting — normal path
		return
	default:
		return // EOF or a real error
	}
}

func (s *Server) plainEcho(conn net.Conn, peeked []byte, rec ConnRecord) {
	pc := &prefixConn{Conn: conn, prefix: peeked}
	s.record(rec)
	greeting, _ := json.Marshal(map[string]any{
		"observed_src": rec.Src, "mss": rec.MSS, "options": rec.Options, "tls": false,
	})
	if _, err := pc.Write(append(greeting, '\n')); err != nil {
		return
	}
	_, _ = io.Copy(pc, pc)
}

func isTimeout(err error) bool {
	ne, ok := err.(net.Error)
	return ok && ne.Timeout()
}

// handleTLS captures the full ClientHello record, computes JA4, completes the
// handshake, then greets with the ClientHello (raw + JA4) and echoes over TLS.
func (s *Server) handleTLS(conn net.Conn, first []byte, rec ConnRecord) {
	// Read the rest of the record header (version[2], length[2]) and the body.
	hdr := make([]byte, 4)
	if _, err := io.ReadFull(conn, hdr); err != nil {
		return
	}
	recLen := int(hdr[2])<<8 | int(hdr[3])
	body := make([]byte, recLen)
	if _, err := io.ReadFull(conn, body); err != nil {
		return
	}
	full := append(append(append([]byte{}, first...), hdr...), body...)

	rec.TLS = true
	if h, ok := parseClientHello(full); ok {
		rec.JA4 = ja4(h)
	}

	// Replay the captured ClientHello into the TLS server.
	cfg := s.TLSConfig.Clone()
	cfg.NextProtos = append([]string{"elt-echo"}, cfg.NextProtos...)
	tconn := tls.Server(&prefixConn{Conn: conn, prefix: full}, cfg)
	if err := tconn.Handshake(); err != nil {
		return
	}
	rec.ALPN = tconn.ConnectionState().NegotiatedProtocol
	s.record(rec)

	greeting, _ := json.Marshal(map[string]any{
		"observed_src":    rec.Src,
		"mss":             rec.MSS,
		"options":         rec.Options,
		"tls":             true,
		"alpn":            rec.ALPN,
		"ja4":             rec.JA4,
		"clienthello_b64": base64.StdEncoding.EncodeToString(full),
	})
	if _, err := tconn.Write(append(greeting, '\n')); err != nil {
		return
	}
	_, _ = io.Copy(tconn, tconn)
}