Add prefix

StrongMonkey · StrongMonkey · commit 881d984e7fd3 · 2025-09-05T16:14:18.000-07:00
Signed-off-by: Daishan Peng &lt;daishan@acorn.io&gt;
diff --git a/cmd/root.go b/cmd/root.go
@@ -35,8 +35,9 @@ type RootCmd struct {
 	EncryptionKey string `name:"encryption-key" env:"ENCRYPTION_KEY" usage:"Base64-encoded 32-byte AES-256 key for encrypting sensitive data (optional)"`
 
 	// Server configuration
-	Port string `name:"port" env:"PORT" usage:"Port to run the server on" default:"8080"`
-	Host string `name:"host" env:"HOST" usage:"Host to bind the server to" default:"localhost"`
+	Port        string `name:"port" env:"PORT" usage:"Port to run the server on" default:"8080"`
+	Host        string `name:"host" env:"HOST" usage:"Host to bind the server to" default:"localhost"`
+	RoutePrefix string `name:"route-prefix" env:"ROUTE_PREFIX" usage:"Optional prefix for all routes (e.g., '/oauth2')"`
 
 	// Logging
 	Verbose bool `name:"verbose,v" usage:"Enable verbose logging"`
@@ -69,6 +70,7 @@ func (c *RootCmd) Run(cobraCmd *cobra.Command, args []string) error {
 		MCPServerURL:      c.MCPServerURL,
 		EncryptionKey:     c.EncryptionKey,
 		Mode:              c.Mode,
+		RoutePrefix:       c.RoutePrefix,
 	}
 
 	// Validate configuration
diff --git a/main.go b/main.go
@@ -1,13 +1,15 @@
 package main
 
 import (
+	"fmt"
 	"os"
 
 	"github.com/obot-platform/mcp-oauth-proxy/cmd"
 )
 
 func main() {
 	if err := cmd.Execute(); err != nil {
+		fmt.Println(err)
 		os.Exit(1)
 	}
 }
diff --git a/pkg/oauth/authorize/authorize.go b/pkg/oauth/authorize/authorize.go
@@ -24,15 +24,17 @@ type Handler struct {
 	scopesSupported []string
 	clientID        string
 	clientSecret    string
+	routePrefix     string
 }
 
-func NewHandler(db AuthorizationStore, provider providers.Provider, scopesSupported []string, clientID, clientSecret string) http.Handler {
+func NewHandler(db AuthorizationStore, provider providers.Provider, scopesSupported []string, clientID, clientSecret, routePrefix string) http.Handler {
 	return &Handler{
 		db:              db,
 		provider:        provider,
 		scopesSupported: scopesSupported,
 		clientID:        clientID,
 		clientSecret:    clientSecret,
+		routePrefix:     routePrefix,
 	}
 }
 
@@ -139,7 +141,7 @@ func (p *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	redirectURI := fmt.Sprintf("%s/callback", handlerutils.GetBaseURL(r))
+	redirectURI := fmt.Sprintf("%s%s/callback", handlerutils.GetBaseURL(r), p.routePrefix)
 
 	// Generate authorization URL with the provider
 	authURL := p.provider.GetAuthorizationURL(
diff --git a/pkg/oauth/callback/callback.go b/pkg/oauth/callback/callback.go
@@ -30,21 +30,23 @@ type Handler struct {
 	clientID      string
 	clientSecret  string
 	mcpUIManager  MCPUIManager
+	routePrefix   string
 }
 
 // MCPUIManager interface for generating JWT tokens
 type MCPUIManager interface {
 	GenerateMCPUICodeForDownstream(bearerToken, refreshToken string) (string, error)
 }
 
-func NewHandler(db Store, provider providers.Provider, encryptionKey []byte, clientID, clientSecret string, mcpUIManager MCPUIManager) http.Handler {
+func NewHandler(db Store, provider providers.Provider, encryptionKey []byte, clientID, clientSecret, routePrefix string, mcpUIManager MCPUIManager) http.Handler {
 	return &Handler{
 		db:            db,
 		provider:      provider,
 		encryptionKey: encryptionKey,
 		clientID:      clientID,
 		clientSecret:  clientSecret,
 		mcpUIManager:  mcpUIManager,
+		routePrefix:   routePrefix,
 	}
 }
 
@@ -150,7 +152,7 @@ func (p *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}()
 
 	// Get provider credentials
-	redirectURI := fmt.Sprintf("%s/callback", handlerutils.GetBaseURL(r))
+	redirectURI := fmt.Sprintf("%s%s/callback", handlerutils.GetBaseURL(r), p.routePrefix)
 
 	// Exchange code for tokens
 	tokenInfo, err := p.provider.ExchangeCodeForToken(r.Context(), code, p.clientID, p.clientSecret, redirectURI)
diff --git a/pkg/proxy/proxy.go b/pkg/proxy/proxy.go
@@ -192,27 +192,30 @@ func (p *OAuthProxy) SetupRoutes(mux *http.ServeMux) {
 		log.Fatalf("Failed to get provider: %v", err)
 	}
 
-	authorizeHandler := authorize.NewHandler(p.db, provider, p.metadata.ScopesSupported, p.GetOAuthClientID(), p.GetOAuthClientSecret())
+	authorizeHandler := authorize.NewHandler(p.db, provider, p.metadata.ScopesSupported, p.GetOAuthClientID(), p.GetOAuthClientSecret(), p.config.RoutePrefix)
 	tokenHandler := token.NewHandler(p.db)
-	callbackHandler := callback.NewHandler(p.db, provider, p.encryptionKey, p.GetOAuthClientID(), p.GetOAuthClientSecret(), p.mcpUIManager)
+	callbackHandler := callback.NewHandler(p.db, provider, p.encryptionKey, p.GetOAuthClientID(), p.GetOAuthClientSecret(), p.config.RoutePrefix, p.mcpUIManager)
 	revokeHandler := revoke.NewHandler(p.db)
 	tokenValidator := validate.NewTokenValidator(p.tokenManager, p.encryptionKey, p.db, provider, p.GetOAuthClientID(), p.GetOAuthClientSecret(), p.metadata.ScopesSupported)
 
-	mux.HandleFunc("GET /health", p.withCORS(p.healthHandler))
+	// Get route prefix from config
+	prefix := p.config.RoutePrefix
+
+	mux.HandleFunc("GET "+prefix+"/health", p.withCORS(p.healthHandler))
 
 	// OAuth endpoints
-	mux.HandleFunc("GET /authorize", p.withCORS(p.withRateLimit(authorizeHandler)))
-	mux.HandleFunc("GET /callback", p.withCORS(p.withRateLimit(callbackHandler)))
-	mux.HandleFunc("POST /token", p.withCORS(p.withRateLimit(tokenHandler)))
-	mux.HandleFunc("POST /revoke", p.withCORS(p.withRateLimit(revokeHandler)))
-	mux.HandleFunc("POST /register", p.withCORS(p.withRateLimit(register.NewHandler(p.db))))
+	mux.HandleFunc("GET "+prefix+"/authorize", p.withCORS(p.withRateLimit(authorizeHandler)))
+	mux.HandleFunc("GET "+prefix+"/callback", p.withCORS(p.withRateLimit(callbackHandler)))
+	mux.HandleFunc("POST "+prefix+"/token", p.withCORS(p.withRateLimit(tokenHandler)))
+	mux.HandleFunc("POST "+prefix+"/revoke", p.withCORS(p.withRateLimit(revokeHandler)))
+	mux.HandleFunc("POST "+prefix+"/register", p.withCORS(p.withRateLimit(register.NewHandler(p.db))))
 
 	// Metadata endpoints
 	mux.HandleFunc("GET /.well-known/oauth-authorization-server", p.withCORS(p.oauthMetadataHandler))
 	mux.HandleFunc("GET /.well-known/oauth-protected-resource", p.withCORS(p.protectedResourceMetadataHandler))
 
 	// Protect everything else
-	mux.HandleFunc("/{path...}", p.withCORS(p.withRateLimit(tokenValidator.WithTokenValidation(p.mcpProxyHandler))))
+	mux.HandleFunc(prefix+"/{path...}", p.withCORS(p.withRateLimit(tokenValidator.WithTokenValidation(p.mcpProxyHandler))))
 }
 
 // GetHandler returns an http.Handler for the OAuth proxy
@@ -270,21 +273,22 @@ func (p *OAuthProxy) healthHandler(w http.ResponseWriter, r *http.Request) {
 
 func (p *OAuthProxy) oauthMetadataHandler(w http.ResponseWriter, r *http.Request) {
 	baseURL := handlerutils.GetBaseURL(r)
+	prefix := p.config.RoutePrefix
 
 	// Create dynamic metadata based on the request
 	metadata := &types.OAuthMetadata{
 		Issuer:                                   baseURL,
 		ServiceDocumentation:                     p.metadata.ServiceDocumentation,
-		AuthorizationEndpoint:                    fmt.Sprintf("%s/authorize", baseURL),
+		AuthorizationEndpoint:                    fmt.Sprintf("%s%s/authorize", baseURL, prefix),
 		ResponseTypesSupported:                   p.metadata.ResponseTypesSupported,
 		CodeChallengeMethodsSupported:            p.metadata.CodeChallengeMethodsSupported,
-		TokenEndpoint:                            fmt.Sprintf("%s/token", baseURL),
+		TokenEndpoint:                            fmt.Sprintf("%s%s/token", baseURL, prefix),
 		TokenEndpointAuthMethodsSupported:        p.metadata.TokenEndpointAuthMethodsSupported,
 		GrantTypesSupported:                      p.metadata.GrantTypesSupported,
 		ScopesSupported:                          p.metadata.ScopesSupported,
-		RevocationEndpoint:                       fmt.Sprintf("%s/revoke", baseURL),
+		RevocationEndpoint:                       fmt.Sprintf("%s%s/revoke", baseURL, prefix),
 		RevocationEndpointAuthMethodsSupported:   p.metadata.RevocationEndpointAuthMethodsSupported,
-		RegistrationEndpoint:                     fmt.Sprintf("%s/register", baseURL),
+		RegistrationEndpoint:                     fmt.Sprintf("%s%s/register", baseURL, prefix),
 		RegistrationEndpointAuthMethodsSupported: p.metadata.RegistrationEndpointAuthMethodsSupported,
 	}
 
@@ -293,9 +297,12 @@ func (p *OAuthProxy) oauthMetadataHandler(w http.ResponseWriter, r *http.Request
 
 func (p *OAuthProxy) protectedResourceMetadataHandler(w http.ResponseWriter, r *http.Request) {
 	baseURL := handlerutils.GetBaseURL(r)
+	prefix := p.config.RoutePrefix
+	resourceURL := baseURL + prefix
+
 	metadata := types.OAuthProtectedResourceMetadata{
-		Resource:              baseURL,
-		AuthorizationServers:  []string{baseURL},
+		Resource:              resourceURL,
+		AuthorizationServers:  []string{baseURL + prefix},
 		Scopes:                p.metadata.ScopesSupported,
 		ResourceName:          p.resourceName,
 		ResourceDocumentation: p.metadata.ServiceDocumentation,
diff --git a/pkg/types/types.go b/pkg/types/types.go
@@ -15,6 +15,7 @@ type Config struct {
 	EncryptionKey     string
 	MCPServerURL      string
 	Mode              string
+	RoutePrefix       string
 }
 
 // TokenData represents stored token data for OAuth 2.1 compliance

Original file line number	Diff line number	Diff line change
`@@ -1,13 +1,15 @@`
`1`	`1`	`package main`
`2`	`2`
`3`	`3`	`import (`
	`4`	`+ "fmt"`
`4`	`5`	`"os"`
`5`	`6`
`6`	`7`	`"github.com/obot-platform/mcp-oauth-proxy/cmd"`
`7`	`8`	`)`
`8`	`9`
`9`	`10`	`func main() {`
`10`	`11`	`if err := cmd.Execute(); err != nil {`
	`12`	`+ fmt.Println(err)`
`11`	`13`	`os.Exit(1)`
`12`	`14`	`}`
`13`	`15`	`}`
Original file line number	Diff line number	Diff line change
`@@ -24,15 +24,17 @@ type Handler struct {`
`24`	`24`	`scopesSupported []string`
`25`	`25`	`clientID string`
`26`	`26`	`clientSecret string`
	`27`	`+ routePrefix string`
`27`	`28`	`}`
`28`	`29`
`29`		`-func NewHandler(db AuthorizationStore, provider providers.Provider, scopesSupported []string, clientID, clientSecret string) http.Handler {`
	`30`	`+func NewHandler(db AuthorizationStore, provider providers.Provider, scopesSupported []string, clientID, clientSecret, routePrefix string) http.Handler {`
`30`	`31`	`return &Handler{`
`31`	`32`	`db: db,`
`32`	`33`	`provider: provider,`
`33`	`34`	`scopesSupported: scopesSupported,`
`34`	`35`	`clientID: clientID,`
`35`	`36`	`clientSecret: clientSecret,`
	`37`	`+ routePrefix: routePrefix,`
`36`	`38`	`}`
`37`	`39`	`}`
`38`	`40`
`@@ -139,7 +141,7 @@ func (p Handler) ServeHTTP(w http.ResponseWriter, r http.Request) {`
`139`	`141`	`return`
`140`	`142`	`}`
`141`	`143`
`142`		`- redirectURI := fmt.Sprintf("%s/callback", handlerutils.GetBaseURL(r))`
	`144`	`+ redirectURI := fmt.Sprintf("%s%s/callback", handlerutils.GetBaseURL(r), p.routePrefix)`
`143`	`145`
`144`	`146`	`// Generate authorization URL with the provider`
`145`	`147`	`authURL := p.provider.GetAuthorizationURL(`
Original file line number	Diff line number	Diff line change
`@@ -15,6 +15,7 @@ type Config struct {`
`15`	`15`	`EncryptionKey string`
`16`	`16`	`MCPServerURL string`
`17`	`17`	`Mode string`
	`18`	`+ RoutePrefix string`
`18`	`19`	`}`
`19`	`20`
`20`	`21`	`// TokenData represents stored token data for OAuth 2.1 compliance`