initial commit

internal/proxy/whoami_test.go

@@ -0,0 +1,618 @@
+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])
+}