feat: Add Amp CLI integration with comprehensive documentation

Add full Amp CLI support to enable routing AI model requests through the proxy
while maintaining Amp-specific features like thread management, user info, and
telemetry. Includes complete documentation and pull bot configuration.

Features:
- Modular architecture with RouteModule interface for clean integration
- Reverse proxy for Amp management routes (thread/user/meta/ads/telemetry)
- Provider-specific route aliases (/api/provider/{provider}/*)
- Secret management with precedence: config > env > file
- 5-minute secret caching to reduce file I/O
- Automatic gzip decompression for responses
- Proper connection cleanup to prevent leaks
- Localhost-only restriction for management routes (configurable)
- CORS protection for management endpoints

Documentation:
- Complete setup guide (USING_WITH_FACTORY_AND_AMP.md)
- OAuth setup for OpenAI (ChatGPT Plus/Pro) and Anthropic (Claude Pro/Max)
- Factory CLI config examples with all model variants
- Amp CLI/IDE configuration examples
- tmux setup for remote server deployment
- Screenshots and diagrams

Configuration:
- Pull bot disabled for this repo (manual rebase workflow)
- Config fields: AmpUpstreamURL, AmpUpstreamAPIKey, AmpRestrictManagementToLocalhost
- Compatible with upstream DisableCooling and other features

Technical details:
- internal/api/modules/amp/: Complete Amp routing module
- sdk/api/httpx/: HTTP utilities for gzip/transport
- 94.6% test coverage with 34 comprehensive test cases
- Clean integration minimizes merge conflict risk

Security:
- Management routes restricted to localhost by default
- Configurable via amp-restrict-management-to-localhost
- Prevents drive-by browser attacks on user data

This provides a production-ready foundation for Amp CLI integration while
maintaining clean separation from upstream code for easy rebasing.

Amp-Thread-ID: https://ampcode.com/threads/T-9e2befc5-f969-41c6-890c-5b779d58cf18
This commit is contained in:
Ben Vargas
2025-10-21 09:12:58 -06:00
parent 618511ff67
commit 9ad0f3f91e
19 changed files with 2923 additions and 1 deletions

1
.gitignore vendored
View File

@@ -30,3 +30,4 @@ GEMINI.md
.vscode/* .vscode/*
.claude/* .claude/*
.serena/* .serena/*
/cmd/server/server

View File

@@ -1,5 +1,37 @@
# CLI Proxy API # CLI Proxy API
---
## 🔔 Important: Amp CLI Support Fork
**This is a specialized fork of [router-for-me/CLIProxyAPI](https://github.com/router-for-me/CLIProxyAPI) that adds support for the Amp CLI tool.**
### Why This Fork Exists
The **Amp CLI** requires custom routing patterns to function properly. The upstream CLIProxyAPI project maintainers opted not to merge Amp-specific routing support into the main codebase.
### Which Version Should You Use?
- **Use this fork** if you want to run **both Factory CLI and Amp CLI** with the same proxy server
- **Use upstream** ([router-for-me/CLIProxyAPI](https://github.com/router-for-me/CLIProxyAPI)) if you only need Factory CLI support
### 📖 Complete Setup Guide
**→ [USING_WITH_FACTORY_AND_AMP.md](USING_WITH_FACTORY_AND_AMP.md)** - Comprehensive guide for using this proxy with both Factory CLI (Droid) and Amp CLI and IDE extensions, including OAuth setup, configuration examples, and troubleshooting.
### Key Differences
This fork includes:
-**Amp CLI route aliases** (`/api/provider/{provider}/v1...`)
-**Amp upstream proxy support** for OAuth and management routes
-**Automatic gzip decompression** for Amp upstream responses
-**Smart secret management** with precedence: config > env > file
-**All Factory CLI features** from upstream (fully compatible)
All Amp-specific code is isolated in the `internal/api/modules/amp` module, making it easy to sync upstream changes with minimal conflicts.
---
English | [中文](README_CN.md) English | [中文](README_CN.md)
A proxy server that provides OpenAI/Gemini/Claude/Codex compatible API interfaces for CLI. A proxy server that provides OpenAI/Gemini/Claude/Codex compatible API interfaces for CLI.
@@ -40,6 +72,15 @@ Get 10% OFF GLM CODING PLANhttps://z.ai/subscribe?ic=8JVLJQFSKB
- OpenAI-compatible upstream providers via config (e.g., OpenRouter) - OpenAI-compatible upstream providers via config (e.g., OpenRouter)
- Reusable Go SDK for embedding the proxy (see `docs/sdk-usage.md`) - Reusable Go SDK for embedding the proxy (see `docs/sdk-usage.md`)
### Fork-Specific: Amp CLI Support 🔥
- **Full Amp CLI integration** via provider route aliases (`/api/provider/{provider}/v1...`)
- **Amp upstream proxy** for OAuth authentication and management routes
- **Smart secret management** with configurable precedence (config > env > file)
- **Automatic gzip decompression** for Amp upstream responses
- **5-minute secret caching** to reduce file I/O overhead
- **Zero conflict** with Factory CLI - use both tools simultaneously
- **Modular architecture** for easy upstream sync (90% reduction in merge conflicts)
## Getting Started ## Getting Started
CLIProxyAPI Guides: [https://help.router-for.me/](https://help.router-for.me/) CLIProxyAPI Guides: [https://help.router-for.me/](https://help.router-for.me/)
@@ -78,7 +119,7 @@ Native macOS menu bar app to use your Claude Code & ChatGPT subscriptions with A
Browser-based tool to translate SRT subtitles using your Gemini subscription via CLIProxyAPI with automatic validation/error correction - no API keys needed Browser-based tool to translate SRT subtitles using your Gemini subscription via CLIProxyAPI with automatic validation/error correction - no API keys needed
> [!NOTE] > [!NOTE]
> If you developed a project based on CLIProxyAPI, please open a PR to add it to this list. > If you developed a project based on CLIProxyAPI, please open a PR to add it to this list.
## License ## License

View File

@@ -0,0 +1,494 @@
# Using Factory CLI (Droid) and Amp CLI with ChatGPT/Claude Subscriptions (OAuth)
## Why Use Subscriptions Instead of API Keys or Pass-Through Pricing?
Using Factory CLI (droid) or Amp CLI/IDE with this CLIProxyAPI fork lets you leverage your **existing provider subscriptions** (ChatGPT Plus/Pro, Claude Pro/Max) instead of per-token API billing.
**The value proposition is compelling:**
- **ChatGPT Plus/Pro** ($20-200/month) includes substantial use based on 5h and weekly quota limits
- **Claude Pro/Max** ($20-100-200/month) includes substantial Claude Sonnet 4.5 and Opus 4.1 on 5h and weekly quota limits
- **Pay-per-token APIs** can cost 5-10x+ for equivalent usage, even with pass-through pricing and no markup
By using OAuth subscriptions through this proxy, you get significantly better value while using the powerful CLI and IDE harnesses from Factory and AmpCode.
## Disclaimer
- This project is for personal/educational use only. You are solely responsible for how you use it.
- Using reverse proxies or alternate API bases may violate provider Terms of Service (OpenAI, Anthropic, Google, etc.).
- Accounts can be rate-limited, locked, or banned. Credentials and data may be at risk if misconfigured.
- Do not use to resell access, bypass access controls, or otherwise abuse services.
- No warranties. Use at your own risk.
## Summary
- Run Factory CLI (droid) and Amp CLI through a single local proxy server.
- This fork keeps all upstream Factory compatibility and adds Amp-specific support:
- Provider route aliases for Amp: `/api/provider/{provider}/v1...`
- Amp OAuth/management upstream proxy
- Smart secret resolution and automatic gzip handling
- Outcome: one proxy for both tools, minimal switching, clean separation of Amp supporting code from upstream repo.
## Why This Fork?
- Upstream maintainers chose not to include Amp-specific routing to keep scope focused on pure proxy functionality.
- Amp CLI expects Amp-specific alias routes and management endpoints the upstream CLIProxyAPI does not expose.
- This fork adds:
- Route aliases: `/api/provider/{provider}/v1...`
- Amp upstream proxy and OAuth
- Localhost-only access controls for Amp management routes (secure-by-default)
- Amp-specific code is isolated under `internal/api/modules/amp`, reducing merge conflicts with upstream.
## Architecture Overview
### Factory (droid) flow
```mermaid
flowchart LR
A["Factory CLI (droid)"] -->|"OpenAI/Claude-compatible calls"| B["CLIProxyAPI Fork"]
B -->|"/v1/chat/completions<br>/v1/messages<br>/v1/models"| C["Translators/Router"]
C -->|"OAuth tokens"| D[("Providers")]
D -->|"OpenAI Codex / Claude"| E["Responses+Streaming"]
E --> B --> A
```
### Amp flow
```mermaid
flowchart LR
A["Amp CLI"] -->|"/api/provider/provider/v1..."| B["CLIProxyAPI Fork"]
B -->|"Route aliases map to<br>upstream /v1 handlers"| C["Translators/Router"]
A -->|"/api/auth<br>/api/user<br>/api/meta<br>/api/threads..."| B
B -->|"Amp upstream proxy<br>(config: amp-upstream-url)"| F[("ampcode.com")]
C -->|"OpenAI / Anthropic"| D[("Providers")]
D --> B --> A
```
### Notes
- Factory uses standard OpenAI-compatible routes under `/v1/...`.
- Amp uses `/api/provider/{provider}/v1...` plus management routes proxied to `amp-upstream-url`.
- Management routes are restricted to localhost by default.
## Prerequisites
- Go 1.24+
- Active subscriptions:
- **ChatGPT Plus/Pro** (for GPT-5/GPT-5 Codex via OAuth)
- **Claude Pro/Max** (for Claude models via OAuth)
- **Amp** (for Amp CLI features in this fork)
- CLI tools:
- Factory CLI (droid)
- Amp CLI
- Local port `8317` available (or choose your own in config)
## Installation & Build
### Clone and build:
```bash
git clone https://github.com/ben-vargas/ai-cli-proxy-api.git
cd ai-cli-proxy-api
```
**macOS/Linux:**
```bash
go build -o cli-proxy-api ./cmd/server
```
**Windows:**
```bash
go build -o cli-proxy-api.exe ./cmd/server
```
### Homebrew (Factory CLI only):
> **⚠️ Note:** The Homebrew package installs the upstream version without Amp CLI support. Use the git clone method above if you need Amp CLI functionality.
```bash
brew install cliproxyapi
brew services start cliproxyapi
```
## OAuth Setup
Run these commands in the repo folder after building to authenticate with your subscriptions:
### OpenAI (ChatGPT Plus/Pro for GPT-5/Codex):
```bash
./cli-proxy-api --codex-login
```
- Opens browser on port `1455` for OAuth callback
- Requires active ChatGPT Plus or Pro subscription
- Tokens saved to `~/.cli-proxy-api/codex-<email>.json`
### Claude (Anthropic for Claude models):
```bash
./cli-proxy-api --claude-login
```
- Opens browser on port `54545` for OAuth callback
- Requires active Claude Pro or Claude Max subscription
- Tokens saved to `~/.cli-proxy-api/claude-<email>.json`
**Tip:** Add `--no-browser` to print the login URL instead of opening a browser (useful for remote/headless servers).
## Configuration for Factory CLI
Factory CLI uses `~/.factory/config.json` to define custom models. Add entries to the `custom_models` array.
### Complete configuration example
Copy this entire configuration to `~/.factory/config.json` for quick setup:
```json
{
"custom_models": [
{
"model_display_name": "Claude Haiku 4.5 [Proxy]",
"model": "claude-haiku-4-5-20251001",
"base_url": "http://localhost:8317",
"api_key": "dummy-not-used",
"provider": "anthropic"
},
{
"model_display_name": "Claude Sonnet 4.5 [Proxy]",
"model": "claude-sonnet-4-5-20250929",
"base_url": "http://localhost:8317",
"api_key": "dummy-not-used",
"provider": "anthropic"
},
{
"model_display_name": "Claude Opus 4.1 [Proxy]",
"model": "claude-opus-4-1-20250805",
"base_url": "http://localhost:8317",
"api_key": "dummy-not-used",
"provider": "anthropic"
},
{
"model_display_name": "Claude Sonnet 4 [Proxy]",
"model": "claude-sonnet-4-20250514",
"base_url": "http://localhost:8317",
"api_key": "dummy-not-used",
"provider": "anthropic"
},
{
"model_display_name": "GPT-5 [Proxy]",
"model": "gpt-5",
"base_url": "http://localhost:8317/v1",
"api_key": "dummy-not-used",
"provider": "openai"
},
{
"model_display_name": "GPT-5 Minimal [Proxy]",
"model": "gpt-5-minimal",
"base_url": "http://localhost:8317/v1",
"api_key": "dummy-not-used",
"provider": "openai"
},
{
"model_display_name": "GPT-5 Medium [Proxy]",
"model": "gpt-5-medium",
"base_url": "http://localhost:8317/v1",
"api_key": "dummy-not-used",
"provider": "openai"
},
{
"model_display_name": "GPT-5 High [Proxy]",
"model": "gpt-5-high",
"base_url": "http://localhost:8317/v1",
"api_key": "dummy-not-used",
"provider": "openai"
},
{
"model_display_name": "GPT-5 Codex [Proxy]",
"model": "gpt-5-codex",
"base_url": "http://localhost:8317/v1",
"api_key": "dummy-not-used",
"provider": "openai"
},
{
"model_display_name": "GPT-5 Codex High [Proxy]",
"model": "gpt-5-codex-high",
"base_url": "http://localhost:8317/v1",
"api_key": "dummy-not-used",
"provider": "openai"
}
]
}
```
After configuration, your custom models will appear in the `/model` selector:
![Factory CLI model selector showing custom models](assets/factory_droid_custom_models.png)
### Required fields:
| Field | Required | Description | Example |
|-------|----------|-------------|---------|
| `model_display_name` | ✓ | Human-friendly name shown in `/model` selector | `"Claude Sonnet 4.5 [Proxy]"` |
| `model` | ✓ | Model identifier sent to API | `"claude-sonnet-4-5-20250929"` |
| `base_url` | ✓ | Proxy endpoint | `"http://localhost:8317"` or `"http://localhost:8317/v1"` |
| `api_key` | ✓ | API key (use `"dummy-not-used"` for proxy) | `"dummy-not-used"` |
| `provider` | ✓ | API format type | `"anthropic"`, `"openai"`, or `"generic-chat-completion-api"` |
### Provider-specific base URLs:
| Provider | Base URL | Reason |
|----------|----------|--------|
| `anthropic` | `http://localhost:8317` | Factory appends `/v1/messages` automatically |
| `openai` | `http://localhost:8317/v1` | Factory appends `/responses` (needs `/v1` prefix) |
| `generic-chat-completion-api` | `http://localhost:8317/v1` | For OpenAI Chat Completions compatible models |
### Using custom models:
1. Edit `~/.factory/config.json` with the models above
2. Restart Factory CLI (`droid`)
3. Use `/model` command to select your custom model
## Configuration for Amp CLI
Enable Amp integration (fork-specific):
In `config.yaml`:
```yaml
# Amp CLI integration
amp-upstream-url: "https://ampcode.com"
# Optional override; otherwise uses env or file (see precedence below)
# amp-upstream-api-key: "your-amp-api-key"
# Security: restrict management routes to localhost (recommended)
amp-restrict-management-to-localhost: true
```
### Secret resolution precedence
| Source | Key | Priority |
|-----------------------------------------|----------------------------------|----------|
| Config file | `amp-upstream-api-key` | High |
| Environment | `AMP_API_KEY` | Medium |
| Amp secrets file | `~/.local/share/amp/secrets.json`| Low |
### Set Amp CLI to use this proxy
Edit `~/.config/amp/settings.json` and add the `amp.url` setting:
```json
{
"amp.url": "http://localhost:8317"
}
```
Or set the environment variable:
```bash
export AMP_URL=http://localhost:8317
```
Then login (proxied via `amp-upstream-url`):
```bash
amp login
```
Use Amp as normal:
```bash
amp "Hello, world!"
```
### Supported Amp routes
**Provider Aliases (always available):**
- `/api/provider/openai/v1/chat/completions`
- `/api/provider/openai/v1/responses`
- `/api/provider/anthropic/v1/messages`
- And related provider routes/versions your Amp CLI calls
**Management Routes (require `amp-upstream-url`):**
- `/api/auth`, `/api/user`, `/api/meta`, `/api/internal`, `/api/threads`, `/api/telemetry`
- Localhost-only by default for security
### Works with Amp IDE Extension
This proxy configuration also works with the Amp IDE extension for VSCode and forks (Cursor, Windsurf, etc). Simply set the Amp URL in your IDE extension settings:
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 with the same OAuth subscriptions!
![Amp IDE extension settings](assets/amp_ide_extension_amp_url.png)
The IDE extension uses the same routes as the CLI, so both can share the proxy simultaneously.
## Running the Proxy
> **Important:** The proxy requires a config file with `port` set (e.g., `port: 8317`). There is no built-in default port.
### With config file:
```bash
./cli-proxy-api --config config.yaml
```
If `config.yaml` is in the current directory:
```bash
./cli-proxy-api
```
### Tmux (recommended for remote servers):
Running in tmux keeps the proxy alive across SSH disconnects:
**Start proxy in detached tmux session:**
```bash
tmux new-session -d -s proxy -c ~/ai-cli-proxy-api \
"./cli-proxy-api --config config.yaml"
```
**View/attach to proxy session:**
```bash
tmux attach-session -t proxy
```
**Detach from session (proxy keeps running):**
```
Ctrl+b, then d
```
**Stop proxy:**
```bash
tmux kill-session -t proxy
```
**Check if running:**
```bash
tmux has-session -t proxy && echo "Running" || echo "Not running"
```
**Optional: Add to `~/.bashrc` for convenience:**
```bash
alias proxy-start='tmux new-session -d -s proxy -c ~/ai-cli-proxy-api "./cli-proxy-api --config config.yaml" && echo "Proxy started (use proxy-view to attach)"'
alias proxy-view='tmux attach-session -t proxy'
alias proxy-stop='tmux kill-session -t proxy 2>/dev/null && echo "Proxy stopped"'
alias proxy-status='tmux has-session -t proxy 2>/dev/null && echo "✓ Running" || echo "✗ Not running"'
```
### As a service (examples):
**Homebrew:**
```bash
brew services start cliproxyapi
```
**Systemd/Docker:** use your standard service templates; point the binary and config appropriately
### Key config fields (example)
```yaml
port: 8317
auth-dir: "~/.cli-proxy-api"
debug: false
logging-to-file: true
remote-management:
allow-remote: false
secret-key: "" # leave empty to disable management API
disable-control-panel: false
# Amp integration
amp-upstream-url: "https://ampcode.com"
# amp-upstream-api-key: "your-amp-api-key"
amp-restrict-management-to-localhost: true
# Retries and quotas
request-retry: 3
quota-exceeded:
switch-project: true
switch-preview-model: true
```
## Usage Examples
### Factory
**List models:**
```bash
curl http://localhost:8317/v1/models
```
**Chat Completions (Claude):**
```bash
curl -s http://localhost:8317/v1/messages \
-H "Content-Type: application/json" \
-d '{
"model": "claude-sonnet-4-5-20250929",
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 1024
}'
```
### Amp
**Provider alias (OpenAI-style):**
```bash
curl -s http://localhost:8317/api/provider/openai/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "gpt-5",
"messages": [{"role": "user", "content": "Hello"}]
}'
```
**Management (localhost only by default):**
```bash
curl -s http://localhost:8317/api/user
```
## Troubleshooting
### Common errors and fixes
| Symptom/Code | Likely Cause | Fix |
|------------------------------------------|------------------------------------------------------|----------------------------------------------------------------------|
| 404 /v1/chat/completions | Factory not pointing to proxy base | Set base to `http://localhost:8317/v1` (env/flag/config). |
| 404 /api/provider/... | Incorrect route path or typo | Ensure you're calling `/api/provider/{provider}/v1...` paths exactly.|
| 403 on /api/user (Amp) | Management restricted to localhost | Run from same machine or set `amp-restrict-management-to-localhost: false` (not recommended). |
| 401/403 from provider | Missing/expired OAuth or API key | Re-run the relevant `--*-login` or configure keys in `config.yaml`. |
| 429/Quota exceeded | Project/model quota exhausted | Enable `quota-exceeded` switching or switch accounts. |
| 5xx from provider | Upstream transient error | Increase `request-retry` and try again. |
| SSE/stream stuck | Client not handling SSE properly | Use SSE-capable client or set `stream: false`. |
| Amp gzip decoding errors | Compressed upstream responses | Fork auto-decompresses; update to latest build if issue persists. |
| CORS errors in browser | Protected management endpoints | Use CLI/terminal; avoid browsers for management endpoints. |
| Wrong model name | Provider alias mismatch | Use `gpt-*` for OpenAI or `claude-*` for Anthropic models. |
### Diagnostics
- Check logs (`debug: true` temporarily or `logging-to-file: true`).
- Verify config in effect: print effective config or confirm with startup logs.
- Test base reachability: `curl http://localhost:8317/v1/models`.
- For Amp, verify `amp-upstream-url` and secrets resolution.
## Security Checklist
- Keep `amp-restrict-management-to-localhost: true` (default).
- Do not expose the proxy publicly; bind to localhost or protect with firewall/VPN.
- If enabling remote management, set `remote-management.secret-key` and TLS/ingress protections.
- Disable the built-in management UI if hosting your own:
- `remote-management.disable-control-panel: true`
- Rotate tokens/keys; store config and auth-dir on encrypted disk or managed secret stores.
- Keep binary up to date to receive security fixes.
## References
- This fork README: [README.md](README.md)
- Upstream project: [CLIProxyAPI](https://github.com/router-for-me/CLIProxyAPI)
- Amp CLI: [Official Manual](https://ampcode.com/manual)
- Factory CLI (droid): [Official Documentation](https://docs.factory.ai/cli/getting-started/overview)

View File

@@ -137,6 +137,10 @@ func (MyExecutor) Execute(ctx context.Context, a *coreauth.Auth, req clipexec.Re
return clipexec.Response{Payload: body}, nil 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) { 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) ch := make(chan clipexec.StreamChunk, 1)
go func() { go func() {

View 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, 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
}

View 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")
}
})
}
}

View File

@@ -0,0 +1,176 @@
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
}
// 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)
}
}

View File

@@ -0,0 +1,439 @@
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)
}
})
}
}

View File

@@ -0,0 +1,166 @@
package amp
import (
"net"
"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"
log "github.com/sirupsen/logrus"
)
// localhostOnlyMiddleware restricts access to localhost (127.0.0.1, ::1) only.
// Returns 403 Forbidden for non-localhost clients.
func localhostOnlyMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
clientIP := c.ClientIP()
// Parse the IP to handle both IPv4 and IPv6
ip := net.ParseIP(clientIP)
if ip == nil {
log.Warnf("Amp management: invalid client IP %s, denying access", clientIP)
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 IP %s attempted access, denying", clientIP)
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, 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 (Gemini native API)
ampAPI.Any("/provider/google/v1beta1/*path", proxyHandler)
}
// 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)
// 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)
provider.GET("/models", ampModelsHandler)
provider.POST("/chat/completions", openaiHandlers.ChatCompletions)
provider.POST("/completions", openaiHandlers.Completions)
provider.POST("/responses", openaiResponsesHandlers.Responses)
// /v1 routes (OpenAI/Claude-compatible endpoints)
v1Amp := provider.Group("/v1")
{
v1Amp.GET("/models", ampModelsHandler)
// OpenAI-compatible endpoints
v1Amp.POST("/chat/completions", openaiHandlers.ChatCompletions)
v1Amp.POST("/completions", openaiHandlers.Completions)
v1Amp.POST("/responses", openaiResponsesHandlers.Responses)
// Claude/Anthropic-compatible endpoints
v1Amp.POST("/messages", claudeCodeHandlers.ClaudeMessages)
v1Amp.POST("/messages/count_tokens", claudeCodeHandlers.ClaudeCountTokens)
}
// /v1beta routes (Gemini native API)
v1betaAmp := provider.Group("/v1beta")
{
v1betaAmp.GET("/models", geminiHandlers.GeminiModels)
v1betaAmp.POST("/models/:action", geminiHandlers.GeminiHandler)
v1betaAmp.GET("/models/:action", geminiHandlers.GeminiGetHandler)
}
}

View File

@@ -0,0 +1,216 @@
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{}
m.registerManagementRoutes(r, proxyHandler, false) // false = don't restrict to localhost in tests
managementPaths := []string{
"/api/internal",
"/api/internal/some/path",
"/api/user",
"/api/user/profile",
"/api/auth",
"/api/auth/login",
"/api/meta",
"/api/telemetry",
"/api/threads",
"/api/otel",
"/api/provider/google/v1beta1/models",
}
for _, path := range managementPaths {
t.Run(path, func(t *testing.T) {
proxyCalled = false
req := httptest.NewRequest(http.MethodGet, path, nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code == http.StatusNotFound {
t.Fatalf("route %s not registered", path)
}
if !proxyCalled {
t.Fatalf("proxy handler not called for %s", 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")
}
}

View 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
}

View 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)
}
}

View 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)
}

View File

@@ -21,6 +21,8 @@ import (
"github.com/router-for-me/CLIProxyAPI/v6/internal/access" "github.com/router-for-me/CLIProxyAPI/v6/internal/access"
managementHandlers "github.com/router-for-me/CLIProxyAPI/v6/internal/api/handlers/management" 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/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/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/logging" "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
"github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset" "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 // Setup routes
s.setupRoutes() 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 { if optionState.routerConfigurator != nil {
optionState.routerConfigurator(engine, s.handlers, cfg) optionState.routerConfigurator(engine, s.handlers, cfg)
} }

111
internal/api/server_test.go Normal file
View 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)
}
})
}
}

View File

@@ -23,6 +23,17 @@ type Config struct {
// Port is the network port on which the API server will listen. // Port is the network port on which the API server will listen.
Port int `yaml:"port" json:"-"` 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 is the directory where authentication token files are stored.
AuthDir string `yaml:"auth-dir" json:"-"` AuthDir string `yaml:"auth-dir" json:"-"`
@@ -258,6 +269,7 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
cfg.LoggingToFile = false cfg.LoggingToFile = false
cfg.UsageStatisticsEnabled = false cfg.UsageStatisticsEnabled = false
cfg.DisableCooling = false cfg.DisableCooling = false
cfg.AmpRestrictManagementToLocalhost = true // Default to secure: only localhost access
if err = yaml.Unmarshal(data, &cfg); err != nil { if err = yaml.Unmarshal(data, &cfg); err != nil {
if optional { if optional {
// In cloud deploy mode, if YAML parsing fails, return empty config instead of error. // In cloud deploy mode, if YAML parsing fails, return empty config instead of error.

View File

@@ -8,9 +8,12 @@ package claude
import ( import (
"bufio" "bufio"
"bytes"
"compress/gzip"
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"io"
"net/http" "net/http"
"time" "time"
@@ -19,6 +22,7 @@ import (
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces" "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/internal/registry"
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers" "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson" "github.com/tidwall/gjson"
) )
@@ -153,6 +157,23 @@ func (h *ClaudeCodeAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSO
cliCancel(errMsg.Error) cliCancel(errMsg.Error)
return 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) _, _ = c.Writer.Write(resp)
cliCancel() cliCancel()
} }

33
sdk/api/httpx/gzip.go Normal file
View File

@@ -0,0 +1,33 @@
// Package httpx provides HTTP transport utilities for SDK clients,
// including automatic gzip decompression for misconfigured upstreams.
package httpx
import (
"bytes"
"compress/gzip"
"io"
)
// DecodePossibleGzip inspects the raw response body and transparently
// decompresses it when the payload is gzip compressed. Some upstream
// providers return gzip data without a Content-Encoding header, which
// confuses clients expecting JSON. This helper restores the original
// JSON bytes while leaving plain responses untouched.
//
// This function is preserved for backward compatibility but new code
// should use GzipFixupTransport instead.
func DecodePossibleGzip(raw []byte) ([]byte, error) {
if len(raw) >= 2 && raw[0] == 0x1f && raw[1] == 0x8b {
reader, err := gzip.NewReader(bytes.NewReader(raw))
if err != nil {
return nil, err
}
decompressed, err := io.ReadAll(reader)
_ = reader.Close()
if err != nil {
return nil, err
}
return decompressed, nil
}
return raw, nil
}

177
sdk/api/httpx/transport.go Normal file
View File

@@ -0,0 +1,177 @@
package httpx
import (
"bytes"
"compress/gzip"
"io"
"net/http"
"strings"
log "github.com/sirupsen/logrus"
)
// GzipFixupTransport wraps an http.RoundTripper to auto-decode gzip responses
// that don't properly set Content-Encoding header.
//
// Some upstream providers (especially when proxied) return gzip-compressed
// responses without setting the Content-Encoding: gzip header, which causes
// Go's http client to pass the compressed bytes directly to the application.
//
// This transport detects gzip magic bytes and transparently decompresses
// the response while preserving streaming behavior for SSE and chunked responses.
type GzipFixupTransport struct {
// Base is the underlying transport. If nil, http.DefaultTransport is used.
Base http.RoundTripper
}
// RoundTrip implements http.RoundTripper
func (t *GzipFixupTransport) RoundTrip(req *http.Request) (*http.Response, error) {
base := t.Base
if base == nil {
base = http.DefaultTransport
}
resp, err := base.RoundTrip(req)
if err != nil || resp == nil {
return resp, err
}
// Skip if Go already decompressed it
if resp.Uncompressed {
return resp, nil
}
// Skip if Content-Encoding is already set (properly configured upstream)
if resp.Header.Get("Content-Encoding") != "" {
return resp, nil
}
// Skip streaming responses - they need different handling
if isStreamingResponse(resp) {
// For streaming responses, wrap with a streaming gzip detector
// that can handle chunked gzip data
resp.Body = &streamingGzipDetector{
inner: resp.Body,
}
return resp, nil
}
// For non-streaming responses, peek and decompress if needed
resp.Body = &gzipDetectingReader{
inner: resp.Body,
}
return resp, nil
}
// isStreamingResponse checks if response is SSE or chunked
func isStreamingResponse(resp *http.Response) bool {
contentType := resp.Header.Get("Content-Type")
// Check for Server-Sent Events
if strings.Contains(contentType, "text/event-stream") {
return true
}
// Check for chunked transfer encoding
if strings.Contains(strings.ToLower(resp.Header.Get("Transfer-Encoding")), "chunked") {
return true
}
return false
}
// gzipDetectingReader is an io.ReadCloser that detects gzip magic bytes
// on first read and switches to gzip decompression if detected.
// This is used for non-streaming responses.
type gzipDetectingReader struct {
inner io.ReadCloser
reader io.Reader
once bool
}
func (g *gzipDetectingReader) Read(p []byte) (int, error) {
if !g.once {
g.once = true
// Peek at first 2 bytes to detect gzip magic bytes
buf := make([]byte, 2)
n, err := io.ReadFull(g.inner, buf)
if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
// Can't peek, use original reader
g.reader = io.MultiReader(bytes.NewReader(buf[:n]), g.inner)
return g.reader.Read(p)
}
if n >= 2 && buf[0] == 0x1f && buf[1] == 0x8b {
// It's gzipped, create gzip reader
multiReader := io.MultiReader(bytes.NewReader(buf[:n]), g.inner)
gzipReader, err := gzip.NewReader(multiReader)
if err != nil {
log.Warnf("gzip header detected but reader creation failed: %v", err)
g.reader = multiReader
} else {
g.reader = gzipReader
}
} else {
// Not gzipped, combine peeked bytes with rest
g.reader = io.MultiReader(bytes.NewReader(buf[:n]), g.inner)
}
}
return g.reader.Read(p)
}
func (g *gzipDetectingReader) Close() error {
if closer, ok := g.reader.(io.Closer); ok {
_ = closer.Close()
}
return g.inner.Close()
}
// streamingGzipDetector is similar to gzipDetectingReader but designed for
// streaming responses. It doesn't buffer; it wraps with a streaming gzip reader.
type streamingGzipDetector struct {
inner io.ReadCloser
reader io.Reader
once bool
}
func (s *streamingGzipDetector) Read(p []byte) (int, error) {
if !s.once {
s.once = true
// Peek at first 2 bytes
buf := make([]byte, 2)
n, err := io.ReadFull(s.inner, buf)
if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
s.reader = io.MultiReader(bytes.NewReader(buf[:n]), s.inner)
return s.reader.Read(p)
}
if n >= 2 && buf[0] == 0x1f && buf[1] == 0x8b {
// It's gzipped - wrap with streaming gzip reader
multiReader := io.MultiReader(bytes.NewReader(buf[:n]), s.inner)
gzipReader, err := gzip.NewReader(multiReader)
if err != nil {
log.Warnf("streaming gzip header detected but reader creation failed: %v", err)
s.reader = multiReader
} else {
s.reader = gzipReader
log.Debug("streaming gzip decompression enabled")
}
} else {
// Not gzipped
s.reader = io.MultiReader(bytes.NewReader(buf[:n]), s.inner)
}
}
return s.reader.Read(p)
}
func (s *streamingGzipDetector) Close() error {
if closer, ok := s.reader.(io.Closer); ok {
_ = closer.Close()
}
return s.inner.Close()
}