mirror of
https://mirror.skon.top/github.com/router-for-me/CLIProxyAPI
synced 2026-04-21 01:10:14 +08:00
Merge pull request #287 from ben-vargas/feat-amp-cli-module
Amp CLI Integration Module
This commit is contained in:
Luis Pater
GitHub
12
README.md
12
README.md
@@ -25,6 +25,7 @@ Get 10% OFF GLM CODING PLAN：https://z.ai/subscribe?ic=8JVLJQFSKB
|
||||
- Claude Code support via OAuth login
|
||||
- Qwen Code support via OAuth login
|
||||
- iFlow support via OAuth login
|
||||
- Amp CLI and IDE extensions support with provider routing
|
||||
- Streaming and non-streaming responses
|
||||
- Function calling/tools support
|
||||
- Multimodal input support (text and images)
|
||||
@@ -48,6 +49,17 @@ CLIProxyAPI Guides: [https://help.router-for.me/](https://help.router-for.me/)
|
||||
|
||||
see [MANAGEMENT_API.md](https://help.router-for.me/management/api)
|
||||
|
||||
## Amp CLI Support
|
||||
|
||||
CLIProxyAPI includes integrated support for [Amp CLI](https://ampcode.com) and Amp IDE extensions, enabling you to use your Google/ChatGPT/Claude OAuth subscriptions with Amp's coding tools:
|
||||
|
||||
- Provider route aliases for Amp's API patterns (`/api/provider/{provider}/v1...`)
|
||||
- Management proxy for OAuth authentication and account features
|
||||
- Smart model fallback with automatic routing
|
||||
- Security-first design with localhost-only management endpoints
|
||||
|
||||
**→ [Complete Amp CLI Integration Guide](docs/amp-cli-integration.md)**
|
||||
|
||||
## SDK Docs
|
||||
|
||||
- Usage: [docs/sdk-usage.md](docs/sdk-usage.md)
|
||||
|
||||
392
docs/amp-cli-integration.md
Normal file
392
docs/amp-cli-integration.md
Normal file
@@ -0,0 +1,392 @@
|
||||
# Amp CLI Integration Guide
|
||||
|
||||
This guide explains how to use CLIProxyAPI with Amp CLI and Amp IDE extensions, enabling you to use your existing Google/ChatGPT/Claude subscriptions (via OAuth) with Amp's CLI.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Which Providers Should You Authenticate?](#which-providers-should-you-authenticate)
|
||||
- [Architecture](#architecture)
|
||||
- [Configuration](#configuration)
|
||||
- [Setup](#setup)
|
||||
- [Usage](#usage)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
|
||||
## Overview
|
||||
|
||||
The Amp CLI integration adds specialized routing to support Amp's API patterns while maintaining full compatibility with all existing CLIProxyAPI features. This allows you to use both traditional CLIProxyAPI features and Amp CLI with the same proxy server.
|
||||
|
||||
### Key Features
|
||||
|
||||
- **Provider route aliases**: Maps Amp's `/api/provider/{provider}/v1...` patterns to CLIProxyAPI handlers
|
||||
- **Management proxy**: Forwards OAuth and account management requests to Amp's control plane
|
||||
- **Smart fallback**: Automatically routes unconfigured models to ampcode.com
|
||||
- **Secret management**: Configurable precedence (config > env > file) with 5-minute caching
|
||||
- **Security-first**: Management routes restricted to localhost by default
|
||||
- **Automatic gzip handling**: Decompresses responses from Amp upstream
|
||||
|
||||
### What You Can Do
|
||||
|
||||
- Use Amp CLI with your Google account (Gemini 3 Pro Preview, Gemini 2.5 Pro, Gemini 2.5 Flash)
|
||||
- Use Amp CLI with your ChatGPT Plus/Pro subscription (GPT-5, GPT-5 Codex models)
|
||||
- Use Amp CLI with your Claude Pro/Max subscription (Claude Sonnet 4.5, Opus 4.1)
|
||||
- Use Amp IDE extensions (VS Code, Cursor, Windsurf, etc.) with the same proxy
|
||||
- Run multiple CLI tools (Factory + Amp) through one proxy server
|
||||
- Route unconfigured models automatically through ampcode.com
|
||||
|
||||
### Which Providers Should You Authenticate?
|
||||
|
||||
**Important**: The providers you need to authenticate depend on which models and features your installed version of Amp currently uses. Amp employs different providers for various agent modes and specialized subagents:
|
||||
|
||||
- **Smart mode**: Uses Google/Gemini models (Gemini 3 Pro)
|
||||
- **Rush mode**: Uses Anthropic/Claude models (Claude Haiku 4.5)
|
||||
- **Oracle subagent**: Uses OpenAI/GPT models (GPT-5 medium reasoning)
|
||||
- **Librarian subagent**: Uses Anthropic/Claude models (Claude Sonnet 4.5)
|
||||
- **Search subagent**: Uses Anthropic/Claude models (Claude Haiku 4.5)
|
||||
- **Review feature**: Uses Google/Gemini models (Gemini 2.5 Flash-Lite)
|
||||
|
||||
For the most current information about which models Amp uses, see the **[Amp Models Documentation](https://ampcode.com/models)**.
|
||||
|
||||
#### Fallback Behavior
|
||||
|
||||
CLIProxyAPI uses a smart fallback system:
|
||||
|
||||
1. **Provider authenticated locally** (`--login`, `--codex-login`, `--claude-login`):
|
||||
- Requests use **your OAuth subscription** (ChatGPT Plus/Pro, Claude Pro/Max, Google account)
|
||||
- You benefit from your subscription's included usage quotas
|
||||
- No Amp credits consumed
|
||||
|
||||
2. **Provider NOT authenticated locally**:
|
||||
- Requests automatically forward to **ampcode.com**
|
||||
- Uses Amp's backend provider connections
|
||||
- **Requires Amp credits** if the provider is paid (OpenAI, Anthropic paid tiers)
|
||||
- May result in errors if Amp credit balance is insufficient
|
||||
|
||||
**Recommendation**: Authenticate all providers you have subscriptions for to maximize value and minimize Amp credit usage. If you don't have subscriptions to all providers Amp uses, ensure you have sufficient Amp credits available for fallback requests.
|
||||
|
||||
## Architecture
|
||||
|
||||
### Request Flow
|
||||
|
||||
```
|
||||
Amp CLI/IDE
|
||||
↓
|
||||
├─ Provider API requests (/api/provider/{provider}/v1/...)
|
||||
│ ↓
|
||||
│ ├─ Model configured locally?
|
||||
│ │ YES → Use local OAuth tokens (OpenAI/Claude/Gemini handlers)
|
||||
│ │ NO → Forward to ampcode.com (reverse proxy)
|
||||
│ ↓
|
||||
│ Response
|
||||
│
|
||||
└─ Management requests (/api/auth, /api/user, /api/threads, ...)
|
||||
↓
|
||||
├─ Localhost check (security)
|
||||
↓
|
||||
└─ Reverse proxy to ampcode.com
|
||||
↓
|
||||
Response (auto-decompressed if gzipped)
|
||||
```
|
||||
|
||||
### Components
|
||||
|
||||
The Amp integration is implemented as a modular routing module (`internal/api/modules/amp/`) with these components:
|
||||
|
||||
1. **Route Aliases** (`routes.go`): Maps Amp-style paths to standard handlers
|
||||
2. **Reverse Proxy** (`proxy.go`): Forwards management requests to ampcode.com
|
||||
3. **Fallback Handler** (`fallback_handlers.go`): Routes unconfigured models to ampcode.com
|
||||
4. **Secret Management** (`secret.go`): Multi-source API key resolution with caching
|
||||
5. **Main Module** (`amp.go`): Orchestrates registration and configuration
|
||||
|
||||
## Configuration
|
||||
|
||||
### Basic Configuration
|
||||
|
||||
Add these fields to your `config.yaml`:
|
||||
|
||||
```yaml
|
||||
# Amp upstream control plane (required for management routes)
|
||||
amp-upstream-url: "https://ampcode.com"
|
||||
|
||||
# Optional: Override API key (otherwise uses env or file)
|
||||
# amp-upstream-api-key: "your-amp-api-key"
|
||||
|
||||
# Security: restrict management routes to localhost (recommended)
|
||||
amp-restrict-management-to-localhost: true
|
||||
```
|
||||
|
||||
### Secret Resolution Precedence
|
||||
|
||||
The Amp module resolves API keys using this precedence order:
|
||||
|
||||
| Source | Key | Priority | Cache |
|
||||
|--------|-----|----------|-------|
|
||||
| Config file | `amp-upstream-api-key` | High | No |
|
||||
| Environment | `AMP_API_KEY` | Medium | No |
|
||||
| Amp secrets file | `~/.local/share/amp/secrets.json` | Low | 5 min |
|
||||
|
||||
**Recommendation**: Use the Amp secrets file (lowest precedence) for normal usage. This file is automatically managed by `amp login`.
|
||||
|
||||
### Security Settings
|
||||
|
||||
**`amp-restrict-management-to-localhost`** (default: `true`)
|
||||
|
||||
When enabled, management routes (`/api/auth`, `/api/user`, `/api/threads`, etc.) only accept connections from localhost (127.0.0.1, ::1). This prevents:
|
||||
- Drive-by browser attacks
|
||||
- Remote access to management endpoints
|
||||
- CORS-based attacks
|
||||
- Header spoofing attacks (e.g., `X-Forwarded-For: 127.0.0.1`)
|
||||
|
||||
#### How It Works
|
||||
|
||||
This restriction uses the **actual TCP connection address** (`RemoteAddr`), not HTTP headers like `X-Forwarded-For`. This prevents header spoofing attacks but has important implications:
|
||||
|
||||
- ✅ **Works for direct connections**: Running CLIProxyAPI directly on your machine or server
|
||||
- ⚠️ **May not work behind reverse proxies**: If deploying behind nginx, Cloudflare, or other proxies, the connection will appear to come from the proxy's IP, not localhost
|
||||
|
||||
#### Reverse Proxy Deployments
|
||||
|
||||
If you need to run CLIProxyAPI behind a reverse proxy (nginx, Caddy, Cloudflare Tunnel, etc.):
|
||||
|
||||
1. **Disable the localhost restriction**:
|
||||
```yaml
|
||||
amp-restrict-management-to-localhost: false
|
||||
```
|
||||
|
||||
2. **Use alternative security measures**:
|
||||
- Firewall rules restricting access to management routes
|
||||
- Proxy-level authentication (HTTP Basic Auth, OAuth)
|
||||
- Network-level isolation (VPN, Tailscale, Cloudflare Access)
|
||||
- Bind CLIProxyAPI to `127.0.0.1` only and access via SSH tunnel
|
||||
|
||||
3. **Example nginx configuration** (blocks external access to management routes):
|
||||
```nginx
|
||||
location /api/auth { deny all; }
|
||||
location /api/user { deny all; }
|
||||
location /api/threads { deny all; }
|
||||
location /api/internal { deny all; }
|
||||
```
|
||||
|
||||
**Important**: Only disable `amp-restrict-management-to-localhost` if you understand the security implications and have other protections in place.
|
||||
|
||||
## Setup
|
||||
|
||||
### 1. Configure CLIProxyAPI
|
||||
|
||||
Create or edit `config.yaml`:
|
||||
|
||||
```yaml
|
||||
port: 8317
|
||||
auth-dir: "~/.cli-proxy-api"
|
||||
|
||||
# Amp integration
|
||||
amp-upstream-url: "https://ampcode.com"
|
||||
amp-restrict-management-to-localhost: true
|
||||
|
||||
# Other standard settings...
|
||||
debug: false
|
||||
logging-to-file: true
|
||||
```
|
||||
|
||||
### 2. Authenticate with Providers
|
||||
|
||||
Run OAuth login for the providers you want to use:
|
||||
|
||||
**Google Account (Gemini 2.5 Pro, Gemini 2.5 Flash, Gemini 3 Pro Preview):**
|
||||
```bash
|
||||
./cli-proxy-api --login
|
||||
```
|
||||
|
||||
**ChatGPT Plus/Pro (GPT-5, GPT-5 Codex):**
|
||||
```bash
|
||||
./cli-proxy-api --codex-login
|
||||
```
|
||||
|
||||
**Claude Pro/Max (Claude Sonnet 4.5, Opus 4.1):**
|
||||
```bash
|
||||
./cli-proxy-api --claude-login
|
||||
```
|
||||
|
||||
Tokens are saved to:
|
||||
- Gemini: `~/.cli-proxy-api/gemini-<email>.json`
|
||||
- OpenAI Codex: `~/.cli-proxy-api/codex-<email>.json`
|
||||
- Claude: `~/.cli-proxy-api/claude-<email>.json`
|
||||
|
||||
### 3. Start the Proxy
|
||||
|
||||
```bash
|
||||
./cli-proxy-api --config config.yaml
|
||||
```
|
||||
|
||||
Or run in background with tmux (recommended for remote servers):
|
||||
|
||||
```bash
|
||||
tmux new-session -d -s proxy "./cli-proxy-api --config config.yaml"
|
||||
```
|
||||
|
||||
### 4. Configure Amp CLI
|
||||
|
||||
#### Option A: Settings File
|
||||
|
||||
Edit `~/.config/amp/settings.json`:
|
||||
|
||||
```json
|
||||
{
|
||||
"amp.url": "http://localhost:8317"
|
||||
}
|
||||
```
|
||||
|
||||
#### Option B: Environment Variable
|
||||
|
||||
```bash
|
||||
export AMP_URL=http://localhost:8317
|
||||
```
|
||||
|
||||
### 5. Login and Use Amp
|
||||
|
||||
Login through the proxy (proxied to ampcode.com):
|
||||
|
||||
```bash
|
||||
amp login
|
||||
```
|
||||
|
||||
Use Amp as normal:
|
||||
|
||||
```bash
|
||||
amp "Write a hello world program in Python"
|
||||
```
|
||||
|
||||
### 6. (Optional) Configure Amp IDE Extension
|
||||
|
||||
The proxy also works with Amp IDE extensions for VS Code, Cursor, Windsurf, etc.
|
||||
|
||||
1. Open Amp extension settings in your IDE
|
||||
2. Set **Amp URL** to `http://localhost:8317`
|
||||
3. Login with your Amp account
|
||||
4. Start using Amp in your IDE
|
||||
|
||||
Both CLI and IDE can use the proxy simultaneously.
|
||||
|
||||
## Usage
|
||||
|
||||
### Supported Routes
|
||||
|
||||
#### Provider Aliases (Always Available)
|
||||
|
||||
These routes work even without `amp-upstream-url` configured:
|
||||
|
||||
- `/api/provider/openai/v1/chat/completions`
|
||||
- `/api/provider/openai/v1/responses`
|
||||
- `/api/provider/anthropic/v1/messages`
|
||||
- `/api/provider/google/v1beta/models/:action`
|
||||
|
||||
Amp CLI calls these routes with your OAuth-authenticated models configured in CLIProxyAPI.
|
||||
|
||||
#### Management Routes (Require `amp-upstream-url`)
|
||||
|
||||
These routes are proxied to ampcode.com:
|
||||
|
||||
- `/api/auth` - Authentication
|
||||
- `/api/user` - User profile
|
||||
- `/api/meta` - Metadata
|
||||
- `/api/threads` - Conversation threads
|
||||
- `/api/telemetry` - Usage telemetry
|
||||
- `/api/internal` - Internal APIs
|
||||
|
||||
**Security**: Restricted to localhost by default.
|
||||
|
||||
### Model Fallback Behavior
|
||||
|
||||
When Amp requests a model:
|
||||
|
||||
1. **Check local configuration**: Does CLIProxyAPI have OAuth tokens for this model's provider?
|
||||
2. **If YES**: Route to local handler (use your OAuth subscription)
|
||||
3. **If NO**: Forward to ampcode.com (use Amp's default routing)
|
||||
|
||||
This enables seamless mixed usage:
|
||||
- Models you've configured (Gemini, ChatGPT, Claude) → Your OAuth subscriptions
|
||||
- Models you haven't configured → Amp's default providers
|
||||
|
||||
### Example API Calls
|
||||
|
||||
**Chat completion with local OAuth:**
|
||||
```bash
|
||||
curl http://localhost:8317/api/provider/openai/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "gpt-5",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
}'
|
||||
```
|
||||
|
||||
**Management endpoint (localhost only):**
|
||||
```bash
|
||||
curl http://localhost:8317/api/user
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Common Issues
|
||||
|
||||
| Symptom | Likely Cause | Fix |
|
||||
|---------|--------------|-----|
|
||||
| 404 on `/api/provider/...` | Incorrect route path | Ensure exact path: `/api/provider/{provider}/v1...` |
|
||||
| 403 on `/api/user` | Non-localhost request | Run from same machine or disable `amp-restrict-management-to-localhost` (not recommended) |
|
||||
| 401/403 from provider | Missing/expired OAuth | Re-run `--codex-login` or `--claude-login` |
|
||||
| Amp gzip errors | Response decompression issue | Update to latest build; auto-decompression should handle this |
|
||||
| Models not using proxy | Wrong Amp URL | Verify `amp.url` setting or `AMP_URL` environment variable |
|
||||
| CORS errors | Protected management endpoint | Use CLI/terminal, not browser |
|
||||
|
||||
### Diagnostics
|
||||
|
||||
**Check proxy logs:**
|
||||
```bash
|
||||
# If logging-to-file: true
|
||||
tail -f logs/requests.log
|
||||
|
||||
# If running in tmux
|
||||
tmux attach-session -t proxy
|
||||
```
|
||||
|
||||
**Enable debug mode** (temporarily):
|
||||
```yaml
|
||||
debug: true
|
||||
```
|
||||
|
||||
**Test basic connectivity:**
|
||||
```bash
|
||||
# Check if proxy is running
|
||||
curl http://localhost:8317/v1/models
|
||||
|
||||
# Check Amp-specific route
|
||||
curl http://localhost:8317/api/provider/openai/v1/models
|
||||
```
|
||||
|
||||
**Verify Amp configuration:**
|
||||
```bash
|
||||
# Check if Amp is using proxy
|
||||
amp config get amp.url
|
||||
|
||||
# Or check environment
|
||||
echo $AMP_URL
|
||||
```
|
||||
|
||||
### Security Checklist
|
||||
|
||||
- ✅ Keep `amp-restrict-management-to-localhost: true` (default)
|
||||
- ✅ Don't expose proxy publicly (bind to localhost or use firewall/VPN)
|
||||
- ✅ Use the Amp secrets file (`~/.local/share/amp/secrets.json`) managed by `amp login`
|
||||
- ✅ Rotate OAuth tokens periodically by re-running login commands
|
||||
- ✅ Store config and auth-dir on encrypted disk if handling sensitive data
|
||||
- ✅ Keep proxy binary up to date for security fixes
|
||||
|
||||
## Additional Resources
|
||||
|
||||
- [CLIProxyAPI Main Documentation](https://help.router-for.me/)
|
||||
- [Amp CLI Official Manual](https://ampcode.com/manual)
|
||||
- [Management API Reference](https://help.router-for.me/management/api)
|
||||
- [SDK Documentation](sdk-usage.md)
|
||||
|
||||
## Disclaimer
|
||||
|
||||
This integration is for personal/educational use. Using reverse proxies or alternate API bases may violate provider Terms of Service. You are solely responsible for how you use this software. Accounts may be rate-limited, locked, or banned. No warranties. Use at your own risk.
|
||||
@@ -137,6 +137,10 @@ func (MyExecutor) Execute(ctx context.Context, a *coreauth.Auth, req clipexec.Re
|
||||
return clipexec.Response{Payload: body}, nil
|
||||
}
|
||||
|
||||
func (MyExecutor) CountTokens(context.Context, *coreauth.Auth, clipexec.Request, clipexec.Options) (clipexec.Response, error) {
|
||||
return clipexec.Response{}, errors.New("count tokens not implemented")
|
||||
}
|
||||
|
||||
func (MyExecutor) ExecuteStream(ctx context.Context, a *coreauth.Auth, req clipexec.Request, opts clipexec.Options) (<-chan clipexec.StreamChunk, error) {
|
||||
ch := make(chan clipexec.StreamChunk, 1)
|
||||
go func() {
|
||||
@@ -146,10 +150,6 @@ func (MyExecutor) ExecuteStream(ctx context.Context, a *coreauth.Auth, req clipe
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
func (MyExecutor) CountTokens(ctx context.Context, a *coreauth.Auth, req clipexec.Request, opts clipexec.Options) (clipexec.Response, error) {
|
||||
return clipexec.Response{}, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (MyExecutor) Refresh(ctx context.Context, a *coreauth.Auth) (*coreauth.Auth, error) {
|
||||
return a, nil
|
||||
}
|
||||
|
||||
185
internal/api/modules/amp/amp.go
Normal file
185
internal/api/modules/amp/amp.go
Normal file
@@ -0,0 +1,185 @@
|
||||
// Package amp implements the Amp CLI routing module, providing OAuth-based
|
||||
// integration with Amp CLI for ChatGPT and Anthropic subscriptions.
|
||||
package amp
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http/httputil"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// Option configures the AmpModule.
|
||||
type Option func(*AmpModule)
|
||||
|
||||
// AmpModule implements the RouteModuleV2 interface for Amp CLI integration.
|
||||
// It provides:
|
||||
// - Reverse proxy to Amp control plane for OAuth/management
|
||||
// - Provider-specific route aliases (/api/provider/{provider}/...)
|
||||
// - Automatic gzip decompression for misconfigured upstreams
|
||||
type AmpModule struct {
|
||||
secretSource SecretSource
|
||||
proxy *httputil.ReverseProxy
|
||||
accessManager *sdkaccess.Manager
|
||||
authMiddleware_ gin.HandlerFunc
|
||||
enabled bool
|
||||
registerOnce sync.Once
|
||||
}
|
||||
|
||||
// New creates a new Amp routing module with the given options.
|
||||
// This is the preferred constructor using the Option pattern.
|
||||
//
|
||||
// Example:
|
||||
//
|
||||
// ampModule := amp.New(
|
||||
// amp.WithAccessManager(accessManager),
|
||||
// amp.WithAuthMiddleware(authMiddleware),
|
||||
// amp.WithSecretSource(customSecret),
|
||||
// )
|
||||
func New(opts ...Option) *AmpModule {
|
||||
m := &AmpModule{
|
||||
secretSource: nil, // Will be created on demand if not provided
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(m)
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// NewLegacy creates a new Amp routing module using the legacy constructor signature.
|
||||
// This is provided for backwards compatibility.
|
||||
//
|
||||
// DEPRECATED: Use New with options instead.
|
||||
func NewLegacy(accessManager *sdkaccess.Manager, authMiddleware gin.HandlerFunc) *AmpModule {
|
||||
return New(
|
||||
WithAccessManager(accessManager),
|
||||
WithAuthMiddleware(authMiddleware),
|
||||
)
|
||||
}
|
||||
|
||||
// WithSecretSource sets a custom secret source for the module.
|
||||
func WithSecretSource(source SecretSource) Option {
|
||||
return func(m *AmpModule) {
|
||||
m.secretSource = source
|
||||
}
|
||||
}
|
||||
|
||||
// WithAccessManager sets the access manager for the module.
|
||||
func WithAccessManager(am *sdkaccess.Manager) Option {
|
||||
return func(m *AmpModule) {
|
||||
m.accessManager = am
|
||||
}
|
||||
}
|
||||
|
||||
// WithAuthMiddleware sets the authentication middleware for provider routes.
|
||||
func WithAuthMiddleware(middleware gin.HandlerFunc) Option {
|
||||
return func(m *AmpModule) {
|
||||
m.authMiddleware_ = middleware
|
||||
}
|
||||
}
|
||||
|
||||
// Name returns the module identifier
|
||||
func (m *AmpModule) Name() string {
|
||||
return "amp-routing"
|
||||
}
|
||||
|
||||
// Register sets up Amp routes if configured.
|
||||
// This implements the RouteModuleV2 interface with Context.
|
||||
// Routes are registered only once via sync.Once for idempotent behavior.
|
||||
func (m *AmpModule) Register(ctx modules.Context) error {
|
||||
upstreamURL := strings.TrimSpace(ctx.Config.AmpUpstreamURL)
|
||||
|
||||
// Determine auth middleware (from module or context)
|
||||
auth := m.getAuthMiddleware(ctx)
|
||||
|
||||
// Use registerOnce to ensure routes are only registered once
|
||||
var regErr error
|
||||
m.registerOnce.Do(func() {
|
||||
// Always register provider aliases - these work without an upstream
|
||||
m.registerProviderAliases(ctx.Engine, ctx.BaseHandler, auth)
|
||||
|
||||
// If no upstream URL, skip proxy routes but provider aliases are still available
|
||||
if upstreamURL == "" {
|
||||
log.Debug("Amp upstream proxy disabled (no upstream URL configured)")
|
||||
log.Debug("Amp provider alias routes registered")
|
||||
m.enabled = false
|
||||
return
|
||||
}
|
||||
|
||||
// Create secret source with precedence: config > env > file
|
||||
// Cache secrets for 5 minutes to reduce file I/O
|
||||
if m.secretSource == nil {
|
||||
m.secretSource = NewMultiSourceSecret(ctx.Config.AmpUpstreamAPIKey, 0 /* default 5min */)
|
||||
}
|
||||
|
||||
// Create reverse proxy with gzip handling via ModifyResponse
|
||||
proxy, err := createReverseProxy(upstreamURL, m.secretSource)
|
||||
if err != nil {
|
||||
regErr = fmt.Errorf("failed to create amp proxy: %w", err)
|
||||
return
|
||||
}
|
||||
|
||||
m.proxy = proxy
|
||||
m.enabled = true
|
||||
|
||||
// Register management proxy routes (requires upstream)
|
||||
// Restrict to localhost by default for security (prevents drive-by browser attacks)
|
||||
handler := proxyHandler(proxy)
|
||||
m.registerManagementRoutes(ctx.Engine, ctx.BaseHandler, handler, ctx.Config.AmpRestrictManagementToLocalhost)
|
||||
|
||||
log.Infof("Amp upstream proxy enabled for: %s", upstreamURL)
|
||||
log.Debug("Amp provider alias routes registered")
|
||||
})
|
||||
|
||||
return regErr
|
||||
}
|
||||
|
||||
// getAuthMiddleware returns the authentication middleware, preferring the
|
||||
// module's configured middleware, then the context middleware, then a fallback.
|
||||
func (m *AmpModule) getAuthMiddleware(ctx modules.Context) gin.HandlerFunc {
|
||||
if m.authMiddleware_ != nil {
|
||||
return m.authMiddleware_
|
||||
}
|
||||
if ctx.AuthMiddleware != nil {
|
||||
return ctx.AuthMiddleware
|
||||
}
|
||||
// Fallback: no authentication (should not happen in production)
|
||||
log.Warn("Amp module: no auth middleware provided, allowing all requests")
|
||||
return func(c *gin.Context) {
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// OnConfigUpdated handles configuration updates.
|
||||
// Currently requires restart for URL changes (could be enhanced for dynamic updates).
|
||||
func (m *AmpModule) OnConfigUpdated(cfg *config.Config) error {
|
||||
if !m.enabled {
|
||||
log.Debug("Amp routing not enabled, skipping config update")
|
||||
return nil
|
||||
}
|
||||
|
||||
upstreamURL := strings.TrimSpace(cfg.AmpUpstreamURL)
|
||||
if upstreamURL == "" {
|
||||
log.Warn("Amp upstream URL removed from config, restart required to disable")
|
||||
return nil
|
||||
}
|
||||
|
||||
// If API key changed, invalidate the cache
|
||||
if m.secretSource != nil {
|
||||
if ms, ok := m.secretSource.(*MultiSourceSecret); ok {
|
||||
ms.InvalidateCache()
|
||||
log.Debug("Amp secret cache invalidated due to config update")
|
||||
}
|
||||
}
|
||||
|
||||
log.Debug("Amp config updated (restart required for URL changes)")
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
303
internal/api/modules/amp/amp_test.go
Normal file
303
internal/api/modules/amp/amp_test.go
Normal file
@@ -0,0 +1,303 @@
|
||||
package amp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
|
||||
)
|
||||
|
||||
func TestAmpModule_Name(t *testing.T) {
|
||||
m := New()
|
||||
if m.Name() != "amp-routing" {
|
||||
t.Fatalf("want amp-routing, got %s", m.Name())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmpModule_New(t *testing.T) {
|
||||
accessManager := sdkaccess.NewManager()
|
||||
authMiddleware := func(c *gin.Context) { c.Next() }
|
||||
|
||||
m := NewLegacy(accessManager, authMiddleware)
|
||||
|
||||
if m.accessManager != accessManager {
|
||||
t.Fatal("accessManager not set")
|
||||
}
|
||||
if m.authMiddleware_ == nil {
|
||||
t.Fatal("authMiddleware not set")
|
||||
}
|
||||
if m.enabled {
|
||||
t.Fatal("enabled should be false initially")
|
||||
}
|
||||
if m.proxy != nil {
|
||||
t.Fatal("proxy should be nil initially")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmpModule_Register_WithUpstream(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
// Fake upstream to ensure URL is valid
|
||||
upstream := httptest.NewServer(nil)
|
||||
defer upstream.Close()
|
||||
|
||||
accessManager := sdkaccess.NewManager()
|
||||
base := &handlers.BaseAPIHandler{}
|
||||
|
||||
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
||||
|
||||
cfg := &config.Config{
|
||||
AmpUpstreamURL: upstream.URL,
|
||||
AmpUpstreamAPIKey: "test-key",
|
||||
}
|
||||
|
||||
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
||||
if err := m.Register(ctx); err != nil {
|
||||
t.Fatalf("register error: %v", err)
|
||||
}
|
||||
|
||||
if !m.enabled {
|
||||
t.Fatal("module should be enabled with upstream URL")
|
||||
}
|
||||
if m.proxy == nil {
|
||||
t.Fatal("proxy should be initialized")
|
||||
}
|
||||
if m.secretSource == nil {
|
||||
t.Fatal("secretSource should be initialized")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmpModule_Register_WithoutUpstream(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
accessManager := sdkaccess.NewManager()
|
||||
base := &handlers.BaseAPIHandler{}
|
||||
|
||||
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
||||
|
||||
cfg := &config.Config{
|
||||
AmpUpstreamURL: "", // No upstream
|
||||
}
|
||||
|
||||
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
||||
if err := m.Register(ctx); err != nil {
|
||||
t.Fatalf("register should not error without upstream: %v", err)
|
||||
}
|
||||
|
||||
if m.enabled {
|
||||
t.Fatal("module should be disabled without upstream URL")
|
||||
}
|
||||
if m.proxy != nil {
|
||||
t.Fatal("proxy should not be initialized without upstream")
|
||||
}
|
||||
|
||||
// But provider aliases should still be registered
|
||||
req := httptest.NewRequest("GET", "/api/provider/openai/models", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code == 404 {
|
||||
t.Fatal("provider aliases should be registered even without upstream")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmpModule_Register_InvalidUpstream(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
accessManager := sdkaccess.NewManager()
|
||||
base := &handlers.BaseAPIHandler{}
|
||||
|
||||
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
||||
|
||||
cfg := &config.Config{
|
||||
AmpUpstreamURL: "://invalid-url",
|
||||
}
|
||||
|
||||
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
||||
if err := m.Register(ctx); err == nil {
|
||||
t.Fatal("expected error for invalid upstream URL")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmpModule_OnConfigUpdated_CacheInvalidation(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
p := filepath.Join(tmpDir, "secrets.json")
|
||||
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":"v1"}`), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
m := &AmpModule{enabled: true}
|
||||
ms := NewMultiSourceSecretWithPath("", p, time.Minute)
|
||||
m.secretSource = ms
|
||||
|
||||
// Warm the cache
|
||||
if _, err := ms.Get(context.Background()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if ms.cache == nil {
|
||||
t.Fatal("expected cache to be set")
|
||||
}
|
||||
|
||||
// Update config - should invalidate cache
|
||||
if err := m.OnConfigUpdated(&config.Config{AmpUpstreamURL: "http://x"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if ms.cache != nil {
|
||||
t.Fatal("expected cache to be invalidated")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmpModule_OnConfigUpdated_NotEnabled(t *testing.T) {
|
||||
m := &AmpModule{enabled: false}
|
||||
|
||||
// Should not error or panic when disabled
|
||||
if err := m.OnConfigUpdated(&config.Config{}); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmpModule_OnConfigUpdated_URLRemoved(t *testing.T) {
|
||||
m := &AmpModule{enabled: true}
|
||||
ms := NewMultiSourceSecret("", 0)
|
||||
m.secretSource = ms
|
||||
|
||||
// Config update with empty URL - should log warning but not error
|
||||
cfg := &config.Config{AmpUpstreamURL: ""}
|
||||
|
||||
if err := m.OnConfigUpdated(cfg); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmpModule_OnConfigUpdated_NonMultiSourceSecret(t *testing.T) {
|
||||
// Test that OnConfigUpdated doesn't panic with StaticSecretSource
|
||||
m := &AmpModule{enabled: true}
|
||||
m.secretSource = NewStaticSecretSource("static-key")
|
||||
|
||||
cfg := &config.Config{AmpUpstreamURL: "http://example.com"}
|
||||
|
||||
// Should not error or panic
|
||||
if err := m.OnConfigUpdated(cfg); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmpModule_AuthMiddleware_Fallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
// Create module with no auth middleware
|
||||
m := &AmpModule{authMiddleware_: nil}
|
||||
|
||||
// Get the fallback middleware via getAuthMiddleware
|
||||
ctx := modules.Context{Engine: r, AuthMiddleware: nil}
|
||||
middleware := m.getAuthMiddleware(ctx)
|
||||
|
||||
if middleware == nil {
|
||||
t.Fatal("getAuthMiddleware should return a fallback, not nil")
|
||||
}
|
||||
|
||||
// Test that it works
|
||||
called := false
|
||||
r.GET("/test", middleware, func(c *gin.Context) {
|
||||
called = true
|
||||
c.String(200, "ok")
|
||||
})
|
||||
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if !called {
|
||||
t.Fatal("fallback middleware should allow requests through")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmpModule_SecretSource_FromConfig(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
upstream := httptest.NewServer(nil)
|
||||
defer upstream.Close()
|
||||
|
||||
accessManager := sdkaccess.NewManager()
|
||||
base := &handlers.BaseAPIHandler{}
|
||||
|
||||
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
||||
|
||||
// Config with explicit API key
|
||||
cfg := &config.Config{
|
||||
AmpUpstreamURL: upstream.URL,
|
||||
AmpUpstreamAPIKey: "config-key",
|
||||
}
|
||||
|
||||
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
||||
if err := m.Register(ctx); err != nil {
|
||||
t.Fatalf("register error: %v", err)
|
||||
}
|
||||
|
||||
// Secret source should be MultiSourceSecret with config key
|
||||
if m.secretSource == nil {
|
||||
t.Fatal("secretSource should be set")
|
||||
}
|
||||
|
||||
// Verify it returns the config key
|
||||
key, err := m.secretSource.Get(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Get error: %v", err)
|
||||
}
|
||||
if key != "config-key" {
|
||||
t.Fatalf("want config-key, got %s", key)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmpModule_ProviderAliasesAlwaysRegistered(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
scenarios := []struct {
|
||||
name string
|
||||
configURL string
|
||||
}{
|
||||
{"with_upstream", "http://example.com"},
|
||||
{"without_upstream", ""},
|
||||
}
|
||||
|
||||
for _, scenario := range scenarios {
|
||||
t.Run(scenario.name, func(t *testing.T) {
|
||||
r := gin.New()
|
||||
accessManager := sdkaccess.NewManager()
|
||||
base := &handlers.BaseAPIHandler{}
|
||||
|
||||
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
||||
|
||||
cfg := &config.Config{AmpUpstreamURL: scenario.configURL}
|
||||
|
||||
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
||||
if err := m.Register(ctx); err != nil && scenario.configURL != "" {
|
||||
t.Fatalf("register error: %v", err)
|
||||
}
|
||||
|
||||
// Provider aliases should always be available
|
||||
req := httptest.NewRequest("GET", "/api/provider/openai/models", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code == 404 {
|
||||
t.Fatal("provider aliases should be registered")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
129
internal/api/modules/amp/fallback_handlers.go
Normal file
129
internal/api/modules/amp/fallback_handlers.go
Normal file
@@ -0,0 +1,129 @@
|
||||
package amp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http/httputil"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// FallbackHandler wraps a standard handler with fallback logic to ampcode.com
|
||||
// when the model's provider is not available in CLIProxyAPI
|
||||
type FallbackHandler struct {
|
||||
getProxy func() *httputil.ReverseProxy
|
||||
}
|
||||
|
||||
// NewFallbackHandler creates a new fallback handler wrapper
|
||||
// The getProxy function allows lazy evaluation of the proxy (useful when proxy is created after routes)
|
||||
func NewFallbackHandler(getProxy func() *httputil.ReverseProxy) *FallbackHandler {
|
||||
return &FallbackHandler{
|
||||
getProxy: getProxy,
|
||||
}
|
||||
}
|
||||
|
||||
// WrapHandler wraps a gin.HandlerFunc with fallback logic
|
||||
// If the model's provider is not configured in CLIProxyAPI, it forwards to ampcode.com
|
||||
func (fh *FallbackHandler) WrapHandler(handler gin.HandlerFunc) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// Read the request body to extract the model name
|
||||
bodyBytes, err := io.ReadAll(c.Request.Body)
|
||||
if err != nil {
|
||||
log.Errorf("amp fallback: failed to read request body: %v", err)
|
||||
handler(c)
|
||||
return
|
||||
}
|
||||
|
||||
// Restore the body for the handler to read
|
||||
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
|
||||
|
||||
// Try to extract model from request body or URL path (for Gemini)
|
||||
modelName := extractModelFromRequest(bodyBytes, c)
|
||||
if modelName == "" {
|
||||
// Can't determine model, proceed with normal handler
|
||||
handler(c)
|
||||
return
|
||||
}
|
||||
|
||||
// Normalize model (handles Gemini thinking suffixes)
|
||||
normalizedModel, _ := util.NormalizeGeminiThinkingModel(modelName)
|
||||
|
||||
// Check if we have providers for this model
|
||||
providers := util.GetProviderName(normalizedModel)
|
||||
|
||||
if len(providers) == 0 {
|
||||
// No providers configured - check if we have a proxy for fallback
|
||||
proxy := fh.getProxy()
|
||||
if proxy != nil {
|
||||
// Fallback to ampcode.com
|
||||
log.Infof("amp fallback: model %s has no configured provider, forwarding to ampcode.com", modelName)
|
||||
|
||||
// Restore body again for the proxy
|
||||
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
|
||||
|
||||
// Forward to ampcode.com
|
||||
proxy.ServeHTTP(c.Writer, c.Request)
|
||||
return
|
||||
}
|
||||
|
||||
// No proxy available, let the normal handler return the error
|
||||
log.Debugf("amp fallback: model %s has no configured provider and no proxy available", modelName)
|
||||
}
|
||||
|
||||
// Providers available or no proxy for fallback, restore body and use normal handler
|
||||
// Filter Anthropic-Beta header to remove features requiring special subscription
|
||||
// This is needed when using local providers (bypassing the Amp proxy)
|
||||
if betaHeader := c.Request.Header.Get("Anthropic-Beta"); betaHeader != "" {
|
||||
filtered := filterBetaFeatures(betaHeader, "context-1m-2025-08-07")
|
||||
if filtered != "" {
|
||||
c.Request.Header.Set("Anthropic-Beta", filtered)
|
||||
} else {
|
||||
c.Request.Header.Del("Anthropic-Beta")
|
||||
}
|
||||
}
|
||||
|
||||
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
|
||||
handler(c)
|
||||
}
|
||||
}
|
||||
|
||||
// extractModelFromRequest attempts to extract the model name from various request formats
|
||||
func extractModelFromRequest(body []byte, c *gin.Context) string {
|
||||
// First try to parse from JSON body (OpenAI, Claude, etc.)
|
||||
var payload map[string]interface{}
|
||||
if err := json.Unmarshal(body, &payload); err == nil {
|
||||
// Check common model field names
|
||||
if model, ok := payload["model"].(string); ok {
|
||||
return model
|
||||
}
|
||||
}
|
||||
|
||||
// For Gemini requests, model is in the URL path
|
||||
// Standard format: /models/{model}:generateContent -> :action parameter
|
||||
if action := c.Param("action"); action != "" {
|
||||
// Split by colon to get model name (e.g., "gemini-pro:generateContent" -> "gemini-pro")
|
||||
parts := strings.Split(action, ":")
|
||||
if len(parts) > 0 && parts[0] != "" {
|
||||
return parts[0]
|
||||
}
|
||||
}
|
||||
|
||||
// AMP CLI format: /publishers/google/models/{model}:method -> *path parameter
|
||||
// Example: /publishers/google/models/gemini-3-pro-preview:streamGenerateContent
|
||||
if path := c.Param("path"); path != "" {
|
||||
// Look for /models/{model}:method pattern
|
||||
if idx := strings.Index(path, "/models/"); idx >= 0 {
|
||||
modelPart := path[idx+8:] // Skip "/models/"
|
||||
// Split by colon to get model name
|
||||
if colonIdx := strings.Index(modelPart, ":"); colonIdx > 0 {
|
||||
return modelPart[:colonIdx]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
45
internal/api/modules/amp/gemini_bridge.go
Normal file
45
internal/api/modules/amp/gemini_bridge.go
Normal file
@@ -0,0 +1,45 @@
|
||||
package amp
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/gemini"
|
||||
)
|
||||
|
||||
// createGeminiBridgeHandler creates a handler that bridges AMP CLI's non-standard Gemini paths
|
||||
// to our standard Gemini handler by rewriting the request context.
|
||||
//
|
||||
// AMP CLI format: /publishers/google/models/gemini-3-pro-preview:streamGenerateContent
|
||||
// Standard format: /models/gemini-3-pro-preview:streamGenerateContent
|
||||
//
|
||||
// This extracts the model+method from the AMP path and sets it as the :action parameter
|
||||
// so the standard Gemini handler can process it.
|
||||
func createGeminiBridgeHandler(geminiHandler *gemini.GeminiAPIHandler) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// Get the full path from the catch-all parameter
|
||||
path := c.Param("path")
|
||||
|
||||
// Extract model:method from AMP CLI path format
|
||||
// Example: /publishers/google/models/gemini-3-pro-preview:streamGenerateContent
|
||||
if idx := strings.Index(path, "/models/"); idx >= 0 {
|
||||
// Extract everything after "/models/"
|
||||
actionPart := path[idx+8:] // Skip "/models/"
|
||||
|
||||
// Set this as the :action parameter that the Gemini handler expects
|
||||
c.Params = append(c.Params, gin.Param{
|
||||
Key: "action",
|
||||
Value: actionPart,
|
||||
})
|
||||
|
||||
// Call the standard Gemini handler
|
||||
geminiHandler.GeminiHandler(c)
|
||||
return
|
||||
}
|
||||
|
||||
// If we can't parse the path, return 400
|
||||
c.JSON(400, gin.H{
|
||||
"error": "Invalid Gemini API path format",
|
||||
})
|
||||
}
|
||||
}
|
||||
195
internal/api/modules/amp/proxy.go
Normal file
195
internal/api/modules/amp/proxy.go
Normal file
@@ -0,0 +1,195 @@
|
||||
package amp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httputil"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// readCloser wraps a reader and forwards Close to a separate closer.
|
||||
// Used to restore peeked bytes while preserving upstream body Close behavior.
|
||||
type readCloser struct {
|
||||
r io.Reader
|
||||
c io.Closer
|
||||
}
|
||||
|
||||
func (rc *readCloser) Read(p []byte) (int, error) { return rc.r.Read(p) }
|
||||
func (rc *readCloser) Close() error { return rc.c.Close() }
|
||||
|
||||
// createReverseProxy creates a reverse proxy handler for Amp upstream
|
||||
// with automatic gzip decompression via ModifyResponse
|
||||
func createReverseProxy(upstreamURL string, secretSource SecretSource) (*httputil.ReverseProxy, error) {
|
||||
parsed, err := url.Parse(upstreamURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid amp upstream url: %w", err)
|
||||
}
|
||||
|
||||
proxy := httputil.NewSingleHostReverseProxy(parsed)
|
||||
originalDirector := proxy.Director
|
||||
|
||||
// Modify outgoing requests to inject API key and fix routing
|
||||
proxy.Director = func(req *http.Request) {
|
||||
originalDirector(req)
|
||||
req.Host = parsed.Host
|
||||
|
||||
// Preserve correlation headers for debugging
|
||||
if req.Header.Get("X-Request-ID") == "" {
|
||||
// Could generate one here if needed
|
||||
}
|
||||
|
||||
// Note: We do NOT filter Anthropic-Beta headers in the proxy path
|
||||
// Users going through ampcode.com proxy are paying for the service and should get all features
|
||||
// including 1M context window (context-1m-2025-08-07)
|
||||
|
||||
// Inject API key from secret source (precedence: config > env > file)
|
||||
if key, err := secretSource.Get(req.Context()); err == nil && key != "" {
|
||||
req.Header.Set("X-Api-Key", key)
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", key))
|
||||
} else if err != nil {
|
||||
log.Warnf("amp secret source error (continuing without auth): %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Modify incoming responses to handle gzip without Content-Encoding
|
||||
// This addresses the same issue as inline handler gzip handling, but at the proxy level
|
||||
proxy.ModifyResponse = func(resp *http.Response) error {
|
||||
// Only process successful responses
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Skip if already marked as gzip (Content-Encoding set)
|
||||
if resp.Header.Get("Content-Encoding") != "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Skip streaming responses (SSE, chunked)
|
||||
if isStreamingResponse(resp) {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Save reference to original upstream body for proper cleanup
|
||||
originalBody := resp.Body
|
||||
|
||||
// Peek at first 2 bytes to detect gzip magic bytes
|
||||
header := make([]byte, 2)
|
||||
n, _ := io.ReadFull(originalBody, header)
|
||||
|
||||
// Check for gzip magic bytes (0x1f 0x8b)
|
||||
// If n < 2, we didn't get enough bytes, so it's not gzip
|
||||
if n >= 2 && header[0] == 0x1f && header[1] == 0x8b {
|
||||
// It's gzip - read the rest of the body
|
||||
rest, err := io.ReadAll(originalBody)
|
||||
if err != nil {
|
||||
// Restore what we read and return original body (preserve Close behavior)
|
||||
resp.Body = &readCloser{
|
||||
r: io.MultiReader(bytes.NewReader(header[:n]), originalBody),
|
||||
c: originalBody,
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Reconstruct complete gzipped data
|
||||
gzippedData := append(header[:n], rest...)
|
||||
|
||||
// Decompress
|
||||
gzipReader, err := gzip.NewReader(bytes.NewReader(gzippedData))
|
||||
if err != nil {
|
||||
log.Warnf("amp proxy: gzip header detected but decompress failed: %v", err)
|
||||
// Close original body and return in-memory copy
|
||||
_ = originalBody.Close()
|
||||
resp.Body = io.NopCloser(bytes.NewReader(gzippedData))
|
||||
return nil
|
||||
}
|
||||
|
||||
decompressed, err := io.ReadAll(gzipReader)
|
||||
_ = gzipReader.Close()
|
||||
if err != nil {
|
||||
log.Warnf("amp proxy: gzip decompress error: %v", err)
|
||||
// Close original body and return in-memory copy
|
||||
_ = originalBody.Close()
|
||||
resp.Body = io.NopCloser(bytes.NewReader(gzippedData))
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close original body since we're replacing with in-memory decompressed content
|
||||
_ = originalBody.Close()
|
||||
|
||||
// Replace body with decompressed content
|
||||
resp.Body = io.NopCloser(bytes.NewReader(decompressed))
|
||||
resp.ContentLength = int64(len(decompressed))
|
||||
|
||||
// Update headers to reflect decompressed state
|
||||
resp.Header.Del("Content-Encoding") // No longer compressed
|
||||
resp.Header.Del("Content-Length") // Remove stale compressed length
|
||||
resp.Header.Set("Content-Length", strconv.FormatInt(resp.ContentLength, 10)) // Set decompressed length
|
||||
|
||||
log.Debugf("amp proxy: decompressed gzip response (%d -> %d bytes)", len(gzippedData), len(decompressed))
|
||||
} else {
|
||||
// Not gzip - restore peeked bytes while preserving Close behavior
|
||||
// Handle edge cases: n might be 0, 1, or 2 depending on EOF
|
||||
resp.Body = &readCloser{
|
||||
r: io.MultiReader(bytes.NewReader(header[:n]), originalBody),
|
||||
c: originalBody,
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Error handler for proxy failures
|
||||
proxy.ErrorHandler = func(rw http.ResponseWriter, req *http.Request, err error) {
|
||||
log.Errorf("amp upstream proxy error for %s %s: %v", req.Method, req.URL.Path, err)
|
||||
rw.Header().Set("Content-Type", "application/json")
|
||||
rw.WriteHeader(http.StatusBadGateway)
|
||||
_, _ = rw.Write([]byte(`{"error":"amp_upstream_proxy_error","message":"Failed to reach Amp upstream"}`))
|
||||
}
|
||||
|
||||
return proxy, nil
|
||||
}
|
||||
|
||||
// isStreamingResponse detects if the response is streaming (SSE only)
|
||||
// Note: We only treat text/event-stream as streaming. Chunked transfer encoding
|
||||
// is a transport-level detail and doesn't mean we can't decompress the full response.
|
||||
// Many JSON APIs use chunked encoding for normal responses.
|
||||
func isStreamingResponse(resp *http.Response) bool {
|
||||
contentType := resp.Header.Get("Content-Type")
|
||||
|
||||
// Only Server-Sent Events are true streaming responses
|
||||
if strings.Contains(contentType, "text/event-stream") {
|
||||
return true
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
// proxyHandler converts httputil.ReverseProxy to gin.HandlerFunc
|
||||
func proxyHandler(proxy *httputil.ReverseProxy) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
proxy.ServeHTTP(c.Writer, c.Request)
|
||||
}
|
||||
}
|
||||
|
||||
// filterBetaFeatures removes a specific beta feature from comma-separated list
|
||||
func filterBetaFeatures(header, featureToRemove string) string {
|
||||
features := strings.Split(header, ",")
|
||||
filtered := make([]string, 0, len(features))
|
||||
|
||||
for _, feature := range features {
|
||||
trimmed := strings.TrimSpace(feature)
|
||||
if trimmed != "" && trimmed != featureToRemove {
|
||||
filtered = append(filtered, trimmed)
|
||||
}
|
||||
}
|
||||
|
||||
return strings.Join(filtered, ",")
|
||||
}
|
||||
500
internal/api/modules/amp/proxy_test.go
Normal file
500
internal/api/modules/amp/proxy_test.go
Normal file
@@ -0,0 +1,500 @@
|
||||
package amp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// Helper: compress data with gzip
|
||||
func gzipBytes(b []byte) []byte {
|
||||
var buf bytes.Buffer
|
||||
zw := gzip.NewWriter(&buf)
|
||||
zw.Write(b)
|
||||
zw.Close()
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
// Helper: create a mock http.Response
|
||||
func mkResp(status int, hdr http.Header, body []byte) *http.Response {
|
||||
if hdr == nil {
|
||||
hdr = http.Header{}
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: status,
|
||||
Header: hdr,
|
||||
Body: io.NopCloser(bytes.NewReader(body)),
|
||||
ContentLength: int64(len(body)),
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateReverseProxy_ValidURL(t *testing.T) {
|
||||
proxy, err := createReverseProxy("http://example.com", NewStaticSecretSource("key"))
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error, got: %v", err)
|
||||
}
|
||||
if proxy == nil {
|
||||
t.Fatal("expected proxy to be created")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateReverseProxy_InvalidURL(t *testing.T) {
|
||||
_, err := createReverseProxy("://invalid", NewStaticSecretSource("key"))
|
||||
if err == nil {
|
||||
t.Fatal("expected error for invalid URL")
|
||||
}
|
||||
}
|
||||
|
||||
func TestModifyResponse_GzipScenarios(t *testing.T) {
|
||||
proxy, err := createReverseProxy("http://example.com", NewStaticSecretSource("k"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
goodJSON := []byte(`{"ok":true}`)
|
||||
good := gzipBytes(goodJSON)
|
||||
truncated := good[:10]
|
||||
corrupted := append([]byte{0x1f, 0x8b}, []byte("notgzip")...)
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
header http.Header
|
||||
body []byte
|
||||
status int
|
||||
wantBody []byte
|
||||
wantCE string
|
||||
}{
|
||||
{
|
||||
name: "decompresses_valid_gzip_no_header",
|
||||
header: http.Header{},
|
||||
body: good,
|
||||
status: 200,
|
||||
wantBody: goodJSON,
|
||||
wantCE: "",
|
||||
},
|
||||
{
|
||||
name: "skips_when_ce_present",
|
||||
header: http.Header{"Content-Encoding": []string{"gzip"}},
|
||||
body: good,
|
||||
status: 200,
|
||||
wantBody: good,
|
||||
wantCE: "gzip",
|
||||
},
|
||||
{
|
||||
name: "passes_truncated_unchanged",
|
||||
header: http.Header{},
|
||||
body: truncated,
|
||||
status: 200,
|
||||
wantBody: truncated,
|
||||
wantCE: "",
|
||||
},
|
||||
{
|
||||
name: "passes_corrupted_unchanged",
|
||||
header: http.Header{},
|
||||
body: corrupted,
|
||||
status: 200,
|
||||
wantBody: corrupted,
|
||||
wantCE: "",
|
||||
},
|
||||
{
|
||||
name: "non_gzip_unchanged",
|
||||
header: http.Header{},
|
||||
body: []byte("plain"),
|
||||
status: 200,
|
||||
wantBody: []byte("plain"),
|
||||
wantCE: "",
|
||||
},
|
||||
{
|
||||
name: "empty_body",
|
||||
header: http.Header{},
|
||||
body: []byte{},
|
||||
status: 200,
|
||||
wantBody: []byte{},
|
||||
wantCE: "",
|
||||
},
|
||||
{
|
||||
name: "single_byte_body",
|
||||
header: http.Header{},
|
||||
body: []byte{0x1f},
|
||||
status: 200,
|
||||
wantBody: []byte{0x1f},
|
||||
wantCE: "",
|
||||
},
|
||||
{
|
||||
name: "skips_non_2xx_status",
|
||||
header: http.Header{},
|
||||
body: good,
|
||||
status: 404,
|
||||
wantBody: good,
|
||||
wantCE: "",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
resp := mkResp(tc.status, tc.header, tc.body)
|
||||
if err := proxy.ModifyResponse(resp); err != nil {
|
||||
t.Fatalf("ModifyResponse error: %v", err)
|
||||
}
|
||||
got, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
t.Fatalf("ReadAll error: %v", err)
|
||||
}
|
||||
if !bytes.Equal(got, tc.wantBody) {
|
||||
t.Fatalf("body mismatch:\nwant: %q\ngot: %q", tc.wantBody, got)
|
||||
}
|
||||
if ce := resp.Header.Get("Content-Encoding"); ce != tc.wantCE {
|
||||
t.Fatalf("Content-Encoding: want %q, got %q", tc.wantCE, ce)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestModifyResponse_UpdatesContentLengthHeader(t *testing.T) {
|
||||
proxy, err := createReverseProxy("http://example.com", NewStaticSecretSource("k"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
goodJSON := []byte(`{"message":"test response"}`)
|
||||
gzipped := gzipBytes(goodJSON)
|
||||
|
||||
// Simulate upstream response with gzip body AND Content-Length header
|
||||
// (this is the scenario the bot flagged - stale Content-Length after decompression)
|
||||
resp := mkResp(200, http.Header{
|
||||
"Content-Length": []string{fmt.Sprintf("%d", len(gzipped))}, // Compressed size
|
||||
}, gzipped)
|
||||
|
||||
if err := proxy.ModifyResponse(resp); err != nil {
|
||||
t.Fatalf("ModifyResponse error: %v", err)
|
||||
}
|
||||
|
||||
// Verify body is decompressed
|
||||
got, _ := io.ReadAll(resp.Body)
|
||||
if !bytes.Equal(got, goodJSON) {
|
||||
t.Fatalf("body should be decompressed, got: %q, want: %q", got, goodJSON)
|
||||
}
|
||||
|
||||
// Verify Content-Length header is updated to decompressed size
|
||||
wantCL := fmt.Sprintf("%d", len(goodJSON))
|
||||
gotCL := resp.Header.Get("Content-Length")
|
||||
if gotCL != wantCL {
|
||||
t.Fatalf("Content-Length header mismatch: want %q (decompressed), got %q", wantCL, gotCL)
|
||||
}
|
||||
|
||||
// Verify struct field also matches
|
||||
if resp.ContentLength != int64(len(goodJSON)) {
|
||||
t.Fatalf("resp.ContentLength mismatch: want %d, got %d", len(goodJSON), resp.ContentLength)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModifyResponse_SkipsStreamingResponses(t *testing.T) {
|
||||
proxy, err := createReverseProxy("http://example.com", NewStaticSecretSource("k"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
goodJSON := []byte(`{"ok":true}`)
|
||||
gzipped := gzipBytes(goodJSON)
|
||||
|
||||
t.Run("sse_skips_decompression", func(t *testing.T) {
|
||||
resp := mkResp(200, http.Header{"Content-Type": []string{"text/event-stream"}}, gzipped)
|
||||
if err := proxy.ModifyResponse(resp); err != nil {
|
||||
t.Fatalf("ModifyResponse error: %v", err)
|
||||
}
|
||||
// SSE should NOT be decompressed
|
||||
got, _ := io.ReadAll(resp.Body)
|
||||
if !bytes.Equal(got, gzipped) {
|
||||
t.Fatal("SSE response should not be decompressed")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestModifyResponse_DecompressesChunkedJSON(t *testing.T) {
|
||||
proxy, err := createReverseProxy("http://example.com", NewStaticSecretSource("k"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
goodJSON := []byte(`{"ok":true}`)
|
||||
gzipped := gzipBytes(goodJSON)
|
||||
|
||||
t.Run("chunked_json_decompresses", func(t *testing.T) {
|
||||
// Chunked JSON responses (like thread APIs) should be decompressed
|
||||
resp := mkResp(200, http.Header{"Transfer-Encoding": []string{"chunked"}}, gzipped)
|
||||
if err := proxy.ModifyResponse(resp); err != nil {
|
||||
t.Fatalf("ModifyResponse error: %v", err)
|
||||
}
|
||||
// Should decompress because it's not SSE
|
||||
got, _ := io.ReadAll(resp.Body)
|
||||
if !bytes.Equal(got, goodJSON) {
|
||||
t.Fatalf("chunked JSON should be decompressed, got: %q, want: %q", got, goodJSON)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestReverseProxy_InjectsHeaders(t *testing.T) {
|
||||
gotHeaders := make(chan http.Header, 1)
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotHeaders <- r.Header.Clone()
|
||||
w.WriteHeader(200)
|
||||
w.Write([]byte(`ok`))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
proxy, err := createReverseProxy(upstream.URL, NewStaticSecretSource("secret"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
proxy.ServeHTTP(w, r)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
res, err := http.Get(srv.URL + "/test")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
res.Body.Close()
|
||||
|
||||
hdr := <-gotHeaders
|
||||
if hdr.Get("X-Api-Key") != "secret" {
|
||||
t.Fatalf("X-Api-Key missing or wrong, got: %q", hdr.Get("X-Api-Key"))
|
||||
}
|
||||
if hdr.Get("Authorization") != "Bearer secret" {
|
||||
t.Fatalf("Authorization missing or wrong, got: %q", hdr.Get("Authorization"))
|
||||
}
|
||||
}
|
||||
|
||||
func TestReverseProxy_EmptySecret(t *testing.T) {
|
||||
gotHeaders := make(chan http.Header, 1)
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotHeaders <- r.Header.Clone()
|
||||
w.WriteHeader(200)
|
||||
w.Write([]byte(`ok`))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
proxy, err := createReverseProxy(upstream.URL, NewStaticSecretSource(""))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
proxy.ServeHTTP(w, r)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
res, err := http.Get(srv.URL + "/test")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
res.Body.Close()
|
||||
|
||||
hdr := <-gotHeaders
|
||||
// Should NOT inject headers when secret is empty
|
||||
if hdr.Get("X-Api-Key") != "" {
|
||||
t.Fatalf("X-Api-Key should not be set, got: %q", hdr.Get("X-Api-Key"))
|
||||
}
|
||||
if authVal := hdr.Get("Authorization"); authVal != "" && authVal != "Bearer " {
|
||||
t.Fatalf("Authorization should not be set, got: %q", authVal)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReverseProxy_ErrorHandler(t *testing.T) {
|
||||
// Point proxy to a non-routable address to trigger error
|
||||
proxy, err := createReverseProxy("http://127.0.0.1:1", NewStaticSecretSource(""))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
proxy.ServeHTTP(w, r)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
res, err := http.Get(srv.URL + "/any")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
body, _ := io.ReadAll(res.Body)
|
||||
res.Body.Close()
|
||||
|
||||
if res.StatusCode != http.StatusBadGateway {
|
||||
t.Fatalf("want 502, got %d", res.StatusCode)
|
||||
}
|
||||
if !bytes.Contains(body, []byte(`"amp_upstream_proxy_error"`)) {
|
||||
t.Fatalf("unexpected body: %s", body)
|
||||
}
|
||||
if ct := res.Header.Get("Content-Type"); ct != "application/json" {
|
||||
t.Fatalf("content-type: want application/json, got %s", ct)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReverseProxy_FullRoundTrip_Gzip(t *testing.T) {
|
||||
// Upstream returns gzipped JSON without Content-Encoding header
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(200)
|
||||
w.Write(gzipBytes([]byte(`{"upstream":"ok"}`)))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
proxy, err := createReverseProxy(upstream.URL, NewStaticSecretSource("key"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
proxy.ServeHTTP(w, r)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
res, err := http.Get(srv.URL + "/test")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
body, _ := io.ReadAll(res.Body)
|
||||
res.Body.Close()
|
||||
|
||||
expected := []byte(`{"upstream":"ok"}`)
|
||||
if !bytes.Equal(body, expected) {
|
||||
t.Fatalf("want decompressed JSON, got: %s", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReverseProxy_FullRoundTrip_PlainJSON(t *testing.T) {
|
||||
// Upstream returns plain JSON
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(200)
|
||||
w.Write([]byte(`{"plain":"json"}`))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
proxy, err := createReverseProxy(upstream.URL, NewStaticSecretSource("key"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
proxy.ServeHTTP(w, r)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
res, err := http.Get(srv.URL + "/test")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
body, _ := io.ReadAll(res.Body)
|
||||
res.Body.Close()
|
||||
|
||||
expected := []byte(`{"plain":"json"}`)
|
||||
if !bytes.Equal(body, expected) {
|
||||
t.Fatalf("want plain JSON unchanged, got: %s", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsStreamingResponse(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
header http.Header
|
||||
want bool
|
||||
}{
|
||||
{
|
||||
name: "sse",
|
||||
header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "chunked_not_streaming",
|
||||
header: http.Header{"Transfer-Encoding": []string{"chunked"}},
|
||||
want: false, // Chunked is transport-level, not streaming
|
||||
},
|
||||
{
|
||||
name: "normal_json",
|
||||
header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "empty",
|
||||
header: http.Header{},
|
||||
want: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
resp := &http.Response{Header: tc.header}
|
||||
got := isStreamingResponse(resp)
|
||||
if got != tc.want {
|
||||
t.Fatalf("want %v, got %v", tc.want, got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFilterBetaFeatures(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
header string
|
||||
featureToRemove string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "Remove context-1m from middle",
|
||||
header: "fine-grained-tool-streaming-2025-05-14,context-1m-2025-08-07,oauth-2025-04-20",
|
||||
featureToRemove: "context-1m-2025-08-07",
|
||||
expected: "fine-grained-tool-streaming-2025-05-14,oauth-2025-04-20",
|
||||
},
|
||||
{
|
||||
name: "Remove context-1m from start",
|
||||
header: "context-1m-2025-08-07,fine-grained-tool-streaming-2025-05-14",
|
||||
featureToRemove: "context-1m-2025-08-07",
|
||||
expected: "fine-grained-tool-streaming-2025-05-14",
|
||||
},
|
||||
{
|
||||
name: "Remove context-1m from end",
|
||||
header: "fine-grained-tool-streaming-2025-05-14,context-1m-2025-08-07",
|
||||
featureToRemove: "context-1m-2025-08-07",
|
||||
expected: "fine-grained-tool-streaming-2025-05-14",
|
||||
},
|
||||
{
|
||||
name: "Feature not present",
|
||||
header: "fine-grained-tool-streaming-2025-05-14,oauth-2025-04-20",
|
||||
featureToRemove: "context-1m-2025-08-07",
|
||||
expected: "fine-grained-tool-streaming-2025-05-14,oauth-2025-04-20",
|
||||
},
|
||||
{
|
||||
name: "Only feature to remove",
|
||||
header: "context-1m-2025-08-07",
|
||||
featureToRemove: "context-1m-2025-08-07",
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "Empty header",
|
||||
header: "",
|
||||
featureToRemove: "context-1m-2025-08-07",
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "Header with spaces",
|
||||
header: "fine-grained-tool-streaming-2025-05-14, context-1m-2025-08-07 , oauth-2025-04-20",
|
||||
featureToRemove: "context-1m-2025-08-07",
|
||||
expected: "fine-grained-tool-streaming-2025-05-14,oauth-2025-04-20",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := filterBetaFeatures(tt.header, tt.featureToRemove)
|
||||
if result != tt.expected {
|
||||
t.Errorf("filterBetaFeatures() = %q, want %q", result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
222
internal/api/modules/amp/routes.go
Normal file
222
internal/api/modules/amp/routes.go
Normal file
@@ -0,0 +1,222 @@
|
||||
package amp
|
||||
|
||||
import (
|
||||
"net"
|
||||
"net/http/httputil"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/claude"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/gemini"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/openai"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// localhostOnlyMiddleware restricts access to localhost (127.0.0.1, ::1) only.
|
||||
// Returns 403 Forbidden for non-localhost clients.
|
||||
//
|
||||
// Security: Uses RemoteAddr (actual TCP connection) instead of ClientIP() to prevent
|
||||
// header spoofing attacks via X-Forwarded-For or similar headers. This means the
|
||||
// middleware will not work correctly behind reverse proxies - users deploying behind
|
||||
// nginx/Cloudflare should disable this feature and use firewall rules instead.
|
||||
func localhostOnlyMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// Use actual TCP connection address (RemoteAddr) to prevent header spoofing
|
||||
// This cannot be forged by X-Forwarded-For or other client-controlled headers
|
||||
remoteAddr := c.Request.RemoteAddr
|
||||
|
||||
// RemoteAddr format is "IP:port" or "[IPv6]:port", extract just the IP
|
||||
host, _, err := net.SplitHostPort(remoteAddr)
|
||||
if err != nil {
|
||||
// Try parsing as raw IP (shouldn't happen with standard HTTP, but be defensive)
|
||||
host = remoteAddr
|
||||
}
|
||||
|
||||
// Parse the IP to handle both IPv4 and IPv6
|
||||
ip := net.ParseIP(host)
|
||||
if ip == nil {
|
||||
log.Warnf("Amp management: invalid RemoteAddr %s, denying access", remoteAddr)
|
||||
c.AbortWithStatusJSON(403, gin.H{
|
||||
"error": "Access denied: management routes restricted to localhost",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Check if IP is loopback (127.0.0.1 or ::1)
|
||||
if !ip.IsLoopback() {
|
||||
log.Warnf("Amp management: non-localhost connection from %s attempted access, denying", remoteAddr)
|
||||
c.AbortWithStatusJSON(403, gin.H{
|
||||
"error": "Access denied: management routes restricted to localhost",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// noCORSMiddleware disables CORS for management routes to prevent browser-based attacks.
|
||||
// This overwrites any global CORS headers set by the server.
|
||||
func noCORSMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// Remove CORS headers to prevent cross-origin access from browsers
|
||||
c.Header("Access-Control-Allow-Origin", "")
|
||||
c.Header("Access-Control-Allow-Methods", "")
|
||||
c.Header("Access-Control-Allow-Headers", "")
|
||||
c.Header("Access-Control-Allow-Credentials", "")
|
||||
|
||||
// For OPTIONS preflight, deny with 403
|
||||
if c.Request.Method == "OPTIONS" {
|
||||
c.AbortWithStatus(403)
|
||||
return
|
||||
}
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// registerManagementRoutes registers Amp management proxy routes
|
||||
// These routes proxy through to the Amp control plane for OAuth, user management, etc.
|
||||
// If restrictToLocalhost is true, routes will only accept connections from 127.0.0.1/::1.
|
||||
func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *handlers.BaseAPIHandler, proxyHandler gin.HandlerFunc, restrictToLocalhost bool) {
|
||||
ampAPI := engine.Group("/api")
|
||||
|
||||
// Always disable CORS for management routes to prevent browser-based attacks
|
||||
ampAPI.Use(noCORSMiddleware())
|
||||
|
||||
// Apply localhost-only restriction if configured
|
||||
if restrictToLocalhost {
|
||||
ampAPI.Use(localhostOnlyMiddleware())
|
||||
log.Info("Amp management routes restricted to localhost only (CORS disabled)")
|
||||
} else {
|
||||
log.Warn("⚠️ Amp management routes are NOT restricted to localhost - this is insecure!")
|
||||
}
|
||||
|
||||
// Management routes - these are proxied directly to Amp upstream
|
||||
ampAPI.Any("/internal", proxyHandler)
|
||||
ampAPI.Any("/internal/*path", proxyHandler)
|
||||
ampAPI.Any("/user", proxyHandler)
|
||||
ampAPI.Any("/user/*path", proxyHandler)
|
||||
ampAPI.Any("/auth", proxyHandler)
|
||||
ampAPI.Any("/auth/*path", proxyHandler)
|
||||
ampAPI.Any("/meta", proxyHandler)
|
||||
ampAPI.Any("/meta/*path", proxyHandler)
|
||||
ampAPI.Any("/ads", proxyHandler)
|
||||
ampAPI.Any("/telemetry", proxyHandler)
|
||||
ampAPI.Any("/telemetry/*path", proxyHandler)
|
||||
ampAPI.Any("/threads", proxyHandler)
|
||||
ampAPI.Any("/threads/*path", proxyHandler)
|
||||
ampAPI.Any("/otel", proxyHandler)
|
||||
ampAPI.Any("/otel/*path", proxyHandler)
|
||||
|
||||
// Google v1beta1 passthrough with OAuth fallback
|
||||
// AMP CLI uses non-standard paths like /publishers/google/models/...
|
||||
// We bridge these to our standard Gemini handler to enable local OAuth.
|
||||
// If no local OAuth is available, falls back to ampcode.com proxy.
|
||||
geminiHandlers := gemini.NewGeminiAPIHandler(baseHandler)
|
||||
geminiBridge := createGeminiBridgeHandler(geminiHandlers)
|
||||
geminiV1Beta1Fallback := NewFallbackHandler(func() *httputil.ReverseProxy {
|
||||
return m.proxy
|
||||
})
|
||||
geminiV1Beta1Handler := geminiV1Beta1Fallback.WrapHandler(geminiBridge)
|
||||
|
||||
// Route POST model calls through Gemini bridge when a local provider exists, otherwise proxy.
|
||||
// All other methods (e.g., GET model listing) always proxy to upstream to preserve Amp CLI behavior.
|
||||
ampAPI.Any("/provider/google/v1beta1/*path", func(c *gin.Context) {
|
||||
if c.Request.Method == "POST" {
|
||||
// Attempt to extract the model name from the AMP-style path
|
||||
if path := c.Param("path"); strings.Contains(path, "/models/") {
|
||||
modelPart := path[strings.Index(path, "/models/")+len("/models/"):]
|
||||
if colonIdx := strings.Index(modelPart, ":"); colonIdx > 0 {
|
||||
modelPart = modelPart[:colonIdx]
|
||||
}
|
||||
if modelPart != "" {
|
||||
normalized, _ := util.NormalizeGeminiThinkingModel(modelPart)
|
||||
// Only handle locally when we have a provider; otherwise fall back to proxy
|
||||
if providers := util.GetProviderName(normalized); len(providers) > 0 {
|
||||
geminiV1Beta1Handler(c)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Non-POST or no local provider available -> proxy upstream
|
||||
proxyHandler(c)
|
||||
})
|
||||
}
|
||||
|
||||
// registerProviderAliases registers /api/provider/{provider}/... routes
|
||||
// These allow Amp CLI to route requests like:
|
||||
//
|
||||
// /api/provider/openai/v1/chat/completions
|
||||
// /api/provider/anthropic/v1/messages
|
||||
// /api/provider/google/v1beta/models
|
||||
func (m *AmpModule) registerProviderAliases(engine *gin.Engine, baseHandler *handlers.BaseAPIHandler, auth gin.HandlerFunc) {
|
||||
// Create handler instances for different providers
|
||||
openaiHandlers := openai.NewOpenAIAPIHandler(baseHandler)
|
||||
geminiHandlers := gemini.NewGeminiAPIHandler(baseHandler)
|
||||
claudeCodeHandlers := claude.NewClaudeCodeAPIHandler(baseHandler)
|
||||
openaiResponsesHandlers := openai.NewOpenAIResponsesAPIHandler(baseHandler)
|
||||
|
||||
// Create fallback handler wrapper that forwards to ampcode.com when provider not found
|
||||
// Uses lazy evaluation to access proxy (which is created after routes are registered)
|
||||
fallbackHandler := NewFallbackHandler(func() *httputil.ReverseProxy {
|
||||
return m.proxy
|
||||
})
|
||||
|
||||
// Provider-specific routes under /api/provider/:provider
|
||||
ampProviders := engine.Group("/api/provider")
|
||||
if auth != nil {
|
||||
ampProviders.Use(auth)
|
||||
}
|
||||
|
||||
provider := ampProviders.Group("/:provider")
|
||||
|
||||
// Dynamic models handler - routes to appropriate provider based on path parameter
|
||||
ampModelsHandler := func(c *gin.Context) {
|
||||
providerName := strings.ToLower(c.Param("provider"))
|
||||
|
||||
switch providerName {
|
||||
case "anthropic":
|
||||
claudeCodeHandlers.ClaudeModels(c)
|
||||
case "google":
|
||||
geminiHandlers.GeminiModels(c)
|
||||
default:
|
||||
// Default to OpenAI-compatible (works for openai, groq, cerebras, etc.)
|
||||
openaiHandlers.OpenAIModels(c)
|
||||
}
|
||||
}
|
||||
|
||||
// Root-level routes (for providers that omit /v1, like groq/cerebras)
|
||||
// Wrap handlers with fallback logic to forward to ampcode.com when provider not found
|
||||
provider.GET("/models", ampModelsHandler) // Models endpoint doesn't need fallback (no body to check)
|
||||
provider.POST("/chat/completions", fallbackHandler.WrapHandler(openaiHandlers.ChatCompletions))
|
||||
provider.POST("/completions", fallbackHandler.WrapHandler(openaiHandlers.Completions))
|
||||
provider.POST("/responses", fallbackHandler.WrapHandler(openaiResponsesHandlers.Responses))
|
||||
|
||||
// /v1 routes (OpenAI/Claude-compatible endpoints)
|
||||
v1Amp := provider.Group("/v1")
|
||||
{
|
||||
v1Amp.GET("/models", ampModelsHandler) // Models endpoint doesn't need fallback
|
||||
|
||||
// OpenAI-compatible endpoints with fallback
|
||||
v1Amp.POST("/chat/completions", fallbackHandler.WrapHandler(openaiHandlers.ChatCompletions))
|
||||
v1Amp.POST("/completions", fallbackHandler.WrapHandler(openaiHandlers.Completions))
|
||||
v1Amp.POST("/responses", fallbackHandler.WrapHandler(openaiResponsesHandlers.Responses))
|
||||
|
||||
// Claude/Anthropic-compatible endpoints with fallback
|
||||
v1Amp.POST("/messages", fallbackHandler.WrapHandler(claudeCodeHandlers.ClaudeMessages))
|
||||
v1Amp.POST("/messages/count_tokens", fallbackHandler.WrapHandler(claudeCodeHandlers.ClaudeCountTokens))
|
||||
}
|
||||
|
||||
// /v1beta routes (Gemini native API)
|
||||
// Note: Gemini handler extracts model from URL path, so fallback logic needs special handling
|
||||
v1betaAmp := provider.Group("/v1beta")
|
||||
{
|
||||
v1betaAmp.GET("/models", geminiHandlers.GeminiModels)
|
||||
v1betaAmp.POST("/models/:action", fallbackHandler.WrapHandler(geminiHandlers.GeminiHandler))
|
||||
v1betaAmp.GET("/models/:action", geminiHandlers.GeminiGetHandler)
|
||||
}
|
||||
}
|
||||
301
internal/api/modules/amp/routes_test.go
Normal file
301
internal/api/modules/amp/routes_test.go
Normal file
@@ -0,0 +1,301 @@
|
||||
package amp
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
|
||||
)
|
||||
|
||||
func TestRegisterManagementRoutes(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
// Spy to track if proxy handler was called
|
||||
proxyCalled := false
|
||||
proxyHandler := func(c *gin.Context) {
|
||||
proxyCalled = true
|
||||
c.String(200, "proxied")
|
||||
}
|
||||
|
||||
m := &AmpModule{}
|
||||
base := &handlers.BaseAPIHandler{}
|
||||
m.registerManagementRoutes(r, base, proxyHandler, false) // false = don't restrict to localhost in tests
|
||||
|
||||
managementPaths := []struct {
|
||||
path string
|
||||
method string
|
||||
}{
|
||||
{"/api/internal", http.MethodGet},
|
||||
{"/api/internal/some/path", http.MethodGet},
|
||||
{"/api/user", http.MethodGet},
|
||||
{"/api/user/profile", http.MethodGet},
|
||||
{"/api/auth", http.MethodGet},
|
||||
{"/api/auth/login", http.MethodGet},
|
||||
{"/api/meta", http.MethodGet},
|
||||
{"/api/telemetry", http.MethodGet},
|
||||
{"/api/threads", http.MethodGet},
|
||||
{"/api/otel", http.MethodGet},
|
||||
// Google v1beta1 bridge should still proxy non-model requests (GET) and allow POST
|
||||
{"/api/provider/google/v1beta1/models", http.MethodGet},
|
||||
{"/api/provider/google/v1beta1/models", http.MethodPost},
|
||||
}
|
||||
|
||||
for _, path := range managementPaths {
|
||||
t.Run(path.path, func(t *testing.T) {
|
||||
proxyCalled = false
|
||||
req := httptest.NewRequest(path.method, path.path, nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code == http.StatusNotFound {
|
||||
t.Fatalf("route %s not registered", path.path)
|
||||
}
|
||||
if !proxyCalled {
|
||||
t.Fatalf("proxy handler not called for %s", path.path)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterProviderAliases_AllProvidersRegistered(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
// Minimal base handler setup (no need to initialize, just check routing)
|
||||
base := &handlers.BaseAPIHandler{}
|
||||
|
||||
// Track if auth middleware was called
|
||||
authCalled := false
|
||||
authMiddleware := func(c *gin.Context) {
|
||||
authCalled = true
|
||||
c.Header("X-Auth", "ok")
|
||||
// Abort with success to avoid calling the actual handler (which needs full setup)
|
||||
c.AbortWithStatus(http.StatusOK)
|
||||
}
|
||||
|
||||
m := &AmpModule{authMiddleware_: authMiddleware}
|
||||
m.registerProviderAliases(r, base, authMiddleware)
|
||||
|
||||
paths := []struct {
|
||||
path string
|
||||
method string
|
||||
}{
|
||||
{"/api/provider/openai/models", http.MethodGet},
|
||||
{"/api/provider/anthropic/models", http.MethodGet},
|
||||
{"/api/provider/google/models", http.MethodGet},
|
||||
{"/api/provider/groq/models", http.MethodGet},
|
||||
{"/api/provider/openai/chat/completions", http.MethodPost},
|
||||
{"/api/provider/anthropic/v1/messages", http.MethodPost},
|
||||
{"/api/provider/google/v1beta/models", http.MethodGet},
|
||||
}
|
||||
|
||||
for _, tc := range paths {
|
||||
t.Run(tc.path, func(t *testing.T) {
|
||||
authCalled = false
|
||||
req := httptest.NewRequest(tc.method, tc.path, nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code == http.StatusNotFound {
|
||||
t.Fatalf("route %s %s not registered", tc.method, tc.path)
|
||||
}
|
||||
if !authCalled {
|
||||
t.Fatalf("auth middleware not executed for %s", tc.path)
|
||||
}
|
||||
if w.Header().Get("X-Auth") != "ok" {
|
||||
t.Fatalf("auth middleware header not set for %s", tc.path)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterProviderAliases_DynamicModelsHandler(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
base := &handlers.BaseAPIHandler{}
|
||||
|
||||
m := &AmpModule{authMiddleware_: func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) }}
|
||||
m.registerProviderAliases(r, base, func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) })
|
||||
|
||||
providers := []string{"openai", "anthropic", "google", "groq", "cerebras"}
|
||||
|
||||
for _, provider := range providers {
|
||||
t.Run(provider, func(t *testing.T) {
|
||||
path := "/api/provider/" + provider + "/models"
|
||||
req := httptest.NewRequest(http.MethodGet, path, nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
// Should not 404
|
||||
if w.Code == http.StatusNotFound {
|
||||
t.Fatalf("models route not found for provider: %s", provider)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterProviderAliases_V1Routes(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
base := &handlers.BaseAPIHandler{}
|
||||
|
||||
m := &AmpModule{authMiddleware_: func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) }}
|
||||
m.registerProviderAliases(r, base, func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) })
|
||||
|
||||
v1Paths := []struct {
|
||||
path string
|
||||
method string
|
||||
}{
|
||||
{"/api/provider/openai/v1/models", http.MethodGet},
|
||||
{"/api/provider/openai/v1/chat/completions", http.MethodPost},
|
||||
{"/api/provider/openai/v1/completions", http.MethodPost},
|
||||
{"/api/provider/anthropic/v1/messages", http.MethodPost},
|
||||
{"/api/provider/anthropic/v1/messages/count_tokens", http.MethodPost},
|
||||
}
|
||||
|
||||
for _, tc := range v1Paths {
|
||||
t.Run(tc.path, func(t *testing.T) {
|
||||
req := httptest.NewRequest(tc.method, tc.path, nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code == http.StatusNotFound {
|
||||
t.Fatalf("v1 route %s %s not registered", tc.method, tc.path)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterProviderAliases_V1BetaRoutes(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
base := &handlers.BaseAPIHandler{}
|
||||
|
||||
m := &AmpModule{authMiddleware_: func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) }}
|
||||
m.registerProviderAliases(r, base, func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) })
|
||||
|
||||
v1betaPaths := []struct {
|
||||
path string
|
||||
method string
|
||||
}{
|
||||
{"/api/provider/google/v1beta/models", http.MethodGet},
|
||||
{"/api/provider/google/v1beta/models/generateContent", http.MethodPost},
|
||||
}
|
||||
|
||||
for _, tc := range v1betaPaths {
|
||||
t.Run(tc.path, func(t *testing.T) {
|
||||
req := httptest.NewRequest(tc.method, tc.path, nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code == http.StatusNotFound {
|
||||
t.Fatalf("v1beta route %s %s not registered", tc.method, tc.path)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterProviderAliases_NoAuthMiddleware(t *testing.T) {
|
||||
// Test that routes still register even if auth middleware is nil (fallback behavior)
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
base := &handlers.BaseAPIHandler{}
|
||||
|
||||
m := &AmpModule{authMiddleware_: nil} // No auth middleware
|
||||
m.registerProviderAliases(r, base, func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) })
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/provider/openai/models", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
// Should still work (with fallback no-op auth)
|
||||
if w.Code == http.StatusNotFound {
|
||||
t.Fatal("routes should register even without auth middleware")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalhostOnlyMiddleware_PreventsSpoofing(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
// Apply localhost-only middleware
|
||||
r.Use(localhostOnlyMiddleware())
|
||||
r.GET("/test", func(c *gin.Context) {
|
||||
c.String(http.StatusOK, "ok")
|
||||
})
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
remoteAddr string
|
||||
forwardedFor string
|
||||
expectedStatus int
|
||||
description string
|
||||
}{
|
||||
{
|
||||
name: "spoofed_header_remote_connection",
|
||||
remoteAddr: "192.168.1.100:12345",
|
||||
forwardedFor: "127.0.0.1",
|
||||
expectedStatus: http.StatusForbidden,
|
||||
description: "Spoofed X-Forwarded-For header should be ignored",
|
||||
},
|
||||
{
|
||||
name: "real_localhost_ipv4",
|
||||
remoteAddr: "127.0.0.1:54321",
|
||||
forwardedFor: "",
|
||||
expectedStatus: http.StatusOK,
|
||||
description: "Real localhost IPv4 connection should work",
|
||||
},
|
||||
{
|
||||
name: "real_localhost_ipv6",
|
||||
remoteAddr: "[::1]:54321",
|
||||
forwardedFor: "",
|
||||
expectedStatus: http.StatusOK,
|
||||
description: "Real localhost IPv6 connection should work",
|
||||
},
|
||||
{
|
||||
name: "remote_ipv4",
|
||||
remoteAddr: "203.0.113.42:8080",
|
||||
forwardedFor: "",
|
||||
expectedStatus: http.StatusForbidden,
|
||||
description: "Remote IPv4 connection should be blocked",
|
||||
},
|
||||
{
|
||||
name: "remote_ipv6",
|
||||
remoteAddr: "[2001:db8::1]:9090",
|
||||
forwardedFor: "",
|
||||
expectedStatus: http.StatusForbidden,
|
||||
description: "Remote IPv6 connection should be blocked",
|
||||
},
|
||||
{
|
||||
name: "spoofed_localhost_ipv6",
|
||||
remoteAddr: "203.0.113.42:8080",
|
||||
forwardedFor: "::1",
|
||||
expectedStatus: http.StatusForbidden,
|
||||
description: "Spoofed X-Forwarded-For with IPv6 localhost should be ignored",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.RemoteAddr = tt.remoteAddr
|
||||
if tt.forwardedFor != "" {
|
||||
req.Header.Set("X-Forwarded-For", tt.forwardedFor)
|
||||
}
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != tt.expectedStatus {
|
||||
t.Errorf("%s: expected status %d, got %d", tt.description, tt.expectedStatus, w.Code)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
155
internal/api/modules/amp/secret.go
Normal file
155
internal/api/modules/amp/secret.go
Normal file
@@ -0,0 +1,155 @@
|
||||
package amp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SecretSource provides Amp API keys with configurable precedence and caching
|
||||
type SecretSource interface {
|
||||
Get(ctx context.Context) (string, error)
|
||||
}
|
||||
|
||||
// cachedSecret holds a secret value with expiration
|
||||
type cachedSecret struct {
|
||||
value string
|
||||
expiresAt time.Time
|
||||
}
|
||||
|
||||
// MultiSourceSecret implements precedence-based secret lookup:
|
||||
// 1. Explicit config value (highest priority)
|
||||
// 2. Environment variable AMP_API_KEY
|
||||
// 3. File-based secret (lowest priority)
|
||||
type MultiSourceSecret struct {
|
||||
explicitKey string
|
||||
envKey string
|
||||
filePath string
|
||||
cacheTTL time.Duration
|
||||
|
||||
mu sync.RWMutex
|
||||
cache *cachedSecret
|
||||
}
|
||||
|
||||
// NewMultiSourceSecret creates a secret source with precedence and caching
|
||||
func NewMultiSourceSecret(explicitKey string, cacheTTL time.Duration) *MultiSourceSecret {
|
||||
if cacheTTL == 0 {
|
||||
cacheTTL = 5 * time.Minute // Default 5 minute cache
|
||||
}
|
||||
|
||||
home, _ := os.UserHomeDir()
|
||||
filePath := filepath.Join(home, ".local", "share", "amp", "secrets.json")
|
||||
|
||||
return &MultiSourceSecret{
|
||||
explicitKey: strings.TrimSpace(explicitKey),
|
||||
envKey: "AMP_API_KEY",
|
||||
filePath: filePath,
|
||||
cacheTTL: cacheTTL,
|
||||
}
|
||||
}
|
||||
|
||||
// NewMultiSourceSecretWithPath creates a secret source with a custom file path (for testing)
|
||||
func NewMultiSourceSecretWithPath(explicitKey string, filePath string, cacheTTL time.Duration) *MultiSourceSecret {
|
||||
if cacheTTL == 0 {
|
||||
cacheTTL = 5 * time.Minute
|
||||
}
|
||||
|
||||
return &MultiSourceSecret{
|
||||
explicitKey: strings.TrimSpace(explicitKey),
|
||||
envKey: "AMP_API_KEY",
|
||||
filePath: filePath,
|
||||
cacheTTL: cacheTTL,
|
||||
}
|
||||
}
|
||||
|
||||
// Get retrieves the Amp API key using precedence: config > env > file
|
||||
// Results are cached for cacheTTL duration to avoid excessive file reads
|
||||
func (s *MultiSourceSecret) Get(ctx context.Context) (string, error) {
|
||||
// Precedence 1: Explicit config key (highest priority, no caching needed)
|
||||
if s.explicitKey != "" {
|
||||
return s.explicitKey, nil
|
||||
}
|
||||
|
||||
// Precedence 2: Environment variable
|
||||
if envValue := strings.TrimSpace(os.Getenv(s.envKey)); envValue != "" {
|
||||
return envValue, nil
|
||||
}
|
||||
|
||||
// Precedence 3: File-based secret (lowest priority, cached)
|
||||
// Check cache first
|
||||
s.mu.RLock()
|
||||
if s.cache != nil && time.Now().Before(s.cache.expiresAt) {
|
||||
value := s.cache.value
|
||||
s.mu.RUnlock()
|
||||
return value, nil
|
||||
}
|
||||
s.mu.RUnlock()
|
||||
|
||||
// Cache miss or expired - read from file
|
||||
key, err := s.readFromFile()
|
||||
if err != nil {
|
||||
// Cache empty result to avoid repeated file reads on missing files
|
||||
s.updateCache("")
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Cache the result
|
||||
s.updateCache(key)
|
||||
return key, nil
|
||||
}
|
||||
|
||||
// readFromFile reads the Amp API key from the secrets file
|
||||
func (s *MultiSourceSecret) readFromFile() (string, error) {
|
||||
content, err := os.ReadFile(s.filePath)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return "", nil // Missing file is not an error, just no key available
|
||||
}
|
||||
return "", fmt.Errorf("failed to read amp secrets from %s: %w", s.filePath, err)
|
||||
}
|
||||
|
||||
var secrets map[string]string
|
||||
if err := json.Unmarshal(content, &secrets); err != nil {
|
||||
return "", fmt.Errorf("failed to parse amp secrets from %s: %w", s.filePath, err)
|
||||
}
|
||||
|
||||
key := strings.TrimSpace(secrets["apiKey@https://ampcode.com/"])
|
||||
return key, nil
|
||||
}
|
||||
|
||||
// updateCache updates the cached secret value
|
||||
func (s *MultiSourceSecret) updateCache(value string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.cache = &cachedSecret{
|
||||
value: value,
|
||||
expiresAt: time.Now().Add(s.cacheTTL),
|
||||
}
|
||||
}
|
||||
|
||||
// InvalidateCache clears the cached secret, forcing a fresh read on next Get
|
||||
func (s *MultiSourceSecret) InvalidateCache() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.cache = nil
|
||||
}
|
||||
|
||||
// StaticSecretSource returns a fixed API key (for testing)
|
||||
type StaticSecretSource struct {
|
||||
key string
|
||||
}
|
||||
|
||||
// NewStaticSecretSource creates a secret source with a fixed key
|
||||
func NewStaticSecretSource(key string) *StaticSecretSource {
|
||||
return &StaticSecretSource{key: strings.TrimSpace(key)}
|
||||
}
|
||||
|
||||
// Get returns the static API key
|
||||
func (s *StaticSecretSource) Get(ctx context.Context) (string, error) {
|
||||
return s.key, nil
|
||||
}
|
||||
280
internal/api/modules/amp/secret_test.go
Normal file
280
internal/api/modules/amp/secret_test.go
Normal file
@@ -0,0 +1,280 @@
|
||||
package amp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestMultiSourceSecret_PrecedenceOrder(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
configKey string
|
||||
envKey string
|
||||
fileJSON string
|
||||
want string
|
||||
}{
|
||||
{"config_wins", "cfg", "env", `{"apiKey@https://ampcode.com/":"file"}`, "cfg"},
|
||||
{"env_wins_when_no_cfg", "", "env", `{"apiKey@https://ampcode.com/":"file"}`, "env"},
|
||||
{"file_when_no_cfg_env", "", "", `{"apiKey@https://ampcode.com/":"file"}`, "file"},
|
||||
{"empty_cfg_trims_then_env", " ", "env", `{"apiKey@https://ampcode.com/":"file"}`, "env"},
|
||||
{"empty_env_then_file", "", " ", `{"apiKey@https://ampcode.com/":"file"}`, "file"},
|
||||
{"missing_file_returns_empty", "", "", "", ""},
|
||||
{"all_empty_returns_empty", " ", " ", `{"apiKey@https://ampcode.com/":" "}`, ""},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
tc := tc // capture range variable
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
secretsPath := filepath.Join(tmpDir, "secrets.json")
|
||||
|
||||
if tc.fileJSON != "" {
|
||||
if err := os.WriteFile(secretsPath, []byte(tc.fileJSON), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
t.Setenv("AMP_API_KEY", tc.envKey)
|
||||
|
||||
s := NewMultiSourceSecretWithPath(tc.configKey, secretsPath, 100*time.Millisecond)
|
||||
got, err := s.Get(ctx)
|
||||
if err != nil && tc.fileJSON != "" && json.Valid([]byte(tc.fileJSON)) {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got != tc.want {
|
||||
t.Fatalf("want %q, got %q", tc.want, got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultiSourceSecret_CacheBehavior(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tmpDir := t.TempDir()
|
||||
p := filepath.Join(tmpDir, "secrets.json")
|
||||
|
||||
// Initial value
|
||||
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":"v1"}`), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
s := NewMultiSourceSecretWithPath("", p, 50*time.Millisecond)
|
||||
|
||||
// First read - should return v1
|
||||
got1, err := s.Get(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("Get failed: %v", err)
|
||||
}
|
||||
if got1 != "v1" {
|
||||
t.Fatalf("expected v1, got %s", got1)
|
||||
}
|
||||
|
||||
// Change file; within TTL we should still see v1 (cached)
|
||||
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":"v2"}`), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got2, _ := s.Get(ctx)
|
||||
if got2 != "v1" {
|
||||
t.Fatalf("cache hit expected v1, got %s", got2)
|
||||
}
|
||||
|
||||
// After TTL expires, should see v2
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
got3, _ := s.Get(ctx)
|
||||
if got3 != "v2" {
|
||||
t.Fatalf("cache miss expected v2, got %s", got3)
|
||||
}
|
||||
|
||||
// Invalidate forces re-read immediately
|
||||
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":"v3"}`), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.InvalidateCache()
|
||||
got4, _ := s.Get(ctx)
|
||||
if got4 != "v3" {
|
||||
t.Fatalf("invalidate expected v3, got %s", got4)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultiSourceSecret_FileHandling(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
t.Run("missing_file_no_error", func(t *testing.T) {
|
||||
s := NewMultiSourceSecretWithPath("", "/nonexistent/path/secrets.json", 100*time.Millisecond)
|
||||
got, err := s.Get(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error for missing file, got: %v", err)
|
||||
}
|
||||
if got != "" {
|
||||
t.Fatalf("expected empty string, got %q", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid_json", func(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
p := filepath.Join(tmpDir, "secrets.json")
|
||||
if err := os.WriteFile(p, []byte(`{invalid json`), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
s := NewMultiSourceSecretWithPath("", p, 100*time.Millisecond)
|
||||
_, err := s.Get(ctx)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for invalid JSON")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("missing_key_in_json", func(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
p := filepath.Join(tmpDir, "secrets.json")
|
||||
if err := os.WriteFile(p, []byte(`{"other":"value"}`), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
s := NewMultiSourceSecretWithPath("", p, 100*time.Millisecond)
|
||||
got, err := s.Get(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got != "" {
|
||||
t.Fatalf("expected empty string for missing key, got %q", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty_key_value", func(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
p := filepath.Join(tmpDir, "secrets.json")
|
||||
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":" "}`), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
s := NewMultiSourceSecretWithPath("", p, 100*time.Millisecond)
|
||||
got, _ := s.Get(ctx)
|
||||
if got != "" {
|
||||
t.Fatalf("expected empty after trim, got %q", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestMultiSourceSecret_Concurrency(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
p := filepath.Join(tmpDir, "secrets.json")
|
||||
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":"concurrent"}`), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
s := NewMultiSourceSecretWithPath("", p, 5*time.Second)
|
||||
ctx := context.Background()
|
||||
|
||||
// Spawn many goroutines calling Get concurrently
|
||||
const goroutines = 50
|
||||
const iterations = 100
|
||||
|
||||
var wg sync.WaitGroup
|
||||
errors := make(chan error, goroutines)
|
||||
|
||||
for i := 0; i < goroutines; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < iterations; j++ {
|
||||
val, err := s.Get(ctx)
|
||||
if err != nil {
|
||||
errors <- err
|
||||
return
|
||||
}
|
||||
if val != "concurrent" {
|
||||
errors <- err
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
close(errors)
|
||||
|
||||
for err := range errors {
|
||||
t.Errorf("concurrency error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStaticSecretSource(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
t.Run("returns_provided_key", func(t *testing.T) {
|
||||
s := NewStaticSecretSource("test-key-123")
|
||||
got, err := s.Get(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got != "test-key-123" {
|
||||
t.Fatalf("want test-key-123, got %q", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("trims_whitespace", func(t *testing.T) {
|
||||
s := NewStaticSecretSource(" test-key ")
|
||||
got, err := s.Get(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got != "test-key" {
|
||||
t.Fatalf("want test-key, got %q", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty_string", func(t *testing.T) {
|
||||
s := NewStaticSecretSource("")
|
||||
got, err := s.Get(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got != "" {
|
||||
t.Fatalf("want empty string, got %q", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestMultiSourceSecret_CacheEmptyResult(t *testing.T) {
|
||||
// Test that missing file results are cached to avoid repeated file reads
|
||||
tmpDir := t.TempDir()
|
||||
p := filepath.Join(tmpDir, "nonexistent.json")
|
||||
|
||||
s := NewMultiSourceSecretWithPath("", p, 100*time.Millisecond)
|
||||
ctx := context.Background()
|
||||
|
||||
// First call - file doesn't exist, should cache empty result
|
||||
got1, err := s.Get(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error for missing file, got: %v", err)
|
||||
}
|
||||
if got1 != "" {
|
||||
t.Fatalf("expected empty string, got %q", got1)
|
||||
}
|
||||
|
||||
// Create the file now
|
||||
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":"new-value"}`), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Second call - should still return empty (cached), not read the new file
|
||||
got2, _ := s.Get(ctx)
|
||||
if got2 != "" {
|
||||
t.Fatalf("cache should return empty, got %q", got2)
|
||||
}
|
||||
|
||||
// After TTL expires, should see the new value
|
||||
time.Sleep(110 * time.Millisecond)
|
||||
got3, _ := s.Get(ctx)
|
||||
if got3 != "new-value" {
|
||||
t.Fatalf("after cache expiry, expected new-value, got %q", got3)
|
||||
}
|
||||
}
|
||||
92
internal/api/modules/modules.go
Normal file
92
internal/api/modules/modules.go
Normal file
@@ -0,0 +1,92 @@
|
||||
// Package modules provides a pluggable routing module system for extending
|
||||
// the API server with optional features without modifying core routing logic.
|
||||
package modules
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
|
||||
)
|
||||
|
||||
// Context encapsulates the dependencies exposed to routing modules during
|
||||
// registration. Modules can use the Gin engine to attach routes, the shared
|
||||
// BaseAPIHandler for constructing SDK-specific handlers, and the resolved
|
||||
// authentication middleware for protecting routes that require API keys.
|
||||
type Context struct {
|
||||
Engine *gin.Engine
|
||||
BaseHandler *handlers.BaseAPIHandler
|
||||
Config *config.Config
|
||||
AuthMiddleware gin.HandlerFunc
|
||||
}
|
||||
|
||||
// RouteModule represents a pluggable routing module that can register routes
|
||||
// and handle configuration updates independently of the core server.
|
||||
//
|
||||
// DEPRECATED: Use RouteModuleV2 for new modules. This interface is kept for
|
||||
// backwards compatibility and will be removed in a future version.
|
||||
type RouteModule interface {
|
||||
// Name returns a human-readable identifier for the module
|
||||
Name() string
|
||||
|
||||
// Register sets up routes and handlers for this module.
|
||||
// It receives the Gin engine, base handlers, and current configuration.
|
||||
// Returns an error if registration fails (errors are logged but don't stop the server).
|
||||
Register(engine *gin.Engine, baseHandler *handlers.BaseAPIHandler, cfg *config.Config) error
|
||||
|
||||
// OnConfigUpdated is called when the configuration is reloaded.
|
||||
// Modules can respond to configuration changes here.
|
||||
// Returns an error if the update cannot be applied.
|
||||
OnConfigUpdated(cfg *config.Config) error
|
||||
}
|
||||
|
||||
// RouteModuleV2 represents a pluggable bundle of routes that can integrate with
|
||||
// the API server without modifying its core routing logic. Implementations can
|
||||
// attach routes during Register and react to configuration updates via
|
||||
// OnConfigUpdated.
|
||||
//
|
||||
// This is the preferred interface for new modules. It uses Context for cleaner
|
||||
// dependency injection and supports idempotent registration.
|
||||
type RouteModuleV2 interface {
|
||||
// Name returns a unique identifier for logging and diagnostics.
|
||||
Name() string
|
||||
|
||||
// Register wires the module's routes into the provided Gin engine. Modules
|
||||
// should treat multiple calls as idempotent and avoid duplicate route
|
||||
// registration when invoked more than once.
|
||||
Register(ctx Context) error
|
||||
|
||||
// OnConfigUpdated notifies the module when the server configuration changes
|
||||
// via hot reload. Implementations can refresh cached state or emit warnings.
|
||||
OnConfigUpdated(cfg *config.Config) error
|
||||
}
|
||||
|
||||
// RegisterModule is a helper that registers a module using either the V1 or V2
|
||||
// interface. This allows gradual migration from V1 to V2 without breaking
|
||||
// existing modules.
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// ctx := modules.Context{
|
||||
// Engine: engine,
|
||||
// BaseHandler: baseHandler,
|
||||
// Config: cfg,
|
||||
// AuthMiddleware: authMiddleware,
|
||||
// }
|
||||
// if err := modules.RegisterModule(ctx, ampModule); err != nil {
|
||||
// log.Errorf("Failed to register module: %v", err)
|
||||
// }
|
||||
func RegisterModule(ctx Context, mod interface{}) error {
|
||||
// Try V2 interface first (preferred)
|
||||
if v2, ok := mod.(RouteModuleV2); ok {
|
||||
return v2.Register(ctx)
|
||||
}
|
||||
|
||||
// Fall back to V1 interface for backwards compatibility
|
||||
if v1, ok := mod.(RouteModule); ok {
|
||||
return v1.Register(ctx.Engine, ctx.BaseHandler, ctx.Config)
|
||||
}
|
||||
|
||||
return fmt.Errorf("unsupported module type %T (must implement RouteModule or RouteModuleV2)", mod)
|
||||
}
|
||||
@@ -21,6 +21,8 @@ import (
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/access"
|
||||
managementHandlers "github.com/router-for-me/CLIProxyAPI/v6/internal/api/handlers/management"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/middleware"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules"
|
||||
ampmodule "github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules/amp"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset"
|
||||
@@ -261,6 +263,20 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
|
||||
|
||||
// Setup routes
|
||||
s.setupRoutes()
|
||||
|
||||
// Register Amp module using V2 interface with Context
|
||||
ampModule := ampmodule.NewLegacy(accessManager, AuthMiddleware(accessManager))
|
||||
ctx := modules.Context{
|
||||
Engine: engine,
|
||||
BaseHandler: s.handlers,
|
||||
Config: cfg,
|
||||
AuthMiddleware: AuthMiddleware(accessManager),
|
||||
}
|
||||
if err := modules.RegisterModule(ctx, ampModule); err != nil {
|
||||
log.Errorf("Failed to register Amp module: %v", err)
|
||||
}
|
||||
|
||||
// Apply additional router configurators from options
|
||||
if optionState.routerConfigurator != nil {
|
||||
optionState.routerConfigurator(engine, s.handlers, cfg)
|
||||
}
|
||||
|
||||
111
internal/api/server_test.go
Normal file
111
internal/api/server_test.go
Normal file
@@ -0,0 +1,111 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
gin "github.com/gin-gonic/gin"
|
||||
proxyconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
|
||||
)
|
||||
|
||||
func newTestServer(t *testing.T) *Server {
|
||||
t.Helper()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tmpDir := t.TempDir()
|
||||
authDir := filepath.Join(tmpDir, "auth")
|
||||
if err := os.MkdirAll(authDir, 0o700); err != nil {
|
||||
t.Fatalf("failed to create auth dir: %v", err)
|
||||
}
|
||||
|
||||
cfg := &proxyconfig.Config{
|
||||
SDKConfig: sdkconfig.SDKConfig{
|
||||
APIKeys: []string{"test-key"},
|
||||
},
|
||||
Port: 0,
|
||||
AuthDir: authDir,
|
||||
Debug: true,
|
||||
LoggingToFile: false,
|
||||
UsageStatisticsEnabled: false,
|
||||
}
|
||||
|
||||
authManager := auth.NewManager(nil, nil, nil)
|
||||
accessManager := sdkaccess.NewManager()
|
||||
|
||||
configPath := filepath.Join(tmpDir, "config.yaml")
|
||||
return NewServer(cfg, authManager, accessManager, configPath)
|
||||
}
|
||||
|
||||
func TestAmpProviderModelRoutes(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
path string
|
||||
wantStatus int
|
||||
wantContains string
|
||||
}{
|
||||
{
|
||||
name: "openai root models",
|
||||
path: "/api/provider/openai/models",
|
||||
wantStatus: http.StatusOK,
|
||||
wantContains: `"object":"list"`,
|
||||
},
|
||||
{
|
||||
name: "groq root models",
|
||||
path: "/api/provider/groq/models",
|
||||
wantStatus: http.StatusOK,
|
||||
wantContains: `"object":"list"`,
|
||||
},
|
||||
{
|
||||
name: "openai models",
|
||||
path: "/api/provider/openai/v1/models",
|
||||
wantStatus: http.StatusOK,
|
||||
wantContains: `"object":"list"`,
|
||||
},
|
||||
{
|
||||
name: "anthropic models",
|
||||
path: "/api/provider/anthropic/v1/models",
|
||||
wantStatus: http.StatusOK,
|
||||
wantContains: `"data"`,
|
||||
},
|
||||
{
|
||||
name: "google models v1",
|
||||
path: "/api/provider/google/v1/models",
|
||||
wantStatus: http.StatusOK,
|
||||
wantContains: `"models"`,
|
||||
},
|
||||
{
|
||||
name: "google models v1beta",
|
||||
path: "/api/provider/google/v1beta/models",
|
||||
wantStatus: http.StatusOK,
|
||||
wantContains: `"models"`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
tc := tc
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
server := newTestServer(t)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, tc.path, nil)
|
||||
req.Header.Set("Authorization", "Bearer test-key")
|
||||
|
||||
rr := httptest.NewRecorder()
|
||||
server.engine.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Code != tc.wantStatus {
|
||||
t.Fatalf("unexpected status code for %s: got %d want %d; body=%s", tc.path, rr.Code, tc.wantStatus, rr.Body.String())
|
||||
}
|
||||
if body := rr.Body.String(); !strings.Contains(body, tc.wantContains) {
|
||||
t.Fatalf("response body for %s missing %q: %s", tc.path, tc.wantContains, body)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -23,6 +23,17 @@ type Config struct {
|
||||
// Port is the network port on which the API server will listen.
|
||||
Port int `yaml:"port" json:"-"`
|
||||
|
||||
// AmpUpstreamURL defines the upstream Amp control plane used for non-provider calls.
|
||||
AmpUpstreamURL string `yaml:"amp-upstream-url" json:"amp-upstream-url"`
|
||||
|
||||
// AmpUpstreamAPIKey optionally overrides the Authorization header when proxying Amp upstream calls.
|
||||
AmpUpstreamAPIKey string `yaml:"amp-upstream-api-key" json:"amp-upstream-api-key"`
|
||||
|
||||
// AmpRestrictManagementToLocalhost restricts Amp management routes (/api/user, /api/threads, etc.)
|
||||
// to only accept connections from localhost (127.0.0.1, ::1). When true, prevents drive-by
|
||||
// browser attacks and remote access to management endpoints. Default: true (recommended).
|
||||
AmpRestrictManagementToLocalhost bool `yaml:"amp-restrict-management-to-localhost" json:"amp-restrict-management-to-localhost"`
|
||||
|
||||
// AuthDir is the directory where authentication token files are stored.
|
||||
AuthDir string `yaml:"auth-dir" json:"-"`
|
||||
|
||||
@@ -258,6 +269,7 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
|
||||
cfg.LoggingToFile = false
|
||||
cfg.UsageStatisticsEnabled = false
|
||||
cfg.DisableCooling = false
|
||||
cfg.AmpRestrictManagementToLocalhost = true // Default to secure: only localhost access
|
||||
if err = yaml.Unmarshal(data, &cfg); err != nil {
|
||||
if optional {
|
||||
// In cloud deploy mode, if YAML parsing fails, return empty config instead of error.
|
||||
|
||||
@@ -62,6 +62,23 @@ func ParseGeminiThinkingSuffix(model string) (string, *int, *bool, bool) {
|
||||
return base, &budgetValue, nil, true
|
||||
}
|
||||
|
||||
func NormalizeGeminiThinkingModel(modelName string) (string, map[string]any) {
|
||||
baseModel, budget, include, matched := ParseGeminiThinkingSuffix(modelName)
|
||||
if !matched {
|
||||
return baseModel, nil
|
||||
}
|
||||
metadata := map[string]any{
|
||||
GeminiOriginalModelMetadataKey: modelName,
|
||||
}
|
||||
if budget != nil {
|
||||
metadata[GeminiThinkingBudgetMetadataKey] = *budget
|
||||
}
|
||||
if include != nil {
|
||||
metadata[GeminiIncludeThoughtsMetadataKey] = *include
|
||||
}
|
||||
return baseModel, metadata
|
||||
}
|
||||
|
||||
func ApplyGeminiThinkingConfig(body []byte, budget *int, includeThoughts *bool) []byte {
|
||||
if budget == nil && includeThoughts == nil {
|
||||
return body
|
||||
|
||||
@@ -8,9 +8,12 @@ package claude
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
@@ -19,6 +22,7 @@ import (
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
@@ -153,6 +157,23 @@ func (h *ClaudeCodeAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSO
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
|
||||
// Decompress gzipped responses - Claude API sometimes returns gzip without Content-Encoding header
|
||||
// This fixes title generation and other non-streaming responses that arrive compressed
|
||||
if len(resp) >= 2 && resp[0] == 0x1f && resp[1] == 0x8b {
|
||||
gzReader, err := gzip.NewReader(bytes.NewReader(resp))
|
||||
if err != nil {
|
||||
log.Warnf("failed to decompress gzipped Claude response: %v", err)
|
||||
} else {
|
||||
defer gzReader.Close()
|
||||
if decompressed, err := io.ReadAll(gzReader); err != nil {
|
||||
log.Warnf("failed to read decompressed Claude response: %v", err)
|
||||
} else {
|
||||
resp = decompressed
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
_, _ = c.Writer.Write(resp)
|
||||
cliCancel()
|
||||
}
|
||||
|
||||
@@ -359,20 +359,7 @@ func cloneBytes(src []byte) []byte {
|
||||
}
|
||||
|
||||
func normalizeModelMetadata(modelName string) (string, map[string]any) {
|
||||
baseModel, budget, include, matched := util.ParseGeminiThinkingSuffix(modelName)
|
||||
if !matched {
|
||||
return baseModel, nil
|
||||
}
|
||||
metadata := map[string]any{
|
||||
util.GeminiOriginalModelMetadataKey: modelName,
|
||||
}
|
||||
if budget != nil {
|
||||
metadata[util.GeminiThinkingBudgetMetadataKey] = *budget
|
||||
}
|
||||
if include != nil {
|
||||
metadata[util.GeminiIncludeThoughtsMetadataKey] = *include
|
||||
}
|
||||
return baseModel, metadata
|
||||
return util.NormalizeGeminiThinkingModel(modelName)
|
||||
}
|
||||
|
||||
func cloneMetadata(src map[string]any) map[string]any {
|
||||
|
||||
Reference in New Issue
Block a user