设置 LLM 认证代理

"""执行 OAuth2 客户端凭据令牌交换的 ext_authz 服务。

作为边车（或独立服务）运行在主认证代理组件旁。
对于每个 ext_authz 检查请求，它返回一个缓存的 OAuth 访问令牌，
当令牌过期时，从配置的令牌端点进行刷新。

环境变量：
  OAUTH_TOKEN_URL    – 令牌端点（例如 https://login.example.com/oauth/token）
  OAUTH_CLIENT_ID    – 客户端凭据授权的客户端 ID
  OAUTH_CLIENT_SECRET– 客户端凭据授权的客户端密钥
  OAUTH_SCOPE        – （可选）请求的空格分隔作用域
  LISTEN_PORT        – （可选）监听端口，默认 10002
"""

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

# ---------------------------------------------------------------------------
# 配置
# ---------------------------------------------------------------------------
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"))

# 在令牌实际过期前多少秒刷新它。
EXPIRY_BUFFER_SECONDS = 30

# ---------------------------------------------------------------------------
# 令牌缓存（线程安全）
# ---------------------------------------------------------------------------
_lock = threading.Lock()
_cached_token: str | None = None
_token_expiry: float = 0  # 历元秒数


def _fetch_token() -> tuple[str, float]:
    """执行客户端凭据授权并返回 (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:
    """返回一个有效的访问令牌，必要时进行刷新。"""
    global _cached_token, _token_expiry
    with _lock:
        if _cached_token and time.time() < _token_expiry:
            return _cached_token
    # 在锁外获取，这样其他请求不会因 I/O 而阻塞。
    token, expiry = _fetch_token()
    with _lock:
        _cached_token = token
        _token_expiry = expiry
    print(f"已刷新 OAuth 令牌（{int(expiry - time.time())}秒后过期）", flush=True)
    return token


# ---------------------------------------------------------------------------
# ext_authz HTTP 处理器
# ---------------------------------------------------------------------------
class Handler(BaseHTTPRequestHandler):
    def do_any(self):
        try:
            token = get_token()
        except Exception as exc:
            print(f"获取 OAuth 令牌失败: {exc}", flush=True)
            self.send_response(500)
            self.send_header("Content-Type", "text/plain")
            self.end_headers()
            self.wfile.write(b"OAuth 令牌交换失败")
            return

        self.send_response(200)
        # 根据需要替换头名称 - 此头将被转发给上游 LLM 提供商/网关。
        self.send_header("Authorization", f"Bearer {token}")
        self.end_headers()

    # 处理 Envoy 为 ext_authz 检查可能发送的所有方法。
    do_GET = do_POST = do_PUT = do_DELETE = do_PATCH = do_HEAD = do_OPTIONS = do_any

    def log_message(self, format, *args):
        # 更简洁的日志 - 仅打印错误。
        pass


if __name__ == "__main__":
    server = HTTPServer(("0.0.0.0", LISTEN_PORT), Handler)
    print(f"ext-authz-oauth 正在监听 :{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: 替换为您的认证逻辑，例如将 JWT 交换为特定于提供商的
            // 令牌，调用密钥管理器等。
            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("请求体解析失败，直接透传: %v", err)
            return resp
        }

        // TODO: 替换为您的转换逻辑。
        // 此示例将 OpenAI 格式的请求体包装在自定义信封中。
        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("请求体序列化失败，直接透传: %v", err)
            return resp
        }

        // 重要：更新 content-length 以匹配新的请求体大小。
        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("监听失败: %v", err)
        }
        s := grpc.NewServer()
        ext_proc.RegisterExternalProcessorServer(s, &server{})
        log.Println("transformer 正在监听 :50051")
        if err := s.Serve(lis); err != nil {
            log.Fatalf("服务启动失败: %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