mirror of
https://github.com/router-for-me/CLIProxyAPIPlus.git
synced 2026-03-12 08:43:58 +00:00
Compare commits
20 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f8aba62860 | ||
|
|
a7eeb06f3d | ||
|
|
9426be7a5c | ||
|
|
4a135f1986 | ||
|
|
c4c02f4ad0 | ||
|
|
b87b9b455f | ||
|
|
db03ae9663 | ||
|
|
969ff6bb68 | ||
|
|
e24af6e545 | ||
|
|
bceecfb2e3 | ||
|
|
48dd987867 | ||
|
|
6a2906e3e5 | ||
|
|
d72886c801 | ||
|
|
6efba3d829 | ||
|
|
6b60bdd139 | ||
|
|
897c40bed8 | ||
|
|
373ea8d7e4 | ||
|
|
b5de004c01 | ||
|
|
94ec772521 | ||
|
|
e216d26731 |
@@ -20,25 +20,6 @@ func (h *Handler) GetConfig(c *gin.Context) {
|
||||
c.JSON(200, &cfgCopy)
|
||||
}
|
||||
|
||||
func (h *Handler) GetConfigYAML(c *gin.Context) {
|
||||
data, err := os.ReadFile(h.configFilePath)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "read_failed", "message": err.Error()})
|
||||
return
|
||||
}
|
||||
var node yaml.Node
|
||||
if err = yaml.Unmarshal(data, &node); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "parse_failed", "message": err.Error()})
|
||||
return
|
||||
}
|
||||
c.Header("Content-Type", "application/yaml; charset=utf-8")
|
||||
c.Header("Vary", "format, Accept")
|
||||
enc := yaml.NewEncoder(c.Writer)
|
||||
enc.SetIndent(2)
|
||||
_ = enc.Encode(&node)
|
||||
_ = enc.Close()
|
||||
}
|
||||
|
||||
func WriteConfig(path string, data []byte) error {
|
||||
data = config.NormalizeCommentIndentation(data)
|
||||
f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
|
||||
@@ -110,9 +91,9 @@ func (h *Handler) PutConfigYAML(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "changed": []string{"config"}})
|
||||
}
|
||||
|
||||
// GetConfigFile returns the raw config.yaml file bytes without re-encoding.
|
||||
// GetConfigYAML returns the raw config.yaml file bytes without re-encoding.
|
||||
// It preserves comments and original formatting/styles.
|
||||
func (h *Handler) GetConfigFile(c *gin.Context) {
|
||||
func (h *Handler) GetConfigYAML(c *gin.Context) {
|
||||
data, err := os.ReadFile(h.configFilePath)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
|
||||
@@ -104,52 +104,6 @@ func (h *Handler) deleteFromStringList(c *gin.Context, target *[]string, after f
|
||||
c.JSON(400, gin.H{"error": "missing index or value"})
|
||||
}
|
||||
|
||||
func sanitizeStringSlice(in []string) []string {
|
||||
out := make([]string, 0, len(in))
|
||||
for i := range in {
|
||||
if trimmed := strings.TrimSpace(in[i]); trimmed != "" {
|
||||
out = append(out, trimmed)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func geminiKeyStringsFromConfig(cfg *config.Config) []string {
|
||||
if cfg == nil || len(cfg.GeminiKey) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]string, 0, len(cfg.GeminiKey))
|
||||
for i := range cfg.GeminiKey {
|
||||
if key := strings.TrimSpace(cfg.GeminiKey[i].APIKey); key != "" {
|
||||
out = append(out, key)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (h *Handler) applyLegacyKeys(keys []string) {
|
||||
if h == nil || h.cfg == nil {
|
||||
return
|
||||
}
|
||||
sanitized := sanitizeStringSlice(keys)
|
||||
existing := make(map[string]config.GeminiKey, len(h.cfg.GeminiKey))
|
||||
for _, entry := range h.cfg.GeminiKey {
|
||||
if key := strings.TrimSpace(entry.APIKey); key != "" {
|
||||
existing[key] = entry
|
||||
}
|
||||
}
|
||||
newList := make([]config.GeminiKey, 0, len(sanitized))
|
||||
for _, key := range sanitized {
|
||||
if entry, ok := existing[key]; ok {
|
||||
newList = append(newList, entry)
|
||||
} else {
|
||||
newList = append(newList, config.GeminiKey{APIKey: key})
|
||||
}
|
||||
}
|
||||
h.cfg.GeminiKey = newList
|
||||
h.cfg.SanitizeGeminiKeys()
|
||||
}
|
||||
|
||||
// api-keys
|
||||
func (h *Handler) GetAPIKeys(c *gin.Context) { c.JSON(200, gin.H{"api-keys": h.cfg.APIKeys}) }
|
||||
func (h *Handler) PutAPIKeys(c *gin.Context) {
|
||||
@@ -165,24 +119,6 @@ func (h *Handler) DeleteAPIKeys(c *gin.Context) {
|
||||
h.deleteFromStringList(c, &h.cfg.APIKeys, func() { h.cfg.Access.Providers = nil })
|
||||
}
|
||||
|
||||
// generative-language-api-key
|
||||
func (h *Handler) GetGlKeys(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"generative-language-api-key": geminiKeyStringsFromConfig(h.cfg)})
|
||||
}
|
||||
func (h *Handler) PutGlKeys(c *gin.Context) {
|
||||
h.putStringList(c, func(v []string) {
|
||||
h.applyLegacyKeys(v)
|
||||
}, nil)
|
||||
}
|
||||
func (h *Handler) PatchGlKeys(c *gin.Context) {
|
||||
target := append([]string(nil), geminiKeyStringsFromConfig(h.cfg)...)
|
||||
h.patchStringList(c, &target, func() { h.applyLegacyKeys(target) })
|
||||
}
|
||||
func (h *Handler) DeleteGlKeys(c *gin.Context) {
|
||||
target := append([]string(nil), geminiKeyStringsFromConfig(h.cfg)...)
|
||||
h.deleteFromStringList(c, &target, func() { h.applyLegacyKeys(target) })
|
||||
}
|
||||
|
||||
// gemini-api-key: []GeminiKey
|
||||
func (h *Handler) GetGeminiKeys(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"gemini-api-key": h.cfg.GeminiKey})
|
||||
|
||||
@@ -27,11 +27,20 @@ type Option func(*AmpModule)
|
||||
type AmpModule struct {
|
||||
secretSource SecretSource
|
||||
proxy *httputil.ReverseProxy
|
||||
proxyMu sync.RWMutex // protects proxy for hot-reload
|
||||
accessManager *sdkaccess.Manager
|
||||
authMiddleware_ gin.HandlerFunc
|
||||
modelMapper *DefaultModelMapper
|
||||
enabled bool
|
||||
registerOnce sync.Once
|
||||
|
||||
// restrictToLocalhost controls localhost-only access for management routes (hot-reloadable)
|
||||
restrictToLocalhost bool
|
||||
restrictMu sync.RWMutex
|
||||
|
||||
// configMu protects lastConfig for partial reload comparison
|
||||
configMu sync.RWMutex
|
||||
lastConfig *config.AmpCode
|
||||
}
|
||||
|
||||
// New creates a new Amp routing module with the given options.
|
||||
@@ -107,6 +116,13 @@ func (m *AmpModule) Register(ctx modules.Context) error {
|
||||
// Initialize model mapper from config (for routing unavailable models to alternatives)
|
||||
m.modelMapper = NewModelMapper(settings.ModelMappings)
|
||||
|
||||
// Store initial config for partial reload comparison
|
||||
settingsCopy := settings
|
||||
m.lastConfig = &settingsCopy
|
||||
|
||||
// Initialize localhost restriction setting (hot-reloadable)
|
||||
m.setRestrictToLocalhost(settings.RestrictManagementToLocalhost)
|
||||
|
||||
// Always register provider aliases - these work without an upstream
|
||||
m.registerProviderAliases(ctx.Engine, ctx.BaseHandler, auth)
|
||||
|
||||
@@ -131,13 +147,12 @@ func (m *AmpModule) Register(ctx modules.Context) error {
|
||||
return
|
||||
}
|
||||
|
||||
m.proxy = proxy
|
||||
m.setProxy(proxy)
|
||||
m.enabled = true
|
||||
|
||||
// Register management proxy routes (requires upstream)
|
||||
// Restrict to localhost by default for security (prevents drive-by browser attacks)
|
||||
handler := proxyHandler(proxy)
|
||||
m.registerManagementRoutes(ctx.Engine, ctx.BaseHandler, handler, settings.RestrictManagementToLocalhost)
|
||||
// Uses dynamic middleware that checks m.IsRestrictedToLocalhost() for hot-reload support
|
||||
m.registerManagementRoutes(ctx.Engine, ctx.BaseHandler)
|
||||
|
||||
log.Infof("amp upstream proxy enabled for: %s", upstreamURL)
|
||||
log.Debug("amp provider alias routes registered")
|
||||
@@ -162,44 +177,165 @@ func (m *AmpModule) getAuthMiddleware(ctx modules.Context) gin.HandlerFunc {
|
||||
}
|
||||
}
|
||||
|
||||
// OnConfigUpdated handles configuration updates.
|
||||
// Currently requires restart for URL changes (could be enhanced for dynamic updates).
|
||||
// OnConfigUpdated handles configuration updates with partial reload support.
|
||||
// Only updates components that have actually changed to avoid unnecessary work.
|
||||
// Supports hot-reload for: model-mappings, upstream-api-key, upstream-url, restrict-management-to-localhost.
|
||||
func (m *AmpModule) OnConfigUpdated(cfg *config.Config) error {
|
||||
settings := cfg.AmpCode
|
||||
newSettings := cfg.AmpCode
|
||||
|
||||
// Update model mappings (hot-reload supported)
|
||||
if m.modelMapper != nil {
|
||||
m.modelMapper.UpdateMappings(settings.ModelMappings)
|
||||
if m.enabled {
|
||||
log.Infof("amp config updated: reloading %d model mapping(s)", len(settings.ModelMappings))
|
||||
}
|
||||
} else if m.enabled {
|
||||
log.Warnf("amp model mapper not initialized, skipping model mapping update")
|
||||
}
|
||||
// Get previous config for comparison
|
||||
m.configMu.RLock()
|
||||
oldSettings := m.lastConfig
|
||||
m.configMu.RUnlock()
|
||||
|
||||
if !m.enabled {
|
||||
return nil
|
||||
}
|
||||
// Track what changed for logging
|
||||
var changes []string
|
||||
|
||||
upstreamURL := strings.TrimSpace(settings.UpstreamURL)
|
||||
if upstreamURL == "" {
|
||||
log.Warn("amp upstream URL removed from config, restart required to disable")
|
||||
return nil
|
||||
}
|
||||
|
||||
// If API key changed, invalidate the cache
|
||||
if m.secretSource != nil {
|
||||
if ms, ok := m.secretSource.(*MultiSourceSecret); ok {
|
||||
ms.InvalidateCache()
|
||||
log.Debug("amp secret cache invalidated due to config update")
|
||||
// Check model mappings change
|
||||
modelMappingsChanged := m.hasModelMappingsChanged(oldSettings, &newSettings)
|
||||
if modelMappingsChanged {
|
||||
if m.modelMapper != nil {
|
||||
m.modelMapper.UpdateMappings(newSettings.ModelMappings)
|
||||
changes = append(changes, "model-mappings")
|
||||
if m.enabled {
|
||||
log.Infof("amp config partial reload: model mappings updated (%d entries)", len(newSettings.ModelMappings))
|
||||
}
|
||||
} else if m.enabled {
|
||||
log.Warnf("amp model mapper not initialized, skipping model mapping update")
|
||||
}
|
||||
}
|
||||
|
||||
log.Debug("amp config updated (restart required for URL changes)")
|
||||
if m.enabled {
|
||||
// Check upstream URL change - now supports hot-reload
|
||||
newUpstreamURL := strings.TrimSpace(newSettings.UpstreamURL)
|
||||
oldUpstreamURL := ""
|
||||
if oldSettings != nil {
|
||||
oldUpstreamURL = strings.TrimSpace(oldSettings.UpstreamURL)
|
||||
}
|
||||
|
||||
if newUpstreamURL == "" && oldUpstreamURL != "" {
|
||||
log.Warn("amp upstream URL removed from config, proxy has been disabled")
|
||||
m.setProxy(nil)
|
||||
changes = append(changes, "upstream-url(disabled)")
|
||||
} else if newUpstreamURL != oldUpstreamURL && newUpstreamURL != "" {
|
||||
// Recreate proxy with new URL
|
||||
proxy, err := createReverseProxy(newUpstreamURL, m.secretSource)
|
||||
if err != nil {
|
||||
log.Errorf("amp config: failed to create proxy for new upstream URL %s: %v", newUpstreamURL, err)
|
||||
} else {
|
||||
m.setProxy(proxy)
|
||||
changes = append(changes, "upstream-url")
|
||||
log.Infof("amp config partial reload: upstream URL updated (%s -> %s)", oldUpstreamURL, newUpstreamURL)
|
||||
}
|
||||
}
|
||||
|
||||
// Check API key change
|
||||
apiKeyChanged := m.hasAPIKeyChanged(oldSettings, &newSettings)
|
||||
if apiKeyChanged {
|
||||
if m.secretSource != nil {
|
||||
if ms, ok := m.secretSource.(*MultiSourceSecret); ok {
|
||||
ms.UpdateExplicitKey(newSettings.UpstreamAPIKey)
|
||||
ms.InvalidateCache()
|
||||
changes = append(changes, "upstream-api-key")
|
||||
log.Debug("amp config partial reload: secret cache invalidated")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check restrict-management-to-localhost change - now supports hot-reload
|
||||
if oldSettings != nil && oldSettings.RestrictManagementToLocalhost != newSettings.RestrictManagementToLocalhost {
|
||||
m.setRestrictToLocalhost(newSettings.RestrictManagementToLocalhost)
|
||||
changes = append(changes, "restrict-management-to-localhost")
|
||||
if newSettings.RestrictManagementToLocalhost {
|
||||
log.Infof("amp config partial reload: management routes now restricted to localhost")
|
||||
} else {
|
||||
log.Warnf("amp config partial reload: management routes now accessible from any IP - this is insecure!")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Store current config for next comparison
|
||||
m.configMu.Lock()
|
||||
settingsCopy := newSettings // copy struct
|
||||
m.lastConfig = &settingsCopy
|
||||
m.configMu.Unlock()
|
||||
|
||||
// Log summary if any changes detected
|
||||
if len(changes) > 0 {
|
||||
log.Debugf("amp config partial reload completed: %v", changes)
|
||||
} else {
|
||||
log.Debug("amp config checked: no changes detected")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// hasModelMappingsChanged compares old and new model mappings.
|
||||
func (m *AmpModule) hasModelMappingsChanged(old *config.AmpCode, new *config.AmpCode) bool {
|
||||
if old == nil {
|
||||
return len(new.ModelMappings) > 0
|
||||
}
|
||||
|
||||
if len(old.ModelMappings) != len(new.ModelMappings) {
|
||||
return true
|
||||
}
|
||||
|
||||
// Build map for efficient comparison
|
||||
oldMap := make(map[string]string, len(old.ModelMappings))
|
||||
for _, mapping := range old.ModelMappings {
|
||||
oldMap[strings.TrimSpace(mapping.From)] = strings.TrimSpace(mapping.To)
|
||||
}
|
||||
|
||||
for _, mapping := range new.ModelMappings {
|
||||
from := strings.TrimSpace(mapping.From)
|
||||
to := strings.TrimSpace(mapping.To)
|
||||
if oldTo, exists := oldMap[from]; !exists || oldTo != to {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
// hasAPIKeyChanged compares old and new API keys.
|
||||
func (m *AmpModule) hasAPIKeyChanged(old *config.AmpCode, new *config.AmpCode) bool {
|
||||
oldKey := ""
|
||||
if old != nil {
|
||||
oldKey = strings.TrimSpace(old.UpstreamAPIKey)
|
||||
}
|
||||
newKey := strings.TrimSpace(new.UpstreamAPIKey)
|
||||
return oldKey != newKey
|
||||
}
|
||||
|
||||
// GetModelMapper returns the model mapper instance (for testing/debugging).
|
||||
func (m *AmpModule) GetModelMapper() *DefaultModelMapper {
|
||||
return m.modelMapper
|
||||
}
|
||||
|
||||
// getProxy returns the current proxy instance (thread-safe for hot-reload).
|
||||
func (m *AmpModule) getProxy() *httputil.ReverseProxy {
|
||||
m.proxyMu.RLock()
|
||||
defer m.proxyMu.RUnlock()
|
||||
return m.proxy
|
||||
}
|
||||
|
||||
// setProxy updates the proxy instance (thread-safe for hot-reload).
|
||||
func (m *AmpModule) setProxy(proxy *httputil.ReverseProxy) {
|
||||
m.proxyMu.Lock()
|
||||
defer m.proxyMu.Unlock()
|
||||
m.proxy = proxy
|
||||
}
|
||||
|
||||
// IsRestrictedToLocalhost returns whether management routes are restricted to localhost.
|
||||
func (m *AmpModule) IsRestrictedToLocalhost() bool {
|
||||
m.restrictMu.RLock()
|
||||
defer m.restrictMu.RUnlock()
|
||||
return m.restrictToLocalhost
|
||||
}
|
||||
|
||||
// setRestrictToLocalhost updates the localhost restriction setting.
|
||||
func (m *AmpModule) setRestrictToLocalhost(restrict bool) {
|
||||
m.restrictMu.Lock()
|
||||
defer m.restrictMu.Unlock()
|
||||
m.restrictToLocalhost = restrict
|
||||
}
|
||||
|
||||
@@ -152,9 +152,9 @@ func TestModelMapper_UpdateMappings_SkipsInvalid(t *testing.T) {
|
||||
mapper := NewModelMapper(nil)
|
||||
|
||||
mapper.UpdateMappings([]config.AmpModelMapping{
|
||||
{From: "", To: "model-b"}, // Invalid: empty from
|
||||
{From: "model-a", To: ""}, // Invalid: empty to
|
||||
{From: " ", To: "model-b"}, // Invalid: whitespace from
|
||||
{From: "", To: "model-b"}, // Invalid: empty from
|
||||
{From: "model-a", To: ""}, // Invalid: empty to
|
||||
{From: " ", To: "model-b"}, // Invalid: whitespace from
|
||||
{From: "model-c", To: "model-d"}, // Valid
|
||||
})
|
||||
|
||||
|
||||
@@ -14,15 +14,16 @@ import (
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// localhostOnlyMiddleware restricts access to localhost (127.0.0.1, ::1) only.
|
||||
// Returns 403 Forbidden for non-localhost clients.
|
||||
//
|
||||
// Security: Uses RemoteAddr (actual TCP connection) instead of ClientIP() to prevent
|
||||
// header spoofing attacks via X-Forwarded-For or similar headers. This means the
|
||||
// middleware will not work correctly behind reverse proxies - users deploying behind
|
||||
// nginx/Cloudflare should disable this feature and use firewall rules instead.
|
||||
func localhostOnlyMiddleware() gin.HandlerFunc {
|
||||
// localhostOnlyMiddleware returns a middleware that dynamically checks the module's
|
||||
// localhost restriction setting. This allows hot-reload of the restriction without restarting.
|
||||
func (m *AmpModule) localhostOnlyMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// Check current setting (hot-reloadable)
|
||||
if !m.IsRestrictedToLocalhost() {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
// Use actual TCP connection address (RemoteAddr) to prevent header spoofing
|
||||
// This cannot be forged by X-Forwarded-For or other client-controlled headers
|
||||
remoteAddr := c.Request.RemoteAddr
|
||||
@@ -79,21 +80,32 @@ func noCORSMiddleware() gin.HandlerFunc {
|
||||
|
||||
// registerManagementRoutes registers Amp management proxy routes
|
||||
// These routes proxy through to the Amp control plane for OAuth, user management, etc.
|
||||
// If restrictToLocalhost is true, routes will only accept connections from 127.0.0.1/::1.
|
||||
func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *handlers.BaseAPIHandler, proxyHandler gin.HandlerFunc, restrictToLocalhost bool) {
|
||||
// Uses dynamic middleware and proxy getter for hot-reload support.
|
||||
func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *handlers.BaseAPIHandler) {
|
||||
ampAPI := engine.Group("/api")
|
||||
|
||||
// Always disable CORS for management routes to prevent browser-based attacks
|
||||
ampAPI.Use(noCORSMiddleware())
|
||||
|
||||
// Apply localhost-only restriction if configured
|
||||
if restrictToLocalhost {
|
||||
ampAPI.Use(localhostOnlyMiddleware())
|
||||
// Apply dynamic localhost-only restriction (hot-reloadable via m.IsRestrictedToLocalhost())
|
||||
ampAPI.Use(m.localhostOnlyMiddleware())
|
||||
|
||||
if m.IsRestrictedToLocalhost() {
|
||||
log.Info("amp management routes restricted to localhost only (CORS disabled)")
|
||||
} else {
|
||||
log.Warn("amp management routes are NOT restricted to localhost - this is insecure!")
|
||||
}
|
||||
|
||||
// Dynamic proxy handler that uses m.getProxy() for hot-reload support
|
||||
proxyHandler := func(c *gin.Context) {
|
||||
proxy := m.getProxy()
|
||||
if proxy == nil {
|
||||
c.JSON(503, gin.H{"error": "amp upstream proxy not available"})
|
||||
return
|
||||
}
|
||||
proxy.ServeHTTP(c.Writer, c.Request)
|
||||
}
|
||||
|
||||
// Management routes - these are proxied directly to Amp upstream
|
||||
ampAPI.Any("/internal", proxyHandler)
|
||||
ampAPI.Any("/internal/*path", proxyHandler)
|
||||
@@ -110,15 +122,20 @@ func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *ha
|
||||
ampAPI.Any("/threads/*path", proxyHandler)
|
||||
ampAPI.Any("/otel", proxyHandler)
|
||||
ampAPI.Any("/otel/*path", proxyHandler)
|
||||
ampAPI.Any("/tab", proxyHandler)
|
||||
ampAPI.Any("/tab/*path", proxyHandler)
|
||||
|
||||
// Root-level routes that AMP CLI expects without /api prefix
|
||||
// These need the same security middleware as the /api/* routes
|
||||
rootMiddleware := []gin.HandlerFunc{noCORSMiddleware()}
|
||||
if restrictToLocalhost {
|
||||
rootMiddleware = append(rootMiddleware, localhostOnlyMiddleware())
|
||||
}
|
||||
// These need the same security middleware as the /api/* routes (dynamic for hot-reload)
|
||||
rootMiddleware := []gin.HandlerFunc{noCORSMiddleware(), m.localhostOnlyMiddleware()}
|
||||
engine.GET("/threads.rss", append(rootMiddleware, proxyHandler)...)
|
||||
|
||||
// Root-level auth routes for CLI login flow
|
||||
// Amp uses multiple auth routes: /auth/cli-login, /auth/callback, /auth/sign-in, /auth/logout
|
||||
// We proxy all /auth/* to support the complete OAuth flow
|
||||
engine.Any("/auth", append(rootMiddleware, proxyHandler)...)
|
||||
engine.Any("/auth/*path", append(rootMiddleware, proxyHandler)...)
|
||||
|
||||
// Google v1beta1 passthrough with OAuth fallback
|
||||
// AMP CLI uses non-standard paths like /publishers/google/models/...
|
||||
// We bridge these to our standard Gemini handler to enable local OAuth.
|
||||
@@ -126,7 +143,7 @@ func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *ha
|
||||
geminiHandlers := gemini.NewGeminiAPIHandler(baseHandler)
|
||||
geminiBridge := createGeminiBridgeHandler(geminiHandlers)
|
||||
geminiV1Beta1Fallback := NewFallbackHandler(func() *httputil.ReverseProxy {
|
||||
return m.proxy
|
||||
return m.getProxy()
|
||||
})
|
||||
geminiV1Beta1Handler := geminiV1Beta1Fallback.WrapHandler(geminiBridge)
|
||||
|
||||
@@ -169,10 +186,10 @@ func (m *AmpModule) registerProviderAliases(engine *gin.Engine, baseHandler *han
|
||||
openaiResponsesHandlers := openai.NewOpenAIResponsesAPIHandler(baseHandler)
|
||||
|
||||
// Create fallback handler wrapper that forwards to ampcode.com when provider not found
|
||||
// Uses lazy evaluation to access proxy (which is created after routes are registered)
|
||||
// Uses m.getProxy() for hot-reload support (proxy can be updated at runtime)
|
||||
// Also includes model mapping support for routing unavailable models to alternatives
|
||||
fallbackHandler := NewFallbackHandlerWithMapper(func() *httputil.ReverseProxy {
|
||||
return m.proxy
|
||||
return m.getProxy()
|
||||
}, m.modelMapper)
|
||||
|
||||
// Provider-specific routes under /api/provider/:provider
|
||||
|
||||
@@ -13,16 +13,26 @@ func TestRegisterManagementRoutes(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
// Spy to track if proxy handler was called
|
||||
proxyCalled := false
|
||||
proxyHandler := func(c *gin.Context) {
|
||||
proxyCalled = true
|
||||
c.String(200, "proxied")
|
||||
// Create module with proxy for testing
|
||||
m := &AmpModule{
|
||||
restrictToLocalhost: false, // disable localhost restriction for tests
|
||||
}
|
||||
|
||||
m := &AmpModule{}
|
||||
// Create a mock proxy that tracks calls
|
||||
proxyCalled := false
|
||||
mockProxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
proxyCalled = true
|
||||
w.WriteHeader(200)
|
||||
w.Write([]byte("proxied"))
|
||||
}))
|
||||
defer mockProxy.Close()
|
||||
|
||||
// Create real proxy to mock server
|
||||
proxy, _ := createReverseProxy(mockProxy.URL, NewStaticSecretSource(""))
|
||||
m.setProxy(proxy)
|
||||
|
||||
base := &handlers.BaseAPIHandler{}
|
||||
m.registerManagementRoutes(r, base, proxyHandler, false) // false = don't restrict to localhost in tests
|
||||
m.registerManagementRoutes(r, base)
|
||||
|
||||
managementPaths := []struct {
|
||||
path string
|
||||
@@ -39,6 +49,11 @@ func TestRegisterManagementRoutes(t *testing.T) {
|
||||
{"/api/threads", http.MethodGet},
|
||||
{"/threads.rss", http.MethodGet}, // Root-level route (no /api prefix)
|
||||
{"/api/otel", http.MethodGet},
|
||||
{"/api/tab", http.MethodGet},
|
||||
{"/api/tab/some/path", http.MethodGet},
|
||||
{"/auth", http.MethodGet}, // Root-level auth route
|
||||
{"/auth/cli-login", http.MethodGet}, // CLI login flow
|
||||
{"/auth/callback", http.MethodGet}, // OAuth callback
|
||||
// Google v1beta1 bridge should still proxy non-model requests (GET) and allow POST
|
||||
{"/api/provider/google/v1beta1/models", http.MethodGet},
|
||||
{"/api/provider/google/v1beta1/models", http.MethodPost},
|
||||
@@ -226,8 +241,13 @@ func TestLocalhostOnlyMiddleware_PreventsSpoofing(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
// Apply localhost-only middleware
|
||||
r.Use(localhostOnlyMiddleware())
|
||||
// Create module with localhost restriction enabled
|
||||
m := &AmpModule{
|
||||
restrictToLocalhost: true,
|
||||
}
|
||||
|
||||
// Apply dynamic localhost-only middleware
|
||||
r.Use(m.localhostOnlyMiddleware())
|
||||
r.GET("/test", func(c *gin.Context) {
|
||||
c.String(http.StatusOK, "ok")
|
||||
})
|
||||
@@ -300,3 +320,53 @@ func TestLocalhostOnlyMiddleware_PreventsSpoofing(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalhostOnlyMiddleware_HotReload(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
|
||||
// Create module with localhost restriction initially enabled
|
||||
m := &AmpModule{
|
||||
restrictToLocalhost: true,
|
||||
}
|
||||
|
||||
// Apply dynamic localhost-only middleware
|
||||
r.Use(m.localhostOnlyMiddleware())
|
||||
r.GET("/test", func(c *gin.Context) {
|
||||
c.String(http.StatusOK, "ok")
|
||||
})
|
||||
|
||||
// Test 1: Remote IP should be blocked when restriction is enabled
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.RemoteAddr = "192.168.1.100:12345"
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Errorf("Expected 403 when restriction enabled, got %d", w.Code)
|
||||
}
|
||||
|
||||
// Test 2: Hot-reload - disable restriction
|
||||
m.setRestrictToLocalhost(false)
|
||||
|
||||
req = httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.RemoteAddr = "192.168.1.100:12345"
|
||||
w = httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("Expected 200 after disabling restriction, got %d", w.Code)
|
||||
}
|
||||
|
||||
// Test 3: Hot-reload - re-enable restriction
|
||||
m.setRestrictToLocalhost(true)
|
||||
|
||||
req = httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.RemoteAddr = "192.168.1.100:12345"
|
||||
w = httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Errorf("Expected 403 after re-enabling restriction, got %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -139,6 +139,17 @@ func (s *MultiSourceSecret) InvalidateCache() {
|
||||
s.cache = nil
|
||||
}
|
||||
|
||||
// UpdateExplicitKey refreshes the config-provided key and clears cache.
|
||||
func (s *MultiSourceSecret) UpdateExplicitKey(key string) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.explicitKey = strings.TrimSpace(key)
|
||||
s.cache = nil
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// StaticSecretSource returns a fixed API key (for testing)
|
||||
type StaticSecretSource struct {
|
||||
key string
|
||||
|
||||
@@ -470,8 +470,8 @@ func (s *Server) registerManagementRoutes() {
|
||||
{
|
||||
mgmt.GET("/usage", s.mgmt.GetUsageStatistics)
|
||||
mgmt.GET("/config", s.mgmt.GetConfig)
|
||||
mgmt.GET("/config.yaml", s.mgmt.GetConfigYAML)
|
||||
mgmt.PUT("/config.yaml", s.mgmt.PutConfigYAML)
|
||||
mgmt.GET("/config.yaml", s.mgmt.GetConfigFile)
|
||||
|
||||
mgmt.GET("/debug", s.mgmt.GetDebug)
|
||||
mgmt.PUT("/debug", s.mgmt.PutDebug)
|
||||
@@ -503,11 +503,6 @@ func (s *Server) registerManagementRoutes() {
|
||||
mgmt.PATCH("/api-keys", s.mgmt.PatchAPIKeys)
|
||||
mgmt.DELETE("/api-keys", s.mgmt.DeleteAPIKeys)
|
||||
|
||||
mgmt.GET("/generative-language-api-key", s.mgmt.GetGlKeys)
|
||||
mgmt.PUT("/generative-language-api-key", s.mgmt.PutGlKeys)
|
||||
mgmt.PATCH("/generative-language-api-key", s.mgmt.PatchGlKeys)
|
||||
mgmt.DELETE("/generative-language-api-key", s.mgmt.DeleteGlKeys)
|
||||
|
||||
mgmt.GET("/gemini-api-key", s.mgmt.GetGeminiKeys)
|
||||
mgmt.PUT("/gemini-api-key", s.mgmt.PutGeminiKeys)
|
||||
mgmt.PATCH("/gemini-api-key", s.mgmt.PatchGeminiKey)
|
||||
|
||||
@@ -38,7 +38,13 @@ func (m *LogFormatter) Format(entry *log.Entry) ([]byte, error) {
|
||||
|
||||
timestamp := entry.Time.Format("2006-01-02 15:04:05")
|
||||
message := strings.TrimRight(entry.Message, "\r\n")
|
||||
formatted := fmt.Sprintf("[%s] [%s] [%s:%d] %s\n", timestamp, entry.Level, filepath.Base(entry.Caller.File), entry.Caller.Line, message)
|
||||
|
||||
var formatted string
|
||||
if entry.Caller != nil {
|
||||
formatted = fmt.Sprintf("[%s] [%s] [%s:%d] %s\n", timestamp, entry.Level, filepath.Base(entry.Caller.File), entry.Caller.Line, message)
|
||||
} else {
|
||||
formatted = fmt.Sprintf("[%s] [%s] %s\n", timestamp, entry.Level, message)
|
||||
}
|
||||
buffer.WriteString(formatted)
|
||||
|
||||
return buffer.Bytes(), nil
|
||||
|
||||
@@ -961,6 +961,7 @@ func GetIFlowModels() []*ModelInfo {
|
||||
{ID: "glm-4.6", DisplayName: "GLM-4.6", Description: "Zhipu GLM 4.6 general model", Created: 1759190400},
|
||||
{ID: "kimi-k2", DisplayName: "Kimi-K2", Description: "Moonshot Kimi K2 general model", Created: 1752192000},
|
||||
{ID: "kimi-k2-thinking", DisplayName: "Kimi-K2-Thinking", Description: "Moonshot Kimi K2 general model", Created: 1762387200},
|
||||
{ID: "deepseek-v3.2-chat", DisplayName: "DeepSeek-V3.2", Description: "DeepSeek V3.2", Created: 1764576000},
|
||||
{ID: "deepseek-v3.2", DisplayName: "DeepSeek-V3.2-Exp", Description: "DeepSeek V3.2 experimental", Created: 1759104000},
|
||||
{ID: "deepseek-v3.1", DisplayName: "DeepSeek-V3.1-Terminus", Description: "DeepSeek V3.1 Terminus", Created: 1756339200},
|
||||
{ID: "deepseek-r1", DisplayName: "DeepSeek-R1", Description: "DeepSeek reasoning model R1", Created: 1737331200},
|
||||
|
||||
@@ -734,6 +734,8 @@ func modelName2Alias(modelName string) string {
|
||||
return "gemini-claude-sonnet-4-5"
|
||||
case "claude-sonnet-4-5-thinking":
|
||||
return "gemini-claude-sonnet-4-5-thinking"
|
||||
case "claude-opus-4-5-thinking":
|
||||
return "gemini-claude-opus-4-5-thinking"
|
||||
case "chat_20706", "chat_23310", "gemini-2.5-flash-thinking", "gemini-3-pro-low", "gemini-2.5-pro":
|
||||
return ""
|
||||
default:
|
||||
@@ -753,6 +755,8 @@ func alias2ModelName(modelName string) string {
|
||||
return "claude-sonnet-4-5"
|
||||
case "gemini-claude-sonnet-4-5-thinking":
|
||||
return "claude-sonnet-4-5-thinking"
|
||||
case "gemini-claude-opus-4-5-thinking":
|
||||
return "claude-opus-4-5-thinking"
|
||||
default:
|
||||
return modelName
|
||||
}
|
||||
|
||||
@@ -667,7 +667,7 @@ func cliPreviewFallbackOrder(model string) []string {
|
||||
case "gemini-2.5-pro":
|
||||
return []string{
|
||||
// "gemini-2.5-pro-preview-05-06",
|
||||
"gemini-2.5-pro-preview-06-05",
|
||||
// "gemini-2.5-pro-preview-06-05",
|
||||
}
|
||||
case "gemini-2.5-flash":
|
||||
return []string{
|
||||
|
||||
@@ -570,6 +570,35 @@ func summarizeExcludedModels(list []string) excludedModelsSummary {
|
||||
}
|
||||
}
|
||||
|
||||
type ampModelMappingsSummary struct {
|
||||
hash string
|
||||
count int
|
||||
}
|
||||
|
||||
func summarizeAmpModelMappings(mappings []config.AmpModelMapping) ampModelMappingsSummary {
|
||||
if len(mappings) == 0 {
|
||||
return ampModelMappingsSummary{}
|
||||
}
|
||||
entries := make([]string, 0, len(mappings))
|
||||
for _, mapping := range mappings {
|
||||
from := strings.TrimSpace(mapping.From)
|
||||
to := strings.TrimSpace(mapping.To)
|
||||
if from == "" && to == "" {
|
||||
continue
|
||||
}
|
||||
entries = append(entries, from+"->"+to)
|
||||
}
|
||||
if len(entries) == 0 {
|
||||
return ampModelMappingsSummary{}
|
||||
}
|
||||
sort.Strings(entries)
|
||||
sum := sha256.Sum256([]byte(strings.Join(entries, "|")))
|
||||
return ampModelMappingsSummary{
|
||||
hash: hex.EncodeToString(sum[:]),
|
||||
count: len(entries),
|
||||
}
|
||||
}
|
||||
|
||||
func summarizeOAuthExcludedModels(entries map[string][]string) map[string]excludedModelsSummary {
|
||||
if len(entries) == 0 {
|
||||
return nil
|
||||
@@ -1762,6 +1791,31 @@ func buildConfigChangeDetails(oldCfg, newCfg *config.Config) []string {
|
||||
}
|
||||
}
|
||||
|
||||
// AmpCode settings (redacted where needed)
|
||||
oldAmpURL := strings.TrimSpace(oldCfg.AmpCode.UpstreamURL)
|
||||
newAmpURL := strings.TrimSpace(newCfg.AmpCode.UpstreamURL)
|
||||
if oldAmpURL != newAmpURL {
|
||||
changes = append(changes, fmt.Sprintf("ampcode.upstream-url: %s -> %s", oldAmpURL, newAmpURL))
|
||||
}
|
||||
oldAmpKey := strings.TrimSpace(oldCfg.AmpCode.UpstreamAPIKey)
|
||||
newAmpKey := strings.TrimSpace(newCfg.AmpCode.UpstreamAPIKey)
|
||||
switch {
|
||||
case oldAmpKey == "" && newAmpKey != "":
|
||||
changes = append(changes, "ampcode.upstream-api-key: added")
|
||||
case oldAmpKey != "" && newAmpKey == "":
|
||||
changes = append(changes, "ampcode.upstream-api-key: removed")
|
||||
case oldAmpKey != newAmpKey:
|
||||
changes = append(changes, "ampcode.upstream-api-key: updated")
|
||||
}
|
||||
if oldCfg.AmpCode.RestrictManagementToLocalhost != newCfg.AmpCode.RestrictManagementToLocalhost {
|
||||
changes = append(changes, fmt.Sprintf("ampcode.restrict-management-to-localhost: %t -> %t", oldCfg.AmpCode.RestrictManagementToLocalhost, newCfg.AmpCode.RestrictManagementToLocalhost))
|
||||
}
|
||||
oldMappings := summarizeAmpModelMappings(oldCfg.AmpCode.ModelMappings)
|
||||
newMappings := summarizeAmpModelMappings(newCfg.AmpCode.ModelMappings)
|
||||
if oldMappings.hash != newMappings.hash {
|
||||
changes = append(changes, fmt.Sprintf("ampcode.model-mappings: updated (%d -> %d entries)", oldMappings.count, newMappings.count))
|
||||
}
|
||||
|
||||
if entries, _ := diffOAuthExcludedModelChanges(oldCfg.OAuthExcludedModels, newCfg.OAuthExcludedModels); len(entries) > 0 {
|
||||
changes = append(changes, entries...)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user