Set up the LLM auth proxy

"""ext_authz service that performs an OAuth2 client-credentials token exchange.

Runs as a sidecar (or standalone service) alongside the main auth-proxy component.
On each ext_authz check request it returns a cached OAuth access token,
refreshing it from the configured token endpoint when expired.

Environment variables:
  OAUTH_TOKEN_URL    – Token endpoint (e.g. https://login.example.com/oauth/token)
  OAUTH_CLIENT_ID    – Client ID for the credentials grant
  OAUTH_CLIENT_SECRET– Client secret for the credentials grant
  OAUTH_SCOPE        – (optional) Space-separated scopes to request
  LISTEN_PORT        – (optional) Port to listen on, default 10002
"""

from http.server import HTTPServer, BaseHTTPRequestHandler
import json
import os
import sys
import threading
import time
import urllib.request
import urllib.parse

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
TOKEN_URL = os.environ["OAUTH_TOKEN_URL"]
CLIENT_ID = os.environ["OAUTH_CLIENT_ID"]
CLIENT_SECRET = os.environ["OAUTH_CLIENT_SECRET"]
SCOPE = os.environ.get("OAUTH_SCOPE", "")
LISTEN_PORT = int(os.environ.get("LISTEN_PORT", "10002"))

# Refresh the token this many seconds before it actually expires.
EXPIRY_BUFFER_SECONDS = 30

# ---------------------------------------------------------------------------
# Token cache (thread-safe)
# ---------------------------------------------------------------------------
_lock = threading.Lock()
_cached_token: str | None = None
_token_expiry: float = 0  # epoch seconds


def _fetch_token() -> tuple[str, float]:
    """Perform a client_credentials grant and return (access_token, expiry_epoch)."""
    data = urllib.parse.urlencode({
        "grant_type": "client_credentials",
        "client_id": CLIENT_ID,
        "client_secret": CLIENT_SECRET,
        **({"scope": SCOPE} if SCOPE else {}),
    }).encode()

    req = urllib.request.Request(
        TOKEN_URL,
        data=data,
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=10) as resp:
        body = json.loads(resp.read())

    access_token = body["access_token"]
    expires_in = int(body.get("expires_in", 3600))
    expiry = time.time() + expires_in - EXPIRY_BUFFER_SECONDS
    return access_token, expiry


def get_token() -> str:
    """Return a valid access token, refreshing if necessary."""
    global _cached_token, _token_expiry
    with _lock:
        if _cached_token and time.time() < _token_expiry:
            return _cached_token
    # Fetch outside the lock so other requests aren't blocked on I/O.
    token, expiry = _fetch_token()
    with _lock:
        _cached_token = token
        _token_expiry = expiry
    print(f"Refreshed OAuth token (expires in {int(expiry - time.time())}s)", flush=True)
    return token


# ---------------------------------------------------------------------------
# ext_authz HTTP handler
# ---------------------------------------------------------------------------
class Handler(BaseHTTPRequestHandler):
    def do_any(self):
        try:
            token = get_token()
        except Exception as exc:
            print(f"OAuth token fetch failed: {exc}", flush=True)
            self.send_response(500)
            self.send_header("Content-Type", "text/plain")
            self.end_headers()
            self.wfile.write(b"OAuth token exchange failed")
            return

        self.send_response(200)
        # Replace the header name as needed - this header will be forwarded to the upstream LLM provider / gateway.
        self.send_header("Authorization", f"Bearer {token}")
        self.end_headers()

    # Handle every method Envoy might send for ext_authz checks.
    do_GET = do_POST = do_PUT = do_DELETE = do_PATCH = do_HEAD = do_OPTIONS = do_any

    def log_message(self, format, *args):
        # Quieter logs — only print errors.
        pass


if __name__ == "__main__":
    server = HTTPServer(("0.0.0.0", LISTEN_PORT), Handler)
    print(f"ext-authz-oauth listening on :{LISTEN_PORT}", flush=True)
    print(f"  token_url={TOKEN_URL} client_id=<redacted>", flush=True)
    server.serve_forever()

apiVersion: v1
kind: ConfigMap
metadata:
  name: transformer-source
data:
  main.go: |
    package main

    import (
        "encoding/json"
        "fmt"
        "io"
        "log"
        "net"
        "strings"

        core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
        ext_proc "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
        "google.golang.org/grpc"
    )

    type server struct {
        ext_proc.UnimplementedExternalProcessorServer
    }

    func (s *server) Process(stream ext_proc.ExternalProcessor_ProcessServer) error {
        for {
            req, err := stream.Recv()
            if err == io.EOF {
                return nil
            }
            if err != nil {
                return err
            }

            var resp *ext_proc.ProcessingResponse
            switch v := req.Request.(type) {
            case *ext_proc.ProcessingRequest_RequestHeaders:
                resp = handleRequestHeaders(v.RequestHeaders)
            case *ext_proc.ProcessingRequest_RequestBody:
                resp = handleRequestBody(v.RequestBody)
            default:
                resp = &ext_proc.ProcessingResponse{}
            }

            if err := stream.Send(resp); err != nil {
                return err
            }
        }
    }

    func handleRequestHeaders(headers *ext_proc.HttpHeaders) *ext_proc.ProcessingResponse {
        var jwtValue string
        for _, h := range headers.Headers.Headers {
            if strings.EqualFold(h.Key, "x-langsmith-llm-auth") {
                if len(h.RawValue) > 0 {
                    jwtValue = string(h.RawValue)
                } else {
                    jwtValue = h.Value
                }
                break
            }
        }

        resp := &ext_proc.ProcessingResponse{
            Response: &ext_proc.ProcessingResponse_RequestHeaders{
                RequestHeaders: &ext_proc.HeadersResponse{},
            },
        }

        if jwtValue != "" {
            // TODO: Replace with your auth logic, e.g. exchange JWT for a
            // provider-specific token, call a secrets manager, etc.
            providerKey := "Bearer your-provider-key"

            headerResp := resp.GetRequestHeaders()
            headerResp.Response = &ext_proc.CommonResponse{
                HeaderMutation: &ext_proc.HeaderMutation{
                    SetHeaders: []*core.HeaderValueOption{
                        {
                            Header: &core.HeaderValue{
                                Key:      "Authorization",
                                RawValue: []byte(providerKey),
                            },
                        },
                    },
                },
            }
        }
        return resp
    }

    func handleRequestBody(body *ext_proc.HttpBody) *ext_proc.ProcessingResponse {
        resp := &ext_proc.ProcessingResponse{
            Response: &ext_proc.ProcessingResponse_RequestBody{
                RequestBody: &ext_proc.BodyResponse{},
            },
        }

        var original map[string]interface{}
        if err := json.Unmarshal(body.Body, &original); err != nil {
            log.Printf("Body parse failed, passing through: %v", err)
            return resp
        }

        // TODO: Replace with your transformation logic.
        // This example wraps the OpenAI-format body in a custom envelope.
        transformed := map[string]interface{}{
            "custom_model":    original["model"],
            "custom_messages": original["messages"],
            "metadata":        map[string]string{"source": "langsmith"},
        }

        newBody, err := json.Marshal(transformed)
        if err != nil {
            log.Printf("Body marshal failed, passing through: %v", err)
            return resp
        }

        // IMPORTANT: update content-length to match the new body size.
        bodyResp := resp.GetRequestBody()
        bodyResp.Response = &ext_proc.CommonResponse{
            Status: ext_proc.CommonResponse_CONTINUE_AND_REPLACE,
            HeaderMutation: &ext_proc.HeaderMutation{
                SetHeaders: []*core.HeaderValueOption{
                    {
                        Header: &core.HeaderValue{
                            Key:      "content-length",
                            RawValue: []byte(fmt.Sprintf("%d", len(newBody))),
                        },
                    },
                },
            },
            BodyMutation: &ext_proc.BodyMutation{
                Mutation: &ext_proc.BodyMutation_Body{
                    Body: newBody,
                },
            },
        }
        return resp
    }

    func main() {
        lis, err := net.Listen("tcp", ":50051")
        if err != nil {
            log.Fatalf("failed to listen: %v", err)
        }
        s := grpc.NewServer()
        ext_proc.RegisterExternalProcessorServer(s, &server{})
        log.Println("transformer listening on :50051")
        if err := s.Serve(lis); err != nil {
            log.Fatalf("failed to serve: %v", err)
        }
    }
  go.mod: |
    module transformer

    go 1.23

    require (
        github.com/envoyproxy/go-control-plane/envoy v1.32.4
        google.golang.org/grpc v1.72.1
    )

apiVersion: apps/v1
kind: Deployment
metadata:
  name: transformer
  labels:
    app: transformer
spec:
  replicas: 1
  selector:
    matchLabels:
      app: transformer
  template:
    metadata:
      labels:
        app: transformer
    spec:
      initContainers:
        - name: build
          image: golang:1.23
          command: ["sh", "-c"]
          args:
            - |
              cp /src/main.go /src/go.mod /build/ &&
              cd /build &&
              go mod tidy &&
              CGO_ENABLED=0 go build -o /build/transformer ./main.go
          volumeMounts:
            - name: source
              mountPath: /src
              readOnly: true
            - name: binary
              mountPath: /build
      containers:
        - name: transformer
          image: gcr.io/distroless/static-debian12:nonroot
          command: ["/app/transformer"]
          ports:
            - containerPort: 50051
          volumeMounts:
            - name: binary
              mountPath: /app
              readOnly: true
      volumes:
        - name: source
          configMap:
            name: transformer-source
        - name: binary
          emptyDir: {}
---
apiVersion: v1
kind: Service
metadata:
  name: transformer
  labels:
    app: transformer
spec:
  selector:
    app: transformer
  ports:
    - port: 50051
      targetPort: 50051
      protocol: TCP