package proxy

import (
	"bufio"
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"io"
	"math/big"
	"net"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/sandbox0-ai/sandbox0/infra-operator/api/config"
	"github.com/sandbox0-ai/sandbox0/manager/pkg/apis/sandbox0/v1alpha1"
	"github.com/sandbox0-ai/sandbox0/netd/pkg/policy"
	"github.com/sandbox0-ai/sandbox0/pkg/egressauth"
	"go.uber.org/zap"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/health"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/metadata"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

type blockingEgressAuthResolver struct {
	called  chan *egressauth.ResolveRequest
	release <-chan struct{}
	resp    *egressauth.ResolveResponse
}

func (r *blockingEgressAuthResolver) Resolve(_ context.Context, req *egressauth.ResolveRequest) (*egressauth.ResolveResponse, error) {
	if r != nil && r.called != nil && req != nil {
		copied := *req
		select {
		case r.called <- &copied:
		default:
		}
	}
	if r != nil && r.release != nil {
		<-r.release
	}
	if r == nil {
		return nil, nil
	}
	return cloneResolveResponse(r.resp), nil
}

func TestPrefixedConnReadsPrefixBeforeConn(t *testing.T) {
	client, server := net.Pipe()
	defer client.Close()
	defer server.Close()

	go func() {
		_, _ = client.Write([]byte("tail"))
	}()

	conn := newPrefixedConn(server, []byte("head-"))
	buf := make([]byte, len("head-tail"))
	n, err := io.ReadFull(conn, buf)
	if err != nil {
		t.Fatalf("read full: %v", err)
	}
	if got := string(buf[:n]); got != "head-tail" {
		t.Fatalf("prefixed conn read = %q", got)
	}
}

func TestCertificateAuthorityIssuesCertificateForHost(t *testing.T) {
	certPEM, keyPEM, err := newSelfSignedCertificateAuthority("sandbox0-test-root", time.Hour)
	if err != nil {
		t.Fatalf("newSelfSignedCertificateAuthority: %v", err)
	}
	authority, err := newCertificateAuthority(certPEM, keyPEM, time.Hour)
	if err != nil {
		t.Fatalf("newCertificateAuthority: %v", err)
	}
	cert, err := authority.CertificateForHost("api.example.com")
	if err != nil {
		t.Fatalf("CertificateForHost: %v", err)
	}
	if cert == nil || len(cert.Certificate) == 0 {
		t.Fatal("expected issued certificate")
	}
	leaf, err := x509.ParseCertificate(cert.Certificate[0])
	if err != nil {
		t.Fatalf("parse leaf: %v", err)
	}
	if len(leaf.DNSNames) != 1 || leaf.DNSNames[0] != "api.example.com" {
		t.Fatalf("unexpected dns names: %#v", leaf.DNSNames)
	}
}

func TestTLSAdapterInterceptsHTTPSAndInjectsHeaders(t *testing.T) {
	mitmCertPEM, mitmKeyPEM, err := newSelfSignedCertificateAuthority("sandbox0-mitm", time.Hour)
	if err != nil {
		t.Fatalf("new mitm ca: %v", err)
	}
	mitmAuthority, err := newCertificateAuthority(mitmCertPEM, mitmKeyPEM, time.Hour)
	if err != nil {
		t.Fatalf("new mitm authority: %v", err)
	}

	upstreamCAPEM, upstreamCAKeyPEM, err := newSelfSignedCertificateAuthority("sandbox0-upstream", time.Hour)
	if err != nil {
		t.Fatalf("new upstream ca: %v", err)
	}
	upstreamAuthority, err := newCertificateAuthority(upstreamCAPEM, upstreamCAKeyPEM, time.Hour)
	if err != nil {
		t.Fatalf("new upstream authority: %v", err)
	}
	upstreamLeaf, err := upstreamAuthority.CertificateForHost("api.example.com")
	if err != nil {
		t.Fatalf("upstream leaf: %v", err)
	}
	upstreamRootPool := x509.NewCertPool()
	if !upstreamRootPool.AppendCertsFromPEM(upstreamCAPEM) {
		t.Fatal("append upstream ca")
	}

	requestHeaders := make(chan http.Header, 1)
	requestBody := make(chan string, 1)
	upstreamListener, err := tls.Listen("tcp", "127.0.0.1:0", &tls.Config{
		Certificates: []tls.Certificate{*upstreamLeaf},
		NextProtos:   []string{"http/1.1"},
		MinVersion:   tls.VersionTLS12,
	})
	if err != nil {
		t.Fatalf("tls listen upstream: %v", err)
	}
	defer upstreamListener.Close()
	upstreamServer := &http.Server{
		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			body, readErr := io.ReadAll(r.Body)
			if readErr != nil {
				t.Errorf("read upstream body: %v", readErr)
				w.WriteHeader(http.StatusInternalServerError)
				return
			}
			requestHeaders <- r.Header.Clone()
			requestBody <- string(body)
			_, _ = w.Write([]byte("secure-ok"))
		}),
	}
	defer upstreamServer.Close()
	go func() {
		_ = upstreamServer.Serve(upstreamListener)
	}()

	proxyListener, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen proxy: %v", err)
	}
	defer proxyListener.Close()

	server := &Server{
		cfg: &config.NetdConfig{
			ProxyUpstreamTimeout: metav1.Duration{Duration: 2 * time.Second},
		},
		logger:            zap.NewNop(),
		tlsAuthority:      mitmAuthority,
		upstreamTLSConfig: &tls.Config{RootCAs: upstreamRootPool, NextProtos: []string{"http/1.1"}},
	}

	done := make(chan error, 1)
	go func() {
		conn, acceptErr := proxyListener.Accept()
		if acceptErr != nil {
			done <- acceptErr
			return
		}
		defer conn.Close()
		req := &adapterRequest{
			Server:   server,
			Compiled: &policy.CompiledPolicy{SandboxID: "sbx_123", TeamID: "team_123", Mode: v1alpha1.NetworkModeAllowAll},
			SrcIP:    "10.0.0.2",
			DestIP:   upstreamListener.Addr().(*net.TCPAddr).IP,
			DestPort: upstreamListener.Addr().(*net.TCPAddr).Port,
			Host:     "api.example.com",
			Conn:     conn,
			EgressAuth: &egressAuthContext{
				Rule: &policy.CompiledEgressAuthRule{
					Name:    "example-https",
					AuthRef: "example-api",
					TLSMode: v1alpha1.EgressTLSModeTerminateReoriginate,
				},
				Resolved: egressauth.NewHTTPHeadersResolveResponse("example-api", map[string]string{
					"Authorization": "Bearer secure-token",
				}, nil),
			},
		}
		done <- (&tlsAdapter{}).Handle(req)
	}()

	clientRootPool := x509.NewCertPool()
	if !clientRootPool.AppendCertsFromPEM(mitmCertPEM) {
		t.Fatal("append mitm ca")
	}
	rawConn, err := net.Dial("tcp4", proxyListener.Addr().String())
	if err != nil {
		t.Fatalf("dial proxy: %v", err)
	}
	defer rawConn.Close()
	clientTLS := tls.Client(rawConn, &tls.Config{
		ServerName: "api.example.com",
		RootCAs:    clientRootPool,
		NextProtos: []string{"http/1.1"},
		MinVersion: tls.VersionTLS12,
	})
	defer clientTLS.Close()
	if err := clientTLS.Handshake(); err != nil {
		t.Fatalf("handshake client tls: %v", err)
	}
	if got := clientTLS.ConnectionState().NegotiatedProtocol; got != "http/1.1" {
		t.Fatalf("negotiated protocol = %q, want http/1.1", got)
	}

	if _, err := io.WriteString(clientTLS, "POST /v1/test HTTP/1.1\r\nHost: api.example.com\r\nContent-Length: 7\r\n\r\npayload"); err != nil {
		t.Fatalf("write tls request: %v", err)
	}
	resp, err := http.ReadResponse(bufio.NewReader(clientTLS), nil)
	if err != nil {
		t.Fatalf("read https response: %v", err)
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("read https response body: %v", err)
	}
	_ = resp.Body.Close()
	if got := string(body); got != "secure-ok" {
		t.Fatalf("response body = %q", got)
	}
	if err := <-done; err != nil {
		t.Fatalf("tls adapter handle: %v", err)
	}

	headers := <-requestHeaders
	if got := headers.Get("Authorization"); got != "Bearer secure-token" {
		t.Fatalf("authorization header = %q", got)
	}
	if got := <-requestBody; got != "payload" {
		t.Fatalf("request body = %q", got)
	}
}

func TestTLSAdapterDefersHTTPSHeaderResolutionUntilAfterDownstreamHandshake(t *testing.T) {
	mitmCertPEM, mitmKeyPEM, err := newSelfSignedCertificateAuthority("sandbox0-mitm", time.Hour)
	if err != nil {
		t.Fatalf("new mitm ca: %v", err)
	}
	mitmAuthority, err := newCertificateAuthority(mitmCertPEM, mitmKeyPEM, time.Hour)
	if err != nil {
		t.Fatalf("new mitm authority: %v", err)
	}

	upstreamCAPEM, upstreamCAKeyPEM, err := newSelfSignedCertificateAuthority("sandbox0-upstream", time.Hour)
	if err != nil {
		t.Fatalf("new upstream ca: %v", err)
	}
	upstreamAuthority, err := newCertificateAuthority(upstreamCAPEM, upstreamCAKeyPEM, time.Hour)
	if err != nil {
		t.Fatalf("new upstream authority: %v", err)
	}
	upstreamLeaf, err := upstreamAuthority.CertificateForHost("api.example.com")
	if err != nil {
		t.Fatalf("upstream leaf: %v", err)
	}
	upstreamRootPool := x509.NewCertPool()
	if !upstreamRootPool.AppendCertsFromPEM(upstreamCAPEM) {
		t.Fatal("append upstream ca")
	}

	requestHeaders := make(chan http.Header, 1)
	upstreamListener, err := tls.Listen("tcp", "127.0.0.1:0", &tls.Config{
		Certificates: []tls.Certificate{*upstreamLeaf},
		NextProtos:   []string{"http/1.1"},
		MinVersion:   tls.VersionTLS12,
	})
	if err != nil {
		t.Fatalf("tls listen upstream: %v", err)
	}
	defer upstreamListener.Close()
	upstreamServer := &http.Server{
		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			requestHeaders <- r.Header.Clone()
			_, _ = io.Copy(io.Discard, r.Body)
			_, _ = w.Write([]byte("secure-ok"))
		}),
	}
	defer upstreamServer.Close()
	go func() {
		_ = upstreamServer.Serve(upstreamListener)
	}()

	proxyListener, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen proxy: %v", err)
	}
	defer proxyListener.Close()

	releaseResolver := make(chan struct{})
	resolver := &blockingEgressAuthResolver{
		called:  make(chan *egressauth.ResolveRequest, 1),
		release: releaseResolver,
		resp: egressauth.NewHTTPHeadersResolveResponse("example-api", map[string]string{
			"Authorization": "Bearer delayed-token",
		}, nil),
	}
	released := false
	defer func() {
		if !released {
			close(releaseResolver)
		}
	}()

	server := &Server{
		cfg: &config.NetdConfig{
			EgressAuthEnabled:    true,
			ProxyUpstreamTimeout: metav1.Duration{Duration: 2 * time.Second},
		},
		logger:            zap.NewNop(),
		tlsAuthority:      mitmAuthority,
		upstreamTLSConfig: &tls.Config{RootCAs: upstreamRootPool, NextProtos: []string{"http/1.1"}},
		authResolver:      resolver,
		authCache:         newMemoryEgressAuthCache(),
	}
	rule := &policy.CompiledEgressAuthRule{
		Name:     "example-https",
		AuthRef:  "example-api",
		Protocol: v1alpha1.EgressAuthProtocolHTTPS,
		TLSMode:  v1alpha1.EgressTLSModeTerminateReoriginate,
	}

	done := make(chan error, 1)
	go func() {
		conn, acceptErr := proxyListener.Accept()
		if acceptErr != nil {
			done <- acceptErr
			return
		}
		defer conn.Close()
		req := &adapterRequest{
			Server:   server,
			Compiled: &policy.CompiledPolicy{SandboxID: "sbx_123", TeamID: "team_123", Mode: v1alpha1.NetworkModeAllowAll},
			SrcIP:    "10.0.0.2",
			DestIP:   upstreamListener.Addr().(*net.TCPAddr).IP,
			DestPort: upstreamListener.Addr().(*net.TCPAddr).Port,
			Host:     "api.example.com",
			Conn:     conn,
		}
		server.attachEgressAuth(req, trafficDecision{
			Action:          decisionActionUseAdapter,
			Transport:       "tcp",
			Protocol:        "tls",
			MatchedAuthRule: rule,
		})
		done <- (&tlsAdapter{}).Handle(req)
	}()

	clientRootPool := x509.NewCertPool()
	if !clientRootPool.AppendCertsFromPEM(mitmCertPEM) {
		t.Fatal("append mitm ca")
	}
	rawConn, err := net.Dial("tcp4", proxyListener.Addr().String())
	if err != nil {
		t.Fatalf("dial proxy: %v", err)
	}
	defer rawConn.Close()
	if err := rawConn.SetDeadline(time.Now().Add(time.Second)); err != nil {
		t.Fatalf("set handshake deadline: %v", err)
	}
	clientTLS := tls.Client(rawConn, &tls.Config{
		ServerName: "api.example.com",
		RootCAs:    clientRootPool,
		NextProtos: []string{"http/1.1"},
		MinVersion: tls.VersionTLS12,
	})
	defer clientTLS.Close()
	if err := clientTLS.Handshake(); err != nil {
		t.Fatalf("handshake client tls before resolver release: %v", err)
	}
	if err := rawConn.SetDeadline(time.Time{}); err != nil {
		t.Fatalf("clear handshake deadline: %v", err)
	}
	select {
	case req := <-resolver.called:
		t.Fatalf("resolver called before downstream request was read: %#v", req)
	default:
	}

	if _, err := io.WriteString(clientTLS, "POST /v1/test HTTP/1.1\r\nHost: api.example.com\r\nContent-Length: 7\r\n\r\npayload"); err != nil {
		t.Fatalf("write tls request: %v", err)
	}
	select {
	case req := <-resolver.called:
		if req.AuthRef != "example-api" {
			t.Fatalf("resolver auth ref = %q, want example-api", req.AuthRef)
		}
	case <-time.After(time.Second):
		t.Fatal("resolver was not called after downstream request was read")
	}
	close(releaseResolver)
	released = true

	resp, err := http.ReadResponse(bufio.NewReader(clientTLS), nil)
	if err != nil {
		t.Fatalf("read https response: %v", err)
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("read https response body: %v", err)
	}
	_ = resp.Body.Close()
	if got := string(body); got != "secure-ok" {
		t.Fatalf("response body = %q", got)
	}
	if err := <-done; err != nil {
		t.Fatalf("tls adapter handle: %v", err)
	}

	headers := <-requestHeaders
	if got := headers.Get("Authorization"); got != "Bearer delayed-token" {
		t.Fatalf("authorization header = %q", got)
	}
}

func TestTLSAdapterInterceptsHTTPSOverHTTP2AndInjectsHeaders(t *testing.T) {
	mitmCertPEM, mitmKeyPEM, err := newSelfSignedCertificateAuthority("sandbox0-mitm", time.Hour)
	if err != nil {
		t.Fatalf("new mitm ca: %v", err)
	}
	mitmAuthority, err := newCertificateAuthority(mitmCertPEM, mitmKeyPEM, time.Hour)
	if err != nil {
		t.Fatalf("new mitm authority: %v", err)
	}

	upstreamCAPEM, upstreamCAKeyPEM, err := newSelfSignedCertificateAuthority("sandbox0-upstream", time.Hour)
	if err != nil {
		t.Fatalf("new upstream ca: %v", err)
	}
	upstreamAuthority, err := newCertificateAuthority(upstreamCAPEM, upstreamCAKeyPEM, time.Hour)
	if err != nil {
		t.Fatalf("new upstream authority: %v", err)
	}
	upstreamLeaf, err := upstreamAuthority.CertificateForHost("api.example.com")
	if err != nil {
		t.Fatalf("upstream leaf: %v", err)
	}
	upstreamRootPool := x509.NewCertPool()
	if !upstreamRootPool.AppendCertsFromPEM(upstreamCAPEM) {
		t.Fatal("append upstream ca")
	}

	requestHeaders := make(chan http.Header, 1)
	requestBody := make(chan string, 1)
	requestProtocol := make(chan string, 1)
	upstreamListener, err := tls.Listen("tcp", "127.0.0.1:0", &tls.Config{
		Certificates: []tls.Certificate{*upstreamLeaf},
		NextProtos:   []string{"h2"},
		MinVersion:   tls.VersionTLS12,
	})
	if err != nil {
		t.Fatalf("tls listen upstream: %v", err)
	}
	defer upstreamListener.Close()
	upstreamServer := &http.Server{
		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			body, readErr := io.ReadAll(r.Body)
			if readErr != nil {
				t.Errorf("read upstream body: %v", readErr)
				w.WriteHeader(http.StatusInternalServerError)
				return
			}
			requestHeaders <- r.Header.Clone()
			requestBody <- string(body)
			requestProtocol <- r.Proto
			_, _ = w.Write([]byte("h2-secure-ok"))
		}),
	}
	if err := http2.ConfigureServer(upstreamServer, &http2.Server{}); err != nil {
		t.Fatalf("configure upstream http2 server: %v", err)
	}
	defer upstreamServer.Close()
	go func() {
		_ = upstreamServer.Serve(upstreamListener)
	}()

	proxyListener, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen proxy: %v", err)
	}
	defer proxyListener.Close()

	server := &Server{
		cfg: &config.NetdConfig{
			ProxyUpstreamTimeout: metav1.Duration{Duration: 2 * time.Second},
		},
		logger:            zap.NewNop(),
		tlsAuthority:      mitmAuthority,
		upstreamTLSConfig: &tls.Config{RootCAs: upstreamRootPool},
	}

	done := make(chan error, 1)
	go func() {
		conn, acceptErr := proxyListener.Accept()
		if acceptErr != nil {
			done <- acceptErr
			return
		}
		defer conn.Close()
		req := &adapterRequest{
			Server:   server,
			Compiled: &policy.CompiledPolicy{SandboxID: "sbx_123", TeamID: "team_123", Mode: v1alpha1.NetworkModeAllowAll},
			SrcIP:    "10.0.0.2",
			DestIP:   upstreamListener.Addr().(*net.TCPAddr).IP,
			DestPort: upstreamListener.Addr().(*net.TCPAddr).Port,
			Host:     "api.example.com",
			Conn:     conn,
			EgressAuth: &egressAuthContext{
				Rule: &policy.CompiledEgressAuthRule{
					Name:     "example-https",
					AuthRef:  "example-api",
					Protocol: v1alpha1.EgressAuthProtocolHTTPS,
					TLSMode:  v1alpha1.EgressTLSModeTerminateReoriginate,
				},
				Resolved: egressauth.NewHTTPHeadersResolveResponse("example-api", map[string]string{
					"Authorization": "Bearer h2-token",
				}, nil),
			},
		}
		done <- (&tlsAdapter{}).Handle(req)
	}()

	clientRootPool := x509.NewCertPool()
	if !clientRootPool.AppendCertsFromPEM(mitmCertPEM) {
		t.Fatal("append mitm ca")
	}
	clientTLSConfig := &tls.Config{
		ServerName: "api.example.com",
		RootCAs:    clientRootPool,
		NextProtos: []string{"h2"},
		MinVersion: tls.VersionTLS12,
	}
	negotiatedProtocol := make(chan string, 1)
	clientTransport := &http2.Transport{
		TLSClientConfig: clientTLSConfig,
		DialTLSContext: func(ctx context.Context, network, addr string, cfg *tls.Config) (net.Conn, error) {
			_ = network
			_ = addr
			tlsCfg := clientTLSConfig
			if cfg != nil {
				tlsCfg = cfg.Clone()
			}
			rawConn, dialErr := net.Dial("tcp4", proxyListener.Addr().String())
			if dialErr != nil {
				return nil, dialErr
			}
			tlsConn := tls.Client(rawConn, tlsCfg)
			if handshakeErr := tlsConn.HandshakeContext(ctx); handshakeErr != nil {
				_ = rawConn.Close()
				return nil, handshakeErr
			}
			negotiatedProtocol <- tlsConn.ConnectionState().NegotiatedProtocol
			return tlsConn, nil
		},
	}
	defer clientTransport.CloseIdleConnections()
	client := &http.Client{
		Transport: clientTransport,
		Timeout:   5 * time.Second,
	}
	httpReq, err := http.NewRequest(http.MethodPost, "https://api.example.com/v1/test", strings.NewReader("payload"))
	if err != nil {
		t.Fatalf("new http2 request: %v", err)
	}
	resp, err := client.Do(httpReq)
	if err != nil {
		t.Fatalf("http2 client request: %v", err)
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("read http2 response body: %v", err)
	}
	_ = resp.Body.Close()
	if got := string(body); got != "h2-secure-ok" {
		t.Fatalf("response body = %q", got)
	}
	if got := <-negotiatedProtocol; got != "h2" {
		t.Fatalf("negotiated protocol = %q, want h2", got)
	}
	if got := <-requestProtocol; got != "HTTP/2.0" {
		t.Fatalf("upstream protocol = %q, want HTTP/2.0", got)
	}
	headers := <-requestHeaders
	if got := headers.Get("Authorization"); got != "Bearer h2-token" {
		t.Fatalf("authorization header = %q", got)
	}
	if got := <-requestBody; got != "payload" {
		t.Fatalf("request body = %q", got)
	}
	clientTransport.CloseIdleConnections()
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("tls adapter handle: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for tls adapter")
	}
}

func TestTLSAdapterInterceptsGRPCAndInjectsMetadata(t *testing.T) {
	mitmCertPEM, mitmKeyPEM, err := newSelfSignedCertificateAuthority("sandbox0-mitm", time.Hour)
	if err != nil {
		t.Fatalf("new mitm ca: %v", err)
	}
	mitmAuthority, err := newCertificateAuthority(mitmCertPEM, mitmKeyPEM, time.Hour)
	if err != nil {
		t.Fatalf("new mitm authority: %v", err)
	}

	upstreamCAPEM, upstreamCAKeyPEM, err := newSelfSignedCertificateAuthority("sandbox0-upstream", time.Hour)
	if err != nil {
		t.Fatalf("new upstream ca: %v", err)
	}
	upstreamAuthority, err := newCertificateAuthority(upstreamCAPEM, upstreamCAKeyPEM, time.Hour)
	if err != nil {
		t.Fatalf("new upstream authority: %v", err)
	}
	upstreamLeaf, err := upstreamAuthority.CertificateForHost("api.example.com")
	if err != nil {
		t.Fatalf("upstream leaf: %v", err)
	}
	upstreamRootPool := x509.NewCertPool()
	if !upstreamRootPool.AppendCertsFromPEM(upstreamCAPEM) {
		t.Fatal("append upstream ca")
	}

	receivedAuth := make(chan string, 1)
	grpcListener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen grpc upstream: %v", err)
	}
	defer grpcListener.Close()
	grpcServer := grpc.NewServer(
		grpc.Creds(credentials.NewTLS(&tls.Config{
			Certificates: []tls.Certificate{*upstreamLeaf},
			NextProtos:   []string{"h2"},
			MinVersion:   tls.VersionTLS12,
		})),
		grpc.UnaryInterceptor(func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
			md, _ := metadata.FromIncomingContext(ctx)
			values := md.Get("authorization")
			if len(values) > 0 {
				receivedAuth <- values[0]
			} else {
				receivedAuth <- ""
			}
			return handler(ctx, req)
		}),
	)
	defer grpcServer.Stop()
	healthServer := health.NewServer()
	healthServer.SetServingStatus("", healthpb.HealthCheckResponse_SERVING)
	healthpb.RegisterHealthServer(grpcServer, healthServer)
	go func() {
		_ = grpcServer.Serve(grpcListener)
	}()

	proxyListener, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen proxy: %v", err)
	}
	defer proxyListener.Close()

	server := &Server{
		cfg: &config.NetdConfig{
			ProxyUpstreamTimeout: metav1.Duration{Duration: 3 * time.Second},
		},
		logger:            zap.NewNop(),
		tlsAuthority:      mitmAuthority,
		upstreamTLSConfig: &tls.Config{RootCAs: upstreamRootPool, NextProtos: []string{"h2"}},
	}

	done := make(chan error, 1)
	go func() {
		conn, acceptErr := proxyListener.Accept()
		if acceptErr != nil {
			done <- acceptErr
			return
		}
		defer conn.Close()
		req := &adapterRequest{
			Server:   server,
			Compiled: &policy.CompiledPolicy{SandboxID: "sbx_123", TeamID: "team_123", Mode: v1alpha1.NetworkModeAllowAll},
			SrcIP:    "10.0.0.2",
			DestIP:   grpcListener.Addr().(*net.TCPAddr).IP,
			DestPort: grpcListener.Addr().(*net.TCPAddr).Port,
			Host:     "api.example.com",
			Conn:     conn,
			EgressAuth: &egressAuthContext{
				Rule: &policy.CompiledEgressAuthRule{
					Name:     "example-grpc",
					AuthRef:  "example-api",
					Protocol: v1alpha1.EgressAuthProtocolGRPC,
					TLSMode:  v1alpha1.EgressTLSModeTerminateReoriginate,
				},
				Resolved: &egressauth.ResolveResponse{
					AuthRef: "example-api",
					Headers: map[string]string{
						"Authorization": "Bearer grpc-token",
					},
				},
			},
		}
		done <- (&tlsAdapter{}).Handle(req)
	}()

	clientRootPool := x509.NewCertPool()
	if !clientRootPool.AppendCertsFromPEM(mitmCertPEM) {
		t.Fatal("append mitm ca")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	clientConn, err := grpc.DialContext(
		ctx,
		proxyListener.Addr().String(),
		grpc.WithBlock(),
		grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{
			ServerName: "api.example.com",
			RootCAs:    clientRootPool,
			NextProtos: []string{"h2"},
			MinVersion: tls.VersionTLS12,
		})),
	)
	if err != nil {
		t.Fatalf("dial grpc proxy: %v", err)
	}

	client := healthpb.NewHealthClient(clientConn)
	resp, err := client.Check(ctx, &healthpb.HealthCheckRequest{})
	if err != nil {
		t.Fatalf("grpc health check: %v", err)
	}
	if resp.GetStatus() != healthpb.HealthCheckResponse_SERVING {
		t.Fatalf("unexpected health status %v", resp.GetStatus())
	}
	if got := <-receivedAuth; got != "Bearer grpc-token" {
		t.Fatalf("authorization metadata = %q", got)
	}
	if err := clientConn.Close(); err != nil {
		t.Fatalf("close grpc client: %v", err)
	}
	if err := <-done; err != nil {
		t.Fatalf("tls adapter handle: %v", err)
	}
}

func TestTLSAdapterFailOpenRelaysRawTLSWhenMITMCAIsMissing(t *testing.T) {
	upstreamCAPEM, upstreamCAKeyPEM, err := newSelfSignedCertificateAuthority("sandbox0-upstream", time.Hour)
	if err != nil {
		t.Fatalf("new upstream ca: %v", err)
	}
	upstreamAuthority, err := newCertificateAuthority(upstreamCAPEM, upstreamCAKeyPEM, time.Hour)
	if err != nil {
		t.Fatalf("new upstream authority: %v", err)
	}
	upstreamLeaf, err := upstreamAuthority.CertificateForHost("api.example.com")
	if err != nil {
		t.Fatalf("upstream leaf: %v", err)
	}

	upstreamListener, err := tls.Listen("tcp", "127.0.0.1:0", &tls.Config{
		Certificates: []tls.Certificate{*upstreamLeaf},
		NextProtos:   []string{"http/1.1"},
		MinVersion:   tls.VersionTLS12,
	})
	if err != nil {
		t.Fatalf("tls listen upstream: %v", err)
	}
	defer upstreamListener.Close()
	upstreamServer := &http.Server{
		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			_, _ = w.Write([]byte("passthrough-ok"))
		}),
	}
	defer upstreamServer.Close()
	go func() {
		_ = upstreamServer.Serve(upstreamListener)
	}()

	proxyListener, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen proxy: %v", err)
	}
	defer proxyListener.Close()

	server := &Server{
		cfg: &config.NetdConfig{
			ProxyUpstreamTimeout: metav1.Duration{Duration: 2 * time.Second},
		},
		logger: zap.NewNop(),
	}

	done := make(chan error, 1)
	go func() {
		conn, acceptErr := proxyListener.Accept()
		if acceptErr != nil {
			done <- acceptErr
			return
		}
		defer conn.Close()
		req := &adapterRequest{
			Server:   server,
			Compiled: &policy.CompiledPolicy{SandboxID: "sbx_123", TeamID: "team_123", Mode: v1alpha1.NetworkModeAllowAll},
			SrcIP:    "10.0.0.2",
			DestIP:   upstreamListener.Addr().(*net.TCPAddr).IP,
			DestPort: upstreamListener.Addr().(*net.TCPAddr).Port,
			Host:     "api.example.com",
			Conn:     conn,
			EgressAuth: &egressAuthContext{
				Rule: &policy.CompiledEgressAuthRule{
					Name:          "example-https",
					AuthRef:       "example-api",
					TLSMode:       v1alpha1.EgressTLSModeTerminateReoriginate,
					FailurePolicy: v1alpha1.EgressAuthFailurePolicyFailOpen,
				},
				FailurePolicy: string(v1alpha1.EgressAuthFailurePolicyFailOpen),
				ResolveError:  io.EOF,
				BypassReason:  "resolve_error",
			},
		}
		done <- (&tlsAdapter{}).Handle(req)
	}()

	clientRootPool := x509.NewCertPool()
	if !clientRootPool.AppendCertsFromPEM(upstreamCAPEM) {
		t.Fatal("append upstream ca")
	}
	rawConn, err := net.Dial("tcp4", proxyListener.Addr().String())
	if err != nil {
		t.Fatalf("dial proxy: %v", err)
	}
	defer rawConn.Close()
	clientTLS := tls.Client(rawConn, &tls.Config{
		ServerName: "api.example.com",
		RootCAs:    clientRootPool,
		NextProtos: []string{"http/1.1"},
		MinVersion: tls.VersionTLS12,
	})
	defer clientTLS.Close()

	if _, err := io.WriteString(clientTLS, "GET /v1/test HTTP/1.1\r\nHost: api.example.com\r\nConnection: close\r\n\r\n"); err != nil {
		t.Fatalf("write tls request: %v", err)
	}
	resp, err := http.ReadResponse(bufio.NewReader(clientTLS), nil)
	if err != nil {
		t.Fatalf("read https response: %v", err)
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("read https response body: %v", err)
	}
	_ = resp.Body.Close()
	if got := string(body); got != "passthrough-ok" {
		t.Fatalf("response body = %q", got)
	}
	if err := clientTLS.Close(); err != nil {
		t.Fatalf("close client tls: %v", err)
	}
	if err := <-done; err != nil {
		t.Fatalf("tls adapter handle: %v", err)
	}
}

func TestTLSAdapterReturnsErrorWhenMITMCAIsMissing(t *testing.T) {
	server := &Server{
		cfg: &config.NetdConfig{
			ProxyUpstreamTimeout: metav1.Duration{Duration: time.Second},
		},
		logger: zap.NewNop(),
	}
	clientConn, upstreamConn := net.Pipe()
	defer clientConn.Close()
	defer upstreamConn.Close()

	err := (&tlsAdapter{}).Handle(&adapterRequest{
		Server:   upstreamConnServer(server),
		Conn:     upstreamConn,
		Host:     "api.example.com",
		DestIP:   net.ParseIP("127.0.0.1"),
		DestPort: 443,
		EgressAuth: &egressAuthContext{
			Rule: &policy.CompiledEgressAuthRule{
				Name:    "example-https",
				AuthRef: "example-api",
				TLSMode: v1alpha1.EgressTLSModeTerminateReoriginate,
			},
			Resolved: &egressauth.ResolveResponse{AuthRef: "example-api"},
		},
	})
	if err == nil || !strings.Contains(err.Error(), "authority") {
		t.Fatalf("expected tls authority error, got %v", err)
	}
}

func TestTLSAdapterTerminatesTLSAndPresentsClientCertificate(t *testing.T) {
	mitmCertPEM, mitmKeyPEM, err := newSelfSignedCertificateAuthority("sandbox0-mitm", time.Hour)
	if err != nil {
		t.Fatalf("new mitm ca: %v", err)
	}
	mitmAuthority, err := newCertificateAuthority(mitmCertPEM, mitmKeyPEM, time.Hour)
	if err != nil {
		t.Fatalf("new mitm authority: %v", err)
	}

	upstreamCAPEM, upstreamCAKeyPEM, err := newSelfSignedCertificateAuthority("sandbox0-upstream", time.Hour)
	if err != nil {
		t.Fatalf("new upstream ca: %v", err)
	}
	upstreamAuthority, err := newCertificateAuthority(upstreamCAPEM, upstreamCAKeyPEM, time.Hour)
	if err != nil {
		t.Fatalf("new upstream authority: %v", err)
	}
	upstreamLeaf, err := upstreamAuthority.CertificateForHost("db.example.com")
	if err != nil {
		t.Fatalf("upstream leaf: %v", err)
	}
	upstreamRootPool := x509.NewCertPool()
	if !upstreamRootPool.AppendCertsFromPEM(upstreamCAPEM) {
		t.Fatal("append upstream ca")
	}

	clientCertPEM, clientKeyPEM, err := newSelfSignedClientCertificate("sandbox0-client", time.Hour)
	if err != nil {
		t.Fatalf("new client cert: %v", err)
	}
	clientCAPool := x509.NewCertPool()
	if !clientCAPool.AppendCertsFromPEM(clientCertPEM) {
		t.Fatal("append client ca")
	}

	upstreamListener, err := tls.Listen("tcp", "127.0.0.1:0", &tls.Config{
		Certificates: []tls.Certificate{*upstreamLeaf},
		ClientAuth:   tls.RequireAndVerifyClientCert,
		ClientCAs:    clientCAPool,
		MinVersion:   tls.VersionTLS12,
	})
	if err != nil {
		t.Fatalf("tls listen upstream: %v", err)
	}
	defer upstreamListener.Close()

	received := make(chan string, 1)
	upstreamDone := make(chan error, 1)
	go func() {
		conn, acceptErr := upstreamListener.Accept()
		if acceptErr != nil {
			upstreamDone <- acceptErr
			return
		}
		defer conn.Close()
		tlsConn := conn.(*tls.Conn)
		if err := tlsConn.Handshake(); err != nil {
			upstreamDone <- err
			return
		}
		state := tlsConn.ConnectionState()
		if len(state.PeerCertificates) == 0 {
			upstreamDone <- io.EOF
			return
		}
		buf := make([]byte, len("hello over tls"))
		if _, err := io.ReadFull(tlsConn, buf); err != nil {
			upstreamDone <- err
			return
		}
		received <- string(buf)
		_, err = tlsConn.Write([]byte("secure reply"))
		upstreamDone <- err
	}()

	proxyListener, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen proxy: %v", err)
	}
	defer proxyListener.Close()

	server := &Server{
		cfg: &config.NetdConfig{
			ProxyUpstreamTimeout: metav1.Duration{Duration: 3 * time.Second},
		},
		logger:            zap.NewNop(),
		tlsAuthority:      mitmAuthority,
		upstreamTLSConfig: &tls.Config{RootCAs: upstreamRootPool},
	}

	done := make(chan error, 1)
	go func() {
		conn, acceptErr := proxyListener.Accept()
		if acceptErr != nil {
			done <- acceptErr
			return
		}
		defer conn.Close()
		req := &adapterRequest{
			Server:   server,
			Compiled: &policy.CompiledPolicy{SandboxID: "sbx_123", TeamID: "team_123", Mode: v1alpha1.NetworkModeAllowAll},
			SrcIP:    "10.0.0.2",
			DestIP:   upstreamListener.Addr().(*net.TCPAddr).IP,
			DestPort: upstreamListener.Addr().(*net.TCPAddr).Port,
			Host:     "db.example.com",
			Conn:     conn,
			EgressAuth: &egressAuthContext{
				Rule: &policy.CompiledEgressAuthRule{
					Name:     "example-mtls",
					AuthRef:  "example-cert",
					Protocol: v1alpha1.EgressAuthProtocolTLS,
					TLSMode:  v1alpha1.EgressTLSModeTerminateReoriginate,
				},
				Resolved: egressauth.NewTLSClientCertificateResolveResponse("example-cert", &egressauth.TLSClientCertificateDirective{
					CertificatePEM: string(clientCertPEM),
					PrivateKeyPEM:  string(clientKeyPEM),
				}, nil),
			},
		}
		done <- (&tlsAdapter{}).Handle(req)
	}()

	clientRootPool := x509.NewCertPool()
	if !clientRootPool.AppendCertsFromPEM(mitmCertPEM) {
		t.Fatal("append mitm ca")
	}
	rawConn, err := net.Dial("tcp4", proxyListener.Addr().String())
	if err != nil {
		t.Fatalf("dial proxy: %v", err)
	}
	defer rawConn.Close()
	clientTLS := tls.Client(rawConn, &tls.Config{
		ServerName: "db.example.com",
		RootCAs:    clientRootPool,
		MinVersion: tls.VersionTLS12,
	})
	defer clientTLS.Close()

	if _, err := io.WriteString(clientTLS, "hello over tls"); err != nil {
		t.Fatalf("write tls payload: %v", err)
	}
	reply := make([]byte, len("secure reply"))
	if _, err := io.ReadFull(clientTLS, reply); err != nil {
		t.Fatalf("read tls reply: %v", err)
	}
	if got := string(reply); got != "secure reply" {
		t.Fatalf("reply = %q", got)
	}
	if err := clientTLS.Close(); err != nil {
		t.Fatalf("close client tls: %v", err)
	}
	if got := <-received; got != "hello over tls" {
		t.Fatalf("upstream payload = %q", got)
	}
	if err := <-upstreamDone; err != nil {
		t.Fatalf("upstream server: %v", err)
	}
	if err := <-done; err != nil {
		t.Fatalf("tls adapter handle: %v", err)
	}
}

func newSelfSignedClientCertificate(commonName string, validity time.Duration) ([]byte, []byte, error) {
	if commonName == "" {
		commonName = "sandbox0-client"
	}
	if validity <= 0 {
		validity = time.Hour
	}
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, nil, err
	}
	serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return nil, nil, err
	}
	now := time.Now().UTC()
	template := &x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName: commonName,
		},
		NotBefore:             now.Add(-time.Minute),
		NotAfter:              now.Add(validity),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}
	der, err := x509.CreateCertificate(rand.Reader, template, template, &privateKey.PublicKey, privateKey)
	if err != nil {
		return nil, nil, err
	}
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)})
	return certPEM, keyPEM, nil
}
