feat: Add Amp CLI integration with comprehensive documentation
Browse filesAdd 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
- .gitignore +1 -0
- README.md +42 -1
- USING_WITH_FACTORY_AND_AMP.md +494 -0
- examples/custom-provider/main.go +4 -0
- internal/api/modules/amp/amp.go +185 -0
- internal/api/modules/amp/amp_test.go +303 -0
- internal/api/modules/amp/proxy.go +176 -0
- internal/api/modules/amp/proxy_test.go +439 -0
- internal/api/modules/amp/routes.go +166 -0
- internal/api/modules/amp/routes_test.go +216 -0
- internal/api/modules/amp/secret.go +155 -0
- internal/api/modules/amp/secret_test.go +280 -0
- internal/api/modules/modules.go +92 -0
- internal/api/server.go +16 -0
- internal/api/server_test.go +111 -0
- internal/config/config.go +12 -0
- sdk/api/handlers/claude/code_handlers.go +21 -0
- sdk/api/httpx/gzip.go +33 -0
- sdk/api/httpx/transport.go +177 -0
|
@@ -30,3 +30,4 @@ GEMINI.md
|
|
| 30 |
.vscode/*
|
| 31 |
.claude/*
|
| 32 |
.serena/*
|
|
|
|
|
|
| 30 |
.vscode/*
|
| 31 |
.claude/*
|
| 32 |
.serena/*
|
| 33 |
+
/cmd/server/server
|
|
@@ -1,5 +1,37 @@
|
|
| 1 |
# CLI Proxy API
|
| 2 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
English | [中文](README_CN.md)
|
| 4 |
|
| 5 |
A proxy server that provides OpenAI/Gemini/Claude/Codex compatible API interfaces for CLI.
|
|
@@ -40,6 +72,15 @@ Get 10% OFF GLM CODING PLAN:https://z.ai/subscribe?ic=8JVLJQFSKB
|
|
| 40 |
- OpenAI-compatible upstream providers via config (e.g., OpenRouter)
|
| 41 |
- Reusable Go SDK for embedding the proxy (see `docs/sdk-usage.md`)
|
| 42 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 43 |
## Getting Started
|
| 44 |
|
| 45 |
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
|
|
| 78 |
|
| 79 |
Browser-based tool to translate SRT subtitles using your Gemini subscription via CLIProxyAPI with automatic validation/error correction - no API keys needed
|
| 80 |
|
| 81 |
-
> [!NOTE]
|
| 82 |
> If you developed a project based on CLIProxyAPI, please open a PR to add it to this list.
|
| 83 |
|
| 84 |
## License
|
|
|
|
| 1 |
# CLI Proxy API
|
| 2 |
|
| 3 |
+
---
|
| 4 |
+
|
| 5 |
+
## 🔔 Important: Amp CLI Support Fork
|
| 6 |
+
|
| 7 |
+
**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.**
|
| 8 |
+
|
| 9 |
+
### Why This Fork Exists
|
| 10 |
+
|
| 11 |
+
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.
|
| 12 |
+
|
| 13 |
+
### Which Version Should You Use?
|
| 14 |
+
|
| 15 |
+
- **Use this fork** if you want to run **both Factory CLI and Amp CLI** with the same proxy server
|
| 16 |
+
- **Use upstream** ([router-for-me/CLIProxyAPI](https://github.com/router-for-me/CLIProxyAPI)) if you only need Factory CLI support
|
| 17 |
+
|
| 18 |
+
### 📖 Complete Setup Guide
|
| 19 |
+
|
| 20 |
+
**→ [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.
|
| 21 |
+
|
| 22 |
+
### Key Differences
|
| 23 |
+
|
| 24 |
+
This fork includes:
|
| 25 |
+
- ✅ **Amp CLI route aliases** (`/api/provider/{provider}/v1...`)
|
| 26 |
+
- ✅ **Amp upstream proxy support** for OAuth and management routes
|
| 27 |
+
- ✅ **Automatic gzip decompression** for Amp upstream responses
|
| 28 |
+
- ✅ **Smart secret management** with precedence: config > env > file
|
| 29 |
+
- ✅ **All Factory CLI features** from upstream (fully compatible)
|
| 30 |
+
|
| 31 |
+
All Amp-specific code is isolated in the `internal/api/modules/amp` module, making it easy to sync upstream changes with minimal conflicts.
|
| 32 |
+
|
| 33 |
+
---
|
| 34 |
+
|
| 35 |
English | [中文](README_CN.md)
|
| 36 |
|
| 37 |
A proxy server that provides OpenAI/Gemini/Claude/Codex compatible API interfaces for CLI.
|
|
|
|
| 72 |
- OpenAI-compatible upstream providers via config (e.g., OpenRouter)
|
| 73 |
- Reusable Go SDK for embedding the proxy (see `docs/sdk-usage.md`)
|
| 74 |
|
| 75 |
+
### Fork-Specific: Amp CLI Support 🔥
|
| 76 |
+
- **Full Amp CLI integration** via provider route aliases (`/api/provider/{provider}/v1...`)
|
| 77 |
+
- **Amp upstream proxy** for OAuth authentication and management routes
|
| 78 |
+
- **Smart secret management** with configurable precedence (config > env > file)
|
| 79 |
+
- **Automatic gzip decompression** for Amp upstream responses
|
| 80 |
+
- **5-minute secret caching** to reduce file I/O overhead
|
| 81 |
+
- **Zero conflict** with Factory CLI - use both tools simultaneously
|
| 82 |
+
- **Modular architecture** for easy upstream sync (90% reduction in merge conflicts)
|
| 83 |
+
|
| 84 |
## Getting Started
|
| 85 |
|
| 86 |
CLIProxyAPI Guides: [https://help.router-for.me/](https://help.router-for.me/)
|
|
|
|
| 119 |
|
| 120 |
Browser-based tool to translate SRT subtitles using your Gemini subscription via CLIProxyAPI with automatic validation/error correction - no API keys needed
|
| 121 |
|
| 122 |
+
> [!NOTE]
|
| 123 |
> If you developed a project based on CLIProxyAPI, please open a PR to add it to this list.
|
| 124 |
|
| 125 |
## License
|
|
@@ -0,0 +1,494 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Using Factory CLI (Droid) and Amp CLI with ChatGPT/Claude Subscriptions (OAuth)
|
| 2 |
+
|
| 3 |
+
|
| 4 |
+
## Why Use Subscriptions Instead of API Keys or Pass-Through Pricing?
|
| 5 |
+
|
| 6 |
+
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.
|
| 7 |
+
|
| 8 |
+
**The value proposition is compelling:**
|
| 9 |
+
- **ChatGPT Plus/Pro** ($20-200/month) includes substantial use based on 5h and weekly quota limits
|
| 10 |
+
- **Claude Pro/Max** ($20-100-200/month) includes substantial Claude Sonnet 4.5 and Opus 4.1 on 5h and weekly quota limits
|
| 11 |
+
- **Pay-per-token APIs** can cost 5-10x+ for equivalent usage, even with pass-through pricing and no markup
|
| 12 |
+
|
| 13 |
+
By using OAuth subscriptions through this proxy, you get significantly better value while using the powerful CLI and IDE harnesses from Factory and AmpCode.
|
| 14 |
+
|
| 15 |
+
## Disclaimer
|
| 16 |
+
|
| 17 |
+
- This project is for personal/educational use only. You are solely responsible for how you use it.
|
| 18 |
+
- Using reverse proxies or alternate API bases may violate provider Terms of Service (OpenAI, Anthropic, Google, etc.).
|
| 19 |
+
- Accounts can be rate-limited, locked, or banned. Credentials and data may be at risk if misconfigured.
|
| 20 |
+
- Do not use to resell access, bypass access controls, or otherwise abuse services.
|
| 21 |
+
- No warranties. Use at your own risk.
|
| 22 |
+
|
| 23 |
+
## Summary
|
| 24 |
+
|
| 25 |
+
- Run Factory CLI (droid) and Amp CLI through a single local proxy server.
|
| 26 |
+
- This fork keeps all upstream Factory compatibility and adds Amp-specific support:
|
| 27 |
+
- Provider route aliases for Amp: `/api/provider/{provider}/v1...`
|
| 28 |
+
- Amp OAuth/management upstream proxy
|
| 29 |
+
- Smart secret resolution and automatic gzip handling
|
| 30 |
+
- Outcome: one proxy for both tools, minimal switching, clean separation of Amp supporting code from upstream repo.
|
| 31 |
+
|
| 32 |
+
## Why This Fork?
|
| 33 |
+
|
| 34 |
+
- Upstream maintainers chose not to include Amp-specific routing to keep scope focused on pure proxy functionality.
|
| 35 |
+
- Amp CLI expects Amp-specific alias routes and management endpoints the upstream CLIProxyAPI does not expose.
|
| 36 |
+
- This fork adds:
|
| 37 |
+
- Route aliases: `/api/provider/{provider}/v1...`
|
| 38 |
+
- Amp upstream proxy and OAuth
|
| 39 |
+
- Localhost-only access controls for Amp management routes (secure-by-default)
|
| 40 |
+
- Amp-specific code is isolated under `internal/api/modules/amp`, reducing merge conflicts with upstream.
|
| 41 |
+
|
| 42 |
+
## Architecture Overview
|
| 43 |
+
|
| 44 |
+
### Factory (droid) flow
|
| 45 |
+
|
| 46 |
+
```mermaid
|
| 47 |
+
flowchart LR
|
| 48 |
+
A["Factory CLI (droid)"] -->|"OpenAI/Claude-compatible calls"| B["CLIProxyAPI Fork"]
|
| 49 |
+
B -->|"/v1/chat/completions<br>/v1/messages<br>/v1/models"| C["Translators/Router"]
|
| 50 |
+
C -->|"OAuth tokens"| D[("Providers")]
|
| 51 |
+
D -->|"OpenAI Codex / Claude"| E["Responses+Streaming"]
|
| 52 |
+
E --> B --> A
|
| 53 |
+
```
|
| 54 |
+
|
| 55 |
+
### Amp flow
|
| 56 |
+
|
| 57 |
+
```mermaid
|
| 58 |
+
flowchart LR
|
| 59 |
+
A["Amp CLI"] -->|"/api/provider/provider/v1..."| B["CLIProxyAPI Fork"]
|
| 60 |
+
B -->|"Route aliases map to<br>upstream /v1 handlers"| C["Translators/Router"]
|
| 61 |
+
A -->|"/api/auth<br>/api/user<br>/api/meta<br>/api/threads..."| B
|
| 62 |
+
B -->|"Amp upstream proxy<br>(config: amp-upstream-url)"| F[("ampcode.com")]
|
| 63 |
+
C -->|"OpenAI / Anthropic"| D[("Providers")]
|
| 64 |
+
D --> B --> A
|
| 65 |
+
```
|
| 66 |
+
|
| 67 |
+
### Notes
|
| 68 |
+
|
| 69 |
+
- Factory uses standard OpenAI-compatible routes under `/v1/...`.
|
| 70 |
+
- Amp uses `/api/provider/{provider}/v1...` plus management routes proxied to `amp-upstream-url`.
|
| 71 |
+
- Management routes are restricted to localhost by default.
|
| 72 |
+
|
| 73 |
+
## Prerequisites
|
| 74 |
+
|
| 75 |
+
- Go 1.24+
|
| 76 |
+
- Active subscriptions:
|
| 77 |
+
- **ChatGPT Plus/Pro** (for GPT-5/GPT-5 Codex via OAuth)
|
| 78 |
+
- **Claude Pro/Max** (for Claude models via OAuth)
|
| 79 |
+
- **Amp** (for Amp CLI features in this fork)
|
| 80 |
+
- CLI tools:
|
| 81 |
+
- Factory CLI (droid)
|
| 82 |
+
- Amp CLI
|
| 83 |
+
- Local port `8317` available (or choose your own in config)
|
| 84 |
+
|
| 85 |
+
## Installation & Build
|
| 86 |
+
|
| 87 |
+
### Clone and build:
|
| 88 |
+
|
| 89 |
+
```bash
|
| 90 |
+
git clone https://github.com/ben-vargas/ai-cli-proxy-api.git
|
| 91 |
+
cd ai-cli-proxy-api
|
| 92 |
+
```
|
| 93 |
+
|
| 94 |
+
**macOS/Linux:**
|
| 95 |
+
```bash
|
| 96 |
+
go build -o cli-proxy-api ./cmd/server
|
| 97 |
+
```
|
| 98 |
+
|
| 99 |
+
**Windows:**
|
| 100 |
+
```bash
|
| 101 |
+
go build -o cli-proxy-api.exe ./cmd/server
|
| 102 |
+
```
|
| 103 |
+
|
| 104 |
+
### Homebrew (Factory CLI only):
|
| 105 |
+
|
| 106 |
+
> **⚠️ Note:** The Homebrew package installs the upstream version without Amp CLI support. Use the git clone method above if you need Amp CLI functionality.
|
| 107 |
+
|
| 108 |
+
```bash
|
| 109 |
+
brew install cliproxyapi
|
| 110 |
+
brew services start cliproxyapi
|
| 111 |
+
```
|
| 112 |
+
|
| 113 |
+
## OAuth Setup
|
| 114 |
+
|
| 115 |
+
Run these commands in the repo folder after building to authenticate with your subscriptions:
|
| 116 |
+
|
| 117 |
+
### OpenAI (ChatGPT Plus/Pro for GPT-5/Codex):
|
| 118 |
+
|
| 119 |
+
```bash
|
| 120 |
+
./cli-proxy-api --codex-login
|
| 121 |
+
```
|
| 122 |
+
|
| 123 |
+
- Opens browser on port `1455` for OAuth callback
|
| 124 |
+
- Requires active ChatGPT Plus or Pro subscription
|
| 125 |
+
- Tokens saved to `~/.cli-proxy-api/codex-<email>.json`
|
| 126 |
+
|
| 127 |
+
### Claude (Anthropic for Claude models):
|
| 128 |
+
|
| 129 |
+
```bash
|
| 130 |
+
./cli-proxy-api --claude-login
|
| 131 |
+
```
|
| 132 |
+
|
| 133 |
+
- Opens browser on port `54545` for OAuth callback
|
| 134 |
+
- Requires active Claude Pro or Claude Max subscription
|
| 135 |
+
- Tokens saved to `~/.cli-proxy-api/claude-<email>.json`
|
| 136 |
+
|
| 137 |
+
**Tip:** Add `--no-browser` to print the login URL instead of opening a browser (useful for remote/headless servers).
|
| 138 |
+
|
| 139 |
+
## Configuration for Factory CLI
|
| 140 |
+
|
| 141 |
+
Factory CLI uses `~/.factory/config.json` to define custom models. Add entries to the `custom_models` array.
|
| 142 |
+
|
| 143 |
+
### Complete configuration example
|
| 144 |
+
|
| 145 |
+
Copy this entire configuration to `~/.factory/config.json` for quick setup:
|
| 146 |
+
|
| 147 |
+
```json
|
| 148 |
+
{
|
| 149 |
+
"custom_models": [
|
| 150 |
+
{
|
| 151 |
+
"model_display_name": "Claude Haiku 4.5 [Proxy]",
|
| 152 |
+
"model": "claude-haiku-4-5-20251001",
|
| 153 |
+
"base_url": "http://localhost:8317",
|
| 154 |
+
"api_key": "dummy-not-used",
|
| 155 |
+
"provider": "anthropic"
|
| 156 |
+
},
|
| 157 |
+
{
|
| 158 |
+
"model_display_name": "Claude Sonnet 4.5 [Proxy]",
|
| 159 |
+
"model": "claude-sonnet-4-5-20250929",
|
| 160 |
+
"base_url": "http://localhost:8317",
|
| 161 |
+
"api_key": "dummy-not-used",
|
| 162 |
+
"provider": "anthropic"
|
| 163 |
+
},
|
| 164 |
+
{
|
| 165 |
+
"model_display_name": "Claude Opus 4.1 [Proxy]",
|
| 166 |
+
"model": "claude-opus-4-1-20250805",
|
| 167 |
+
"base_url": "http://localhost:8317",
|
| 168 |
+
"api_key": "dummy-not-used",
|
| 169 |
+
"provider": "anthropic"
|
| 170 |
+
},
|
| 171 |
+
{
|
| 172 |
+
"model_display_name": "Claude Sonnet 4 [Proxy]",
|
| 173 |
+
"model": "claude-sonnet-4-20250514",
|
| 174 |
+
"base_url": "http://localhost:8317",
|
| 175 |
+
"api_key": "dummy-not-used",
|
| 176 |
+
"provider": "anthropic"
|
| 177 |
+
},
|
| 178 |
+
{
|
| 179 |
+
"model_display_name": "GPT-5 [Proxy]",
|
| 180 |
+
"model": "gpt-5",
|
| 181 |
+
"base_url": "http://localhost:8317/v1",
|
| 182 |
+
"api_key": "dummy-not-used",
|
| 183 |
+
"provider": "openai"
|
| 184 |
+
},
|
| 185 |
+
{
|
| 186 |
+
"model_display_name": "GPT-5 Minimal [Proxy]",
|
| 187 |
+
"model": "gpt-5-minimal",
|
| 188 |
+
"base_url": "http://localhost:8317/v1",
|
| 189 |
+
"api_key": "dummy-not-used",
|
| 190 |
+
"provider": "openai"
|
| 191 |
+
},
|
| 192 |
+
{
|
| 193 |
+
"model_display_name": "GPT-5 Medium [Proxy]",
|
| 194 |
+
"model": "gpt-5-medium",
|
| 195 |
+
"base_url": "http://localhost:8317/v1",
|
| 196 |
+
"api_key": "dummy-not-used",
|
| 197 |
+
"provider": "openai"
|
| 198 |
+
},
|
| 199 |
+
{
|
| 200 |
+
"model_display_name": "GPT-5 High [Proxy]",
|
| 201 |
+
"model": "gpt-5-high",
|
| 202 |
+
"base_url": "http://localhost:8317/v1",
|
| 203 |
+
"api_key": "dummy-not-used",
|
| 204 |
+
"provider": "openai"
|
| 205 |
+
},
|
| 206 |
+
{
|
| 207 |
+
"model_display_name": "GPT-5 Codex [Proxy]",
|
| 208 |
+
"model": "gpt-5-codex",
|
| 209 |
+
"base_url": "http://localhost:8317/v1",
|
| 210 |
+
"api_key": "dummy-not-used",
|
| 211 |
+
"provider": "openai"
|
| 212 |
+
},
|
| 213 |
+
{
|
| 214 |
+
"model_display_name": "GPT-5 Codex High [Proxy]",
|
| 215 |
+
"model": "gpt-5-codex-high",
|
| 216 |
+
"base_url": "http://localhost:8317/v1",
|
| 217 |
+
"api_key": "dummy-not-used",
|
| 218 |
+
"provider": "openai"
|
| 219 |
+
}
|
| 220 |
+
]
|
| 221 |
+
}
|
| 222 |
+
```
|
| 223 |
+
|
| 224 |
+
After configuration, your custom models will appear in the `/model` selector:
|
| 225 |
+
|
| 226 |
+

|
| 227 |
+
|
| 228 |
+
### Required fields:
|
| 229 |
+
|
| 230 |
+
| Field | Required | Description | Example |
|
| 231 |
+
|-------|----------|-------------|---------|
|
| 232 |
+
| `model_display_name` | ✓ | Human-friendly name shown in `/model` selector | `"Claude Sonnet 4.5 [Proxy]"` |
|
| 233 |
+
| `model` | ✓ | Model identifier sent to API | `"claude-sonnet-4-5-20250929"` |
|
| 234 |
+
| `base_url` | ✓ | Proxy endpoint | `"http://localhost:8317"` or `"http://localhost:8317/v1"` |
|
| 235 |
+
| `api_key` | ✓ | API key (use `"dummy-not-used"` for proxy) | `"dummy-not-used"` |
|
| 236 |
+
| `provider` | ✓ | API format type | `"anthropic"`, `"openai"`, or `"generic-chat-completion-api"` |
|
| 237 |
+
|
| 238 |
+
### Provider-specific base URLs:
|
| 239 |
+
|
| 240 |
+
| Provider | Base URL | Reason |
|
| 241 |
+
|----------|----------|--------|
|
| 242 |
+
| `anthropic` | `http://localhost:8317` | Factory appends `/v1/messages` automatically |
|
| 243 |
+
| `openai` | `http://localhost:8317/v1` | Factory appends `/responses` (needs `/v1` prefix) |
|
| 244 |
+
| `generic-chat-completion-api` | `http://localhost:8317/v1` | For OpenAI Chat Completions compatible models |
|
| 245 |
+
|
| 246 |
+
### Using custom models:
|
| 247 |
+
|
| 248 |
+
1. Edit `~/.factory/config.json` with the models above
|
| 249 |
+
2. Restart Factory CLI (`droid`)
|
| 250 |
+
3. Use `/model` command to select your custom model
|
| 251 |
+
|
| 252 |
+
## Configuration for Amp CLI
|
| 253 |
+
|
| 254 |
+
Enable Amp integration (fork-specific):
|
| 255 |
+
|
| 256 |
+
In `config.yaml`:
|
| 257 |
+
|
| 258 |
+
```yaml
|
| 259 |
+
# Amp CLI integration
|
| 260 |
+
amp-upstream-url: "https://ampcode.com"
|
| 261 |
+
|
| 262 |
+
# Optional override; otherwise uses env or file (see precedence below)
|
| 263 |
+
# amp-upstream-api-key: "your-amp-api-key"
|
| 264 |
+
|
| 265 |
+
# Security: restrict management routes to localhost (recommended)
|
| 266 |
+
amp-restrict-management-to-localhost: true
|
| 267 |
+
```
|
| 268 |
+
|
| 269 |
+
### Secret resolution precedence
|
| 270 |
+
|
| 271 |
+
| Source | Key | Priority |
|
| 272 |
+
|-----------------------------------------|----------------------------------|----------|
|
| 273 |
+
| Config file | `amp-upstream-api-key` | High |
|
| 274 |
+
| Environment | `AMP_API_KEY` | Medium |
|
| 275 |
+
| Amp secrets file | `~/.local/share/amp/secrets.json`| Low |
|
| 276 |
+
|
| 277 |
+
### Set Amp CLI to use this proxy
|
| 278 |
+
|
| 279 |
+
Edit `~/.config/amp/settings.json` and add the `amp.url` setting:
|
| 280 |
+
|
| 281 |
+
```json
|
| 282 |
+
{
|
| 283 |
+
"amp.url": "http://localhost:8317"
|
| 284 |
+
}
|
| 285 |
+
```
|
| 286 |
+
|
| 287 |
+
Or set the environment variable:
|
| 288 |
+
|
| 289 |
+
```bash
|
| 290 |
+
export AMP_URL=http://localhost:8317
|
| 291 |
+
```
|
| 292 |
+
|
| 293 |
+
Then login (proxied via `amp-upstream-url`):
|
| 294 |
+
|
| 295 |
+
```bash
|
| 296 |
+
amp login
|
| 297 |
+
```
|
| 298 |
+
|
| 299 |
+
Use Amp as normal:
|
| 300 |
+
|
| 301 |
+
```bash
|
| 302 |
+
amp "Hello, world!"
|
| 303 |
+
```
|
| 304 |
+
|
| 305 |
+
### Supported Amp routes
|
| 306 |
+
|
| 307 |
+
**Provider Aliases (always available):**
|
| 308 |
+
- `/api/provider/openai/v1/chat/completions`
|
| 309 |
+
- `/api/provider/openai/v1/responses`
|
| 310 |
+
- `/api/provider/anthropic/v1/messages`
|
| 311 |
+
- And related provider routes/versions your Amp CLI calls
|
| 312 |
+
|
| 313 |
+
**Management Routes (require `amp-upstream-url`):**
|
| 314 |
+
- `/api/auth`, `/api/user`, `/api/meta`, `/api/internal`, `/api/threads`, `/api/telemetry`
|
| 315 |
+
- Localhost-only by default for security
|
| 316 |
+
|
| 317 |
+
### Works with Amp IDE Extension
|
| 318 |
+
|
| 319 |
+
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:
|
| 320 |
+
|
| 321 |
+
1. Open Amp extension settings in your IDE
|
| 322 |
+
2. Set **Amp URL** to `http://localhost:8317`
|
| 323 |
+
3. Login with your Amp account
|
| 324 |
+
4. Start using Amp in your IDE with the same OAuth subscriptions!
|
| 325 |
+
|
| 326 |
+

|
| 327 |
+
|
| 328 |
+
The IDE extension uses the same routes as the CLI, so both can share the proxy simultaneously.
|
| 329 |
+
|
| 330 |
+
## Running the Proxy
|
| 331 |
+
|
| 332 |
+
> **Important:** The proxy requires a config file with `port` set (e.g., `port: 8317`). There is no built-in default port.
|
| 333 |
+
|
| 334 |
+
### With config file:
|
| 335 |
+
|
| 336 |
+
```bash
|
| 337 |
+
./cli-proxy-api --config config.yaml
|
| 338 |
+
```
|
| 339 |
+
|
| 340 |
+
If `config.yaml` is in the current directory:
|
| 341 |
+
|
| 342 |
+
```bash
|
| 343 |
+
./cli-proxy-api
|
| 344 |
+
```
|
| 345 |
+
|
| 346 |
+
### Tmux (recommended for remote servers):
|
| 347 |
+
|
| 348 |
+
Running in tmux keeps the proxy alive across SSH disconnects:
|
| 349 |
+
|
| 350 |
+
**Start proxy in detached tmux session:**
|
| 351 |
+
```bash
|
| 352 |
+
tmux new-session -d -s proxy -c ~/ai-cli-proxy-api \
|
| 353 |
+
"./cli-proxy-api --config config.yaml"
|
| 354 |
+
```
|
| 355 |
+
|
| 356 |
+
**View/attach to proxy session:**
|
| 357 |
+
```bash
|
| 358 |
+
tmux attach-session -t proxy
|
| 359 |
+
```
|
| 360 |
+
|
| 361 |
+
**Detach from session (proxy keeps running):**
|
| 362 |
+
```
|
| 363 |
+
Ctrl+b, then d
|
| 364 |
+
```
|
| 365 |
+
|
| 366 |
+
**Stop proxy:**
|
| 367 |
+
```bash
|
| 368 |
+
tmux kill-session -t proxy
|
| 369 |
+
```
|
| 370 |
+
|
| 371 |
+
**Check if running:**
|
| 372 |
+
```bash
|
| 373 |
+
tmux has-session -t proxy && echo "Running" || echo "Not running"
|
| 374 |
+
```
|
| 375 |
+
|
| 376 |
+
**Optional: Add to `~/.bashrc` for convenience:**
|
| 377 |
+
```bash
|
| 378 |
+
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)"'
|
| 379 |
+
alias proxy-view='tmux attach-session -t proxy'
|
| 380 |
+
alias proxy-stop='tmux kill-session -t proxy 2>/dev/null && echo "Proxy stopped"'
|
| 381 |
+
alias proxy-status='tmux has-session -t proxy 2>/dev/null && echo "✓ Running" || echo "✗ Not running"'
|
| 382 |
+
```
|
| 383 |
+
|
| 384 |
+
### As a service (examples):
|
| 385 |
+
|
| 386 |
+
**Homebrew:**
|
| 387 |
+
```bash
|
| 388 |
+
brew services start cliproxyapi
|
| 389 |
+
```
|
| 390 |
+
|
| 391 |
+
**Systemd/Docker:** use your standard service templates; point the binary and config appropriately
|
| 392 |
+
|
| 393 |
+
### Key config fields (example)
|
| 394 |
+
|
| 395 |
+
```yaml
|
| 396 |
+
port: 8317
|
| 397 |
+
auth-dir: "~/.cli-proxy-api"
|
| 398 |
+
debug: false
|
| 399 |
+
logging-to-file: true
|
| 400 |
+
|
| 401 |
+
remote-management:
|
| 402 |
+
allow-remote: false
|
| 403 |
+
secret-key: "" # leave empty to disable management API
|
| 404 |
+
disable-control-panel: false
|
| 405 |
+
|
| 406 |
+
# Amp integration
|
| 407 |
+
amp-upstream-url: "https://ampcode.com"
|
| 408 |
+
# amp-upstream-api-key: "your-amp-api-key"
|
| 409 |
+
amp-restrict-management-to-localhost: true
|
| 410 |
+
|
| 411 |
+
# Retries and quotas
|
| 412 |
+
request-retry: 3
|
| 413 |
+
quota-exceeded:
|
| 414 |
+
switch-project: true
|
| 415 |
+
switch-preview-model: true
|
| 416 |
+
```
|
| 417 |
+
|
| 418 |
+
## Usage Examples
|
| 419 |
+
|
| 420 |
+
### Factory
|
| 421 |
+
|
| 422 |
+
**List models:**
|
| 423 |
+
```bash
|
| 424 |
+
curl http://localhost:8317/v1/models
|
| 425 |
+
```
|
| 426 |
+
|
| 427 |
+
**Chat Completions (Claude):**
|
| 428 |
+
```bash
|
| 429 |
+
curl -s http://localhost:8317/v1/messages \
|
| 430 |
+
-H "Content-Type: application/json" \
|
| 431 |
+
-d '{
|
| 432 |
+
"model": "claude-sonnet-4-5-20250929",
|
| 433 |
+
"messages": [{"role": "user", "content": "Hello"}],
|
| 434 |
+
"max_tokens": 1024
|
| 435 |
+
}'
|
| 436 |
+
```
|
| 437 |
+
|
| 438 |
+
### Amp
|
| 439 |
+
|
| 440 |
+
**Provider alias (OpenAI-style):**
|
| 441 |
+
```bash
|
| 442 |
+
curl -s http://localhost:8317/api/provider/openai/v1/chat/completions \
|
| 443 |
+
-H "Content-Type: application/json" \
|
| 444 |
+
-d '{
|
| 445 |
+
"model": "gpt-5",
|
| 446 |
+
"messages": [{"role": "user", "content": "Hello"}]
|
| 447 |
+
}'
|
| 448 |
+
```
|
| 449 |
+
|
| 450 |
+
**Management (localhost only by default):**
|
| 451 |
+
```bash
|
| 452 |
+
curl -s http://localhost:8317/api/user
|
| 453 |
+
```
|
| 454 |
+
|
| 455 |
+
## Troubleshooting
|
| 456 |
+
|
| 457 |
+
### Common errors and fixes
|
| 458 |
+
|
| 459 |
+
| Symptom/Code | Likely Cause | Fix |
|
| 460 |
+
|------------------------------------------|------------------------------------------------------|----------------------------------------------------------------------|
|
| 461 |
+
| 404 /v1/chat/completions | Factory not pointing to proxy base | Set base to `http://localhost:8317/v1` (env/flag/config). |
|
| 462 |
+
| 404 /api/provider/... | Incorrect route path or typo | Ensure you're calling `/api/provider/{provider}/v1...` paths exactly.|
|
| 463 |
+
| 403 on /api/user (Amp) | Management restricted to localhost | Run from same machine or set `amp-restrict-management-to-localhost: false` (not recommended). |
|
| 464 |
+
| 401/403 from provider | Missing/expired OAuth or API key | Re-run the relevant `--*-login` or configure keys in `config.yaml`. |
|
| 465 |
+
| 429/Quota exceeded | Project/model quota exhausted | Enable `quota-exceeded` switching or switch accounts. |
|
| 466 |
+
| 5xx from provider | Upstream transient error | Increase `request-retry` and try again. |
|
| 467 |
+
| SSE/stream stuck | Client not handling SSE properly | Use SSE-capable client or set `stream: false`. |
|
| 468 |
+
| Amp gzip decoding errors | Compressed upstream responses | Fork auto-decompresses; update to latest build if issue persists. |
|
| 469 |
+
| CORS errors in browser | Protected management endpoints | Use CLI/terminal; avoid browsers for management endpoints. |
|
| 470 |
+
| Wrong model name | Provider alias mismatch | Use `gpt-*` for OpenAI or `claude-*` for Anthropic models. |
|
| 471 |
+
|
| 472 |
+
### Diagnostics
|
| 473 |
+
|
| 474 |
+
- Check logs (`debug: true` temporarily or `logging-to-file: true`).
|
| 475 |
+
- Verify config in effect: print effective config or confirm with startup logs.
|
| 476 |
+
- Test base reachability: `curl http://localhost:8317/v1/models`.
|
| 477 |
+
- For Amp, verify `amp-upstream-url` and secrets resolution.
|
| 478 |
+
|
| 479 |
+
## Security Checklist
|
| 480 |
+
|
| 481 |
+
- Keep `amp-restrict-management-to-localhost: true` (default).
|
| 482 |
+
- Do not expose the proxy publicly; bind to localhost or protect with firewall/VPN.
|
| 483 |
+
- If enabling remote management, set `remote-management.secret-key` and TLS/ingress protections.
|
| 484 |
+
- Disable the built-in management UI if hosting your own:
|
| 485 |
+
- `remote-management.disable-control-panel: true`
|
| 486 |
+
- Rotate tokens/keys; store config and auth-dir on encrypted disk or managed secret stores.
|
| 487 |
+
- Keep binary up to date to receive security fixes.
|
| 488 |
+
|
| 489 |
+
## References
|
| 490 |
+
|
| 491 |
+
- This fork README: [README.md](README.md)
|
| 492 |
+
- Upstream project: [CLIProxyAPI](https://github.com/router-for-me/CLIProxyAPI)
|
| 493 |
+
- Amp CLI: [Official Manual](https://ampcode.com/manual)
|
| 494 |
+
- Factory CLI (droid): [Official Documentation](https://docs.factory.ai/cli/getting-started/overview)
|
|
@@ -137,6 +137,10 @@ func (MyExecutor) Execute(ctx context.Context, a *coreauth.Auth, req clipexec.Re
|
|
| 137 |
return clipexec.Response{Payload: body}, nil
|
| 138 |
}
|
| 139 |
|
|
|
|
|
|
|
|
|
|
|
|
|
| 140 |
func (MyExecutor) ExecuteStream(ctx context.Context, a *coreauth.Auth, req clipexec.Request, opts clipexec.Options) (<-chan clipexec.StreamChunk, error) {
|
| 141 |
ch := make(chan clipexec.StreamChunk, 1)
|
| 142 |
go func() {
|
|
|
|
| 137 |
return clipexec.Response{Payload: body}, nil
|
| 138 |
}
|
| 139 |
|
| 140 |
+
func (MyExecutor) CountTokens(context.Context, *coreauth.Auth, clipexec.Request, clipexec.Options) (clipexec.Response, error) {
|
| 141 |
+
return clipexec.Response{}, errors.New("count tokens not implemented")
|
| 142 |
+
}
|
| 143 |
+
|
| 144 |
func (MyExecutor) ExecuteStream(ctx context.Context, a *coreauth.Auth, req clipexec.Request, opts clipexec.Options) (<-chan clipexec.StreamChunk, error) {
|
| 145 |
ch := make(chan clipexec.StreamChunk, 1)
|
| 146 |
go func() {
|
|
@@ -0,0 +1,185 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
// Package amp implements the Amp CLI routing module, providing OAuth-based
|
| 2 |
+
// integration with Amp CLI for ChatGPT and Anthropic subscriptions.
|
| 3 |
+
package amp
|
| 4 |
+
|
| 5 |
+
import (
|
| 6 |
+
"fmt"
|
| 7 |
+
"net/http/httputil"
|
| 8 |
+
"strings"
|
| 9 |
+
"sync"
|
| 10 |
+
|
| 11 |
+
"github.com/gin-gonic/gin"
|
| 12 |
+
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules"
|
| 13 |
+
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
| 14 |
+
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
|
| 15 |
+
log "github.com/sirupsen/logrus"
|
| 16 |
+
)
|
| 17 |
+
|
| 18 |
+
// Option configures the AmpModule.
|
| 19 |
+
type Option func(*AmpModule)
|
| 20 |
+
|
| 21 |
+
// AmpModule implements the RouteModuleV2 interface for Amp CLI integration.
|
| 22 |
+
// It provides:
|
| 23 |
+
// - Reverse proxy to Amp control plane for OAuth/management
|
| 24 |
+
// - Provider-specific route aliases (/api/provider/{provider}/...)
|
| 25 |
+
// - Automatic gzip decompression for misconfigured upstreams
|
| 26 |
+
type AmpModule struct {
|
| 27 |
+
secretSource SecretSource
|
| 28 |
+
proxy *httputil.ReverseProxy
|
| 29 |
+
accessManager *sdkaccess.Manager
|
| 30 |
+
authMiddleware_ gin.HandlerFunc
|
| 31 |
+
enabled bool
|
| 32 |
+
registerOnce sync.Once
|
| 33 |
+
}
|
| 34 |
+
|
| 35 |
+
// New creates a new Amp routing module with the given options.
|
| 36 |
+
// This is the preferred constructor using the Option pattern.
|
| 37 |
+
//
|
| 38 |
+
// Example:
|
| 39 |
+
//
|
| 40 |
+
// ampModule := amp.New(
|
| 41 |
+
// amp.WithAccessManager(accessManager),
|
| 42 |
+
// amp.WithAuthMiddleware(authMiddleware),
|
| 43 |
+
// amp.WithSecretSource(customSecret),
|
| 44 |
+
// )
|
| 45 |
+
func New(opts ...Option) *AmpModule {
|
| 46 |
+
m := &AmpModule{
|
| 47 |
+
secretSource: nil, // Will be created on demand if not provided
|
| 48 |
+
}
|
| 49 |
+
for _, opt := range opts {
|
| 50 |
+
opt(m)
|
| 51 |
+
}
|
| 52 |
+
return m
|
| 53 |
+
}
|
| 54 |
+
|
| 55 |
+
// NewLegacy creates a new Amp routing module using the legacy constructor signature.
|
| 56 |
+
// This is provided for backwards compatibility.
|
| 57 |
+
//
|
| 58 |
+
// DEPRECATED: Use New with options instead.
|
| 59 |
+
func NewLegacy(accessManager *sdkaccess.Manager, authMiddleware gin.HandlerFunc) *AmpModule {
|
| 60 |
+
return New(
|
| 61 |
+
WithAccessManager(accessManager),
|
| 62 |
+
WithAuthMiddleware(authMiddleware),
|
| 63 |
+
)
|
| 64 |
+
}
|
| 65 |
+
|
| 66 |
+
// WithSecretSource sets a custom secret source for the module.
|
| 67 |
+
func WithSecretSource(source SecretSource) Option {
|
| 68 |
+
return func(m *AmpModule) {
|
| 69 |
+
m.secretSource = source
|
| 70 |
+
}
|
| 71 |
+
}
|
| 72 |
+
|
| 73 |
+
// WithAccessManager sets the access manager for the module.
|
| 74 |
+
func WithAccessManager(am *sdkaccess.Manager) Option {
|
| 75 |
+
return func(m *AmpModule) {
|
| 76 |
+
m.accessManager = am
|
| 77 |
+
}
|
| 78 |
+
}
|
| 79 |
+
|
| 80 |
+
// WithAuthMiddleware sets the authentication middleware for provider routes.
|
| 81 |
+
func WithAuthMiddleware(middleware gin.HandlerFunc) Option {
|
| 82 |
+
return func(m *AmpModule) {
|
| 83 |
+
m.authMiddleware_ = middleware
|
| 84 |
+
}
|
| 85 |
+
}
|
| 86 |
+
|
| 87 |
+
// Name returns the module identifier
|
| 88 |
+
func (m *AmpModule) Name() string {
|
| 89 |
+
return "amp-routing"
|
| 90 |
+
}
|
| 91 |
+
|
| 92 |
+
// Register sets up Amp routes if configured.
|
| 93 |
+
// This implements the RouteModuleV2 interface with Context.
|
| 94 |
+
// Routes are registered only once via sync.Once for idempotent behavior.
|
| 95 |
+
func (m *AmpModule) Register(ctx modules.Context) error {
|
| 96 |
+
upstreamURL := strings.TrimSpace(ctx.Config.AmpUpstreamURL)
|
| 97 |
+
|
| 98 |
+
// Determine auth middleware (from module or context)
|
| 99 |
+
auth := m.getAuthMiddleware(ctx)
|
| 100 |
+
|
| 101 |
+
// Use registerOnce to ensure routes are only registered once
|
| 102 |
+
var regErr error
|
| 103 |
+
m.registerOnce.Do(func() {
|
| 104 |
+
// Always register provider aliases - these work without an upstream
|
| 105 |
+
m.registerProviderAliases(ctx.Engine, ctx.BaseHandler, auth)
|
| 106 |
+
|
| 107 |
+
// If no upstream URL, skip proxy routes but provider aliases are still available
|
| 108 |
+
if upstreamURL == "" {
|
| 109 |
+
log.Debug("Amp upstream proxy disabled (no upstream URL configured)")
|
| 110 |
+
log.Debug("Amp provider alias routes registered")
|
| 111 |
+
m.enabled = false
|
| 112 |
+
return
|
| 113 |
+
}
|
| 114 |
+
|
| 115 |
+
// Create secret source with precedence: config > env > file
|
| 116 |
+
// Cache secrets for 5 minutes to reduce file I/O
|
| 117 |
+
if m.secretSource == nil {
|
| 118 |
+
m.secretSource = NewMultiSourceSecret(ctx.Config.AmpUpstreamAPIKey, 0 /* default 5min */)
|
| 119 |
+
}
|
| 120 |
+
|
| 121 |
+
// Create reverse proxy with gzip handling via ModifyResponse
|
| 122 |
+
proxy, err := createReverseProxy(upstreamURL, m.secretSource)
|
| 123 |
+
if err != nil {
|
| 124 |
+
regErr = fmt.Errorf("failed to create amp proxy: %w", err)
|
| 125 |
+
return
|
| 126 |
+
}
|
| 127 |
+
|
| 128 |
+
m.proxy = proxy
|
| 129 |
+
m.enabled = true
|
| 130 |
+
|
| 131 |
+
// Register management proxy routes (requires upstream)
|
| 132 |
+
// Restrict to localhost by default for security (prevents drive-by browser attacks)
|
| 133 |
+
handler := proxyHandler(proxy)
|
| 134 |
+
m.registerManagementRoutes(ctx.Engine, handler, ctx.Config.AmpRestrictManagementToLocalhost)
|
| 135 |
+
|
| 136 |
+
log.Infof("Amp upstream proxy enabled for: %s", upstreamURL)
|
| 137 |
+
log.Debug("Amp provider alias routes registered")
|
| 138 |
+
})
|
| 139 |
+
|
| 140 |
+
return regErr
|
| 141 |
+
}
|
| 142 |
+
|
| 143 |
+
// getAuthMiddleware returns the authentication middleware, preferring the
|
| 144 |
+
// module's configured middleware, then the context middleware, then a fallback.
|
| 145 |
+
func (m *AmpModule) getAuthMiddleware(ctx modules.Context) gin.HandlerFunc {
|
| 146 |
+
if m.authMiddleware_ != nil {
|
| 147 |
+
return m.authMiddleware_
|
| 148 |
+
}
|
| 149 |
+
if ctx.AuthMiddleware != nil {
|
| 150 |
+
return ctx.AuthMiddleware
|
| 151 |
+
}
|
| 152 |
+
// Fallback: no authentication (should not happen in production)
|
| 153 |
+
log.Warn("Amp module: no auth middleware provided, allowing all requests")
|
| 154 |
+
return func(c *gin.Context) {
|
| 155 |
+
c.Next()
|
| 156 |
+
}
|
| 157 |
+
}
|
| 158 |
+
|
| 159 |
+
// OnConfigUpdated handles configuration updates.
|
| 160 |
+
// Currently requires restart for URL changes (could be enhanced for dynamic updates).
|
| 161 |
+
func (m *AmpModule) OnConfigUpdated(cfg *config.Config) error {
|
| 162 |
+
if !m.enabled {
|
| 163 |
+
log.Debug("Amp routing not enabled, skipping config update")
|
| 164 |
+
return nil
|
| 165 |
+
}
|
| 166 |
+
|
| 167 |
+
upstreamURL := strings.TrimSpace(cfg.AmpUpstreamURL)
|
| 168 |
+
if upstreamURL == "" {
|
| 169 |
+
log.Warn("Amp upstream URL removed from config, restart required to disable")
|
| 170 |
+
return nil
|
| 171 |
+
}
|
| 172 |
+
|
| 173 |
+
// If API key changed, invalidate the cache
|
| 174 |
+
if m.secretSource != nil {
|
| 175 |
+
if ms, ok := m.secretSource.(*MultiSourceSecret); ok {
|
| 176 |
+
ms.InvalidateCache()
|
| 177 |
+
log.Debug("Amp secret cache invalidated due to config update")
|
| 178 |
+
}
|
| 179 |
+
}
|
| 180 |
+
|
| 181 |
+
log.Debug("Amp config updated (restart required for URL changes)")
|
| 182 |
+
return nil
|
| 183 |
+
}
|
| 184 |
+
|
| 185 |
+
|
|
@@ -0,0 +1,303 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
package amp
|
| 2 |
+
|
| 3 |
+
import (
|
| 4 |
+
"context"
|
| 5 |
+
"net/http/httptest"
|
| 6 |
+
"os"
|
| 7 |
+
"path/filepath"
|
| 8 |
+
"testing"
|
| 9 |
+
"time"
|
| 10 |
+
|
| 11 |
+
"github.com/gin-gonic/gin"
|
| 12 |
+
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules"
|
| 13 |
+
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
| 14 |
+
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
|
| 15 |
+
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
|
| 16 |
+
)
|
| 17 |
+
|
| 18 |
+
func TestAmpModule_Name(t *testing.T) {
|
| 19 |
+
m := New()
|
| 20 |
+
if m.Name() != "amp-routing" {
|
| 21 |
+
t.Fatalf("want amp-routing, got %s", m.Name())
|
| 22 |
+
}
|
| 23 |
+
}
|
| 24 |
+
|
| 25 |
+
func TestAmpModule_New(t *testing.T) {
|
| 26 |
+
accessManager := sdkaccess.NewManager()
|
| 27 |
+
authMiddleware := func(c *gin.Context) { c.Next() }
|
| 28 |
+
|
| 29 |
+
m := NewLegacy(accessManager, authMiddleware)
|
| 30 |
+
|
| 31 |
+
if m.accessManager != accessManager {
|
| 32 |
+
t.Fatal("accessManager not set")
|
| 33 |
+
}
|
| 34 |
+
if m.authMiddleware_ == nil {
|
| 35 |
+
t.Fatal("authMiddleware not set")
|
| 36 |
+
}
|
| 37 |
+
if m.enabled {
|
| 38 |
+
t.Fatal("enabled should be false initially")
|
| 39 |
+
}
|
| 40 |
+
if m.proxy != nil {
|
| 41 |
+
t.Fatal("proxy should be nil initially")
|
| 42 |
+
}
|
| 43 |
+
}
|
| 44 |
+
|
| 45 |
+
func TestAmpModule_Register_WithUpstream(t *testing.T) {
|
| 46 |
+
gin.SetMode(gin.TestMode)
|
| 47 |
+
r := gin.New()
|
| 48 |
+
|
| 49 |
+
// Fake upstream to ensure URL is valid
|
| 50 |
+
upstream := httptest.NewServer(nil)
|
| 51 |
+
defer upstream.Close()
|
| 52 |
+
|
| 53 |
+
accessManager := sdkaccess.NewManager()
|
| 54 |
+
base := &handlers.BaseAPIHandler{}
|
| 55 |
+
|
| 56 |
+
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
| 57 |
+
|
| 58 |
+
cfg := &config.Config{
|
| 59 |
+
AmpUpstreamURL: upstream.URL,
|
| 60 |
+
AmpUpstreamAPIKey: "test-key",
|
| 61 |
+
}
|
| 62 |
+
|
| 63 |
+
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
| 64 |
+
if err := m.Register(ctx); err != nil {
|
| 65 |
+
t.Fatalf("register error: %v", err)
|
| 66 |
+
}
|
| 67 |
+
|
| 68 |
+
if !m.enabled {
|
| 69 |
+
t.Fatal("module should be enabled with upstream URL")
|
| 70 |
+
}
|
| 71 |
+
if m.proxy == nil {
|
| 72 |
+
t.Fatal("proxy should be initialized")
|
| 73 |
+
}
|
| 74 |
+
if m.secretSource == nil {
|
| 75 |
+
t.Fatal("secretSource should be initialized")
|
| 76 |
+
}
|
| 77 |
+
}
|
| 78 |
+
|
| 79 |
+
func TestAmpModule_Register_WithoutUpstream(t *testing.T) {
|
| 80 |
+
gin.SetMode(gin.TestMode)
|
| 81 |
+
r := gin.New()
|
| 82 |
+
|
| 83 |
+
accessManager := sdkaccess.NewManager()
|
| 84 |
+
base := &handlers.BaseAPIHandler{}
|
| 85 |
+
|
| 86 |
+
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
| 87 |
+
|
| 88 |
+
cfg := &config.Config{
|
| 89 |
+
AmpUpstreamURL: "", // No upstream
|
| 90 |
+
}
|
| 91 |
+
|
| 92 |
+
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
| 93 |
+
if err := m.Register(ctx); err != nil {
|
| 94 |
+
t.Fatalf("register should not error without upstream: %v", err)
|
| 95 |
+
}
|
| 96 |
+
|
| 97 |
+
if m.enabled {
|
| 98 |
+
t.Fatal("module should be disabled without upstream URL")
|
| 99 |
+
}
|
| 100 |
+
if m.proxy != nil {
|
| 101 |
+
t.Fatal("proxy should not be initialized without upstream")
|
| 102 |
+
}
|
| 103 |
+
|
| 104 |
+
// But provider aliases should still be registered
|
| 105 |
+
req := httptest.NewRequest("GET", "/api/provider/openai/models", nil)
|
| 106 |
+
w := httptest.NewRecorder()
|
| 107 |
+
r.ServeHTTP(w, req)
|
| 108 |
+
|
| 109 |
+
if w.Code == 404 {
|
| 110 |
+
t.Fatal("provider aliases should be registered even without upstream")
|
| 111 |
+
}
|
| 112 |
+
}
|
| 113 |
+
|
| 114 |
+
func TestAmpModule_Register_InvalidUpstream(t *testing.T) {
|
| 115 |
+
gin.SetMode(gin.TestMode)
|
| 116 |
+
r := gin.New()
|
| 117 |
+
|
| 118 |
+
accessManager := sdkaccess.NewManager()
|
| 119 |
+
base := &handlers.BaseAPIHandler{}
|
| 120 |
+
|
| 121 |
+
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
| 122 |
+
|
| 123 |
+
cfg := &config.Config{
|
| 124 |
+
AmpUpstreamURL: "://invalid-url",
|
| 125 |
+
}
|
| 126 |
+
|
| 127 |
+
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
| 128 |
+
if err := m.Register(ctx); err == nil {
|
| 129 |
+
t.Fatal("expected error for invalid upstream URL")
|
| 130 |
+
}
|
| 131 |
+
}
|
| 132 |
+
|
| 133 |
+
func TestAmpModule_OnConfigUpdated_CacheInvalidation(t *testing.T) {
|
| 134 |
+
tmpDir := t.TempDir()
|
| 135 |
+
p := filepath.Join(tmpDir, "secrets.json")
|
| 136 |
+
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":"v1"}`), 0600); err != nil {
|
| 137 |
+
t.Fatal(err)
|
| 138 |
+
}
|
| 139 |
+
|
| 140 |
+
m := &AmpModule{enabled: true}
|
| 141 |
+
ms := NewMultiSourceSecretWithPath("", p, time.Minute)
|
| 142 |
+
m.secretSource = ms
|
| 143 |
+
|
| 144 |
+
// Warm the cache
|
| 145 |
+
if _, err := ms.Get(context.Background()); err != nil {
|
| 146 |
+
t.Fatal(err)
|
| 147 |
+
}
|
| 148 |
+
|
| 149 |
+
if ms.cache == nil {
|
| 150 |
+
t.Fatal("expected cache to be set")
|
| 151 |
+
}
|
| 152 |
+
|
| 153 |
+
// Update config - should invalidate cache
|
| 154 |
+
if err := m.OnConfigUpdated(&config.Config{AmpUpstreamURL: "http://x"}); err != nil {
|
| 155 |
+
t.Fatal(err)
|
| 156 |
+
}
|
| 157 |
+
|
| 158 |
+
if ms.cache != nil {
|
| 159 |
+
t.Fatal("expected cache to be invalidated")
|
| 160 |
+
}
|
| 161 |
+
}
|
| 162 |
+
|
| 163 |
+
func TestAmpModule_OnConfigUpdated_NotEnabled(t *testing.T) {
|
| 164 |
+
m := &AmpModule{enabled: false}
|
| 165 |
+
|
| 166 |
+
// Should not error or panic when disabled
|
| 167 |
+
if err := m.OnConfigUpdated(&config.Config{}); err != nil {
|
| 168 |
+
t.Fatalf("unexpected error: %v", err)
|
| 169 |
+
}
|
| 170 |
+
}
|
| 171 |
+
|
| 172 |
+
func TestAmpModule_OnConfigUpdated_URLRemoved(t *testing.T) {
|
| 173 |
+
m := &AmpModule{enabled: true}
|
| 174 |
+
ms := NewMultiSourceSecret("", 0)
|
| 175 |
+
m.secretSource = ms
|
| 176 |
+
|
| 177 |
+
// Config update with empty URL - should log warning but not error
|
| 178 |
+
cfg := &config.Config{AmpUpstreamURL: ""}
|
| 179 |
+
|
| 180 |
+
if err := m.OnConfigUpdated(cfg); err != nil {
|
| 181 |
+
t.Fatalf("unexpected error: %v", err)
|
| 182 |
+
}
|
| 183 |
+
}
|
| 184 |
+
|
| 185 |
+
func TestAmpModule_OnConfigUpdated_NonMultiSourceSecret(t *testing.T) {
|
| 186 |
+
// Test that OnConfigUpdated doesn't panic with StaticSecretSource
|
| 187 |
+
m := &AmpModule{enabled: true}
|
| 188 |
+
m.secretSource = NewStaticSecretSource("static-key")
|
| 189 |
+
|
| 190 |
+
cfg := &config.Config{AmpUpstreamURL: "http://example.com"}
|
| 191 |
+
|
| 192 |
+
// Should not error or panic
|
| 193 |
+
if err := m.OnConfigUpdated(cfg); err != nil {
|
| 194 |
+
t.Fatalf("unexpected error: %v", err)
|
| 195 |
+
}
|
| 196 |
+
}
|
| 197 |
+
|
| 198 |
+
func TestAmpModule_AuthMiddleware_Fallback(t *testing.T) {
|
| 199 |
+
gin.SetMode(gin.TestMode)
|
| 200 |
+
r := gin.New()
|
| 201 |
+
|
| 202 |
+
// Create module with no auth middleware
|
| 203 |
+
m := &AmpModule{authMiddleware_: nil}
|
| 204 |
+
|
| 205 |
+
// Get the fallback middleware via getAuthMiddleware
|
| 206 |
+
ctx := modules.Context{Engine: r, AuthMiddleware: nil}
|
| 207 |
+
middleware := m.getAuthMiddleware(ctx)
|
| 208 |
+
|
| 209 |
+
if middleware == nil {
|
| 210 |
+
t.Fatal("getAuthMiddleware should return a fallback, not nil")
|
| 211 |
+
}
|
| 212 |
+
|
| 213 |
+
// Test that it works
|
| 214 |
+
called := false
|
| 215 |
+
r.GET("/test", middleware, func(c *gin.Context) {
|
| 216 |
+
called = true
|
| 217 |
+
c.String(200, "ok")
|
| 218 |
+
})
|
| 219 |
+
|
| 220 |
+
req := httptest.NewRequest("GET", "/test", nil)
|
| 221 |
+
w := httptest.NewRecorder()
|
| 222 |
+
r.ServeHTTP(w, req)
|
| 223 |
+
|
| 224 |
+
if !called {
|
| 225 |
+
t.Fatal("fallback middleware should allow requests through")
|
| 226 |
+
}
|
| 227 |
+
}
|
| 228 |
+
|
| 229 |
+
func TestAmpModule_SecretSource_FromConfig(t *testing.T) {
|
| 230 |
+
gin.SetMode(gin.TestMode)
|
| 231 |
+
r := gin.New()
|
| 232 |
+
|
| 233 |
+
upstream := httptest.NewServer(nil)
|
| 234 |
+
defer upstream.Close()
|
| 235 |
+
|
| 236 |
+
accessManager := sdkaccess.NewManager()
|
| 237 |
+
base := &handlers.BaseAPIHandler{}
|
| 238 |
+
|
| 239 |
+
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
| 240 |
+
|
| 241 |
+
// Config with explicit API key
|
| 242 |
+
cfg := &config.Config{
|
| 243 |
+
AmpUpstreamURL: upstream.URL,
|
| 244 |
+
AmpUpstreamAPIKey: "config-key",
|
| 245 |
+
}
|
| 246 |
+
|
| 247 |
+
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
| 248 |
+
if err := m.Register(ctx); err != nil {
|
| 249 |
+
t.Fatalf("register error: %v", err)
|
| 250 |
+
}
|
| 251 |
+
|
| 252 |
+
// Secret source should be MultiSourceSecret with config key
|
| 253 |
+
if m.secretSource == nil {
|
| 254 |
+
t.Fatal("secretSource should be set")
|
| 255 |
+
}
|
| 256 |
+
|
| 257 |
+
// Verify it returns the config key
|
| 258 |
+
key, err := m.secretSource.Get(context.Background())
|
| 259 |
+
if err != nil {
|
| 260 |
+
t.Fatalf("Get error: %v", err)
|
| 261 |
+
}
|
| 262 |
+
if key != "config-key" {
|
| 263 |
+
t.Fatalf("want config-key, got %s", key)
|
| 264 |
+
}
|
| 265 |
+
}
|
| 266 |
+
|
| 267 |
+
func TestAmpModule_ProviderAliasesAlwaysRegistered(t *testing.T) {
|
| 268 |
+
gin.SetMode(gin.TestMode)
|
| 269 |
+
|
| 270 |
+
scenarios := []struct {
|
| 271 |
+
name string
|
| 272 |
+
configURL string
|
| 273 |
+
}{
|
| 274 |
+
{"with_upstream", "http://example.com"},
|
| 275 |
+
{"without_upstream", ""},
|
| 276 |
+
}
|
| 277 |
+
|
| 278 |
+
for _, scenario := range scenarios {
|
| 279 |
+
t.Run(scenario.name, func(t *testing.T) {
|
| 280 |
+
r := gin.New()
|
| 281 |
+
accessManager := sdkaccess.NewManager()
|
| 282 |
+
base := &handlers.BaseAPIHandler{}
|
| 283 |
+
|
| 284 |
+
m := NewLegacy(accessManager, func(c *gin.Context) { c.Next() })
|
| 285 |
+
|
| 286 |
+
cfg := &config.Config{AmpUpstreamURL: scenario.configURL}
|
| 287 |
+
|
| 288 |
+
ctx := modules.Context{Engine: r, BaseHandler: base, Config: cfg, AuthMiddleware: func(c *gin.Context) { c.Next() }}
|
| 289 |
+
if err := m.Register(ctx); err != nil && scenario.configURL != "" {
|
| 290 |
+
t.Fatalf("register error: %v", err)
|
| 291 |
+
}
|
| 292 |
+
|
| 293 |
+
// Provider aliases should always be available
|
| 294 |
+
req := httptest.NewRequest("GET", "/api/provider/openai/models", nil)
|
| 295 |
+
w := httptest.NewRecorder()
|
| 296 |
+
r.ServeHTTP(w, req)
|
| 297 |
+
|
| 298 |
+
if w.Code == 404 {
|
| 299 |
+
t.Fatal("provider aliases should be registered")
|
| 300 |
+
}
|
| 301 |
+
})
|
| 302 |
+
}
|
| 303 |
+
}
|
|
@@ -0,0 +1,176 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
package amp
|
| 2 |
+
|
| 3 |
+
import (
|
| 4 |
+
"bytes"
|
| 5 |
+
"compress/gzip"
|
| 6 |
+
"fmt"
|
| 7 |
+
"io"
|
| 8 |
+
"net/http"
|
| 9 |
+
"net/http/httputil"
|
| 10 |
+
"net/url"
|
| 11 |
+
"strconv"
|
| 12 |
+
"strings"
|
| 13 |
+
|
| 14 |
+
"github.com/gin-gonic/gin"
|
| 15 |
+
log "github.com/sirupsen/logrus"
|
| 16 |
+
)
|
| 17 |
+
|
| 18 |
+
// readCloser wraps a reader and forwards Close to a separate closer.
|
| 19 |
+
// Used to restore peeked bytes while preserving upstream body Close behavior.
|
| 20 |
+
type readCloser struct {
|
| 21 |
+
r io.Reader
|
| 22 |
+
c io.Closer
|
| 23 |
+
}
|
| 24 |
+
|
| 25 |
+
func (rc *readCloser) Read(p []byte) (int, error) { return rc.r.Read(p) }
|
| 26 |
+
func (rc *readCloser) Close() error { return rc.c.Close() }
|
| 27 |
+
|
| 28 |
+
// createReverseProxy creates a reverse proxy handler for Amp upstream
|
| 29 |
+
// with automatic gzip decompression via ModifyResponse
|
| 30 |
+
func createReverseProxy(upstreamURL string, secretSource SecretSource) (*httputil.ReverseProxy, error) {
|
| 31 |
+
parsed, err := url.Parse(upstreamURL)
|
| 32 |
+
if err != nil {
|
| 33 |
+
return nil, fmt.Errorf("invalid amp upstream url: %w", err)
|
| 34 |
+
}
|
| 35 |
+
|
| 36 |
+
proxy := httputil.NewSingleHostReverseProxy(parsed)
|
| 37 |
+
originalDirector := proxy.Director
|
| 38 |
+
|
| 39 |
+
// Modify outgoing requests to inject API key and fix routing
|
| 40 |
+
proxy.Director = func(req *http.Request) {
|
| 41 |
+
originalDirector(req)
|
| 42 |
+
req.Host = parsed.Host
|
| 43 |
+
|
| 44 |
+
// Preserve correlation headers for debugging
|
| 45 |
+
if req.Header.Get("X-Request-ID") == "" {
|
| 46 |
+
// Could generate one here if needed
|
| 47 |
+
}
|
| 48 |
+
|
| 49 |
+
// Inject API key from secret source (precedence: config > env > file)
|
| 50 |
+
if key, err := secretSource.Get(req.Context()); err == nil && key != "" {
|
| 51 |
+
req.Header.Set("X-Api-Key", key)
|
| 52 |
+
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", key))
|
| 53 |
+
} else if err != nil {
|
| 54 |
+
log.Warnf("amp secret source error (continuing without auth): %v", err)
|
| 55 |
+
}
|
| 56 |
+
}
|
| 57 |
+
|
| 58 |
+
// Modify incoming responses to handle gzip without Content-Encoding
|
| 59 |
+
// This addresses the same issue as inline handler gzip handling, but at the proxy level
|
| 60 |
+
proxy.ModifyResponse = func(resp *http.Response) error {
|
| 61 |
+
// Only process successful responses
|
| 62 |
+
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
| 63 |
+
return nil
|
| 64 |
+
}
|
| 65 |
+
|
| 66 |
+
// Skip if already marked as gzip (Content-Encoding set)
|
| 67 |
+
if resp.Header.Get("Content-Encoding") != "" {
|
| 68 |
+
return nil
|
| 69 |
+
}
|
| 70 |
+
|
| 71 |
+
// Skip streaming responses (SSE, chunked)
|
| 72 |
+
if isStreamingResponse(resp) {
|
| 73 |
+
return nil
|
| 74 |
+
}
|
| 75 |
+
|
| 76 |
+
// Save reference to original upstream body for proper cleanup
|
| 77 |
+
originalBody := resp.Body
|
| 78 |
+
|
| 79 |
+
// Peek at first 2 bytes to detect gzip magic bytes
|
| 80 |
+
header := make([]byte, 2)
|
| 81 |
+
n, _ := io.ReadFull(originalBody, header)
|
| 82 |
+
|
| 83 |
+
// Check for gzip magic bytes (0x1f 0x8b)
|
| 84 |
+
// If n < 2, we didn't get enough bytes, so it's not gzip
|
| 85 |
+
if n >= 2 && header[0] == 0x1f && header[1] == 0x8b {
|
| 86 |
+
// It's gzip - read the rest of the body
|
| 87 |
+
rest, err := io.ReadAll(originalBody)
|
| 88 |
+
if err != nil {
|
| 89 |
+
// Restore what we read and return original body (preserve Close behavior)
|
| 90 |
+
resp.Body = &readCloser{
|
| 91 |
+
r: io.MultiReader(bytes.NewReader(header[:n]), originalBody),
|
| 92 |
+
c: originalBody,
|
| 93 |
+
}
|
| 94 |
+
return nil
|
| 95 |
+
}
|
| 96 |
+
|
| 97 |
+
// Reconstruct complete gzipped data
|
| 98 |
+
gzippedData := append(header[:n], rest...)
|
| 99 |
+
|
| 100 |
+
// Decompress
|
| 101 |
+
gzipReader, err := gzip.NewReader(bytes.NewReader(gzippedData))
|
| 102 |
+
if err != nil {
|
| 103 |
+
log.Warnf("amp proxy: gzip header detected but decompress failed: %v", err)
|
| 104 |
+
// Close original body and return in-memory copy
|
| 105 |
+
_ = originalBody.Close()
|
| 106 |
+
resp.Body = io.NopCloser(bytes.NewReader(gzippedData))
|
| 107 |
+
return nil
|
| 108 |
+
}
|
| 109 |
+
|
| 110 |
+
decompressed, err := io.ReadAll(gzipReader)
|
| 111 |
+
_ = gzipReader.Close()
|
| 112 |
+
if err != nil {
|
| 113 |
+
log.Warnf("amp proxy: gzip decompress error: %v", err)
|
| 114 |
+
// Close original body and return in-memory copy
|
| 115 |
+
_ = originalBody.Close()
|
| 116 |
+
resp.Body = io.NopCloser(bytes.NewReader(gzippedData))
|
| 117 |
+
return nil
|
| 118 |
+
}
|
| 119 |
+
|
| 120 |
+
// Close original body since we're replacing with in-memory decompressed content
|
| 121 |
+
_ = originalBody.Close()
|
| 122 |
+
|
| 123 |
+
// Replace body with decompressed content
|
| 124 |
+
resp.Body = io.NopCloser(bytes.NewReader(decompressed))
|
| 125 |
+
resp.ContentLength = int64(len(decompressed))
|
| 126 |
+
|
| 127 |
+
// Update headers to reflect decompressed state
|
| 128 |
+
resp.Header.Del("Content-Encoding") // No longer compressed
|
| 129 |
+
resp.Header.Del("Content-Length") // Remove stale compressed length
|
| 130 |
+
resp.Header.Set("Content-Length", strconv.FormatInt(resp.ContentLength, 10)) // Set decompressed length
|
| 131 |
+
|
| 132 |
+
log.Debugf("amp proxy: decompressed gzip response (%d -> %d bytes)", len(gzippedData), len(decompressed))
|
| 133 |
+
} else {
|
| 134 |
+
// Not gzip - restore peeked bytes while preserving Close behavior
|
| 135 |
+
// Handle edge cases: n might be 0, 1, or 2 depending on EOF
|
| 136 |
+
resp.Body = &readCloser{
|
| 137 |
+
r: io.MultiReader(bytes.NewReader(header[:n]), originalBody),
|
| 138 |
+
c: originalBody,
|
| 139 |
+
}
|
| 140 |
+
}
|
| 141 |
+
|
| 142 |
+
return nil
|
| 143 |
+
}
|
| 144 |
+
|
| 145 |
+
// Error handler for proxy failures
|
| 146 |
+
proxy.ErrorHandler = func(rw http.ResponseWriter, req *http.Request, err error) {
|
| 147 |
+
log.Errorf("amp upstream proxy error for %s %s: %v", req.Method, req.URL.Path, err)
|
| 148 |
+
rw.Header().Set("Content-Type", "application/json")
|
| 149 |
+
rw.WriteHeader(http.StatusBadGateway)
|
| 150 |
+
_, _ = rw.Write([]byte(`{"error":"amp_upstream_proxy_error","message":"Failed to reach Amp upstream"}`))
|
| 151 |
+
}
|
| 152 |
+
|
| 153 |
+
return proxy, nil
|
| 154 |
+
}
|
| 155 |
+
|
| 156 |
+
// isStreamingResponse detects if the response is streaming (SSE only)
|
| 157 |
+
// Note: We only treat text/event-stream as streaming. Chunked transfer encoding
|
| 158 |
+
// is a transport-level detail and doesn't mean we can't decompress the full response.
|
| 159 |
+
// Many JSON APIs use chunked encoding for normal responses.
|
| 160 |
+
func isStreamingResponse(resp *http.Response) bool {
|
| 161 |
+
contentType := resp.Header.Get("Content-Type")
|
| 162 |
+
|
| 163 |
+
// Only Server-Sent Events are true streaming responses
|
| 164 |
+
if strings.Contains(contentType, "text/event-stream") {
|
| 165 |
+
return true
|
| 166 |
+
}
|
| 167 |
+
|
| 168 |
+
return false
|
| 169 |
+
}
|
| 170 |
+
|
| 171 |
+
// proxyHandler converts httputil.ReverseProxy to gin.HandlerFunc
|
| 172 |
+
func proxyHandler(proxy *httputil.ReverseProxy) gin.HandlerFunc {
|
| 173 |
+
return func(c *gin.Context) {
|
| 174 |
+
proxy.ServeHTTP(c.Writer, c.Request)
|
| 175 |
+
}
|
| 176 |
+
}
|
|
@@ -0,0 +1,439 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
package amp
|
| 2 |
+
|
| 3 |
+
import (
|
| 4 |
+
"bytes"
|
| 5 |
+
"compress/gzip"
|
| 6 |
+
"fmt"
|
| 7 |
+
"io"
|
| 8 |
+
"net/http"
|
| 9 |
+
"net/http/httptest"
|
| 10 |
+
"testing"
|
| 11 |
+
)
|
| 12 |
+
|
| 13 |
+
// Helper: compress data with gzip
|
| 14 |
+
func gzipBytes(b []byte) []byte {
|
| 15 |
+
var buf bytes.Buffer
|
| 16 |
+
zw := gzip.NewWriter(&buf)
|
| 17 |
+
zw.Write(b)
|
| 18 |
+
zw.Close()
|
| 19 |
+
return buf.Bytes()
|
| 20 |
+
}
|
| 21 |
+
|
| 22 |
+
// Helper: create a mock http.Response
|
| 23 |
+
func mkResp(status int, hdr http.Header, body []byte) *http.Response {
|
| 24 |
+
if hdr == nil {
|
| 25 |
+
hdr = http.Header{}
|
| 26 |
+
}
|
| 27 |
+
return &http.Response{
|
| 28 |
+
StatusCode: status,
|
| 29 |
+
Header: hdr,
|
| 30 |
+
Body: io.NopCloser(bytes.NewReader(body)),
|
| 31 |
+
ContentLength: int64(len(body)),
|
| 32 |
+
}
|
| 33 |
+
}
|
| 34 |
+
|
| 35 |
+
func TestCreateReverseProxy_ValidURL(t *testing.T) {
|
| 36 |
+
proxy, err := createReverseProxy("http://example.com", NewStaticSecretSource("key"))
|
| 37 |
+
if err != nil {
|
| 38 |
+
t.Fatalf("expected no error, got: %v", err)
|
| 39 |
+
}
|
| 40 |
+
if proxy == nil {
|
| 41 |
+
t.Fatal("expected proxy to be created")
|
| 42 |
+
}
|
| 43 |
+
}
|
| 44 |
+
|
| 45 |
+
func TestCreateReverseProxy_InvalidURL(t *testing.T) {
|
| 46 |
+
_, err := createReverseProxy("://invalid", NewStaticSecretSource("key"))
|
| 47 |
+
if err == nil {
|
| 48 |
+
t.Fatal("expected error for invalid URL")
|
| 49 |
+
}
|
| 50 |
+
}
|
| 51 |
+
|
| 52 |
+
func TestModifyResponse_GzipScenarios(t *testing.T) {
|
| 53 |
+
proxy, err := createReverseProxy("http://example.com", NewStaticSecretSource("k"))
|
| 54 |
+
if err != nil {
|
| 55 |
+
t.Fatal(err)
|
| 56 |
+
}
|
| 57 |
+
|
| 58 |
+
goodJSON := []byte(`{"ok":true}`)
|
| 59 |
+
good := gzipBytes(goodJSON)
|
| 60 |
+
truncated := good[:10]
|
| 61 |
+
corrupted := append([]byte{0x1f, 0x8b}, []byte("notgzip")...)
|
| 62 |
+
|
| 63 |
+
cases := []struct {
|
| 64 |
+
name string
|
| 65 |
+
header http.Header
|
| 66 |
+
body []byte
|
| 67 |
+
status int
|
| 68 |
+
wantBody []byte
|
| 69 |
+
wantCE string
|
| 70 |
+
}{
|
| 71 |
+
{
|
| 72 |
+
name: "decompresses_valid_gzip_no_header",
|
| 73 |
+
header: http.Header{},
|
| 74 |
+
body: good,
|
| 75 |
+
status: 200,
|
| 76 |
+
wantBody: goodJSON,
|
| 77 |
+
wantCE: "",
|
| 78 |
+
},
|
| 79 |
+
{
|
| 80 |
+
name: "skips_when_ce_present",
|
| 81 |
+
header: http.Header{"Content-Encoding": []string{"gzip"}},
|
| 82 |
+
body: good,
|
| 83 |
+
status: 200,
|
| 84 |
+
wantBody: good,
|
| 85 |
+
wantCE: "gzip",
|
| 86 |
+
},
|
| 87 |
+
{
|
| 88 |
+
name: "passes_truncated_unchanged",
|
| 89 |
+
header: http.Header{},
|
| 90 |
+
body: truncated,
|
| 91 |
+
status: 200,
|
| 92 |
+
wantBody: truncated,
|
| 93 |
+
wantCE: "",
|
| 94 |
+
},
|
| 95 |
+
{
|
| 96 |
+
name: "passes_corrupted_unchanged",
|
| 97 |
+
header: http.Header{},
|
| 98 |
+
body: corrupted,
|
| 99 |
+
status: 200,
|
| 100 |
+
wantBody: corrupted,
|
| 101 |
+
wantCE: "",
|
| 102 |
+
},
|
| 103 |
+
{
|
| 104 |
+
name: "non_gzip_unchanged",
|
| 105 |
+
header: http.Header{},
|
| 106 |
+
body: []byte("plain"),
|
| 107 |
+
status: 200,
|
| 108 |
+
wantBody: []byte("plain"),
|
| 109 |
+
wantCE: "",
|
| 110 |
+
},
|
| 111 |
+
{
|
| 112 |
+
name: "empty_body",
|
| 113 |
+
header: http.Header{},
|
| 114 |
+
body: []byte{},
|
| 115 |
+
status: 200,
|
| 116 |
+
wantBody: []byte{},
|
| 117 |
+
wantCE: "",
|
| 118 |
+
},
|
| 119 |
+
{
|
| 120 |
+
name: "single_byte_body",
|
| 121 |
+
header: http.Header{},
|
| 122 |
+
body: []byte{0x1f},
|
| 123 |
+
status: 200,
|
| 124 |
+
wantBody: []byte{0x1f},
|
| 125 |
+
wantCE: "",
|
| 126 |
+
},
|
| 127 |
+
{
|
| 128 |
+
name: "skips_non_2xx_status",
|
| 129 |
+
header: http.Header{},
|
| 130 |
+
body: good,
|
| 131 |
+
status: 404,
|
| 132 |
+
wantBody: good,
|
| 133 |
+
wantCE: "",
|
| 134 |
+
},
|
| 135 |
+
}
|
| 136 |
+
|
| 137 |
+
for _, tc := range cases {
|
| 138 |
+
t.Run(tc.name, func(t *testing.T) {
|
| 139 |
+
resp := mkResp(tc.status, tc.header, tc.body)
|
| 140 |
+
if err := proxy.ModifyResponse(resp); err != nil {
|
| 141 |
+
t.Fatalf("ModifyResponse error: %v", err)
|
| 142 |
+
}
|
| 143 |
+
got, err := io.ReadAll(resp.Body)
|
| 144 |
+
if err != nil {
|
| 145 |
+
t.Fatalf("ReadAll error: %v", err)
|
| 146 |
+
}
|
| 147 |
+
if !bytes.Equal(got, tc.wantBody) {
|
| 148 |
+
t.Fatalf("body mismatch:\nwant: %q\ngot: %q", tc.wantBody, got)
|
| 149 |
+
}
|
| 150 |
+
if ce := resp.Header.Get("Content-Encoding"); ce != tc.wantCE {
|
| 151 |
+
t.Fatalf("Content-Encoding: want %q, got %q", tc.wantCE, ce)
|
| 152 |
+
}
|
| 153 |
+
})
|
| 154 |
+
}
|
| 155 |
+
}
|
| 156 |
+
|
| 157 |
+
func TestModifyResponse_UpdatesContentLengthHeader(t *testing.T) {
|
| 158 |
+
proxy, err := createReverseProxy("http://example.com", NewStaticSecretSource("k"))
|
| 159 |
+
if err != nil {
|
| 160 |
+
t.Fatal(err)
|
| 161 |
+
}
|
| 162 |
+
|
| 163 |
+
goodJSON := []byte(`{"message":"test response"}`)
|
| 164 |
+
gzipped := gzipBytes(goodJSON)
|
| 165 |
+
|
| 166 |
+
// Simulate upstream response with gzip body AND Content-Length header
|
| 167 |
+
// (this is the scenario the bot flagged - stale Content-Length after decompression)
|
| 168 |
+
resp := mkResp(200, http.Header{
|
| 169 |
+
"Content-Length": []string{fmt.Sprintf("%d", len(gzipped))}, // Compressed size
|
| 170 |
+
}, gzipped)
|
| 171 |
+
|
| 172 |
+
if err := proxy.ModifyResponse(resp); err != nil {
|
| 173 |
+
t.Fatalf("ModifyResponse error: %v", err)
|
| 174 |
+
}
|
| 175 |
+
|
| 176 |
+
// Verify body is decompressed
|
| 177 |
+
got, _ := io.ReadAll(resp.Body)
|
| 178 |
+
if !bytes.Equal(got, goodJSON) {
|
| 179 |
+
t.Fatalf("body should be decompressed, got: %q, want: %q", got, goodJSON)
|
| 180 |
+
}
|
| 181 |
+
|
| 182 |
+
// Verify Content-Length header is updated to decompressed size
|
| 183 |
+
wantCL := fmt.Sprintf("%d", len(goodJSON))
|
| 184 |
+
gotCL := resp.Header.Get("Content-Length")
|
| 185 |
+
if gotCL != wantCL {
|
| 186 |
+
t.Fatalf("Content-Length header mismatch: want %q (decompressed), got %q", wantCL, gotCL)
|
| 187 |
+
}
|
| 188 |
+
|
| 189 |
+
// Verify struct field also matches
|
| 190 |
+
if resp.ContentLength != int64(len(goodJSON)) {
|
| 191 |
+
t.Fatalf("resp.ContentLength mismatch: want %d, got %d", len(goodJSON), resp.ContentLength)
|
| 192 |
+
}
|
| 193 |
+
}
|
| 194 |
+
|
| 195 |
+
func TestModifyResponse_SkipsStreamingResponses(t *testing.T) {
|
| 196 |
+
proxy, err := createReverseProxy("http://example.com", NewStaticSecretSource("k"))
|
| 197 |
+
if err != nil {
|
| 198 |
+
t.Fatal(err)
|
| 199 |
+
}
|
| 200 |
+
|
| 201 |
+
goodJSON := []byte(`{"ok":true}`)
|
| 202 |
+
gzipped := gzipBytes(goodJSON)
|
| 203 |
+
|
| 204 |
+
t.Run("sse_skips_decompression", func(t *testing.T) {
|
| 205 |
+
resp := mkResp(200, http.Header{"Content-Type": []string{"text/event-stream"}}, gzipped)
|
| 206 |
+
if err := proxy.ModifyResponse(resp); err != nil {
|
| 207 |
+
t.Fatalf("ModifyResponse error: %v", err)
|
| 208 |
+
}
|
| 209 |
+
// SSE should NOT be decompressed
|
| 210 |
+
got, _ := io.ReadAll(resp.Body)
|
| 211 |
+
if !bytes.Equal(got, gzipped) {
|
| 212 |
+
t.Fatal("SSE response should not be decompressed")
|
| 213 |
+
}
|
| 214 |
+
})
|
| 215 |
+
}
|
| 216 |
+
|
| 217 |
+
func TestModifyResponse_DecompressesChunkedJSON(t *testing.T) {
|
| 218 |
+
proxy, err := createReverseProxy("http://example.com", NewStaticSecretSource("k"))
|
| 219 |
+
if err != nil {
|
| 220 |
+
t.Fatal(err)
|
| 221 |
+
}
|
| 222 |
+
|
| 223 |
+
goodJSON := []byte(`{"ok":true}`)
|
| 224 |
+
gzipped := gzipBytes(goodJSON)
|
| 225 |
+
|
| 226 |
+
t.Run("chunked_json_decompresses", func(t *testing.T) {
|
| 227 |
+
// Chunked JSON responses (like thread APIs) should be decompressed
|
| 228 |
+
resp := mkResp(200, http.Header{"Transfer-Encoding": []string{"chunked"}}, gzipped)
|
| 229 |
+
if err := proxy.ModifyResponse(resp); err != nil {
|
| 230 |
+
t.Fatalf("ModifyResponse error: %v", err)
|
| 231 |
+
}
|
| 232 |
+
// Should decompress because it's not SSE
|
| 233 |
+
got, _ := io.ReadAll(resp.Body)
|
| 234 |
+
if !bytes.Equal(got, goodJSON) {
|
| 235 |
+
t.Fatalf("chunked JSON should be decompressed, got: %q, want: %q", got, goodJSON)
|
| 236 |
+
}
|
| 237 |
+
})
|
| 238 |
+
}
|
| 239 |
+
|
| 240 |
+
func TestReverseProxy_InjectsHeaders(t *testing.T) {
|
| 241 |
+
gotHeaders := make(chan http.Header, 1)
|
| 242 |
+
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
| 243 |
+
gotHeaders <- r.Header.Clone()
|
| 244 |
+
w.WriteHeader(200)
|
| 245 |
+
w.Write([]byte(`ok`))
|
| 246 |
+
}))
|
| 247 |
+
defer upstream.Close()
|
| 248 |
+
|
| 249 |
+
proxy, err := createReverseProxy(upstream.URL, NewStaticSecretSource("secret"))
|
| 250 |
+
if err != nil {
|
| 251 |
+
t.Fatal(err)
|
| 252 |
+
}
|
| 253 |
+
|
| 254 |
+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
| 255 |
+
proxy.ServeHTTP(w, r)
|
| 256 |
+
}))
|
| 257 |
+
defer srv.Close()
|
| 258 |
+
|
| 259 |
+
res, err := http.Get(srv.URL + "/test")
|
| 260 |
+
if err != nil {
|
| 261 |
+
t.Fatal(err)
|
| 262 |
+
}
|
| 263 |
+
res.Body.Close()
|
| 264 |
+
|
| 265 |
+
hdr := <-gotHeaders
|
| 266 |
+
if hdr.Get("X-Api-Key") != "secret" {
|
| 267 |
+
t.Fatalf("X-Api-Key missing or wrong, got: %q", hdr.Get("X-Api-Key"))
|
| 268 |
+
}
|
| 269 |
+
if hdr.Get("Authorization") != "Bearer secret" {
|
| 270 |
+
t.Fatalf("Authorization missing or wrong, got: %q", hdr.Get("Authorization"))
|
| 271 |
+
}
|
| 272 |
+
}
|
| 273 |
+
|
| 274 |
+
func TestReverseProxy_EmptySecret(t *testing.T) {
|
| 275 |
+
gotHeaders := make(chan http.Header, 1)
|
| 276 |
+
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
| 277 |
+
gotHeaders <- r.Header.Clone()
|
| 278 |
+
w.WriteHeader(200)
|
| 279 |
+
w.Write([]byte(`ok`))
|
| 280 |
+
}))
|
| 281 |
+
defer upstream.Close()
|
| 282 |
+
|
| 283 |
+
proxy, err := createReverseProxy(upstream.URL, NewStaticSecretSource(""))
|
| 284 |
+
if err != nil {
|
| 285 |
+
t.Fatal(err)
|
| 286 |
+
}
|
| 287 |
+
|
| 288 |
+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
| 289 |
+
proxy.ServeHTTP(w, r)
|
| 290 |
+
}))
|
| 291 |
+
defer srv.Close()
|
| 292 |
+
|
| 293 |
+
res, err := http.Get(srv.URL + "/test")
|
| 294 |
+
if err != nil {
|
| 295 |
+
t.Fatal(err)
|
| 296 |
+
}
|
| 297 |
+
res.Body.Close()
|
| 298 |
+
|
| 299 |
+
hdr := <-gotHeaders
|
| 300 |
+
// Should NOT inject headers when secret is empty
|
| 301 |
+
if hdr.Get("X-Api-Key") != "" {
|
| 302 |
+
t.Fatalf("X-Api-Key should not be set, got: %q", hdr.Get("X-Api-Key"))
|
| 303 |
+
}
|
| 304 |
+
if authVal := hdr.Get("Authorization"); authVal != "" && authVal != "Bearer " {
|
| 305 |
+
t.Fatalf("Authorization should not be set, got: %q", authVal)
|
| 306 |
+
}
|
| 307 |
+
}
|
| 308 |
+
|
| 309 |
+
func TestReverseProxy_ErrorHandler(t *testing.T) {
|
| 310 |
+
// Point proxy to a non-routable address to trigger error
|
| 311 |
+
proxy, err := createReverseProxy("http://127.0.0.1:1", NewStaticSecretSource(""))
|
| 312 |
+
if err != nil {
|
| 313 |
+
t.Fatal(err)
|
| 314 |
+
}
|
| 315 |
+
|
| 316 |
+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
| 317 |
+
proxy.ServeHTTP(w, r)
|
| 318 |
+
}))
|
| 319 |
+
defer srv.Close()
|
| 320 |
+
|
| 321 |
+
res, err := http.Get(srv.URL + "/any")
|
| 322 |
+
if err != nil {
|
| 323 |
+
t.Fatal(err)
|
| 324 |
+
}
|
| 325 |
+
body, _ := io.ReadAll(res.Body)
|
| 326 |
+
res.Body.Close()
|
| 327 |
+
|
| 328 |
+
if res.StatusCode != http.StatusBadGateway {
|
| 329 |
+
t.Fatalf("want 502, got %d", res.StatusCode)
|
| 330 |
+
}
|
| 331 |
+
if !bytes.Contains(body, []byte(`"amp_upstream_proxy_error"`)) {
|
| 332 |
+
t.Fatalf("unexpected body: %s", body)
|
| 333 |
+
}
|
| 334 |
+
if ct := res.Header.Get("Content-Type"); ct != "application/json" {
|
| 335 |
+
t.Fatalf("content-type: want application/json, got %s", ct)
|
| 336 |
+
}
|
| 337 |
+
}
|
| 338 |
+
|
| 339 |
+
func TestReverseProxy_FullRoundTrip_Gzip(t *testing.T) {
|
| 340 |
+
// Upstream returns gzipped JSON without Content-Encoding header
|
| 341 |
+
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
| 342 |
+
w.WriteHeader(200)
|
| 343 |
+
w.Write(gzipBytes([]byte(`{"upstream":"ok"}`)))
|
| 344 |
+
}))
|
| 345 |
+
defer upstream.Close()
|
| 346 |
+
|
| 347 |
+
proxy, err := createReverseProxy(upstream.URL, NewStaticSecretSource("key"))
|
| 348 |
+
if err != nil {
|
| 349 |
+
t.Fatal(err)
|
| 350 |
+
}
|
| 351 |
+
|
| 352 |
+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
| 353 |
+
proxy.ServeHTTP(w, r)
|
| 354 |
+
}))
|
| 355 |
+
defer srv.Close()
|
| 356 |
+
|
| 357 |
+
res, err := http.Get(srv.URL + "/test")
|
| 358 |
+
if err != nil {
|
| 359 |
+
t.Fatal(err)
|
| 360 |
+
}
|
| 361 |
+
body, _ := io.ReadAll(res.Body)
|
| 362 |
+
res.Body.Close()
|
| 363 |
+
|
| 364 |
+
expected := []byte(`{"upstream":"ok"}`)
|
| 365 |
+
if !bytes.Equal(body, expected) {
|
| 366 |
+
t.Fatalf("want decompressed JSON, got: %s", body)
|
| 367 |
+
}
|
| 368 |
+
}
|
| 369 |
+
|
| 370 |
+
func TestReverseProxy_FullRoundTrip_PlainJSON(t *testing.T) {
|
| 371 |
+
// Upstream returns plain JSON
|
| 372 |
+
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
| 373 |
+
w.Header().Set("Content-Type", "application/json")
|
| 374 |
+
w.WriteHeader(200)
|
| 375 |
+
w.Write([]byte(`{"plain":"json"}`))
|
| 376 |
+
}))
|
| 377 |
+
defer upstream.Close()
|
| 378 |
+
|
| 379 |
+
proxy, err := createReverseProxy(upstream.URL, NewStaticSecretSource("key"))
|
| 380 |
+
if err != nil {
|
| 381 |
+
t.Fatal(err)
|
| 382 |
+
}
|
| 383 |
+
|
| 384 |
+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
| 385 |
+
proxy.ServeHTTP(w, r)
|
| 386 |
+
}))
|
| 387 |
+
defer srv.Close()
|
| 388 |
+
|
| 389 |
+
res, err := http.Get(srv.URL + "/test")
|
| 390 |
+
if err != nil {
|
| 391 |
+
t.Fatal(err)
|
| 392 |
+
}
|
| 393 |
+
body, _ := io.ReadAll(res.Body)
|
| 394 |
+
res.Body.Close()
|
| 395 |
+
|
| 396 |
+
expected := []byte(`{"plain":"json"}`)
|
| 397 |
+
if !bytes.Equal(body, expected) {
|
| 398 |
+
t.Fatalf("want plain JSON unchanged, got: %s", body)
|
| 399 |
+
}
|
| 400 |
+
}
|
| 401 |
+
|
| 402 |
+
func TestIsStreamingResponse(t *testing.T) {
|
| 403 |
+
cases := []struct {
|
| 404 |
+
name string
|
| 405 |
+
header http.Header
|
| 406 |
+
want bool
|
| 407 |
+
}{
|
| 408 |
+
{
|
| 409 |
+
name: "sse",
|
| 410 |
+
header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
| 411 |
+
want: true,
|
| 412 |
+
},
|
| 413 |
+
{
|
| 414 |
+
name: "chunked_not_streaming",
|
| 415 |
+
header: http.Header{"Transfer-Encoding": []string{"chunked"}},
|
| 416 |
+
want: false, // Chunked is transport-level, not streaming
|
| 417 |
+
},
|
| 418 |
+
{
|
| 419 |
+
name: "normal_json",
|
| 420 |
+
header: http.Header{"Content-Type": []string{"application/json"}},
|
| 421 |
+
want: false,
|
| 422 |
+
},
|
| 423 |
+
{
|
| 424 |
+
name: "empty",
|
| 425 |
+
header: http.Header{},
|
| 426 |
+
want: false,
|
| 427 |
+
},
|
| 428 |
+
}
|
| 429 |
+
|
| 430 |
+
for _, tc := range cases {
|
| 431 |
+
t.Run(tc.name, func(t *testing.T) {
|
| 432 |
+
resp := &http.Response{Header: tc.header}
|
| 433 |
+
got := isStreamingResponse(resp)
|
| 434 |
+
if got != tc.want {
|
| 435 |
+
t.Fatalf("want %v, got %v", tc.want, got)
|
| 436 |
+
}
|
| 437 |
+
})
|
| 438 |
+
}
|
| 439 |
+
}
|
|
@@ -0,0 +1,166 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
package amp
|
| 2 |
+
|
| 3 |
+
import (
|
| 4 |
+
"net"
|
| 5 |
+
"strings"
|
| 6 |
+
|
| 7 |
+
"github.com/gin-gonic/gin"
|
| 8 |
+
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
|
| 9 |
+
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/claude"
|
| 10 |
+
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/gemini"
|
| 11 |
+
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/openai"
|
| 12 |
+
log "github.com/sirupsen/logrus"
|
| 13 |
+
)
|
| 14 |
+
|
| 15 |
+
// localhostOnlyMiddleware restricts access to localhost (127.0.0.1, ::1) only.
|
| 16 |
+
// Returns 403 Forbidden for non-localhost clients.
|
| 17 |
+
func localhostOnlyMiddleware() gin.HandlerFunc {
|
| 18 |
+
return func(c *gin.Context) {
|
| 19 |
+
clientIP := c.ClientIP()
|
| 20 |
+
|
| 21 |
+
// Parse the IP to handle both IPv4 and IPv6
|
| 22 |
+
ip := net.ParseIP(clientIP)
|
| 23 |
+
if ip == nil {
|
| 24 |
+
log.Warnf("Amp management: invalid client IP %s, denying access", clientIP)
|
| 25 |
+
c.AbortWithStatusJSON(403, gin.H{
|
| 26 |
+
"error": "Access denied: management routes restricted to localhost",
|
| 27 |
+
})
|
| 28 |
+
return
|
| 29 |
+
}
|
| 30 |
+
|
| 31 |
+
// Check if IP is loopback (127.0.0.1 or ::1)
|
| 32 |
+
if !ip.IsLoopback() {
|
| 33 |
+
log.Warnf("Amp management: non-localhost IP %s attempted access, denying", clientIP)
|
| 34 |
+
c.AbortWithStatusJSON(403, gin.H{
|
| 35 |
+
"error": "Access denied: management routes restricted to localhost",
|
| 36 |
+
})
|
| 37 |
+
return
|
| 38 |
+
}
|
| 39 |
+
|
| 40 |
+
c.Next()
|
| 41 |
+
}
|
| 42 |
+
}
|
| 43 |
+
|
| 44 |
+
// noCORSMiddleware disables CORS for management routes to prevent browser-based attacks.
|
| 45 |
+
// This overwrites any global CORS headers set by the server.
|
| 46 |
+
func noCORSMiddleware() gin.HandlerFunc {
|
| 47 |
+
return func(c *gin.Context) {
|
| 48 |
+
// Remove CORS headers to prevent cross-origin access from browsers
|
| 49 |
+
c.Header("Access-Control-Allow-Origin", "")
|
| 50 |
+
c.Header("Access-Control-Allow-Methods", "")
|
| 51 |
+
c.Header("Access-Control-Allow-Headers", "")
|
| 52 |
+
c.Header("Access-Control-Allow-Credentials", "")
|
| 53 |
+
|
| 54 |
+
// For OPTIONS preflight, deny with 403
|
| 55 |
+
if c.Request.Method == "OPTIONS" {
|
| 56 |
+
c.AbortWithStatus(403)
|
| 57 |
+
return
|
| 58 |
+
}
|
| 59 |
+
|
| 60 |
+
c.Next()
|
| 61 |
+
}
|
| 62 |
+
}
|
| 63 |
+
|
| 64 |
+
// registerManagementRoutes registers Amp management proxy routes
|
| 65 |
+
// These routes proxy through to the Amp control plane for OAuth, user management, etc.
|
| 66 |
+
// If restrictToLocalhost is true, routes will only accept connections from 127.0.0.1/::1.
|
| 67 |
+
func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, proxyHandler gin.HandlerFunc, restrictToLocalhost bool) {
|
| 68 |
+
ampAPI := engine.Group("/api")
|
| 69 |
+
|
| 70 |
+
// Always disable CORS for management routes to prevent browser-based attacks
|
| 71 |
+
ampAPI.Use(noCORSMiddleware())
|
| 72 |
+
|
| 73 |
+
// Apply localhost-only restriction if configured
|
| 74 |
+
if restrictToLocalhost {
|
| 75 |
+
ampAPI.Use(localhostOnlyMiddleware())
|
| 76 |
+
log.Info("Amp management routes restricted to localhost only (CORS disabled)")
|
| 77 |
+
} else {
|
| 78 |
+
log.Warn("⚠️ Amp management routes are NOT restricted to localhost - this is insecure!")
|
| 79 |
+
}
|
| 80 |
+
|
| 81 |
+
// Management routes - these are proxied directly to Amp upstream
|
| 82 |
+
ampAPI.Any("/internal", proxyHandler)
|
| 83 |
+
ampAPI.Any("/internal/*path", proxyHandler)
|
| 84 |
+
ampAPI.Any("/user", proxyHandler)
|
| 85 |
+
ampAPI.Any("/user/*path", proxyHandler)
|
| 86 |
+
ampAPI.Any("/auth", proxyHandler)
|
| 87 |
+
ampAPI.Any("/auth/*path", proxyHandler)
|
| 88 |
+
ampAPI.Any("/meta", proxyHandler)
|
| 89 |
+
ampAPI.Any("/meta/*path", proxyHandler)
|
| 90 |
+
ampAPI.Any("/ads", proxyHandler)
|
| 91 |
+
ampAPI.Any("/telemetry", proxyHandler)
|
| 92 |
+
ampAPI.Any("/telemetry/*path", proxyHandler)
|
| 93 |
+
ampAPI.Any("/threads", proxyHandler)
|
| 94 |
+
ampAPI.Any("/threads/*path", proxyHandler)
|
| 95 |
+
ampAPI.Any("/otel", proxyHandler)
|
| 96 |
+
ampAPI.Any("/otel/*path", proxyHandler)
|
| 97 |
+
|
| 98 |
+
// Google v1beta1 passthrough (Gemini native API)
|
| 99 |
+
ampAPI.Any("/provider/google/v1beta1/*path", proxyHandler)
|
| 100 |
+
}
|
| 101 |
+
|
| 102 |
+
// registerProviderAliases registers /api/provider/{provider}/... routes
|
| 103 |
+
// These allow Amp CLI to route requests like:
|
| 104 |
+
//
|
| 105 |
+
// /api/provider/openai/v1/chat/completions
|
| 106 |
+
// /api/provider/anthropic/v1/messages
|
| 107 |
+
// /api/provider/google/v1beta/models
|
| 108 |
+
func (m *AmpModule) registerProviderAliases(engine *gin.Engine, baseHandler *handlers.BaseAPIHandler, auth gin.HandlerFunc) {
|
| 109 |
+
// Create handler instances for different providers
|
| 110 |
+
openaiHandlers := openai.NewOpenAIAPIHandler(baseHandler)
|
| 111 |
+
geminiHandlers := gemini.NewGeminiAPIHandler(baseHandler)
|
| 112 |
+
claudeCodeHandlers := claude.NewClaudeCodeAPIHandler(baseHandler)
|
| 113 |
+
openaiResponsesHandlers := openai.NewOpenAIResponsesAPIHandler(baseHandler)
|
| 114 |
+
|
| 115 |
+
// Provider-specific routes under /api/provider/:provider
|
| 116 |
+
ampProviders := engine.Group("/api/provider")
|
| 117 |
+
if auth != nil {
|
| 118 |
+
ampProviders.Use(auth)
|
| 119 |
+
}
|
| 120 |
+
|
| 121 |
+
provider := ampProviders.Group("/:provider")
|
| 122 |
+
|
| 123 |
+
// Dynamic models handler - routes to appropriate provider based on path parameter
|
| 124 |
+
ampModelsHandler := func(c *gin.Context) {
|
| 125 |
+
providerName := strings.ToLower(c.Param("provider"))
|
| 126 |
+
|
| 127 |
+
switch providerName {
|
| 128 |
+
case "anthropic":
|
| 129 |
+
claudeCodeHandlers.ClaudeModels(c)
|
| 130 |
+
case "google":
|
| 131 |
+
geminiHandlers.GeminiModels(c)
|
| 132 |
+
default:
|
| 133 |
+
// Default to OpenAI-compatible (works for openai, groq, cerebras, etc.)
|
| 134 |
+
openaiHandlers.OpenAIModels(c)
|
| 135 |
+
}
|
| 136 |
+
}
|
| 137 |
+
|
| 138 |
+
// Root-level routes (for providers that omit /v1, like groq/cerebras)
|
| 139 |
+
provider.GET("/models", ampModelsHandler)
|
| 140 |
+
provider.POST("/chat/completions", openaiHandlers.ChatCompletions)
|
| 141 |
+
provider.POST("/completions", openaiHandlers.Completions)
|
| 142 |
+
provider.POST("/responses", openaiResponsesHandlers.Responses)
|
| 143 |
+
|
| 144 |
+
// /v1 routes (OpenAI/Claude-compatible endpoints)
|
| 145 |
+
v1Amp := provider.Group("/v1")
|
| 146 |
+
{
|
| 147 |
+
v1Amp.GET("/models", ampModelsHandler)
|
| 148 |
+
|
| 149 |
+
// OpenAI-compatible endpoints
|
| 150 |
+
v1Amp.POST("/chat/completions", openaiHandlers.ChatCompletions)
|
| 151 |
+
v1Amp.POST("/completions", openaiHandlers.Completions)
|
| 152 |
+
v1Amp.POST("/responses", openaiResponsesHandlers.Responses)
|
| 153 |
+
|
| 154 |
+
// Claude/Anthropic-compatible endpoints
|
| 155 |
+
v1Amp.POST("/messages", claudeCodeHandlers.ClaudeMessages)
|
| 156 |
+
v1Amp.POST("/messages/count_tokens", claudeCodeHandlers.ClaudeCountTokens)
|
| 157 |
+
}
|
| 158 |
+
|
| 159 |
+
// /v1beta routes (Gemini native API)
|
| 160 |
+
v1betaAmp := provider.Group("/v1beta")
|
| 161 |
+
{
|
| 162 |
+
v1betaAmp.GET("/models", geminiHandlers.GeminiModels)
|
| 163 |
+
v1betaAmp.POST("/models/:action", geminiHandlers.GeminiHandler)
|
| 164 |
+
v1betaAmp.GET("/models/:action", geminiHandlers.GeminiGetHandler)
|
| 165 |
+
}
|
| 166 |
+
}
|
|
@@ -0,0 +1,216 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
package amp
|
| 2 |
+
|
| 3 |
+
import (
|
| 4 |
+
"net/http"
|
| 5 |
+
"net/http/httptest"
|
| 6 |
+
"testing"
|
| 7 |
+
|
| 8 |
+
"github.com/gin-gonic/gin"
|
| 9 |
+
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
|
| 10 |
+
)
|
| 11 |
+
|
| 12 |
+
func TestRegisterManagementRoutes(t *testing.T) {
|
| 13 |
+
gin.SetMode(gin.TestMode)
|
| 14 |
+
r := gin.New()
|
| 15 |
+
|
| 16 |
+
// Spy to track if proxy handler was called
|
| 17 |
+
proxyCalled := false
|
| 18 |
+
proxyHandler := func(c *gin.Context) {
|
| 19 |
+
proxyCalled = true
|
| 20 |
+
c.String(200, "proxied")
|
| 21 |
+
}
|
| 22 |
+
|
| 23 |
+
m := &AmpModule{}
|
| 24 |
+
m.registerManagementRoutes(r, proxyHandler, false) // false = don't restrict to localhost in tests
|
| 25 |
+
|
| 26 |
+
managementPaths := []string{
|
| 27 |
+
"/api/internal",
|
| 28 |
+
"/api/internal/some/path",
|
| 29 |
+
"/api/user",
|
| 30 |
+
"/api/user/profile",
|
| 31 |
+
"/api/auth",
|
| 32 |
+
"/api/auth/login",
|
| 33 |
+
"/api/meta",
|
| 34 |
+
"/api/telemetry",
|
| 35 |
+
"/api/threads",
|
| 36 |
+
"/api/otel",
|
| 37 |
+
"/api/provider/google/v1beta1/models",
|
| 38 |
+
}
|
| 39 |
+
|
| 40 |
+
for _, path := range managementPaths {
|
| 41 |
+
t.Run(path, func(t *testing.T) {
|
| 42 |
+
proxyCalled = false
|
| 43 |
+
req := httptest.NewRequest(http.MethodGet, path, nil)
|
| 44 |
+
w := httptest.NewRecorder()
|
| 45 |
+
r.ServeHTTP(w, req)
|
| 46 |
+
|
| 47 |
+
if w.Code == http.StatusNotFound {
|
| 48 |
+
t.Fatalf("route %s not registered", path)
|
| 49 |
+
}
|
| 50 |
+
if !proxyCalled {
|
| 51 |
+
t.Fatalf("proxy handler not called for %s", path)
|
| 52 |
+
}
|
| 53 |
+
})
|
| 54 |
+
}
|
| 55 |
+
}
|
| 56 |
+
|
| 57 |
+
func TestRegisterProviderAliases_AllProvidersRegistered(t *testing.T) {
|
| 58 |
+
gin.SetMode(gin.TestMode)
|
| 59 |
+
r := gin.New()
|
| 60 |
+
|
| 61 |
+
// Minimal base handler setup (no need to initialize, just check routing)
|
| 62 |
+
base := &handlers.BaseAPIHandler{}
|
| 63 |
+
|
| 64 |
+
// Track if auth middleware was called
|
| 65 |
+
authCalled := false
|
| 66 |
+
authMiddleware := func(c *gin.Context) {
|
| 67 |
+
authCalled = true
|
| 68 |
+
c.Header("X-Auth", "ok")
|
| 69 |
+
// Abort with success to avoid calling the actual handler (which needs full setup)
|
| 70 |
+
c.AbortWithStatus(http.StatusOK)
|
| 71 |
+
}
|
| 72 |
+
|
| 73 |
+
m := &AmpModule{authMiddleware_: authMiddleware}
|
| 74 |
+
m.registerProviderAliases(r, base, authMiddleware)
|
| 75 |
+
|
| 76 |
+
paths := []struct {
|
| 77 |
+
path string
|
| 78 |
+
method string
|
| 79 |
+
}{
|
| 80 |
+
{"/api/provider/openai/models", http.MethodGet},
|
| 81 |
+
{"/api/provider/anthropic/models", http.MethodGet},
|
| 82 |
+
{"/api/provider/google/models", http.MethodGet},
|
| 83 |
+
{"/api/provider/groq/models", http.MethodGet},
|
| 84 |
+
{"/api/provider/openai/chat/completions", http.MethodPost},
|
| 85 |
+
{"/api/provider/anthropic/v1/messages", http.MethodPost},
|
| 86 |
+
{"/api/provider/google/v1beta/models", http.MethodGet},
|
| 87 |
+
}
|
| 88 |
+
|
| 89 |
+
for _, tc := range paths {
|
| 90 |
+
t.Run(tc.path, func(t *testing.T) {
|
| 91 |
+
authCalled = false
|
| 92 |
+
req := httptest.NewRequest(tc.method, tc.path, nil)
|
| 93 |
+
w := httptest.NewRecorder()
|
| 94 |
+
r.ServeHTTP(w, req)
|
| 95 |
+
|
| 96 |
+
if w.Code == http.StatusNotFound {
|
| 97 |
+
t.Fatalf("route %s %s not registered", tc.method, tc.path)
|
| 98 |
+
}
|
| 99 |
+
if !authCalled {
|
| 100 |
+
t.Fatalf("auth middleware not executed for %s", tc.path)
|
| 101 |
+
}
|
| 102 |
+
if w.Header().Get("X-Auth") != "ok" {
|
| 103 |
+
t.Fatalf("auth middleware header not set for %s", tc.path)
|
| 104 |
+
}
|
| 105 |
+
})
|
| 106 |
+
}
|
| 107 |
+
}
|
| 108 |
+
|
| 109 |
+
func TestRegisterProviderAliases_DynamicModelsHandler(t *testing.T) {
|
| 110 |
+
gin.SetMode(gin.TestMode)
|
| 111 |
+
r := gin.New()
|
| 112 |
+
|
| 113 |
+
base := &handlers.BaseAPIHandler{}
|
| 114 |
+
|
| 115 |
+
m := &AmpModule{authMiddleware_: func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) }}
|
| 116 |
+
m.registerProviderAliases(r, base, func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) })
|
| 117 |
+
|
| 118 |
+
providers := []string{"openai", "anthropic", "google", "groq", "cerebras"}
|
| 119 |
+
|
| 120 |
+
for _, provider := range providers {
|
| 121 |
+
t.Run(provider, func(t *testing.T) {
|
| 122 |
+
path := "/api/provider/" + provider + "/models"
|
| 123 |
+
req := httptest.NewRequest(http.MethodGet, path, nil)
|
| 124 |
+
w := httptest.NewRecorder()
|
| 125 |
+
r.ServeHTTP(w, req)
|
| 126 |
+
|
| 127 |
+
// Should not 404
|
| 128 |
+
if w.Code == http.StatusNotFound {
|
| 129 |
+
t.Fatalf("models route not found for provider: %s", provider)
|
| 130 |
+
}
|
| 131 |
+
})
|
| 132 |
+
}
|
| 133 |
+
}
|
| 134 |
+
|
| 135 |
+
func TestRegisterProviderAliases_V1Routes(t *testing.T) {
|
| 136 |
+
gin.SetMode(gin.TestMode)
|
| 137 |
+
r := gin.New()
|
| 138 |
+
|
| 139 |
+
base := &handlers.BaseAPIHandler{}
|
| 140 |
+
|
| 141 |
+
m := &AmpModule{authMiddleware_: func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) }}
|
| 142 |
+
m.registerProviderAliases(r, base, func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) })
|
| 143 |
+
|
| 144 |
+
v1Paths := []struct {
|
| 145 |
+
path string
|
| 146 |
+
method string
|
| 147 |
+
}{
|
| 148 |
+
{"/api/provider/openai/v1/models", http.MethodGet},
|
| 149 |
+
{"/api/provider/openai/v1/chat/completions", http.MethodPost},
|
| 150 |
+
{"/api/provider/openai/v1/completions", http.MethodPost},
|
| 151 |
+
{"/api/provider/anthropic/v1/messages", http.MethodPost},
|
| 152 |
+
{"/api/provider/anthropic/v1/messages/count_tokens", http.MethodPost},
|
| 153 |
+
}
|
| 154 |
+
|
| 155 |
+
for _, tc := range v1Paths {
|
| 156 |
+
t.Run(tc.path, func(t *testing.T) {
|
| 157 |
+
req := httptest.NewRequest(tc.method, tc.path, nil)
|
| 158 |
+
w := httptest.NewRecorder()
|
| 159 |
+
r.ServeHTTP(w, req)
|
| 160 |
+
|
| 161 |
+
if w.Code == http.StatusNotFound {
|
| 162 |
+
t.Fatalf("v1 route %s %s not registered", tc.method, tc.path)
|
| 163 |
+
}
|
| 164 |
+
})
|
| 165 |
+
}
|
| 166 |
+
}
|
| 167 |
+
|
| 168 |
+
func TestRegisterProviderAliases_V1BetaRoutes(t *testing.T) {
|
| 169 |
+
gin.SetMode(gin.TestMode)
|
| 170 |
+
r := gin.New()
|
| 171 |
+
|
| 172 |
+
base := &handlers.BaseAPIHandler{}
|
| 173 |
+
|
| 174 |
+
m := &AmpModule{authMiddleware_: func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) }}
|
| 175 |
+
m.registerProviderAliases(r, base, func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) })
|
| 176 |
+
|
| 177 |
+
v1betaPaths := []struct {
|
| 178 |
+
path string
|
| 179 |
+
method string
|
| 180 |
+
}{
|
| 181 |
+
{"/api/provider/google/v1beta/models", http.MethodGet},
|
| 182 |
+
{"/api/provider/google/v1beta/models/generateContent", http.MethodPost},
|
| 183 |
+
}
|
| 184 |
+
|
| 185 |
+
for _, tc := range v1betaPaths {
|
| 186 |
+
t.Run(tc.path, func(t *testing.T) {
|
| 187 |
+
req := httptest.NewRequest(tc.method, tc.path, nil)
|
| 188 |
+
w := httptest.NewRecorder()
|
| 189 |
+
r.ServeHTTP(w, req)
|
| 190 |
+
|
| 191 |
+
if w.Code == http.StatusNotFound {
|
| 192 |
+
t.Fatalf("v1beta route %s %s not registered", tc.method, tc.path)
|
| 193 |
+
}
|
| 194 |
+
})
|
| 195 |
+
}
|
| 196 |
+
}
|
| 197 |
+
|
| 198 |
+
func TestRegisterProviderAliases_NoAuthMiddleware(t *testing.T) {
|
| 199 |
+
// Test that routes still register even if auth middleware is nil (fallback behavior)
|
| 200 |
+
gin.SetMode(gin.TestMode)
|
| 201 |
+
r := gin.New()
|
| 202 |
+
|
| 203 |
+
base := &handlers.BaseAPIHandler{}
|
| 204 |
+
|
| 205 |
+
m := &AmpModule{authMiddleware_: nil} // No auth middleware
|
| 206 |
+
m.registerProviderAliases(r, base, func(c *gin.Context) { c.AbortWithStatus(http.StatusOK) })
|
| 207 |
+
|
| 208 |
+
req := httptest.NewRequest(http.MethodGet, "/api/provider/openai/models", nil)
|
| 209 |
+
w := httptest.NewRecorder()
|
| 210 |
+
r.ServeHTTP(w, req)
|
| 211 |
+
|
| 212 |
+
// Should still work (with fallback no-op auth)
|
| 213 |
+
if w.Code == http.StatusNotFound {
|
| 214 |
+
t.Fatal("routes should register even without auth middleware")
|
| 215 |
+
}
|
| 216 |
+
}
|
|
@@ -0,0 +1,155 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
package amp
|
| 2 |
+
|
| 3 |
+
import (
|
| 4 |
+
"context"
|
| 5 |
+
"encoding/json"
|
| 6 |
+
"fmt"
|
| 7 |
+
"os"
|
| 8 |
+
"path/filepath"
|
| 9 |
+
"strings"
|
| 10 |
+
"sync"
|
| 11 |
+
"time"
|
| 12 |
+
)
|
| 13 |
+
|
| 14 |
+
// SecretSource provides Amp API keys with configurable precedence and caching
|
| 15 |
+
type SecretSource interface {
|
| 16 |
+
Get(ctx context.Context) (string, error)
|
| 17 |
+
}
|
| 18 |
+
|
| 19 |
+
// cachedSecret holds a secret value with expiration
|
| 20 |
+
type cachedSecret struct {
|
| 21 |
+
value string
|
| 22 |
+
expiresAt time.Time
|
| 23 |
+
}
|
| 24 |
+
|
| 25 |
+
// MultiSourceSecret implements precedence-based secret lookup:
|
| 26 |
+
// 1. Explicit config value (highest priority)
|
| 27 |
+
// 2. Environment variable AMP_API_KEY
|
| 28 |
+
// 3. File-based secret (lowest priority)
|
| 29 |
+
type MultiSourceSecret struct {
|
| 30 |
+
explicitKey string
|
| 31 |
+
envKey string
|
| 32 |
+
filePath string
|
| 33 |
+
cacheTTL time.Duration
|
| 34 |
+
|
| 35 |
+
mu sync.RWMutex
|
| 36 |
+
cache *cachedSecret
|
| 37 |
+
}
|
| 38 |
+
|
| 39 |
+
// NewMultiSourceSecret creates a secret source with precedence and caching
|
| 40 |
+
func NewMultiSourceSecret(explicitKey string, cacheTTL time.Duration) *MultiSourceSecret {
|
| 41 |
+
if cacheTTL == 0 {
|
| 42 |
+
cacheTTL = 5 * time.Minute // Default 5 minute cache
|
| 43 |
+
}
|
| 44 |
+
|
| 45 |
+
home, _ := os.UserHomeDir()
|
| 46 |
+
filePath := filepath.Join(home, ".local", "share", "amp", "secrets.json")
|
| 47 |
+
|
| 48 |
+
return &MultiSourceSecret{
|
| 49 |
+
explicitKey: strings.TrimSpace(explicitKey),
|
| 50 |
+
envKey: "AMP_API_KEY",
|
| 51 |
+
filePath: filePath,
|
| 52 |
+
cacheTTL: cacheTTL,
|
| 53 |
+
}
|
| 54 |
+
}
|
| 55 |
+
|
| 56 |
+
// NewMultiSourceSecretWithPath creates a secret source with a custom file path (for testing)
|
| 57 |
+
func NewMultiSourceSecretWithPath(explicitKey string, filePath string, cacheTTL time.Duration) *MultiSourceSecret {
|
| 58 |
+
if cacheTTL == 0 {
|
| 59 |
+
cacheTTL = 5 * time.Minute
|
| 60 |
+
}
|
| 61 |
+
|
| 62 |
+
return &MultiSourceSecret{
|
| 63 |
+
explicitKey: strings.TrimSpace(explicitKey),
|
| 64 |
+
envKey: "AMP_API_KEY",
|
| 65 |
+
filePath: filePath,
|
| 66 |
+
cacheTTL: cacheTTL,
|
| 67 |
+
}
|
| 68 |
+
}
|
| 69 |
+
|
| 70 |
+
// Get retrieves the Amp API key using precedence: config > env > file
|
| 71 |
+
// Results are cached for cacheTTL duration to avoid excessive file reads
|
| 72 |
+
func (s *MultiSourceSecret) Get(ctx context.Context) (string, error) {
|
| 73 |
+
// Precedence 1: Explicit config key (highest priority, no caching needed)
|
| 74 |
+
if s.explicitKey != "" {
|
| 75 |
+
return s.explicitKey, nil
|
| 76 |
+
}
|
| 77 |
+
|
| 78 |
+
// Precedence 2: Environment variable
|
| 79 |
+
if envValue := strings.TrimSpace(os.Getenv(s.envKey)); envValue != "" {
|
| 80 |
+
return envValue, nil
|
| 81 |
+
}
|
| 82 |
+
|
| 83 |
+
// Precedence 3: File-based secret (lowest priority, cached)
|
| 84 |
+
// Check cache first
|
| 85 |
+
s.mu.RLock()
|
| 86 |
+
if s.cache != nil && time.Now().Before(s.cache.expiresAt) {
|
| 87 |
+
value := s.cache.value
|
| 88 |
+
s.mu.RUnlock()
|
| 89 |
+
return value, nil
|
| 90 |
+
}
|
| 91 |
+
s.mu.RUnlock()
|
| 92 |
+
|
| 93 |
+
// Cache miss or expired - read from file
|
| 94 |
+
key, err := s.readFromFile()
|
| 95 |
+
if err != nil {
|
| 96 |
+
// Cache empty result to avoid repeated file reads on missing files
|
| 97 |
+
s.updateCache("")
|
| 98 |
+
return "", err
|
| 99 |
+
}
|
| 100 |
+
|
| 101 |
+
// Cache the result
|
| 102 |
+
s.updateCache(key)
|
| 103 |
+
return key, nil
|
| 104 |
+
}
|
| 105 |
+
|
| 106 |
+
// readFromFile reads the Amp API key from the secrets file
|
| 107 |
+
func (s *MultiSourceSecret) readFromFile() (string, error) {
|
| 108 |
+
content, err := os.ReadFile(s.filePath)
|
| 109 |
+
if err != nil {
|
| 110 |
+
if os.IsNotExist(err) {
|
| 111 |
+
return "", nil // Missing file is not an error, just no key available
|
| 112 |
+
}
|
| 113 |
+
return "", fmt.Errorf("failed to read amp secrets from %s: %w", s.filePath, err)
|
| 114 |
+
}
|
| 115 |
+
|
| 116 |
+
var secrets map[string]string
|
| 117 |
+
if err := json.Unmarshal(content, &secrets); err != nil {
|
| 118 |
+
return "", fmt.Errorf("failed to parse amp secrets from %s: %w", s.filePath, err)
|
| 119 |
+
}
|
| 120 |
+
|
| 121 |
+
key := strings.TrimSpace(secrets["apiKey@https://ampcode.com/"])
|
| 122 |
+
return key, nil
|
| 123 |
+
}
|
| 124 |
+
|
| 125 |
+
// updateCache updates the cached secret value
|
| 126 |
+
func (s *MultiSourceSecret) updateCache(value string) {
|
| 127 |
+
s.mu.Lock()
|
| 128 |
+
defer s.mu.Unlock()
|
| 129 |
+
s.cache = &cachedSecret{
|
| 130 |
+
value: value,
|
| 131 |
+
expiresAt: time.Now().Add(s.cacheTTL),
|
| 132 |
+
}
|
| 133 |
+
}
|
| 134 |
+
|
| 135 |
+
// InvalidateCache clears the cached secret, forcing a fresh read on next Get
|
| 136 |
+
func (s *MultiSourceSecret) InvalidateCache() {
|
| 137 |
+
s.mu.Lock()
|
| 138 |
+
defer s.mu.Unlock()
|
| 139 |
+
s.cache = nil
|
| 140 |
+
}
|
| 141 |
+
|
| 142 |
+
// StaticSecretSource returns a fixed API key (for testing)
|
| 143 |
+
type StaticSecretSource struct {
|
| 144 |
+
key string
|
| 145 |
+
}
|
| 146 |
+
|
| 147 |
+
// NewStaticSecretSource creates a secret source with a fixed key
|
| 148 |
+
func NewStaticSecretSource(key string) *StaticSecretSource {
|
| 149 |
+
return &StaticSecretSource{key: strings.TrimSpace(key)}
|
| 150 |
+
}
|
| 151 |
+
|
| 152 |
+
// Get returns the static API key
|
| 153 |
+
func (s *StaticSecretSource) Get(ctx context.Context) (string, error) {
|
| 154 |
+
return s.key, nil
|
| 155 |
+
}
|
|
@@ -0,0 +1,280 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
package amp
|
| 2 |
+
|
| 3 |
+
import (
|
| 4 |
+
"context"
|
| 5 |
+
"encoding/json"
|
| 6 |
+
"os"
|
| 7 |
+
"path/filepath"
|
| 8 |
+
"sync"
|
| 9 |
+
"testing"
|
| 10 |
+
"time"
|
| 11 |
+
)
|
| 12 |
+
|
| 13 |
+
func TestMultiSourceSecret_PrecedenceOrder(t *testing.T) {
|
| 14 |
+
ctx := context.Background()
|
| 15 |
+
|
| 16 |
+
cases := []struct {
|
| 17 |
+
name string
|
| 18 |
+
configKey string
|
| 19 |
+
envKey string
|
| 20 |
+
fileJSON string
|
| 21 |
+
want string
|
| 22 |
+
}{
|
| 23 |
+
{"config_wins", "cfg", "env", `{"apiKey@https://ampcode.com/":"file"}`, "cfg"},
|
| 24 |
+
{"env_wins_when_no_cfg", "", "env", `{"apiKey@https://ampcode.com/":"file"}`, "env"},
|
| 25 |
+
{"file_when_no_cfg_env", "", "", `{"apiKey@https://ampcode.com/":"file"}`, "file"},
|
| 26 |
+
{"empty_cfg_trims_then_env", " ", "env", `{"apiKey@https://ampcode.com/":"file"}`, "env"},
|
| 27 |
+
{"empty_env_then_file", "", " ", `{"apiKey@https://ampcode.com/":"file"}`, "file"},
|
| 28 |
+
{"missing_file_returns_empty", "", "", "", ""},
|
| 29 |
+
{"all_empty_returns_empty", " ", " ", `{"apiKey@https://ampcode.com/":" "}`, ""},
|
| 30 |
+
}
|
| 31 |
+
|
| 32 |
+
for _, tc := range cases {
|
| 33 |
+
tc := tc // capture range variable
|
| 34 |
+
t.Run(tc.name, func(t *testing.T) {
|
| 35 |
+
tmpDir := t.TempDir()
|
| 36 |
+
secretsPath := filepath.Join(tmpDir, "secrets.json")
|
| 37 |
+
|
| 38 |
+
if tc.fileJSON != "" {
|
| 39 |
+
if err := os.WriteFile(secretsPath, []byte(tc.fileJSON), 0600); err != nil {
|
| 40 |
+
t.Fatal(err)
|
| 41 |
+
}
|
| 42 |
+
}
|
| 43 |
+
|
| 44 |
+
t.Setenv("AMP_API_KEY", tc.envKey)
|
| 45 |
+
|
| 46 |
+
s := NewMultiSourceSecretWithPath(tc.configKey, secretsPath, 100*time.Millisecond)
|
| 47 |
+
got, err := s.Get(ctx)
|
| 48 |
+
if err != nil && tc.fileJSON != "" && json.Valid([]byte(tc.fileJSON)) {
|
| 49 |
+
t.Fatalf("unexpected error: %v", err)
|
| 50 |
+
}
|
| 51 |
+
if got != tc.want {
|
| 52 |
+
t.Fatalf("want %q, got %q", tc.want, got)
|
| 53 |
+
}
|
| 54 |
+
})
|
| 55 |
+
}
|
| 56 |
+
}
|
| 57 |
+
|
| 58 |
+
func TestMultiSourceSecret_CacheBehavior(t *testing.T) {
|
| 59 |
+
ctx := context.Background()
|
| 60 |
+
tmpDir := t.TempDir()
|
| 61 |
+
p := filepath.Join(tmpDir, "secrets.json")
|
| 62 |
+
|
| 63 |
+
// Initial value
|
| 64 |
+
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":"v1"}`), 0600); err != nil {
|
| 65 |
+
t.Fatal(err)
|
| 66 |
+
}
|
| 67 |
+
|
| 68 |
+
s := NewMultiSourceSecretWithPath("", p, 50*time.Millisecond)
|
| 69 |
+
|
| 70 |
+
// First read - should return v1
|
| 71 |
+
got1, err := s.Get(ctx)
|
| 72 |
+
if err != nil {
|
| 73 |
+
t.Fatalf("Get failed: %v", err)
|
| 74 |
+
}
|
| 75 |
+
if got1 != "v1" {
|
| 76 |
+
t.Fatalf("expected v1, got %s", got1)
|
| 77 |
+
}
|
| 78 |
+
|
| 79 |
+
// Change file; within TTL we should still see v1 (cached)
|
| 80 |
+
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":"v2"}`), 0600); err != nil {
|
| 81 |
+
t.Fatal(err)
|
| 82 |
+
}
|
| 83 |
+
got2, _ := s.Get(ctx)
|
| 84 |
+
if got2 != "v1" {
|
| 85 |
+
t.Fatalf("cache hit expected v1, got %s", got2)
|
| 86 |
+
}
|
| 87 |
+
|
| 88 |
+
// After TTL expires, should see v2
|
| 89 |
+
time.Sleep(60 * time.Millisecond)
|
| 90 |
+
got3, _ := s.Get(ctx)
|
| 91 |
+
if got3 != "v2" {
|
| 92 |
+
t.Fatalf("cache miss expected v2, got %s", got3)
|
| 93 |
+
}
|
| 94 |
+
|
| 95 |
+
// Invalidate forces re-read immediately
|
| 96 |
+
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":"v3"}`), 0600); err != nil {
|
| 97 |
+
t.Fatal(err)
|
| 98 |
+
}
|
| 99 |
+
s.InvalidateCache()
|
| 100 |
+
got4, _ := s.Get(ctx)
|
| 101 |
+
if got4 != "v3" {
|
| 102 |
+
t.Fatalf("invalidate expected v3, got %s", got4)
|
| 103 |
+
}
|
| 104 |
+
}
|
| 105 |
+
|
| 106 |
+
func TestMultiSourceSecret_FileHandling(t *testing.T) {
|
| 107 |
+
ctx := context.Background()
|
| 108 |
+
|
| 109 |
+
t.Run("missing_file_no_error", func(t *testing.T) {
|
| 110 |
+
s := NewMultiSourceSecretWithPath("", "/nonexistent/path/secrets.json", 100*time.Millisecond)
|
| 111 |
+
got, err := s.Get(ctx)
|
| 112 |
+
if err != nil {
|
| 113 |
+
t.Fatalf("expected no error for missing file, got: %v", err)
|
| 114 |
+
}
|
| 115 |
+
if got != "" {
|
| 116 |
+
t.Fatalf("expected empty string, got %q", got)
|
| 117 |
+
}
|
| 118 |
+
})
|
| 119 |
+
|
| 120 |
+
t.Run("invalid_json", func(t *testing.T) {
|
| 121 |
+
tmpDir := t.TempDir()
|
| 122 |
+
p := filepath.Join(tmpDir, "secrets.json")
|
| 123 |
+
if err := os.WriteFile(p, []byte(`{invalid json`), 0600); err != nil {
|
| 124 |
+
t.Fatal(err)
|
| 125 |
+
}
|
| 126 |
+
|
| 127 |
+
s := NewMultiSourceSecretWithPath("", p, 100*time.Millisecond)
|
| 128 |
+
_, err := s.Get(ctx)
|
| 129 |
+
if err == nil {
|
| 130 |
+
t.Fatal("expected error for invalid JSON")
|
| 131 |
+
}
|
| 132 |
+
})
|
| 133 |
+
|
| 134 |
+
t.Run("missing_key_in_json", func(t *testing.T) {
|
| 135 |
+
tmpDir := t.TempDir()
|
| 136 |
+
p := filepath.Join(tmpDir, "secrets.json")
|
| 137 |
+
if err := os.WriteFile(p, []byte(`{"other":"value"}`), 0600); err != nil {
|
| 138 |
+
t.Fatal(err)
|
| 139 |
+
}
|
| 140 |
+
|
| 141 |
+
s := NewMultiSourceSecretWithPath("", p, 100*time.Millisecond)
|
| 142 |
+
got, err := s.Get(ctx)
|
| 143 |
+
if err != nil {
|
| 144 |
+
t.Fatalf("unexpected error: %v", err)
|
| 145 |
+
}
|
| 146 |
+
if got != "" {
|
| 147 |
+
t.Fatalf("expected empty string for missing key, got %q", got)
|
| 148 |
+
}
|
| 149 |
+
})
|
| 150 |
+
|
| 151 |
+
t.Run("empty_key_value", func(t *testing.T) {
|
| 152 |
+
tmpDir := t.TempDir()
|
| 153 |
+
p := filepath.Join(tmpDir, "secrets.json")
|
| 154 |
+
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":" "}`), 0600); err != nil {
|
| 155 |
+
t.Fatal(err)
|
| 156 |
+
}
|
| 157 |
+
|
| 158 |
+
s := NewMultiSourceSecretWithPath("", p, 100*time.Millisecond)
|
| 159 |
+
got, _ := s.Get(ctx)
|
| 160 |
+
if got != "" {
|
| 161 |
+
t.Fatalf("expected empty after trim, got %q", got)
|
| 162 |
+
}
|
| 163 |
+
})
|
| 164 |
+
}
|
| 165 |
+
|
| 166 |
+
func TestMultiSourceSecret_Concurrency(t *testing.T) {
|
| 167 |
+
tmpDir := t.TempDir()
|
| 168 |
+
p := filepath.Join(tmpDir, "secrets.json")
|
| 169 |
+
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":"concurrent"}`), 0600); err != nil {
|
| 170 |
+
t.Fatal(err)
|
| 171 |
+
}
|
| 172 |
+
|
| 173 |
+
s := NewMultiSourceSecretWithPath("", p, 5*time.Second)
|
| 174 |
+
ctx := context.Background()
|
| 175 |
+
|
| 176 |
+
// Spawn many goroutines calling Get concurrently
|
| 177 |
+
const goroutines = 50
|
| 178 |
+
const iterations = 100
|
| 179 |
+
|
| 180 |
+
var wg sync.WaitGroup
|
| 181 |
+
errors := make(chan error, goroutines)
|
| 182 |
+
|
| 183 |
+
for i := 0; i < goroutines; i++ {
|
| 184 |
+
wg.Add(1)
|
| 185 |
+
go func() {
|
| 186 |
+
defer wg.Done()
|
| 187 |
+
for j := 0; j < iterations; j++ {
|
| 188 |
+
val, err := s.Get(ctx)
|
| 189 |
+
if err != nil {
|
| 190 |
+
errors <- err
|
| 191 |
+
return
|
| 192 |
+
}
|
| 193 |
+
if val != "concurrent" {
|
| 194 |
+
errors <- err
|
| 195 |
+
return
|
| 196 |
+
}
|
| 197 |
+
}
|
| 198 |
+
}()
|
| 199 |
+
}
|
| 200 |
+
|
| 201 |
+
wg.Wait()
|
| 202 |
+
close(errors)
|
| 203 |
+
|
| 204 |
+
for err := range errors {
|
| 205 |
+
t.Errorf("concurrency error: %v", err)
|
| 206 |
+
}
|
| 207 |
+
}
|
| 208 |
+
|
| 209 |
+
func TestStaticSecretSource(t *testing.T) {
|
| 210 |
+
ctx := context.Background()
|
| 211 |
+
|
| 212 |
+
t.Run("returns_provided_key", func(t *testing.T) {
|
| 213 |
+
s := NewStaticSecretSource("test-key-123")
|
| 214 |
+
got, err := s.Get(ctx)
|
| 215 |
+
if err != nil {
|
| 216 |
+
t.Fatalf("unexpected error: %v", err)
|
| 217 |
+
}
|
| 218 |
+
if got != "test-key-123" {
|
| 219 |
+
t.Fatalf("want test-key-123, got %q", got)
|
| 220 |
+
}
|
| 221 |
+
})
|
| 222 |
+
|
| 223 |
+
t.Run("trims_whitespace", func(t *testing.T) {
|
| 224 |
+
s := NewStaticSecretSource(" test-key ")
|
| 225 |
+
got, err := s.Get(ctx)
|
| 226 |
+
if err != nil {
|
| 227 |
+
t.Fatalf("unexpected error: %v", err)
|
| 228 |
+
}
|
| 229 |
+
if got != "test-key" {
|
| 230 |
+
t.Fatalf("want test-key, got %q", got)
|
| 231 |
+
}
|
| 232 |
+
})
|
| 233 |
+
|
| 234 |
+
t.Run("empty_string", func(t *testing.T) {
|
| 235 |
+
s := NewStaticSecretSource("")
|
| 236 |
+
got, err := s.Get(ctx)
|
| 237 |
+
if err != nil {
|
| 238 |
+
t.Fatalf("unexpected error: %v", err)
|
| 239 |
+
}
|
| 240 |
+
if got != "" {
|
| 241 |
+
t.Fatalf("want empty string, got %q", got)
|
| 242 |
+
}
|
| 243 |
+
})
|
| 244 |
+
}
|
| 245 |
+
|
| 246 |
+
func TestMultiSourceSecret_CacheEmptyResult(t *testing.T) {
|
| 247 |
+
// Test that missing file results are cached to avoid repeated file reads
|
| 248 |
+
tmpDir := t.TempDir()
|
| 249 |
+
p := filepath.Join(tmpDir, "nonexistent.json")
|
| 250 |
+
|
| 251 |
+
s := NewMultiSourceSecretWithPath("", p, 100*time.Millisecond)
|
| 252 |
+
ctx := context.Background()
|
| 253 |
+
|
| 254 |
+
// First call - file doesn't exist, should cache empty result
|
| 255 |
+
got1, err := s.Get(ctx)
|
| 256 |
+
if err != nil {
|
| 257 |
+
t.Fatalf("expected no error for missing file, got: %v", err)
|
| 258 |
+
}
|
| 259 |
+
if got1 != "" {
|
| 260 |
+
t.Fatalf("expected empty string, got %q", got1)
|
| 261 |
+
}
|
| 262 |
+
|
| 263 |
+
// Create the file now
|
| 264 |
+
if err := os.WriteFile(p, []byte(`{"apiKey@https://ampcode.com/":"new-value"}`), 0600); err != nil {
|
| 265 |
+
t.Fatal(err)
|
| 266 |
+
}
|
| 267 |
+
|
| 268 |
+
// Second call - should still return empty (cached), not read the new file
|
| 269 |
+
got2, _ := s.Get(ctx)
|
| 270 |
+
if got2 != "" {
|
| 271 |
+
t.Fatalf("cache should return empty, got %q", got2)
|
| 272 |
+
}
|
| 273 |
+
|
| 274 |
+
// After TTL expires, should see the new value
|
| 275 |
+
time.Sleep(110 * time.Millisecond)
|
| 276 |
+
got3, _ := s.Get(ctx)
|
| 277 |
+
if got3 != "new-value" {
|
| 278 |
+
t.Fatalf("after cache expiry, expected new-value, got %q", got3)
|
| 279 |
+
}
|
| 280 |
+
}
|
|
@@ -0,0 +1,92 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
// Package modules provides a pluggable routing module system for extending
|
| 2 |
+
// the API server with optional features without modifying core routing logic.
|
| 3 |
+
package modules
|
| 4 |
+
|
| 5 |
+
import (
|
| 6 |
+
"fmt"
|
| 7 |
+
|
| 8 |
+
"github.com/gin-gonic/gin"
|
| 9 |
+
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
| 10 |
+
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
|
| 11 |
+
)
|
| 12 |
+
|
| 13 |
+
// Context encapsulates the dependencies exposed to routing modules during
|
| 14 |
+
// registration. Modules can use the Gin engine to attach routes, the shared
|
| 15 |
+
// BaseAPIHandler for constructing SDK-specific handlers, and the resolved
|
| 16 |
+
// authentication middleware for protecting routes that require API keys.
|
| 17 |
+
type Context struct {
|
| 18 |
+
Engine *gin.Engine
|
| 19 |
+
BaseHandler *handlers.BaseAPIHandler
|
| 20 |
+
Config *config.Config
|
| 21 |
+
AuthMiddleware gin.HandlerFunc
|
| 22 |
+
}
|
| 23 |
+
|
| 24 |
+
// RouteModule represents a pluggable routing module that can register routes
|
| 25 |
+
// and handle configuration updates independently of the core server.
|
| 26 |
+
//
|
| 27 |
+
// DEPRECATED: Use RouteModuleV2 for new modules. This interface is kept for
|
| 28 |
+
// backwards compatibility and will be removed in a future version.
|
| 29 |
+
type RouteModule interface {
|
| 30 |
+
// Name returns a human-readable identifier for the module
|
| 31 |
+
Name() string
|
| 32 |
+
|
| 33 |
+
// Register sets up routes and handlers for this module.
|
| 34 |
+
// It receives the Gin engine, base handlers, and current configuration.
|
| 35 |
+
// Returns an error if registration fails (errors are logged but don't stop the server).
|
| 36 |
+
Register(engine *gin.Engine, baseHandler *handlers.BaseAPIHandler, cfg *config.Config) error
|
| 37 |
+
|
| 38 |
+
// OnConfigUpdated is called when the configuration is reloaded.
|
| 39 |
+
// Modules can respond to configuration changes here.
|
| 40 |
+
// Returns an error if the update cannot be applied.
|
| 41 |
+
OnConfigUpdated(cfg *config.Config) error
|
| 42 |
+
}
|
| 43 |
+
|
| 44 |
+
// RouteModuleV2 represents a pluggable bundle of routes that can integrate with
|
| 45 |
+
// the API server without modifying its core routing logic. Implementations can
|
| 46 |
+
// attach routes during Register and react to configuration updates via
|
| 47 |
+
// OnConfigUpdated.
|
| 48 |
+
//
|
| 49 |
+
// This is the preferred interface for new modules. It uses Context for cleaner
|
| 50 |
+
// dependency injection and supports idempotent registration.
|
| 51 |
+
type RouteModuleV2 interface {
|
| 52 |
+
// Name returns a unique identifier for logging and diagnostics.
|
| 53 |
+
Name() string
|
| 54 |
+
|
| 55 |
+
// Register wires the module's routes into the provided Gin engine. Modules
|
| 56 |
+
// should treat multiple calls as idempotent and avoid duplicate route
|
| 57 |
+
// registration when invoked more than once.
|
| 58 |
+
Register(ctx Context) error
|
| 59 |
+
|
| 60 |
+
// OnConfigUpdated notifies the module when the server configuration changes
|
| 61 |
+
// via hot reload. Implementations can refresh cached state or emit warnings.
|
| 62 |
+
OnConfigUpdated(cfg *config.Config) error
|
| 63 |
+
}
|
| 64 |
+
|
| 65 |
+
// RegisterModule is a helper that registers a module using either the V1 or V2
|
| 66 |
+
// interface. This allows gradual migration from V1 to V2 without breaking
|
| 67 |
+
// existing modules.
|
| 68 |
+
//
|
| 69 |
+
// Example usage:
|
| 70 |
+
//
|
| 71 |
+
// ctx := modules.Context{
|
| 72 |
+
// Engine: engine,
|
| 73 |
+
// BaseHandler: baseHandler,
|
| 74 |
+
// Config: cfg,
|
| 75 |
+
// AuthMiddleware: authMiddleware,
|
| 76 |
+
// }
|
| 77 |
+
// if err := modules.RegisterModule(ctx, ampModule); err != nil {
|
| 78 |
+
// log.Errorf("Failed to register module: %v", err)
|
| 79 |
+
// }
|
| 80 |
+
func RegisterModule(ctx Context, mod interface{}) error {
|
| 81 |
+
// Try V2 interface first (preferred)
|
| 82 |
+
if v2, ok := mod.(RouteModuleV2); ok {
|
| 83 |
+
return v2.Register(ctx)
|
| 84 |
+
}
|
| 85 |
+
|
| 86 |
+
// Fall back to V1 interface for backwards compatibility
|
| 87 |
+
if v1, ok := mod.(RouteModule); ok {
|
| 88 |
+
return v1.Register(ctx.Engine, ctx.BaseHandler, ctx.Config)
|
| 89 |
+
}
|
| 90 |
+
|
| 91 |
+
return fmt.Errorf("unsupported module type %T (must implement RouteModule or RouteModuleV2)", mod)
|
| 92 |
+
}
|
|
@@ -21,6 +21,8 @@ import (
|
|
| 21 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/access"
|
| 22 |
managementHandlers "github.com/router-for-me/CLIProxyAPI/v6/internal/api/handlers/management"
|
| 23 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/middleware"
|
|
|
|
|
|
|
| 24 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
| 25 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
|
| 26 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset"
|
|
@@ -261,6 +263,20 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
|
|
| 261 |
|
| 262 |
// Setup routes
|
| 263 |
s.setupRoutes()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 264 |
if optionState.routerConfigurator != nil {
|
| 265 |
optionState.routerConfigurator(engine, s.handlers, cfg)
|
| 266 |
}
|
|
|
|
| 21 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/access"
|
| 22 |
managementHandlers "github.com/router-for-me/CLIProxyAPI/v6/internal/api/handlers/management"
|
| 23 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/middleware"
|
| 24 |
+
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules"
|
| 25 |
+
ampmodule "github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules/amp"
|
| 26 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
| 27 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
|
| 28 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset"
|
|
|
|
| 263 |
|
| 264 |
// Setup routes
|
| 265 |
s.setupRoutes()
|
| 266 |
+
|
| 267 |
+
// Register Amp module using V2 interface with Context
|
| 268 |
+
ampModule := ampmodule.NewLegacy(accessManager, AuthMiddleware(accessManager))
|
| 269 |
+
ctx := modules.Context{
|
| 270 |
+
Engine: engine,
|
| 271 |
+
BaseHandler: s.handlers,
|
| 272 |
+
Config: cfg,
|
| 273 |
+
AuthMiddleware: AuthMiddleware(accessManager),
|
| 274 |
+
}
|
| 275 |
+
if err := modules.RegisterModule(ctx, ampModule); err != nil {
|
| 276 |
+
log.Errorf("Failed to register Amp module: %v", err)
|
| 277 |
+
}
|
| 278 |
+
|
| 279 |
+
// Apply additional router configurators from options
|
| 280 |
if optionState.routerConfigurator != nil {
|
| 281 |
optionState.routerConfigurator(engine, s.handlers, cfg)
|
| 282 |
}
|
|
@@ -0,0 +1,111 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
package api
|
| 2 |
+
|
| 3 |
+
import (
|
| 4 |
+
"net/http"
|
| 5 |
+
"net/http/httptest"
|
| 6 |
+
"os"
|
| 7 |
+
"path/filepath"
|
| 8 |
+
"strings"
|
| 9 |
+
"testing"
|
| 10 |
+
|
| 11 |
+
gin "github.com/gin-gonic/gin"
|
| 12 |
+
proxyconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
| 13 |
+
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
|
| 14 |
+
"github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
| 15 |
+
sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
|
| 16 |
+
)
|
| 17 |
+
|
| 18 |
+
func newTestServer(t *testing.T) *Server {
|
| 19 |
+
t.Helper()
|
| 20 |
+
|
| 21 |
+
gin.SetMode(gin.TestMode)
|
| 22 |
+
|
| 23 |
+
tmpDir := t.TempDir()
|
| 24 |
+
authDir := filepath.Join(tmpDir, "auth")
|
| 25 |
+
if err := os.MkdirAll(authDir, 0o700); err != nil {
|
| 26 |
+
t.Fatalf("failed to create auth dir: %v", err)
|
| 27 |
+
}
|
| 28 |
+
|
| 29 |
+
cfg := &proxyconfig.Config{
|
| 30 |
+
SDKConfig: sdkconfig.SDKConfig{
|
| 31 |
+
APIKeys: []string{"test-key"},
|
| 32 |
+
},
|
| 33 |
+
Port: 0,
|
| 34 |
+
AuthDir: authDir,
|
| 35 |
+
Debug: true,
|
| 36 |
+
LoggingToFile: false,
|
| 37 |
+
UsageStatisticsEnabled: false,
|
| 38 |
+
}
|
| 39 |
+
|
| 40 |
+
authManager := auth.NewManager(nil, nil, nil)
|
| 41 |
+
accessManager := sdkaccess.NewManager()
|
| 42 |
+
|
| 43 |
+
configPath := filepath.Join(tmpDir, "config.yaml")
|
| 44 |
+
return NewServer(cfg, authManager, accessManager, configPath)
|
| 45 |
+
}
|
| 46 |
+
|
| 47 |
+
func TestAmpProviderModelRoutes(t *testing.T) {
|
| 48 |
+
testCases := []struct {
|
| 49 |
+
name string
|
| 50 |
+
path string
|
| 51 |
+
wantStatus int
|
| 52 |
+
wantContains string
|
| 53 |
+
}{
|
| 54 |
+
{
|
| 55 |
+
name: "openai root models",
|
| 56 |
+
path: "/api/provider/openai/models",
|
| 57 |
+
wantStatus: http.StatusOK,
|
| 58 |
+
wantContains: `"object":"list"`,
|
| 59 |
+
},
|
| 60 |
+
{
|
| 61 |
+
name: "groq root models",
|
| 62 |
+
path: "/api/provider/groq/models",
|
| 63 |
+
wantStatus: http.StatusOK,
|
| 64 |
+
wantContains: `"object":"list"`,
|
| 65 |
+
},
|
| 66 |
+
{
|
| 67 |
+
name: "openai models",
|
| 68 |
+
path: "/api/provider/openai/v1/models",
|
| 69 |
+
wantStatus: http.StatusOK,
|
| 70 |
+
wantContains: `"object":"list"`,
|
| 71 |
+
},
|
| 72 |
+
{
|
| 73 |
+
name: "anthropic models",
|
| 74 |
+
path: "/api/provider/anthropic/v1/models",
|
| 75 |
+
wantStatus: http.StatusOK,
|
| 76 |
+
wantContains: `"data"`,
|
| 77 |
+
},
|
| 78 |
+
{
|
| 79 |
+
name: "google models v1",
|
| 80 |
+
path: "/api/provider/google/v1/models",
|
| 81 |
+
wantStatus: http.StatusOK,
|
| 82 |
+
wantContains: `"models"`,
|
| 83 |
+
},
|
| 84 |
+
{
|
| 85 |
+
name: "google models v1beta",
|
| 86 |
+
path: "/api/provider/google/v1beta/models",
|
| 87 |
+
wantStatus: http.StatusOK,
|
| 88 |
+
wantContains: `"models"`,
|
| 89 |
+
},
|
| 90 |
+
}
|
| 91 |
+
|
| 92 |
+
for _, tc := range testCases {
|
| 93 |
+
tc := tc
|
| 94 |
+
t.Run(tc.name, func(t *testing.T) {
|
| 95 |
+
server := newTestServer(t)
|
| 96 |
+
|
| 97 |
+
req := httptest.NewRequest(http.MethodGet, tc.path, nil)
|
| 98 |
+
req.Header.Set("Authorization", "Bearer test-key")
|
| 99 |
+
|
| 100 |
+
rr := httptest.NewRecorder()
|
| 101 |
+
server.engine.ServeHTTP(rr, req)
|
| 102 |
+
|
| 103 |
+
if rr.Code != tc.wantStatus {
|
| 104 |
+
t.Fatalf("unexpected status code for %s: got %d want %d; body=%s", tc.path, rr.Code, tc.wantStatus, rr.Body.String())
|
| 105 |
+
}
|
| 106 |
+
if body := rr.Body.String(); !strings.Contains(body, tc.wantContains) {
|
| 107 |
+
t.Fatalf("response body for %s missing %q: %s", tc.path, tc.wantContains, body)
|
| 108 |
+
}
|
| 109 |
+
})
|
| 110 |
+
}
|
| 111 |
+
}
|
|
@@ -23,6 +23,17 @@ type Config struct {
|
|
| 23 |
// Port is the network port on which the API server will listen.
|
| 24 |
Port int `yaml:"port" json:"-"`
|
| 25 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 26 |
// AuthDir is the directory where authentication token files are stored.
|
| 27 |
AuthDir string `yaml:"auth-dir" json:"-"`
|
| 28 |
|
|
@@ -258,6 +269,7 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
|
|
| 258 |
cfg.LoggingToFile = false
|
| 259 |
cfg.UsageStatisticsEnabled = false
|
| 260 |
cfg.DisableCooling = false
|
|
|
|
| 261 |
if err = yaml.Unmarshal(data, &cfg); err != nil {
|
| 262 |
if optional {
|
| 263 |
// In cloud deploy mode, if YAML parsing fails, return empty config instead of error.
|
|
|
|
| 23 |
// Port is the network port on which the API server will listen.
|
| 24 |
Port int `yaml:"port" json:"-"`
|
| 25 |
|
| 26 |
+
// AmpUpstreamURL defines the upstream Amp control plane used for non-provider calls.
|
| 27 |
+
AmpUpstreamURL string `yaml:"amp-upstream-url" json:"amp-upstream-url"`
|
| 28 |
+
|
| 29 |
+
// AmpUpstreamAPIKey optionally overrides the Authorization header when proxying Amp upstream calls.
|
| 30 |
+
AmpUpstreamAPIKey string `yaml:"amp-upstream-api-key" json:"amp-upstream-api-key"`
|
| 31 |
+
|
| 32 |
+
// AmpRestrictManagementToLocalhost restricts Amp management routes (/api/user, /api/threads, etc.)
|
| 33 |
+
// to only accept connections from localhost (127.0.0.1, ::1). When true, prevents drive-by
|
| 34 |
+
// browser attacks and remote access to management endpoints. Default: true (recommended).
|
| 35 |
+
AmpRestrictManagementToLocalhost bool `yaml:"amp-restrict-management-to-localhost" json:"amp-restrict-management-to-localhost"`
|
| 36 |
+
|
| 37 |
// AuthDir is the directory where authentication token files are stored.
|
| 38 |
AuthDir string `yaml:"auth-dir" json:"-"`
|
| 39 |
|
|
|
|
| 269 |
cfg.LoggingToFile = false
|
| 270 |
cfg.UsageStatisticsEnabled = false
|
| 271 |
cfg.DisableCooling = false
|
| 272 |
+
cfg.AmpRestrictManagementToLocalhost = true // Default to secure: only localhost access
|
| 273 |
if err = yaml.Unmarshal(data, &cfg); err != nil {
|
| 274 |
if optional {
|
| 275 |
// In cloud deploy mode, if YAML parsing fails, return empty config instead of error.
|
|
@@ -8,9 +8,12 @@ package claude
|
|
| 8 |
|
| 9 |
import (
|
| 10 |
"bufio"
|
|
|
|
|
|
|
| 11 |
"context"
|
| 12 |
"encoding/json"
|
| 13 |
"fmt"
|
|
|
|
| 14 |
"net/http"
|
| 15 |
"time"
|
| 16 |
|
|
@@ -19,6 +22,7 @@ import (
|
|
| 19 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
|
| 20 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
| 21 |
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
|
|
|
|
| 22 |
"github.com/tidwall/gjson"
|
| 23 |
)
|
| 24 |
|
|
@@ -153,6 +157,23 @@ func (h *ClaudeCodeAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSO
|
|
| 153 |
cliCancel(errMsg.Error)
|
| 154 |
return
|
| 155 |
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 156 |
_, _ = c.Writer.Write(resp)
|
| 157 |
cliCancel()
|
| 158 |
}
|
|
|
|
| 8 |
|
| 9 |
import (
|
| 10 |
"bufio"
|
| 11 |
+
"bytes"
|
| 12 |
+
"compress/gzip"
|
| 13 |
"context"
|
| 14 |
"encoding/json"
|
| 15 |
"fmt"
|
| 16 |
+
"io"
|
| 17 |
"net/http"
|
| 18 |
"time"
|
| 19 |
|
|
|
|
| 22 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
|
| 23 |
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
| 24 |
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
|
| 25 |
+
log "github.com/sirupsen/logrus"
|
| 26 |
"github.com/tidwall/gjson"
|
| 27 |
)
|
| 28 |
|
|
|
|
| 157 |
cliCancel(errMsg.Error)
|
| 158 |
return
|
| 159 |
}
|
| 160 |
+
|
| 161 |
+
// Decompress gzipped responses - Claude API sometimes returns gzip without Content-Encoding header
|
| 162 |
+
// This fixes title generation and other non-streaming responses that arrive compressed
|
| 163 |
+
if len(resp) >= 2 && resp[0] == 0x1f && resp[1] == 0x8b {
|
| 164 |
+
gzReader, err := gzip.NewReader(bytes.NewReader(resp))
|
| 165 |
+
if err != nil {
|
| 166 |
+
log.Warnf("failed to decompress gzipped Claude response: %v", err)
|
| 167 |
+
} else {
|
| 168 |
+
defer gzReader.Close()
|
| 169 |
+
if decompressed, err := io.ReadAll(gzReader); err != nil {
|
| 170 |
+
log.Warnf("failed to read decompressed Claude response: %v", err)
|
| 171 |
+
} else {
|
| 172 |
+
resp = decompressed
|
| 173 |
+
}
|
| 174 |
+
}
|
| 175 |
+
}
|
| 176 |
+
|
| 177 |
_, _ = c.Writer.Write(resp)
|
| 178 |
cliCancel()
|
| 179 |
}
|
|
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
// Package httpx provides HTTP transport utilities for SDK clients,
|
| 2 |
+
// including automatic gzip decompression for misconfigured upstreams.
|
| 3 |
+
package httpx
|
| 4 |
+
|
| 5 |
+
import (
|
| 6 |
+
"bytes"
|
| 7 |
+
"compress/gzip"
|
| 8 |
+
"io"
|
| 9 |
+
)
|
| 10 |
+
|
| 11 |
+
// DecodePossibleGzip inspects the raw response body and transparently
|
| 12 |
+
// decompresses it when the payload is gzip compressed. Some upstream
|
| 13 |
+
// providers return gzip data without a Content-Encoding header, which
|
| 14 |
+
// confuses clients expecting JSON. This helper restores the original
|
| 15 |
+
// JSON bytes while leaving plain responses untouched.
|
| 16 |
+
//
|
| 17 |
+
// This function is preserved for backward compatibility but new code
|
| 18 |
+
// should use GzipFixupTransport instead.
|
| 19 |
+
func DecodePossibleGzip(raw []byte) ([]byte, error) {
|
| 20 |
+
if len(raw) >= 2 && raw[0] == 0x1f && raw[1] == 0x8b {
|
| 21 |
+
reader, err := gzip.NewReader(bytes.NewReader(raw))
|
| 22 |
+
if err != nil {
|
| 23 |
+
return nil, err
|
| 24 |
+
}
|
| 25 |
+
decompressed, err := io.ReadAll(reader)
|
| 26 |
+
_ = reader.Close()
|
| 27 |
+
if err != nil {
|
| 28 |
+
return nil, err
|
| 29 |
+
}
|
| 30 |
+
return decompressed, nil
|
| 31 |
+
}
|
| 32 |
+
return raw, nil
|
| 33 |
+
}
|
|
@@ -0,0 +1,177 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
package httpx
|
| 2 |
+
|
| 3 |
+
import (
|
| 4 |
+
"bytes"
|
| 5 |
+
"compress/gzip"
|
| 6 |
+
"io"
|
| 7 |
+
"net/http"
|
| 8 |
+
"strings"
|
| 9 |
+
|
| 10 |
+
log "github.com/sirupsen/logrus"
|
| 11 |
+
)
|
| 12 |
+
|
| 13 |
+
// GzipFixupTransport wraps an http.RoundTripper to auto-decode gzip responses
|
| 14 |
+
// that don't properly set Content-Encoding header.
|
| 15 |
+
//
|
| 16 |
+
// Some upstream providers (especially when proxied) return gzip-compressed
|
| 17 |
+
// responses without setting the Content-Encoding: gzip header, which causes
|
| 18 |
+
// Go's http client to pass the compressed bytes directly to the application.
|
| 19 |
+
//
|
| 20 |
+
// This transport detects gzip magic bytes and transparently decompresses
|
| 21 |
+
// the response while preserving streaming behavior for SSE and chunked responses.
|
| 22 |
+
type GzipFixupTransport struct {
|
| 23 |
+
// Base is the underlying transport. If nil, http.DefaultTransport is used.
|
| 24 |
+
Base http.RoundTripper
|
| 25 |
+
}
|
| 26 |
+
|
| 27 |
+
// RoundTrip implements http.RoundTripper
|
| 28 |
+
func (t *GzipFixupTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
| 29 |
+
base := t.Base
|
| 30 |
+
if base == nil {
|
| 31 |
+
base = http.DefaultTransport
|
| 32 |
+
}
|
| 33 |
+
|
| 34 |
+
resp, err := base.RoundTrip(req)
|
| 35 |
+
if err != nil || resp == nil {
|
| 36 |
+
return resp, err
|
| 37 |
+
}
|
| 38 |
+
|
| 39 |
+
// Skip if Go already decompressed it
|
| 40 |
+
if resp.Uncompressed {
|
| 41 |
+
return resp, nil
|
| 42 |
+
}
|
| 43 |
+
|
| 44 |
+
// Skip if Content-Encoding is already set (properly configured upstream)
|
| 45 |
+
if resp.Header.Get("Content-Encoding") != "" {
|
| 46 |
+
return resp, nil
|
| 47 |
+
}
|
| 48 |
+
|
| 49 |
+
// Skip streaming responses - they need different handling
|
| 50 |
+
if isStreamingResponse(resp) {
|
| 51 |
+
// For streaming responses, wrap with a streaming gzip detector
|
| 52 |
+
// that can handle chunked gzip data
|
| 53 |
+
resp.Body = &streamingGzipDetector{
|
| 54 |
+
inner: resp.Body,
|
| 55 |
+
}
|
| 56 |
+
return resp, nil
|
| 57 |
+
}
|
| 58 |
+
|
| 59 |
+
// For non-streaming responses, peek and decompress if needed
|
| 60 |
+
resp.Body = &gzipDetectingReader{
|
| 61 |
+
inner: resp.Body,
|
| 62 |
+
}
|
| 63 |
+
|
| 64 |
+
return resp, nil
|
| 65 |
+
}
|
| 66 |
+
|
| 67 |
+
// isStreamingResponse checks if response is SSE or chunked
|
| 68 |
+
func isStreamingResponse(resp *http.Response) bool {
|
| 69 |
+
contentType := resp.Header.Get("Content-Type")
|
| 70 |
+
|
| 71 |
+
// Check for Server-Sent Events
|
| 72 |
+
if strings.Contains(contentType, "text/event-stream") {
|
| 73 |
+
return true
|
| 74 |
+
}
|
| 75 |
+
|
| 76 |
+
// Check for chunked transfer encoding
|
| 77 |
+
if strings.Contains(strings.ToLower(resp.Header.Get("Transfer-Encoding")), "chunked") {
|
| 78 |
+
return true
|
| 79 |
+
}
|
| 80 |
+
|
| 81 |
+
return false
|
| 82 |
+
}
|
| 83 |
+
|
| 84 |
+
// gzipDetectingReader is an io.ReadCloser that detects gzip magic bytes
|
| 85 |
+
// on first read and switches to gzip decompression if detected.
|
| 86 |
+
// This is used for non-streaming responses.
|
| 87 |
+
type gzipDetectingReader struct {
|
| 88 |
+
inner io.ReadCloser
|
| 89 |
+
reader io.Reader
|
| 90 |
+
once bool
|
| 91 |
+
}
|
| 92 |
+
|
| 93 |
+
func (g *gzipDetectingReader) Read(p []byte) (int, error) {
|
| 94 |
+
if !g.once {
|
| 95 |
+
g.once = true
|
| 96 |
+
|
| 97 |
+
// Peek at first 2 bytes to detect gzip magic bytes
|
| 98 |
+
buf := make([]byte, 2)
|
| 99 |
+
n, err := io.ReadFull(g.inner, buf)
|
| 100 |
+
if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
|
| 101 |
+
// Can't peek, use original reader
|
| 102 |
+
g.reader = io.MultiReader(bytes.NewReader(buf[:n]), g.inner)
|
| 103 |
+
return g.reader.Read(p)
|
| 104 |
+
}
|
| 105 |
+
|
| 106 |
+
if n >= 2 && buf[0] == 0x1f && buf[1] == 0x8b {
|
| 107 |
+
// It's gzipped, create gzip reader
|
| 108 |
+
multiReader := io.MultiReader(bytes.NewReader(buf[:n]), g.inner)
|
| 109 |
+
gzipReader, err := gzip.NewReader(multiReader)
|
| 110 |
+
if err != nil {
|
| 111 |
+
log.Warnf("gzip header detected but reader creation failed: %v", err)
|
| 112 |
+
g.reader = multiReader
|
| 113 |
+
} else {
|
| 114 |
+
g.reader = gzipReader
|
| 115 |
+
}
|
| 116 |
+
} else {
|
| 117 |
+
// Not gzipped, combine peeked bytes with rest
|
| 118 |
+
g.reader = io.MultiReader(bytes.NewReader(buf[:n]), g.inner)
|
| 119 |
+
}
|
| 120 |
+
}
|
| 121 |
+
|
| 122 |
+
return g.reader.Read(p)
|
| 123 |
+
}
|
| 124 |
+
|
| 125 |
+
func (g *gzipDetectingReader) Close() error {
|
| 126 |
+
if closer, ok := g.reader.(io.Closer); ok {
|
| 127 |
+
_ = closer.Close()
|
| 128 |
+
}
|
| 129 |
+
return g.inner.Close()
|
| 130 |
+
}
|
| 131 |
+
|
| 132 |
+
// streamingGzipDetector is similar to gzipDetectingReader but designed for
|
| 133 |
+
// streaming responses. It doesn't buffer; it wraps with a streaming gzip reader.
|
| 134 |
+
type streamingGzipDetector struct {
|
| 135 |
+
inner io.ReadCloser
|
| 136 |
+
reader io.Reader
|
| 137 |
+
once bool
|
| 138 |
+
}
|
| 139 |
+
|
| 140 |
+
func (s *streamingGzipDetector) Read(p []byte) (int, error) {
|
| 141 |
+
if !s.once {
|
| 142 |
+
s.once = true
|
| 143 |
+
|
| 144 |
+
// Peek at first 2 bytes
|
| 145 |
+
buf := make([]byte, 2)
|
| 146 |
+
n, err := io.ReadFull(s.inner, buf)
|
| 147 |
+
if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
|
| 148 |
+
s.reader = io.MultiReader(bytes.NewReader(buf[:n]), s.inner)
|
| 149 |
+
return s.reader.Read(p)
|
| 150 |
+
}
|
| 151 |
+
|
| 152 |
+
if n >= 2 && buf[0] == 0x1f && buf[1] == 0x8b {
|
| 153 |
+
// It's gzipped - wrap with streaming gzip reader
|
| 154 |
+
multiReader := io.MultiReader(bytes.NewReader(buf[:n]), s.inner)
|
| 155 |
+
gzipReader, err := gzip.NewReader(multiReader)
|
| 156 |
+
if err != nil {
|
| 157 |
+
log.Warnf("streaming gzip header detected but reader creation failed: %v", err)
|
| 158 |
+
s.reader = multiReader
|
| 159 |
+
} else {
|
| 160 |
+
s.reader = gzipReader
|
| 161 |
+
log.Debug("streaming gzip decompression enabled")
|
| 162 |
+
}
|
| 163 |
+
} else {
|
| 164 |
+
// Not gzipped
|
| 165 |
+
s.reader = io.MultiReader(bytes.NewReader(buf[:n]), s.inner)
|
| 166 |
+
}
|
| 167 |
+
}
|
| 168 |
+
|
| 169 |
+
return s.reader.Read(p)
|
| 170 |
+
}
|
| 171 |
+
|
| 172 |
+
func (s *streamingGzipDetector) Close() error {
|
| 173 |
+
if closer, ok := s.reader.(io.Closer); ok {
|
| 174 |
+
_ = closer.Close()
|
| 175 |
+
}
|
| 176 |
+
return s.inner.Close()
|
| 177 |
+
}
|