helix-proxy/internal/proxy/whoami_test.go

package proxy

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"helix-proxy/internal/store"

	"golang.org/x/net/websocket"
)

// WhoamiRequest captures what the upstream mock received.
type WhoamiRequest struct {
	Method     string              `json:"method"`
	Path       string              `json:"path"`
	Query      string              `json:"query,omitempty"`
	Host       string              `json:"host"`
	Headers    map[string][]string `json:"headers"`
	Body       string              `json:"body,omitempty"`
	RemoteAddr string              `json:"remoteAddr"`
}

// WhoamiRoute configures a response for a specific path (and optionally method).
type WhoamiRoute struct {
	Method  string
	Status  int
	Body    string
	Headers map[string]string
}

// WhoamiConfig configures whoami upstream behavior.
type WhoamiConfig struct {
	DefaultStatus  int
	DefaultBody    string
	DefaultHeaders map[string]string
	Routes         map[string]WhoamiRoute // key: path or "METHOD path"
	ValidPaths     []string               // exact paths allowed; empty = all
	ValidPrefixes  []string               // path prefixes allowed (checked after ValidPaths)
}

// WhoamiServer is a test upstream that records requests and returns configurable responses.
type WhoamiServer struct {
	Server   *httptest.Server
	Hostname string
	Port     int

	cfg WhoamiConfig

	mu       sync.Mutex
	requests []WhoamiRequest
}

// NewWhoami starts a whoami upstream and registers cleanup on t.
func NewWhoami(t *testing.T, cfg WhoamiConfig) *WhoamiServer {
	t.Helper()
	if cfg.DefaultStatus == 0 {
		cfg.DefaultStatus = http.StatusOK
	}
	w := &WhoamiServer{cfg: cfg}
	w.Server = httptest.NewServer(http.HandlerFunc(w.serve))
	t.Cleanup(w.Server.Close)

	u, err := url.Parse(w.Server.URL)
	if err != nil {
		t.Fatalf("whoami url: %v", err)
	}
	w.Hostname = u.Hostname()
	w.Port, err = strconv.Atoi(u.Port())
	if err != nil {
		t.Fatalf("whoami port: %v", err)
	}
	return w
}

func (w *WhoamiServer) serve(rw http.ResponseWriter, r *http.Request) {
	body, _ := io.ReadAll(r.Body)
	_ = r.Body.Close()

	req := WhoamiRequest{
		Method:     r.Method,
		Path:       r.URL.Path,
		Query:      r.URL.RawQuery,
		Host:       r.Host,
		Headers:    cloneHeaderMap(r.Header),
		Body:       string(body),
		RemoteAddr: r.RemoteAddr,
	}

	w.mu.Lock()
	w.requests = append(w.requests, req)
	w.mu.Unlock()

	if !w.pathAllowed(r.URL.Path) {
		http.Error(rw, "whoami: path not allowed", http.StatusNotFound)
		return
	}

	route, ok := w.matchRoute(r.Method, r.URL.Path)
	status := w.cfg.DefaultStatus
	respBody := w.cfg.DefaultBody
	respHeaders := map[string]string{}
	for k, v := range w.cfg.DefaultHeaders {
		respHeaders[k] = v
	}

	if ok {
		if route.Status != 0 {
			status = route.Status
		}
		if route.Body != "" {
			respBody = route.Body
		}
		for k, v := range route.Headers {
			respHeaders[k] = v
		}
	}

	// If no custom body, echo the captured request as JSON (whoami-style).
	if respBody == "" {
		b, err := json.Marshal(req)
		if err != nil {
			http.Error(rw, "whoami: marshal", http.StatusInternalServerError)
			return
		}
		respBody = string(b)
	}

	rw.Header().Set("Content-Type", "application/json")
	rw.Header().Set("X-Whoami-Method", req.Method)
	rw.Header().Set("X-Whoami-Path", req.Path)
	rw.Header().Set("X-Whoami-Host", req.Host)
	rw.Header().Set("X-Whoami-Query", req.Query)
	for name, vals := range req.Headers {
		rw.Header().Set("X-Whoami-H-"+name, strings.Join(vals, ", "))
	}
	for k, v := range respHeaders {
		rw.Header().Set(k, v)
	}
	rw.WriteHeader(status)
	_, _ = io.WriteString(rw, respBody)
}

func (w *WhoamiServer) pathAllowed(path string) bool {
	if len(w.cfg.ValidPaths) == 0 && len(w.cfg.ValidPrefixes) == 0 {
		return true
	}
	for _, p := range w.cfg.ValidPaths {
		if path == p {
			return true
		}
	}
	for _, p := range w.cfg.ValidPrefixes {
		if strings.HasPrefix(path, p) {
			return true
		}
	}
	return false
}

func (w *WhoamiServer) matchRoute(method, path string) (WhoamiRoute, bool) {
	if len(w.cfg.Routes) == 0 {
		return WhoamiRoute{}, false
	}
	key := strings.ToUpper(method) + " " + path
	if r, ok := w.cfg.Routes[key]; ok {
		return r, true
	}
	if r, ok := w.cfg.Routes[path]; ok {
		return r, true
	}
	return WhoamiRoute{}, false
}

// Requests returns a copy of all recorded upstream requests.
func (w *WhoamiServer) Requests() []WhoamiRequest {
	w.mu.Lock()
	defer w.mu.Unlock()
	out := make([]WhoamiRequest, len(w.requests))
	copy(out, w.requests)
	return out
}

// LastRequest returns the most recent upstream request, if any.
func (w *WhoamiServer) LastRequest() (WhoamiRequest, bool) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if len(w.requests) == 0 {
		return WhoamiRequest{}, false
	}
	return w.requests[len(w.requests)-1], true
}

// ClearRequests resets the recorded request history.
func (w *WhoamiServer) ClearRequests() {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.requests = nil
}

// ParseWhoamiBody decodes a JSON whoami echo body from a proxied response.
func ParseWhoamiBody(body []byte) (WhoamiRequest, error) {
	var req WhoamiRequest
	err := json.Unmarshal(body, &req)
	return req, err
}

// HeaderValue returns the first value for a header name (case-insensitive).
func (r WhoamiRequest) HeaderValue(name string) string {
	for k, vs := range r.Headers {
		if strings.EqualFold(k, name) && len(vs) > 0 {
			return vs[0]
		}
	}
	return ""
}

func cloneHeaderMap(h http.Header) map[string][]string {
	out := make(map[string][]string, len(h))
	for k, vs := range h {
		cp := make([]string, len(vs))
		copy(cp, vs)
		out[k] = cp
	}
	return out
}

// TCPEchoServer accepts connections and echoes a fixed greeting.
type TCPEchoServer struct {
	Listener net.Listener
	Host     string
	Port     int
	Greeting string
}

// NewTCPEcho starts a TCP server that writes Greeting on each accepted connection.
func NewTCPEcho(t *testing.T, greeting string) *TCPEchoServer {
	t.Helper()
	if greeting == "" {
		greeting = "STREAM-ECHO\n"
	}
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("tcp echo listen: %v", err)
	}
	_, portStr, _ := net.SplitHostPort(ln.Addr().String())
	port, _ := strconv.Atoi(portStr)
	s := &TCPEchoServer{Listener: ln, Host: "127.0.0.1", Port: port, Greeting: greeting}

	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				return
			}
			go func(c net.Conn) {
				defer c.Close()
				_, _ = io.WriteString(c, s.Greeting)
			}(conn)
		}
	}()

	t.Cleanup(func() { _ = ln.Close() })
	return s
}

// pickFreePort returns a likely-free TCP port on 127.0.0.1.
func pickFreePort(t *testing.T) int {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("pick port: %v", err)
	}
	_, portStr, _ := net.SplitHostPort(ln.Addr().String())
	_ = ln.Close()
	port, _ := strconv.Atoi(portStr)
	return port
}

// UDPEchoServer responds to each datagram with a fixed payload.
type UDPEchoServer struct {
	Conn     *net.UDPConn
	Host     string
	Port     int
	Response []byte
}

// NewUDPEcho starts a UDP server that writes Response for every datagram received.
func NewUDPEcho(t *testing.T, response string) *UDPEchoServer {
	t.Helper()
	if response == "" {
		response = "UDP-ECHO\n"
	}
	addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("udp resolve: %v", err)
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		t.Fatalf("udp listen: %v", err)
	}
	_, portStr, _ := net.SplitHostPort(conn.LocalAddr().String())
	port, _ := strconv.Atoi(portStr)
	s := &UDPEchoServer{
		Conn:     conn,
		Host:     "127.0.0.1",
		Port:     port,
		Response: []byte(response),
	}

	go func() {
		buf := make([]byte, 4096)
		for {
			n, client, err := conn.ReadFromUDP(buf)
			if err != nil {
				return
			}
			_, _ = conn.WriteToUDP(s.Response, client)
			_ = n
		}
	}()

	t.Cleanup(func() { _ = conn.Close() })
	return s
}

// tryWaitForTLSStream returns nil once a TLS stream listener accepts connections.
func tryWaitForTLSStream(addr, serverName string, trust *tls.Certificate, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	var lastErr error
	pool := x509.NewCertPool()
	if len(trust.Certificate) > 0 {
		if leaf, err := x509.ParseCertificate(trust.Certificate[0]); err == nil {
			pool.AddCert(leaf)
		}
	}
	cfg := &tls.Config{ServerName: serverName, RootCAs: pool, MinVersion: tls.VersionTLS12}
	for time.Now().Before(deadline) {
		conn, err := tls.DialWithDialer(&net.Dialer{Timeout: 200 * time.Millisecond}, "tcp", addr, cfg)
		if err == nil {
			_ = conn.Close()
			return nil
		}
		lastErr = err
		time.Sleep(20 * time.Millisecond)
	}
	return lastErr
}

// waitForTLSStream dials addr with TLS until success or timeout.
func waitForTLSStream(t *testing.T, addr, serverName string, trust *tls.Certificate, timeout time.Duration) {
	t.Helper()
	if err := tryWaitForTLSStream(addr, serverName, trust, timeout); err != nil {
		t.Fatalf("wait for tls stream %s: %v", addr, err)
	}
}

// tryWaitForTCP returns nil once a TCP listener accepts connections.
func tryWaitForTCP(addr string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	var lastErr error
	for time.Now().Before(deadline) {
		conn, err := net.DialTimeout("tcp", addr, 200*time.Millisecond)
		if err == nil {
			_ = conn.Close()
			return nil
		}
		lastErr = err
		time.Sleep(20 * time.Millisecond)
	}
	return lastErr
}

// waitForTCP dials addr until success or timeout (stream listeners start async).
func waitForTCP(t *testing.T, addr string, timeout time.Duration) {
	t.Helper()
	if err := tryWaitForTCP(addr, timeout); err != nil {
		t.Fatalf("wait for tcp %s: %v", addr, err)
	}
}

// tlsConfigForCert builds a client TLS config that trusts the given server certificate.
func tlsConfigForCert(t *testing.T, serverName string, cert *tls.Certificate) *tls.Config {
	t.Helper()
	if len(cert.Certificate) == 0 {
		t.Fatal("tls cert has no certificate chain")
	}
	pool := x509.NewCertPool()
	if !pool.AppendCertsFromPEM(cert.Certificate[0]) {
		// X509KeyPair stores DER; re-parse for the pool
		leaf, err := x509.ParseCertificate(cert.Certificate[0])
		if err != nil {
			t.Fatalf("parse leaf cert: %v", err)
		}
		pool.AddCert(leaf)
	}
	return &tls.Config{
		ServerName: serverName,
		RootCAs:    pool,
		MinVersion: tls.VersionTLS12,
	}
}

// DialTLSStream connects to a TLS-terminated stream listener.
func DialTLSStream(t *testing.T, addr, serverName string, trust *tls.Certificate) net.Conn {
	t.Helper()
	var cfg *tls.Config
	if trust != nil {
		cfg = tlsConfigForCert(t, serverName, trust)
	} else {
		cfg = &tls.Config{
			ServerName:         serverName,
			InsecureSkipVerify: true, //nolint:gosec // test-only self-signed dev certs
			MinVersion:         tls.VersionTLS12,
		}
	}
	conn, err := tls.DialWithDialer(&net.Dialer{Timeout: 2 * time.Second}, "tcp", addr, cfg)
	if err != nil {
		t.Fatalf("tls dial %s: %v", addr, err)
	}
	t.Cleanup(func() { _ = conn.Close() })
	return conn
}

// waitForUDP sends a probe datagram until a response arrives or timeout.
func waitForUDP(t *testing.T, addr string, timeout time.Duration) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	target, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		t.Fatalf("resolve udp %s: %v", addr, err)
	}
	var lastErr error
	for time.Now().Before(deadline) {
		conn, err := net.DialUDP("udp", nil, target)
		if err != nil {
			lastErr = err
			time.Sleep(20 * time.Millisecond)
			continue
		}
		_ = conn.SetDeadline(time.Now().Add(200 * time.Millisecond))
		if _, err := conn.Write([]byte("probe")); err != nil {
			lastErr = err
			_ = conn.Close()
			time.Sleep(20 * time.Millisecond)
			continue
		}
		buf := make([]byte, 8)
		_, err = conn.Read(buf)
		_ = conn.Close()
		if err == nil {
			return
		}
		lastErr = err
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatalf("wait for udp %s: %v", addr, lastErr)
}

// WhoamiWSServer is a websocket upstream that records the upgrade request and echoes messages.
type WhoamiWSServer struct {
	Server   *httptest.Server
	Hostname string
	Port     int
	Path     string

	mu         sync.Mutex
	upgradeReq WhoamiRequest
	messages   []string
}

// NewWhoamiWS starts a websocket echo upstream on path (default /ws).
func NewWhoamiWS(t *testing.T, path string) *WhoamiWSServer {
	t.Helper()
	if path == "" {
		path = "/ws"
	}
	w := &WhoamiWSServer{Path: path}
	mux := http.NewServeMux()
	mux.Handle(path, websocket.Handler(w.handleWS))
	w.Server = httptest.NewServer(mux)
	t.Cleanup(w.Server.Close)

	u, err := url.Parse(w.Server.URL)
	if err != nil {
		t.Fatalf("whoami ws url: %v", err)
	}
	w.Hostname = u.Hostname()
	w.Port, err = strconv.Atoi(u.Port())
	if err != nil {
		t.Fatalf("whoami ws port: %v", err)
	}
	return w
}

func (w *WhoamiWSServer) handleWS(ws *websocket.Conn) {
	req := ws.Request()
	body, _ := io.ReadAll(req.Body)
	_ = req.Body.Close()
	captured := WhoamiRequest{
		Method:     req.Method,
		Path:       req.URL.Path,
		Query:      req.URL.RawQuery,
		Host:       req.Host,
		Headers:    cloneHeaderMap(req.Header),
		Body:       string(body),
		RemoteAddr: req.RemoteAddr,
	}
	w.mu.Lock()
	w.upgradeReq = captured
	w.mu.Unlock()

	for {
		var msg string
		if err := websocket.Message.Receive(ws, &msg); err != nil {
			return
		}
		w.mu.Lock()
		w.messages = append(w.messages, msg)
		w.mu.Unlock()
		if err := websocket.Message.Send(ws, "echo:"+msg); err != nil {
			return
		}
	}
}

// UpgradeRequest returns the HTTP request observed at websocket upgrade time.
func (w *WhoamiWSServer) UpgradeRequest() (WhoamiRequest, bool) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.upgradeReq.Method == "" {
		return WhoamiRequest{}, false
	}
	return w.upgradeReq, true
}

// Messages returns client messages received after upgrade.
func (w *WhoamiWSServer) Messages() []string {
	w.mu.Lock()
	defer w.mu.Unlock()
	out := make([]string, len(w.messages))
	copy(out, w.messages)
	return out
}

// cleanupStreams disables all streams and reloads the engine on test end.
func cleanupStreams(t *testing.T, eng *Engine, st store.Store) {
	t.Helper()
	t.Cleanup(func() {
		for _, s := range st.GetStreams() {
			if !s.Enabled {
				continue
			}
			s.Enabled = false
			_, _ = st.UpdateStream(s)
		}
		eng.ReloadFromStore()
	})
}

// newProxyServer exposes eng.Handler() on an httptest server (needed for websocket hijack).
func newProxyServer(t *testing.T, eng *Engine) *httptest.Server {
	t.Helper()
	srv := httptest.NewServer(eng.Handler())
	t.Cleanup(srv.Close)
	return srv
}

// dialWhoamiWS connects through a proxy front to the websocket upstream path.
// The vhost is embedded in the handshake URL (Host header); TCP dials proxyURL directly.
func dialWhoamiWS(t *testing.T, proxyURL, vhost, wsPath string, hdrs http.Header) *websocket.Conn {
	t.Helper()
	proxyU, err := url.Parse(proxyURL)
	if err != nil {
		t.Fatalf("proxy url: %v", err)
	}
	wsURL := "ws://" + vhost + wsPath
	cfg, err := websocket.NewConfig(wsURL, "http://"+vhost+"/")
	if err != nil {
		t.Fatalf("ws config: %v", err)
	}
	if hdrs != nil {
		cfg.Header = hdrs.Clone()
	}
	conn, err := net.DialTimeout("tcp", proxyU.Host, 2*time.Second)
	if err != nil {
		t.Fatalf("tcp dial %s: %v", proxyU.Host, err)
	}
	ws, err := websocket.NewClient(cfg, conn)
	if err != nil {
		_ = conn.Close()
		t.Fatalf("ws handshake %s via %s: %v", wsURL, proxyU.Host, err)
	}
	t.Cleanup(func() { _ = ws.Close() })
	return ws
}

// ReadStreamGreeting reads up to n bytes from conn before deadline.
func ReadStreamGreeting(t *testing.T, conn net.Conn, n int, timeout time.Duration) string {
	t.Helper()
	_ = conn.SetReadDeadline(time.Now().Add(timeout))
	buf := make([]byte, n)
	got, err := io.ReadAtLeast(conn, buf, 1)
	if err != nil {
		t.Fatalf("stream read: %v", err)
	}
	return string(buf[:got])
}