Skip to content

Commit 881d984

Browse files
StrongMonkeyStrongMonkey
committed
Add prefix
Signed-off-by: Daishan Peng <daishan@acorn.io>
1 parent 8da9ba9 commit 881d984

File tree

6 files changed

+37
-21
lines changed

6 files changed

+37
-21
lines changed

‎cmd/root.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,9 @@ type RootCmd struct {
3535
EncryptionKey string `name:"encryption-key" env:"ENCRYPTION_KEY" usage:"Base64-encoded 32-byte AES-256 key for encrypting sensitive data (optional)"`
3636

3737
// Server configuration
38-
Port string `name:"port" env:"PORT" usage:"Port to run the server on" default:"8080"`
39-
Host string `name:"host" env:"HOST" usage:"Host to bind the server to" default:"localhost"`
38+
Port string `name:"port" env:"PORT" usage:"Port to run the server on" default:"8080"`
39+
Host string `name:"host" env:"HOST" usage:"Host to bind the server to" default:"localhost"`
40+
RoutePrefix string `name:"route-prefix" env:"ROUTE_PREFIX" usage:"Optional prefix for all routes (e.g., '/oauth2')"`
4041

4142
// Logging
4243
Verbose bool `name:"verbose,v" usage:"Enable verbose logging"`
@@ -69,6 +70,7 @@ func (c *RootCmd) Run(cobraCmd *cobra.Command, args []string) error {
6970
MCPServerURL: c.MCPServerURL,
7071
EncryptionKey: c.EncryptionKey,
7172
Mode: c.Mode,
73+
RoutePrefix: c.RoutePrefix,
7274
}
7375

7476
// Validate configuration

‎main.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
package main
22

33
import (
4+
"fmt"
45
"os"
56

67
"github.com/obot-platform/mcp-oauth-proxy/cmd"
78
)
89

910
func main() {
1011
if err := cmd.Execute(); err != nil {
12+
fmt.Println(err)
1113
os.Exit(1)
1214
}
1315
}

‎pkg/oauth/authorize/authorize.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,15 +24,17 @@ type Handler struct {
2424
scopesSupported []string
2525
clientID string
2626
clientSecret string
27+
routePrefix string
2728
}
2829

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 {
3031
return &Handler{
3132
db: db,
3233
provider: provider,
3334
scopesSupported: scopesSupported,
3435
clientID: clientID,
3536
clientSecret: clientSecret,
37+
routePrefix: routePrefix,
3638
}
3739
}
3840

@@ -139,7 +141,7 @@ func (p *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
139141
return
140142
}
141143

142-
redirectURI := fmt.Sprintf("%s/callback", handlerutils.GetBaseURL(r))
144+
redirectURI := fmt.Sprintf("%s%s/callback", handlerutils.GetBaseURL(r), p.routePrefix)
143145

144146
// Generate authorization URL with the provider
145147
authURL := p.provider.GetAuthorizationURL(

‎pkg/oauth/callback/callback.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,21 +30,23 @@ type Handler struct {
3030
clientID string
3131
clientSecret string
3232
mcpUIManager MCPUIManager
33+
routePrefix string
3334
}
3435

3536
// MCPUIManager interface for generating JWT tokens
3637
type MCPUIManager interface {
3738
GenerateMCPUICodeForDownstream(bearerToken, refreshToken string) (string, error)
3839
}
3940

40-
func NewHandler(db Store, provider providers.Provider, encryptionKey []byte, clientID, clientSecret string, mcpUIManager MCPUIManager) http.Handler {
41+
func NewHandler(db Store, provider providers.Provider, encryptionKey []byte, clientID, clientSecret, routePrefix string, mcpUIManager MCPUIManager) http.Handler {
4142
return &Handler{
4243
db: db,
4344
provider: provider,
4445
encryptionKey: encryptionKey,
4546
clientID: clientID,
4647
clientSecret: clientSecret,
4748
mcpUIManager: mcpUIManager,
49+
routePrefix: routePrefix,
4850
}
4951
}
5052

@@ -150,7 +152,7 @@ func (p *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
150152
}()
151153

152154
// Get provider credentials
153-
redirectURI := fmt.Sprintf("%s/callback", handlerutils.GetBaseURL(r))
155+
redirectURI := fmt.Sprintf("%s%s/callback", handlerutils.GetBaseURL(r), p.routePrefix)
154156

155157
// Exchange code for tokens
156158
tokenInfo, err := p.provider.ExchangeCodeForToken(r.Context(), code, p.clientID, p.clientSecret, redirectURI)

‎pkg/proxy/proxy.go‎

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -192,27 +192,30 @@ func (p *OAuthProxy) SetupRoutes(mux *http.ServeMux) {
192192
log.Fatalf("Failed to get provider: %v", err)
193193
}
194194

195-
authorizeHandler := authorize.NewHandler(p.db, provider, p.metadata.ScopesSupported, p.GetOAuthClientID(), p.GetOAuthClientSecret())
195+
authorizeHandler := authorize.NewHandler(p.db, provider, p.metadata.ScopesSupported, p.GetOAuthClientID(), p.GetOAuthClientSecret(), p.config.RoutePrefix)
196196
tokenHandler := token.NewHandler(p.db)
197-
callbackHandler := callback.NewHandler(p.db, provider, p.encryptionKey, p.GetOAuthClientID(), p.GetOAuthClientSecret(), p.mcpUIManager)
197+
callbackHandler := callback.NewHandler(p.db, provider, p.encryptionKey, p.GetOAuthClientID(), p.GetOAuthClientSecret(), p.config.RoutePrefix, p.mcpUIManager)
198198
revokeHandler := revoke.NewHandler(p.db)
199199
tokenValidator := validate.NewTokenValidator(p.tokenManager, p.encryptionKey, p.db, provider, p.GetOAuthClientID(), p.GetOAuthClientSecret(), p.metadata.ScopesSupported)
200200

201-
mux.HandleFunc("GET /health", p.withCORS(p.healthHandler))
201+
// Get route prefix from config
202+
prefix := p.config.RoutePrefix
203+
204+
mux.HandleFunc("GET "+prefix+"/health", p.withCORS(p.healthHandler))
202205

203206
// OAuth endpoints
204-
mux.HandleFunc("GET /authorize", p.withCORS(p.withRateLimit(authorizeHandler)))
205-
mux.HandleFunc("GET /callback", p.withCORS(p.withRateLimit(callbackHandler)))
206-
mux.HandleFunc("POST /token", p.withCORS(p.withRateLimit(tokenHandler)))
207-
mux.HandleFunc("POST /revoke", p.withCORS(p.withRateLimit(revokeHandler)))
208-
mux.HandleFunc("POST /register", p.withCORS(p.withRateLimit(register.NewHandler(p.db))))
207+
mux.HandleFunc("GET "+prefix+"/authorize", p.withCORS(p.withRateLimit(authorizeHandler)))
208+
mux.HandleFunc("GET "+prefix+"/callback", p.withCORS(p.withRateLimit(callbackHandler)))
209+
mux.HandleFunc("POST "+prefix+"/token", p.withCORS(p.withRateLimit(tokenHandler)))
210+
mux.HandleFunc("POST "+prefix+"/revoke", p.withCORS(p.withRateLimit(revokeHandler)))
211+
mux.HandleFunc("POST "+prefix+"/register", p.withCORS(p.withRateLimit(register.NewHandler(p.db))))
209212

210213
// Metadata endpoints
211214
mux.HandleFunc("GET /.well-known/oauth-authorization-server", p.withCORS(p.oauthMetadataHandler))
212215
mux.HandleFunc("GET /.well-known/oauth-protected-resource", p.withCORS(p.protectedResourceMetadataHandler))
213216

214217
// Protect everything else
215-
mux.HandleFunc("/{path...}", p.withCORS(p.withRateLimit(tokenValidator.WithTokenValidation(p.mcpProxyHandler))))
218+
mux.HandleFunc(prefix+"/{path...}", p.withCORS(p.withRateLimit(tokenValidator.WithTokenValidation(p.mcpProxyHandler))))
216219
}
217220

218221
// GetHandler returns an http.Handler for the OAuth proxy
@@ -270,21 +273,22 @@ func (p *OAuthProxy) healthHandler(w http.ResponseWriter, r *http.Request) {
270273

271274
func (p *OAuthProxy) oauthMetadataHandler(w http.ResponseWriter, r *http.Request) {
272275
baseURL := handlerutils.GetBaseURL(r)
276+
prefix := p.config.RoutePrefix
273277

274278
// Create dynamic metadata based on the request
275279
metadata := &types.OAuthMetadata{
276280
Issuer: baseURL,
277281
ServiceDocumentation: p.metadata.ServiceDocumentation,
278-
AuthorizationEndpoint: fmt.Sprintf("%s/authorize", baseURL),
282+
AuthorizationEndpoint: fmt.Sprintf("%s%s/authorize", baseURL, prefix),
279283
ResponseTypesSupported: p.metadata.ResponseTypesSupported,
280284
CodeChallengeMethodsSupported: p.metadata.CodeChallengeMethodsSupported,
281-
TokenEndpoint: fmt.Sprintf("%s/token", baseURL),
285+
TokenEndpoint: fmt.Sprintf("%s%s/token", baseURL, prefix),
282286
TokenEndpointAuthMethodsSupported: p.metadata.TokenEndpointAuthMethodsSupported,
283287
GrantTypesSupported: p.metadata.GrantTypesSupported,
284288
ScopesSupported: p.metadata.ScopesSupported,
285-
RevocationEndpoint: fmt.Sprintf("%s/revoke", baseURL),
289+
RevocationEndpoint: fmt.Sprintf("%s%s/revoke", baseURL, prefix),
286290
RevocationEndpointAuthMethodsSupported: p.metadata.RevocationEndpointAuthMethodsSupported,
287-
RegistrationEndpoint: fmt.Sprintf("%s/register", baseURL),
291+
RegistrationEndpoint: fmt.Sprintf("%s%s/register", baseURL, prefix),
288292
RegistrationEndpointAuthMethodsSupported: p.metadata.RegistrationEndpointAuthMethodsSupported,
289293
}
290294

@@ -293,9 +297,12 @@ func (p *OAuthProxy) oauthMetadataHandler(w http.ResponseWriter, r *http.Request
293297

294298
func (p *OAuthProxy) protectedResourceMetadataHandler(w http.ResponseWriter, r *http.Request) {
295299
baseURL := handlerutils.GetBaseURL(r)
300+
prefix := p.config.RoutePrefix
301+
resourceURL := baseURL + prefix
302+
296303
metadata := types.OAuthProtectedResourceMetadata{
297-
Resource: baseURL,
298-
AuthorizationServers: []string{baseURL},
304+
Resource: resourceURL,
305+
AuthorizationServers: []string{baseURL + prefix},
299306
Scopes: p.metadata.ScopesSupported,
300307
ResourceName: p.resourceName,
301308
ResourceDocumentation: p.metadata.ServiceDocumentation,

‎pkg/types/types.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ type Config struct {
1515
EncryptionKey string
1616
MCPServerURL string
1717
Mode string
18+
RoutePrefix string
1819
}
1920

2021
// TokenData represents stored token data for OAuth 2.1 compliance

0 commit comments

Comments
 (0)