helix-proxy/internal/proxy/integration_test.go

package proxy

import (
	"crypto/tls"
	"encoding/json"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"helix-proxy/internal/certificate"
	"helix-proxy/internal/config"
	"helix-proxy/internal/store"

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

// TestEngine_Integration_FullSetup wires proxy hosts, redirections, dead hosts,
// access lists, certificates, locations, and streams against configurable whoami upstreams.
func TestEngine_Integration_FullSetup(t *testing.T) {
	eng, st := newTestEngine(t)
	defer closeStore(st)

	// --- upstream mocks ---
	mainUpstream := NewWhoami(t, WhoamiConfig{
		Routes: map[string]WhoamiRoute{
			"/health": {Status: 200, Body: `{"status":"ok"}`, Headers: map[string]string{"X-Secret": "hidden"}},
		},
	})
	apiUpstream := NewWhoami(t, WhoamiConfig{
		ValidPrefixes: []string{"/v2"},
	})
	authUpstream := NewWhoami(t, WhoamiConfig{})
	tcpEcho := NewTCPEcho(t, "STREAM-OK\n")

	certRec := issueDevCert(t, st, "secure.example", "stream.example")

	// --- access list ---
	al, _ := st.CreateAccessList(store.AccessList{
		Name:       "integration-acl",
		SatisfyAny: false,
		Clients:    []store.AccessClient{{Address: "all", Directive: "allow"}},
		Items:      []store.AccessItem{{Username: "testuser", Password: "secret"}},
		PassAuth:   false,
	})

	// --- proxy host: headers, location rewrite, access, cert ---
	_, _ = st.CreateProxyHost(store.ProxyHost{
		DomainNames:         []string{"proxy-int.example"},
		ForwardScheme:       "http",
		ForwardHost:         mainUpstream.Hostname,
		ForwardPort:         mainUpstream.Port,
		CertificateID:       certRec.ID,
		AccessListID:        al.ID,
		TrustForwardedProto: true,
		RequestHeaders:      []store.HeaderRule{{Name: "X-Injected-Req", Value: "yes"}},
		ResponseHeaders:     []store.HeaderRule{{Name: "X-Injected-Resp", Value: "yes"}},
		HideHeaders:         []string{"X-Secret"},
		Enabled:             true,
		Locations: []store.Location{{
			Path:          "/api",
			ForwardScheme: "http",
			ForwardHost:   apiUpstream.Hostname,
			ForwardPort:   apiUpstream.Port,
			ForwardPath:   "/v2",
		}},
	})

	// proxy with passAuth=true to verify Authorization reaches upstream
	alPass, _ := st.CreateAccessList(store.AccessList{
		Name:     "pass-auth-acl",
		Clients:  []store.AccessClient{{Address: "all", Directive: "allow"}},
		Items:    []store.AccessItem{{Username: "u", Password: "p"}},
		PassAuth: true,
	})
	_, _ = st.CreateProxyHost(store.ProxyHost{
		DomainNames:   []string{"passauth.example"},
		ForwardScheme: "http",
		ForwardHost:   authUpstream.Hostname,
		ForwardPort:   authUpstream.Port,
		AccessListID:  alPass.ID,
		Enabled:       true,
	})

	// --- redirection host ---
	_, _ = st.CreateRedirectionHost(store.RedirectionHost{
		DomainNames:       []string{"redir-int.example"},
		ForwardHTTPCode:   302,
		ForwardScheme:     "https",
		ForwardDomainName: "target-int.example",
		PreservePath:      true,
		Enabled:           true,
	})

	// --- dead hosts ---
	_, _ = st.CreateDeadHost(store.DeadHost{
		DomainNames: []string{"dead-int.example"},
		Enabled:     true,
		Meta:        map[string]any{"html": "<h1>integration dead</h1>"},
	})
	wwwDir := config.Resolve("www")
	_ = os.MkdirAll(wwwDir, 0o755)
	_ = os.WriteFile(filepath.Join(wwwDir, "404.html"), []byte("<h1>INTEGRATION 404 FILE</h1>"), 0o644)
	_, _ = st.CreateDeadHost(store.DeadHost{DomainNames: []string{"deadfile-int.example"}, Enabled: true})

	// --- default site modes ---
	_ = st.SetSetting("default_site", "404")

	// --- stream (TCP) ---
	t.Setenv("DISABLE_IPV6", "1")
	streamPort := pickFreePort(t)
	_, _ = st.CreateStream(store.Stream{
		IncomingPort:   streamPort,
		ForwardingHost: tcpEcho.Host,
		ForwardingPort: tcpEcho.Port,
		TCPForwarding:  true,
		Enabled:        true,
	})
	cleanupStreams(t, eng, st)

	eng.ReloadFromStore()
	h := eng.Handler()

	// --- redirection ---
	rr := doReq(h, "GET", "redir-int.example", "/foo?bar=1", nil, nil)
	if rr.Code != 302 || !strings.Contains(rr.Header().Get("Location"), "https://target-int.example/foo?bar=1") {
		t.Errorf("redirection: code=%d loc=%q", rr.Code, rr.Header().Get("Location"))
	}

	// --- dead hosts ---
	rr = doReq(h, "GET", "dead-int.example", "/", nil, nil)
	if rr.Code != 404 || !strings.Contains(rr.Body.String(), "integration dead") {
		t.Errorf("dead meta: code=%d body=%q", rr.Code, rr.Body.String())
	}
	rr = doReq(h, "GET", "deadfile-int.example", "/", nil, nil)
	if rr.Code != 404 || !strings.Contains(rr.Body.String(), "INTEGRATION 404 FILE") {
		t.Errorf("dead file: code=%d body=%q", rr.Code, rr.Body.String())
	}

	// --- default 404 ---
	rr = doReq(h, "GET", "unknown-int.example", "/", nil, nil)
	if rr.Code != 404 {
		t.Errorf("default 404: %d", rr.Code)
	}

	// --- certificate SNI ---
	cert, err := eng.GetCertificate(&tls.ClientHelloInfo{ServerName: "secure.example"})
	if err != nil || cert == nil {
		t.Errorf("GetCertificate secure.example: err=%v cert=%v", err, cert)
	}

	// --- proxy: main path with auth + injected headers ---
	authHdr := http.Header{
		"Authorization":     []string{"Basic " + basic("testuser", "secret")},
		"X-Forwarded-For":   []string{"10.0.0.5"},
		"X-Forwarded-Proto": []string{"https"},
	}
	rr = doReq(h, "GET", "proxy-int.example", "/health", authHdr, nil)
	if rr.Code != 200 {
		t.Fatalf("proxy /health: code=%d body=%s", rr.Code, rr.Body.String())
	}
	if rr.Header().Get("X-Injected-Resp") != "yes" {
		t.Error("response header injection missing")
	}
	if rr.Header().Get("X-Secret") != "" {
		t.Error("hide header failed: X-Secret leaked")
	}
	if got, ok := mainUpstream.LastRequest(); !ok {
		t.Fatal("main upstream saw no request")
	} else {
		if got.Path != "/health" {
			t.Errorf("main path: got %q want /health", got.Path)
		}
		if got.HeaderValue("X-Injected-Req") != "yes" {
			t.Errorf("request header injection: got %q", got.HeaderValue("X-Injected-Req"))
		}
		if got.HeaderValue("Authorization") != "" {
			t.Error("passAuth=false should strip Authorization before upstream")
		}
		if got.HeaderValue("X-Forwarded-Proto") != "https" {
			t.Errorf("X-Forwarded-Proto: got %q", got.HeaderValue("X-Forwarded-Proto"))
		}
	}

	// --- location path rewrite ---
	apiUpstream.ClearRequests()
	rr = doReq(h, "GET", "proxy-int.example", "/api/foo", authHdr, nil)
	if rr.Code != 200 {
		t.Fatalf("location /api/foo: code=%d", rr.Code)
	}
	got, ok := apiUpstream.LastRequest()
	if !ok {
		t.Fatal("api upstream saw no request for location")
	}
	if got.Path != "/v2/foo" {
		t.Errorf("location rewrite: upstream path=%q want /v2/foo", got.Path)
	}
	if rr.Header().Get("X-Whoami-Path") != "/v2/foo" {
		t.Errorf("X-Whoami-Path header: %q", rr.Header().Get("X-Whoami-Path"))
	}

	// --- passAuth=true keeps Authorization ---
	authUpstream.ClearRequests()
	rr = doReq(h, "GET", "passauth.example", "/", http.Header{
		"Authorization": []string{"Basic " + basic("u", "p")},
	}, nil)
	if rr.Code != 200 {
		t.Fatalf("passauth proxy: code=%d", rr.Code)
	}
	got, ok = authUpstream.LastRequest()
	if !ok {
		t.Fatal("auth upstream saw no request")
	}
	if got.HeaderValue("Authorization") == "" {
		t.Error("passAuth=true should forward Authorization to upstream")
	}

	// --- whoami JSON echo on default route ---
	mainUpstream.ClearRequests()
	rr = doReq(h, "GET", "proxy-int.example", "/echo-test?q=1", authHdr, nil)
	if rr.Code != 200 {
		t.Fatalf("echo route: code=%d", rr.Code)
	}
	var echoed WhoamiRequest
	if err := json.Unmarshal(rr.Body.Bytes(), &echoed); err != nil {
		t.Fatalf("parse whoami json: %v body=%s", err, rr.Body.String())
	}
	if echoed.Path != "/echo-test" || echoed.Query != "q=1" {
		t.Errorf("whoami json echo: path=%q query=%q", echoed.Path, echoed.Query)
	}

	// --- stream TCP forwarding ---
	waitForTCP(t, net.JoinHostPort("127.0.0.1", itoa(streamPort)), 5*time.Second)
	conn, err := net.DialTimeout("tcp", net.JoinHostPort("127.0.0.1", itoa(streamPort)), 2*time.Second)
	if err != nil {
		t.Fatalf("dial stream port %d: %v", streamPort, err)
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(2 * time.Second))
	if _, err := io.WriteString(conn, "ping\n"); err != nil {
		t.Fatalf("stream write: %v", err)
	}
	buf := make([]byte, 64)
	n, err := conn.Read(buf)
	if err != nil {
		t.Fatalf("stream read: %v", err)
	}
	if !strings.Contains(string(buf[:n]), "STREAM-OK") {
		t.Errorf("stream echo: got %q", string(buf[:n]))
	}

	// reload lifecycle: cleanupStreams t.Cleanup also exercises disable path
	strms := st.GetStreams()
	if len(strms) == 0 {
		t.Fatal("no streams in store")
	}
}

// TestEngine_Integration_Websocket verifies end-to-end websocket upgrade and message echo.
func TestEngine_Integration_Websocket(t *testing.T) {
	eng, st := newTestEngine(t)
	defer closeStore(st)

	wsUp := NewWhoamiWS(t, "/ws")
	_, _ = st.CreateProxyHost(store.ProxyHost{
		DomainNames:           []string{"ws-int.example"},
		ForwardScheme:         "http",
		ForwardHost:           wsUp.Hostname,
		ForwardPort:           wsUp.Port,
		Enabled:               true,
		AllowWebsocketUpgrade: true,
		TrustForwardedProto:   true,
	})
	eng.ReloadFromStore()

	proxy := newProxyServer(t, eng)
	ws := dialWhoamiWS(t, proxy.URL, "ws-int.example", "/ws", http.Header{
		"X-Forwarded-For":   []string{"10.9.8.7"},
		"X-Forwarded-Proto": []string{"https"},
	})

	if err := websocket.Message.Send(ws, "hello"); err != nil {
		t.Fatalf("ws send: %v", err)
	}
	var reply string
	if err := websocket.Message.Receive(ws, &reply); err != nil {
		t.Fatalf("ws recv: %v", err)
	}
	if reply != "echo:hello" {
		t.Errorf("ws echo: got %q", reply)
	}

	got, ok := wsUp.UpgradeRequest()
	if !ok {
		t.Fatal("upstream missed websocket upgrade request")
	}
	if got.Path != "/ws" {
		t.Errorf("upgrade path: %q", got.Path)
	}
	if !strings.Contains(got.HeaderValue("X-Forwarded-For"), "10.9.8.7") {
		t.Errorf("upgrade X-Forwarded-For: %q", got.HeaderValue("X-Forwarded-For"))
	}
	if got.HeaderValue("X-Forwarded-Proto") != "https" {
		t.Errorf("upgrade X-Forwarded-Proto: %q", got.HeaderValue("X-Forwarded-Proto"))
	}
	if got.HeaderValue("Connection") != "Upgrade" || got.HeaderValue("Upgrade") != "websocket" {
		t.Errorf("upgrade headers: conn=%q up=%q", got.HeaderValue("Connection"), got.HeaderValue("Upgrade"))
	}
	msgs := wsUp.Messages()
	if len(msgs) != 1 || msgs[0] != "hello" {
		t.Errorf("upstream messages: %v", msgs)
	}
}

// TestEngine_Integration_Caching exercises cache HIT/MISS through whoami upstreams.
func TestEngine_Integration_Caching(t *testing.T) {
	eng, st := newTestEngine(t)
	defer closeStore(st)

	upstream := NewWhoami(t, WhoamiConfig{
		DefaultHeaders: map[string]string{"Cache-Control": "max-age=3600"},
		Routes: map[string]WhoamiRoute{
			"/asset": {Body: "CACHED-ASSET"},
		},
	})
	_, _ = st.CreateProxyHost(store.ProxyHost{
		DomainNames:    []string{"cache-int.example"},
		ForwardScheme:  "http",
		ForwardHost:    upstream.Hostname,
		ForwardPort:    upstream.Port,
		Enabled:        true,
		CachingEnabled: true,
	})
	eng.ReloadFromStore()
	h := eng.Handler()

	rr1 := doReq(h, "GET", "cache-int.example", "/asset", nil, nil)
	if rr1.Code != 200 || rr1.Header().Get("X-Proxy-Cache") != "MISS" || rr1.Body.String() != "CACHED-ASSET" {
		t.Fatalf("cache MISS: code=%d cache=%s body=%q", rr1.Code, rr1.Header().Get("X-Proxy-Cache"), rr1.Body.String())
	}
	if reqs := upstream.Requests(); len(reqs) != 1 {
		t.Fatalf("upstream requests after MISS: %d", len(reqs))
	}

	rr2 := doReq(h, "GET", "cache-int.example", "/asset", nil, nil)
	if rr2.Header().Get("X-Proxy-Cache") != "HIT" || rr2.Body.String() != "CACHED-ASSET" {
		t.Errorf("cache HIT: cache=%s body=%q", rr2.Header().Get("X-Proxy-Cache"), rr2.Body.String())
	}
	if reqs := upstream.Requests(); len(reqs) != 1 {
		t.Errorf("upstream should not be hit again on HIT; requests=%d", len(reqs))
	}
}

// TestEngine_Integration_ProxyMiddleware exercises ssl_forced, block_exploits, and HSTS
// against a real whoami upstream.
func TestEngine_Integration_ProxyMiddleware(t *testing.T) {
	eng, st := newTestEngine(t)
	defer closeStore(st)

	upstream := NewWhoami(t, WhoamiConfig{})
	_, _ = st.CreateProxyHost(store.ProxyHost{
		DomainNames:         []string{"middleware.example"},
		ForwardScheme:       "http",
		ForwardHost:         upstream.Hostname,
		ForwardPort:         upstream.Port,
		Enabled:             true,
		SslForced:           true,
		BlockExploits:       true,
		HstsEnabled:         true,
		HstsSubdomains:      true,
		TrustForwardedProto: true,
	})
	eng.ReloadFromStore()
	h := eng.Handler()

	// ssl_forced redirects plain http
	rr := doReq(h, "GET", "middleware.example", "/ok", http.Header{
		"X-Forwarded-Proto": []string{"http"},
	}, nil)
	if rr.Code != 301 || !strings.HasPrefix(rr.Header().Get("Location"), "https://") {
		t.Errorf("ssl_forced: code=%d loc=%q", rr.Code, rr.Header().Get("Location"))
	}

	// block_exploits on https path
	req := httptest.NewRequest("GET", "/.env", nil)
	req.Host = "middleware.example"
	req.TLS = &tls.ConnectionState{}
	rr = httptest.NewRecorder()
	h.ServeHTTP(rr, req)
	if rr.Code != 403 || !strings.Contains(rr.Body.String(), "exploits") {
		t.Errorf("block_exploits: code=%d body=%q", rr.Code, rr.Body.String())
	}

	// HSTS + successful whoami proxy on https
	req = httptest.NewRequest("GET", "/ok", nil)
	req.Host = "middleware.example"
	req.TLS = &tls.ConnectionState{}
	rr = httptest.NewRecorder()
	h.ServeHTTP(rr, req)
	if rr.Code != 200 {
		t.Fatalf("proxy ok: code=%d body=%s", rr.Code, rr.Body.String())
	}
	if hs := rr.Header().Get("Strict-Transport-Security"); !strings.Contains(hs, "includeSubDomains") {
		t.Errorf("hsts: %q", hs)
	}
	got, ok := upstream.LastRequest()
	if !ok || got.Path != "/ok" {
		t.Errorf("upstream path: got=%v ok=%v", got.Path, ok)
	}
}

// issueDevCert creates and issues a self-signed certificate in development mode.
func issueDevCert(t *testing.T, st store.Store, domains ...string) store.Certificate {
	t.Helper()
	t.Setenv("PROXY_MODE", "development")
	cm := certificate.NewManager(st)
	rec, _ := st.CreateCertificate(store.Certificate{
		Provider:    "letsencrypt",
		DomainNames: domains,
		Meta:        map[string]any{},
	})
	if err := cm.Issue(rec, cm.GetEmailForCert(rec)); err != nil {
		t.Fatalf("issue cert: %v", err)
	}
	rec, _ = st.GetCertificate(rec.ID)
	return rec
}

// TestEngine_Integration_TLSStream verifies TLS-terminated TCP stream forwarding.
func TestEngine_Integration_TLSStream(t *testing.T) {
	eng, st := newTestEngine(t)
	defer closeStore(st)
	t.Setenv("DISABLE_IPV6", "1")

	certRec := issueDevCert(t, st, "stream-tls.example")
	tcpEcho := NewTCPEcho(t, "STREAM-TLS-OK\n")
	streamPort := 20000 + int(time.Now().UnixNano()%10000)

	_, _ = st.CreateStream(store.Stream{
		IncomingPort:   streamPort,
		ForwardingHost: tcpEcho.Host,
		ForwardingPort: tcpEcho.Port,
		TCPForwarding:  true,
		Enabled:        true,
		CertificateID:  certRec.ID,
	})
	cleanupStreams(t, eng, st)
	eng.ReloadFromStore()

	if eng.getStreamCert(certRec.ID) == nil {
		t.Fatalf("stream cert %d not loaded after reload", certRec.ID)
	}

	trust, err := eng.GetCertificate(&tls.ClientHelloInfo{ServerName: "stream-tls.example"})
	if err != nil {
		t.Fatalf("GetCertificate: %v", err)
	}

	addr := net.JoinHostPort("127.0.0.1", itoa(streamPort))
	var waitErr error
	for attempt := 0; attempt < 3; attempt++ {
		waitErr = tryWaitForTLSStream(addr, "stream-tls.example", trust, 2*time.Second)
		if waitErr == nil {
			break
		}
		eng.ReloadFromStore()
	}
	if waitErr != nil {
		t.Fatalf("tls stream listener on %s: %v", addr, waitErr)
	}

	conn := DialTLSStream(t, addr, "stream-tls.example", trust)
	if _, err := io.WriteString(conn, "ping\n"); err != nil {
		t.Fatalf("tls stream write: %v", err)
	}
	got := ReadStreamGreeting(t, conn, 64, 2*time.Second)
	if !strings.Contains(got, "STREAM-TLS-OK") {
		t.Errorf("tls stream echo: got %q", got)
	}

	// plain TCP to a TLS listener should fail handshake quickly
	raw, err := net.DialTimeout("tcp", addr, time.Second)
	if err != nil {
		t.Fatalf("raw dial: %v", err)
	}
	defer raw.Close()
	_ = raw.SetDeadline(time.Now().Add(time.Second))
	if _, err := io.WriteString(raw, "not-tls\n"); err != nil {
		t.Fatalf("raw write: %v", err)
	}
	buf := make([]byte, 16)
	if _, err := raw.Read(buf); err == nil {
		t.Error("expected plain TCP to TLS stream to fail read, got data")
	}
}

// TestEngine_Integration_UDPStream verifies UDP stream forwarding.
func TestEngine_Integration_UDPStream(t *testing.T) {
	eng, st := newTestEngine(t)
	defer closeStore(st)
	t.Setenv("DISABLE_IPV6", "1") // bind 0.0.0.0 so 127.0.0.1 dials reach the listener

	udpEcho := NewUDPEcho(t, "UDP-STREAM-OK\n")
	streamPort := pickFreePort(t)

	_, _ = st.CreateStream(store.Stream{
		IncomingPort:   streamPort,
		ForwardingHost: udpEcho.Host,
		ForwardingPort: udpEcho.Port,
		UDPForwarding:  true,
		Enabled:        true,
	})
	cleanupStreams(t, eng, st)
	eng.ReloadFromStore()

	addr := net.JoinHostPort("127.0.0.1", itoa(streamPort))
	waitForUDP(t, addr, 2*time.Second)

	target, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		t.Fatalf("resolve stream udp: %v", err)
	}
	conn, err := net.DialUDP("udp", nil, target)
	if err != nil {
		t.Fatalf("dial stream udp: %v", err)
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(2 * time.Second))

	if _, err := conn.Write([]byte("probe")); err != nil {
		t.Fatalf("udp write: %v", err)
	}
	buf := make([]byte, 64)
	n, err := conn.Read(buf)
	if err != nil {
		t.Fatalf("udp read: %v", err)
	}
	if !strings.Contains(string(buf[:n]), "UDP-STREAM-OK") {
		t.Errorf("udp stream echo: got %q", string(buf[:n]))
	}
}