// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// This file contains HTTP tests for the ECP proxy.
package main

import (
	"context"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"io"
	"log"
	"math/big"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"testing"
	"time"
)

const (
	successMessage = "successful response!"
	errorMessage   = "error message!"
)

var (
	certs1 = NewMTLSInMemoryCerts()
	certs2 = NewMTLSInMemoryCerts()
)

func generateCert(ca *x509.Certificate, caKey *rsa.PrivateKey, cn string, isClient bool) (tls.Certificate, []byte, crypto.PrivateKey) {
	template := &x509.Certificate{
		SerialNumber: big.NewInt(time.Now().UnixNano()),
		Subject:      pkix.Name{CommonName: cn},
		NotBefore:    time.Now(),
		NotAfter:     time.Now().Add(24 * time.Hour),
		KeyUsage:     x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		DNSNames:     []string{cn},                       // important for hostname match
		IPAddresses:  []net.IP{net.ParseIP("127.0.0.1")}, // allow IP connect
	}
	if isClient {
		template.ExtKeyUsage = []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth}
		template.DNSNames = nil // clients don’t need SANs
		template.IPAddresses = nil
	}

	key, _ := rsa.GenerateKey(rand.Reader, 2048)
	certDER, _ := x509.CreateCertificate(rand.Reader, template, ca, &key.PublicKey, caKey)

	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})

	tlsCert, _ := tls.X509KeyPair(certPEM, keyPEM)
	return tlsCert, certPEM, key
}

// Struct holding all generated certs/pools
type MTLSInMemoryCerts struct {
	CAPool     *x509.CertPool
	ServerCert tls.Certificate
	ServerKey  crypto.PrivateKey
	ClientCert tls.Certificate
	ClientKey  crypto.PrivateKey
}

// Factory function: creates CA, server, and client certs in memory
func NewMTLSInMemoryCerts() *MTLSInMemoryCerts {
	// 1. Generate CA
	ca := &x509.Certificate{
		SerialNumber:          big.NewInt(1),
		Subject:               pkix.Name{CommonName: "Test CA"},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(24 * time.Hour),
		IsCA:                  true,
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		BasicConstraintsValid: true,
	}
	caKey, _ := rsa.GenerateKey(rand.Reader, 2048)
	caDER, _ := x509.CreateCertificate(rand.Reader, ca, ca, &caKey.PublicKey, caKey)

	caPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: caDER})
	caPool := x509.NewCertPool()
	if !caPool.AppendCertsFromPEM(caPEM) {
		panic("failed to append CA cert")
	}

	// 2. Generate server cert (signed by CA, includes SANs)
	serverCert, _, serverKey := generateCert(ca, caKey, "localhost", false)

	// 3. Generate client cert (signed by CA)
	clientCert, _, clientKey := generateCert(ca, caKey, "TestClient", true)

	return &MTLSInMemoryCerts{
		CAPool:     caPool,
		ServerCert: serverCert,
		ServerKey:  serverKey,
		ClientCert: clientCert,
		ClientKey:  clientKey,
	}
}

func createTLSBackendServer(expectedResponse string, expectedStatusCode int, certs *MTLSInMemoryCerts) *httptest.Server {
	backend := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if expectedStatusCode == http.StatusOK {
			fmt.Fprint(w, expectedResponse)
		} else {
			// Write error
			http.Error(w, expectedResponse, expectedStatusCode)
		}
	}))
	backend.TLS = &tls.Config{
		Certificates: []tls.Certificate{certs.ServerCert},
		ClientCAs:    certs.CAPool,
		ClientAuth:   tls.RequireAndVerifyClientCert,
	}
	backend.Config.ErrorLog = log.New(os.Stderr, "BackendServer: ", log.LstdFlags)
	fmt.Printf("createTLSBackendServer: Backend listening on %s\n", backend.Listener.Addr())
	return backend
}

func createSimpleTLSBackendServer(expectedResponse string, expectedStatusCode int, certs *MTLSInMemoryCerts) *httptest.Server {
	backend := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if expectedStatusCode == http.StatusOK {
			fmt.Fprint(w, expectedResponse)
		} else {
			// Write error
			http.Error(w, expectedResponse, expectedStatusCode)
		}
	}))
	backend.TLS = &tls.Config{
		Certificates: []tls.Certificate{certs.ServerCert},
		// No ClientAuth required for simple TLS
	}
	return backend
}

// Send HTTP requests to ECP Proxy. Validate responses.
func TestECPProxyWithHTTPClient(t *testing.T) {
	// Init MTLS backendServer
	validMTLSBackendServer := createTLSBackendServer(successMessage, http.StatusOK, certs1)
	validMTLSBackendServer.StartTLS()
	defer validMTLSBackendServer.Close()

	// Init MTLS backend server that always returns an (non-mtls related) error
	backendServerReturnsError := createTLSBackendServer(errorMessage, http.StatusInternalServerError, certs1)
	backendServerReturnsError.StartTLS()
	defer backendServerReturnsError.Close()

	// Init plain HTTPS backend server (no mTLS client auth required)
	plainBackendServer := createSimpleTLSBackendServer(successMessage, http.StatusOK, certs1)
	plainBackendServer.StartTLS()
	defer plainBackendServer.Close()

	// Init a passthrough proxy, to test proxy chaining
	passthroughProxyServer := httptest.NewServer(http.HandlerFunc(ProxyHandler))
	defer passthroughProxyServer.Close()

	// To trick the proxy into selecting the ECPTransport (mTLS), the target host must contain ".mtls.".
	// We replace the local IP (127.0.0.1) from our mock servers with "test.mtls.local" in the request headers.
	validMTLSBackendServerHost := strings.Replace(validMTLSBackendServer.Listener.Addr().String(), "127.0.0.1", "test.mtls.local", 1)
	mtlsValidMTLSBackendServerHost := strings.Replace(validMTLSBackendServer.Listener.Addr().String(), "127.0.0.1", "mtls.test.local", 1)
	backendServerReturnsErrorHost := strings.Replace(backendServerReturnsError.Listener.Addr().String(), "127.0.0.1", "test.mtls.local", 1)
	plainBackendServerHost := plainBackendServer.Listener.Addr().String()
	validPassthroughURL := passthroughProxyServer.URL
	invalidPassthroughURL := "1234-bad-server-mtls.com"

	testCases := []struct {
		name                    string
		ecpMTLSCerts            *MTLSInMemoryCerts
		targetHostHeader        string
		passthroughProxyAddress string
		wantStatusCode          int
		expectedResponseBody    string
		wantECPProxyError       bool
	}{
		{
			name:                    "successful request with mtls. prefix",
			ecpMTLSCerts:            certs1,
			targetHostHeader:        mtlsValidMTLSBackendServerHost,
			passthroughProxyAddress: "",
			wantStatusCode:          http.StatusOK,
			expectedResponseBody:    successMessage,
			wantECPProxyError:       false,
		},
		{
			name:                    "successful request",
			ecpMTLSCerts:            certs1,
			targetHostHeader:        validMTLSBackendServerHost,
			passthroughProxyAddress: "",
			wantStatusCode:          http.StatusOK,
			expectedResponseBody:    successMessage,
			wantECPProxyError:       false,
		},
		{
			name:                    "successful plain http request (fallback)",
			ecpMTLSCerts:            certs1,
			targetHostHeader:        plainBackendServerHost,
			passthroughProxyAddress: "",
			wantStatusCode:          http.StatusOK,
			expectedResponseBody:    successMessage,
			wantECPProxyError:       false,
		},
		{
			name:                    "invalid target header (dns failure)",
			ecpMTLSCerts:            certs1,
			targetHostHeader:        "123-not-allowed-mtls.com",
			passthroughProxyAddress: "",
			wantStatusCode:          http.StatusBadGateway,
			wantECPProxyError:       true,
		},
		{
			name:                    "empty target header",
			ecpMTLSCerts:            certs1,
			targetHostHeader:        "",
			passthroughProxyAddress: "",
			wantStatusCode:          http.StatusBadRequest,
			wantECPProxyError:       true,
		},
		{
			name:                    "ecp proxy server is using invalid mtls certs, unable to complete handshake",
			ecpMTLSCerts:            certs2, // these do not match validMTLSBackendServer's mTLS certs (certs1)
			targetHostHeader:        validMTLSBackendServerHost,
			passthroughProxyAddress: "",
			wantStatusCode:          http.StatusBadGateway,
			wantECPProxyError:       true,
		},
		{
			name:                    "backend server returns an error",
			ecpMTLSCerts:            certs1,
			targetHostHeader:        backendServerReturnsErrorHost,
			passthroughProxyAddress: "",
			wantStatusCode:          http.StatusInternalServerError,
			wantECPProxyError:       false, // no ecp internal error, just forward the response
		},
		// tests that use proxy Chaining
		{
			name:                    "successful request, with proxy chain",
			ecpMTLSCerts:            certs1,
			targetHostHeader:        validMTLSBackendServerHost,
			passthroughProxyAddress: validPassthroughURL,
			wantStatusCode:          http.StatusOK,
			expectedResponseBody:    successMessage,
			wantECPProxyError:       false,
		},
		{
			name:                    "successful plain http request with proxy chain",
			ecpMTLSCerts:            certs1,
			targetHostHeader:        plainBackendServerHost,
			passthroughProxyAddress: validPassthroughURL,
			wantStatusCode:          http.StatusOK,
			expectedResponseBody:    successMessage,
			wantECPProxyError:       false,
		},
		{
			name:                    "empty target header with proxy chain",
			ecpMTLSCerts:            certs1,
			targetHostHeader:        "",
			passthroughProxyAddress: validPassthroughURL,
			wantStatusCode:          http.StatusBadRequest,
			wantECPProxyError:       true,
		},
		{
			name:                    "ecp proxy server is using invalid mtls certs, unable to complete handshake with proxy chain",
			ecpMTLSCerts:            certs2, // these do not match validMTLSBackendServer's mTLS certs (certs1)
			targetHostHeader:        validMTLSBackendServerHost,
			passthroughProxyAddress: validPassthroughURL,
			wantStatusCode:          http.StatusBadGateway,
			wantECPProxyError:       true,
		},
		{
			name:                    "backend server returns an error with proxy chain",
			ecpMTLSCerts:            certs1,
			targetHostHeader:        backendServerReturnsErrorHost,
			passthroughProxyAddress: validPassthroughURL,
			wantStatusCode:          http.StatusInternalServerError,
			wantECPProxyError:       false, // no ecp internal error, just forward the response
		},
		{
			name:                    "passthrough proxy doesnt exist",
			ecpMTLSCerts:            certs1,
			targetHostHeader:        validMTLSBackendServerHost,
			passthroughProxyAddress: invalidPassthroughURL,
			wantStatusCode:          http.StatusBadGateway,
			wantECPProxyError:       true,
		},
	}

	for i, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			ecpProxyPort := 8080 + i

			tlsConfig := &tls.Config{
				Certificates:       []tls.Certificate{tc.ecpMTLSCerts.ClientCert},
				InsecureSkipVerify: true,
				RootCAs:            tc.ecpMTLSCerts.CAPool,
				ServerName:         "localhost",
			}

			proxyConfig := &ProxyConfig{
				Port:      ecpProxyPort,
				TlsConfig: tlsConfig,
			}

			if tc.passthroughProxyAddress != "" {
				passthroughProxyServerURL, err := url.Parse(tc.passthroughProxyAddress)
				if err != nil {
					t.Fatalf("Failed to parse passthrough proxy URL: %v", err)
				}
				proxyConfig.UpstreamProxyURL = passthroughProxyServerURL
			}

			ecpTransport := newTransport(proxyConfig.TlsConfig, proxyConfig).(*http.Transport)

			defaultTLSConfig := &tls.Config{
				RootCAs:            tc.ecpMTLSCerts.CAPool,
				InsecureSkipVerify: true,
				ServerName:         "localhost",
			}
			defaultTransport := newTransport(defaultTLSConfig, proxyConfig).(*http.Transport)

			// Because we tricked the proxy with the fake domain "test.mtls.local" above to trigger mTLS,
			// the actual TCP dial will fail DNS resolution. We use a custom DialContext to intercept this
			// fake domain and redirect the connection back to 127.0.0.1 where the mock server is listening.
			customDialContext := func(ctx context.Context, network, addr string) (net.Conn, error) {
				addr = strings.Replace(addr, "test.mtls.local", "127.0.0.1", 1)
				addr = strings.Replace(addr, "mtls.test.local", "127.0.0.1", 1)
				return (&net.Dialer{
					Timeout:   proxyConfig.DialTimeout,
					KeepAlive: proxyConfig.KeepAlivePeriod,
				}).DialContext(ctx, network, addr)
			}
			ecpTransport.DialContext = customDialContext
			defaultTransport.DialContext = customDialContext

			routingTransport := &RoutingTransport{
				ECPTransport:     ecpTransport,
				DefaultTransport: defaultTransport,
			}

			handler := newECPProxyHandler(routingTransport)
			ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
			defer stop()

			go runServer(ctx, proxyConfig, handler)
			// Wait a short time for the server to become ready.
			time.Sleep(100 * time.Millisecond)

			expProxyURL := fmt.Sprintf("http://127.0.0.1:%d/", ecpProxyPort)
			req, err := http.NewRequest("GET", expProxyURL, nil)
			if err != nil {
				t.Fatalf("failed to create request: %v", err)
			}

			req.Header.Set("x-goog-ecpproxy-target-host", tc.targetHostHeader)
			fmt.Printf("TestECPProxyWithHTTPClient: Requesting %s with target %s\n", req.URL, tc.targetHostHeader)
			client := &http.Client{}
			resp, err := client.Do(req)
			if err != nil {
				t.Fatalf("request via proxy failed: %v", err)
			}
			defer resp.Body.Close()

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

			// Assert on status code first, as it's always expected
			if resp.StatusCode != tc.wantStatusCode {
				t.Errorf("expected status code %d, got %d", tc.wantStatusCode, resp.StatusCode)
			}

			if tc.wantStatusCode == http.StatusOK {
				// We expect a plain text success response
				gotBody := string(body)
				if gotBody != tc.expectedResponseBody {
					t.Errorf("expected body %q, got %q", tc.expectedResponseBody, gotBody)
				}
			} else {
				if tc.wantECPProxyError {
					if resp.Header.Get(ecpInternalErrorHeader) == "" {
						t.Errorf("missing error response header")
					}
					var gotError ErrorResponse
					if err := json.Unmarshal(body, &gotError); err != nil {
						t.Fatalf("failed to unmarshal JSON error response: %v", err)
					}
					if gotError.Code != tc.wantStatusCode {
						t.Errorf("incorrect status code %d, got %d", tc.wantStatusCode, gotError.Code)
					}
					if gotError.Error == "" {
						t.Errorf("error message field `Error` is empty")

					}
					if gotError.Message == "" {
						t.Errorf("error message field `Message` is empty")
					}
				} else {
					if resp.Header.Get(ecpInternalErrorHeader) != "" {
						t.Errorf("unexpected error header included")
					}
				}
			}
		})
	}
}