smallprox/proxy_test.go

package main

import (
	"crypto/tls"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
)

func TestParseTLSVersion(t *testing.T) {
	tests := []struct {
		input   string
		want    uint16
		wantErr bool
	}{
		{"1.0", tls.VersionTLS10, false},
		{"TLS1.0", tls.VersionTLS10, false},
		{"1.1", tls.VersionTLS11, false},
		{"1.2", tls.VersionTLS12, false},
		{"tls1.2", tls.VersionTLS12, false},
		{"1.3", tls.VersionTLS13, false},
		{"2.0", 0, true},
		{"invalid", 0, true},
	}

	for _, tt := range tests {
		got, err := parseTLSVersion(tt.input)
		if (err != nil) != tt.wantErr {
			t.Errorf("parseTLSVersion(%q) error = %v, wantErr %v", tt.input, err, tt.wantErr)
			continue
		}
		if got != tt.want {
			t.Errorf("parseTLSVersion(%q) = %x, want %x", tt.input, got, tt.want)
		}
	}
}

func TestParseTargetURL(t *testing.T) {
	tests := []struct {
		input    string
		wantHost string
		wantSch  string
		wantErr  bool
	}{
		{"https://example.com", "example.com", "https", false},
		{"http://localhost:9000", "localhost:9000", "http", false},
		{"example.com:8443", "example.com:8443", "https", false},
		{"", "", "", true},
		{"ftp://example.com", "", "", true},
	}

	for _, tt := range tests {
		u, err := ParseTargetURL(tt.input)
		if (err != nil) != tt.wantErr {
			t.Errorf("ParseTargetURL(%q) error = %v, wantErr %v", tt.input, err, tt.wantErr)
			continue
		}
		if !tt.wantErr {
			if u.Host != tt.wantHost || u.Scheme != tt.wantSch {
				t.Errorf("ParseTargetURL(%q) = %s://%s, want %s://%s", tt.input, u.Scheme, u.Host, tt.wantSch, tt.wantHost)
			}
		}
	}
}

func TestGenerateSelfSignedCert(t *testing.T) {
	cert, err := generateSelfSignedCert([]string{"custom.local", "192.168.1.50"})
	if err != nil {
		t.Fatalf("generateSelfSignedCert() failed: %v", err)
	}

	if len(cert.Certificate) == 0 {
		t.Fatal("generateSelfSignedCert() returned empty certificate chain")
	}
	if cert.PrivateKey == nil {
		t.Fatal("generateSelfSignedCert() returned nil private key")
	}
}

func TestProxy_TLS12_To_TLS13_Backend(t *testing.T) {
	// 1. Setup backend server that ONLY accepts TLS 1.3
	backendHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/test-endpoint" {
			w.Header().Set("X-Backend-Received-Proto", r.Proto)
			w.Header().Set("X-Backend-Received-Host", r.Host)
			w.WriteHeader(http.StatusOK)
			body, _ := io.ReadAll(r.Body)
			w.Write([]byte("backend-response: " + string(body)))
			return
		}
		http.NotFound(w, r)
	})

	backendServer := httptest.NewUnstartedServer(backendHandler)
	backendServer.TLS = &tls.Config{
		MinVersion: tls.VersionTLS13,
		MaxVersion: tls.VersionTLS13,
	}
	backendServer.StartTLS()
	defer backendServer.Close()

	backendURL, err := url.Parse(backendServer.URL)
	if err != nil {
		t.Fatalf("failed to parse backend URL: %v", err)
	}

	// 2. Setup smallprox handler pointing to the TLS 1.3 backend
	proxyHandler := NewProxyHandler(ProxyConfig{
		TargetURL:    backendURL,
		PreserveHost: false,
	})

	proxyCert, err := generateSelfSignedCert(nil)
	if err != nil {
		t.Fatalf("failed to generate proxy cert: %v", err)
	}

	proxyServer := httptest.NewUnstartedServer(proxyHandler)
	// Proxy frontend only accepts TLS 1.2
	proxyServer.TLS = &tls.Config{
		Certificates: []tls.Certificate{proxyCert},
		MinVersion:   tls.VersionTLS12,
		MaxVersion:   tls.VersionTLS12,
	}
	proxyServer.StartTLS()
	defer proxyServer.Close()

	// 3. Client configured strictly for TLS 1.2
	client := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: true,
				MinVersion:         tls.VersionTLS12,
				MaxVersion:         tls.VersionTLS12,
			},
		},
	}

	// Make request through proxy
	req, err := http.NewRequest(http.MethodPost, proxyServer.URL+"/test-endpoint", strings.NewReader("hello-from-tls12-client"))
	if err != nil {
		t.Fatalf("failed to create request: %v", err)
	}

	resp, err := client.Do(req)
	if err != nil {
		t.Fatalf("client request through proxy failed: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		t.Fatalf("expected status 200 OK, got %d", resp.StatusCode)
	}

	if resp.TLS == nil || resp.TLS.Version != tls.VersionTLS12 {
		t.Errorf("expected client TLS version to be TLS 1.2 (0x0303), got %x", resp.TLS.Version)
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("failed to read response body: %v", err)
	}

	expectedBody := "backend-response: hello-from-tls12-client"
	if string(body) != expectedBody {
		t.Errorf("expected body %q, got %q", expectedBody, string(body))
	}

	// Verify Host was rewritten to backend host
	if receivedHost := resp.Header.Get("X-Backend-Received-Host"); receivedHost != backendURL.Host {
		t.Errorf("expected backend to receive host %q, got %q", backendURL.Host, receivedHost)
	}
}

func TestProxy_PreserveHost(t *testing.T) {
	backendHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("X-Received-Host", r.Host)
		w.WriteHeader(http.StatusOK)
	})

	backendServer := httptest.NewServer(backendHandler)
	defer backendServer.Close()

	backendURL, err := url.Parse(backendServer.URL)
	if err != nil {
		t.Fatalf("failed to parse backend URL: %v", err)
	}

	proxyHandler := NewProxyHandler(ProxyConfig{
		TargetURL:    backendURL,
		PreserveHost: true,
	})

	proxyServer := httptest.NewServer(proxyHandler)
	defer proxyServer.Close()

	req, err := http.NewRequest(http.MethodGet, proxyServer.URL+"/test", nil)
	if err != nil {
		t.Fatalf("failed to create request: %v", err)
	}
	req.Host = "custom-domain.net"

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	if got := resp.Header.Get("X-Received-Host"); got != "custom-domain.net" {
		t.Errorf("expected preserved host %q, got %q", "custom-domain.net", got)
	}
}