refactor: stop rewriting paths in proxy (#15)

This change also properly handles redirects from the downstream MCP
server and rewrites them to point to the proxy.

Signed-off-by: Donnie Adams <donnie@obot.ai>

Author: thedadams

README.md

@@ -271,7 +271,7 @@ You can customize the scopes based on your needs. Common additional scopes inclu
          OAUTH_CLIENT_ID: "your-oauth-client-id"
          OAUTH_CLIENT_SECRET: "your-oauth-client-secret"
          OAUTH_AUTHORIZE_URL: "https://your-oauth-provider.com/oauth/authorize"
-         MCP_SERVER_URL: "http://localhost:3000/mcp"
+         MCP_SERVER_URL: "http://localhost:3000"
          ENCRYPTION_KEY: "your-base64-encoded-32-byte-key"
        ports:
          - "8080:8080"
@@ -295,7 +295,7 @@ You can customize the scopes based on your needs. Common additional scopes inclu
          OAUTH_CLIENT_ID: "your-oauth-client-id"
          OAUTH_CLIENT_SECRET: "your-oauth-client-secret"
          OAUTH_AUTHORIZE_URL: "https://your-oauth-provider.com/oauth/authorize"
-         MCP_SERVER_URL: "http://localhost:3000/mcp"
+         MCP_SERVER_URL: "http://localhost:3000"
        volumes:
          - ./data:/app/data # Persist SQLite database
        ports:
@@ -342,7 +342,7 @@ You can customize the scopes based on your needs. Common additional scopes inclu
 
 ### MCP Proxy
 
-- `ANY /mcp/*` - Proxies requests to MCP server with user context headers
+- `ANY /*` - Proxies any request not mentioned above to MCP server with user context headers
 
 ## OAuth Flow
 

docker-compose.yml

@@ -19,7 +19,7 @@ services:
   #     OAUTH_CLIENT_ID: "your-oauth-client-id"
   #     OAUTH_CLIENT_SECRET: "your-oauth-client-secret"
   #     OAUTH_AUTHORIZE_URL: "https://your-oauth-provider.com/oauth/authorize"
-  #     MCP_SERVER_URL: "http://localhost:3000/mcp"
+  #     MCP_SERVER_URL: "http://localhost:3000"
   #   ports:
   #     - "8080:8080"
   #   depends_on:

main.go

@@ -9,7 +9,10 @@ import (
 
 func main() {
 	// Load configuration from environment variables
-	config := proxy.LoadConfigFromEnv()
+	config, err := proxy.LoadConfigFromEnv()
+	if err != nil {
+		log.Fatalf("Failed to load configuration: %v", err)
+	}
 
 	proxy, err := proxy.NewOAuthProxy(config)
 	if err != nil {

main_test.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"log"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -61,7 +62,10 @@ func TestIntegrationFlow(t *testing.T) {
 	}()
 
 	// Create OAuth proxy
-	config := proxy.LoadConfigFromEnv()
+	config, err := proxy.LoadConfigFromEnv()
+	if err != nil {
+		log.Fatalf("Failed to load configuration: %v", err)
+	}
 	oauthProxy, err := proxy.NewOAuthProxy(config)
 	if err != nil {
 		t.Skipf("Skipping test due to database connection error: %v", err)
@@ -99,7 +103,7 @@ func TestIntegrationFlow(t *testing.T) {
 	// Test protected resource metadata
 	t.Run("ProtectedResourceMetadata", func(t *testing.T) {
 		w := httptest.NewRecorder()
-		req := httptest.NewRequest("GET", "/.well-known/oauth-protected-resource/mcp", nil)
+		req := httptest.NewRequest("GET", "/.well-known/oauth-protected-resource", nil)
 		handler.ServeHTTP(w, req)
 
 		assert.Equal(t, http.StatusOK, w.Code)
@@ -165,7 +169,10 @@ func TestOAuthProxyCreation(t *testing.T) {
 	}()
 
 	// Create OAuth proxy
-	config := proxy.LoadConfigFromEnv()
+	config, err := proxy.LoadConfigFromEnv()
+	if err != nil {
+		log.Fatalf("Failed to load configuration: %v", err)
+	}
 	oauthProxy, err := proxy.NewOAuthProxy(config)
 	require.NoError(t, err, "Should be able to create OAuth proxy with valid environment")
 	require.NotNil(t, oauthProxy, "OAuth proxy should not be nil")
@@ -215,7 +222,10 @@ func TestOAuthProxyStart(t *testing.T) {
 	}()
 
 	// Create OAuth proxy
-	config := proxy.LoadConfigFromEnv()
+	config, err := proxy.LoadConfigFromEnv()
+	if err != nil {
+		log.Fatalf("Failed to load configuration: %v", err)
+	}
 	oauthProxy, err := proxy.NewOAuthProxy(config)
 	require.NoError(t, err)
 	defer func() {

pkg/oauth/validate/validatetoken.go

@@ -28,7 +28,7 @@ func (p *TokenValidator) WithTokenValidation(next http.HandlerFunc) http.Handler
 		authHeader := r.Header.Get("Authorization")
 		if authHeader == "" {
 			// Return 401 with proper WWW-Authenticate header
-			resourceMetadataUrl := fmt.Sprintf("%s/.well-known/oauth-protected-resource/mcp", handlerutils.GetBaseURL(r))
+			resourceMetadataUrl := fmt.Sprintf("%s/.well-known/oauth-protected-resource", handlerutils.GetBaseURL(r))
 			wwwAuthValue := fmt.Sprintf(`Bearer error="invalid_token", error_description="Missing Authorization header", resource_metadata="%s"`, resourceMetadataUrl)
 			w.Header().Set("WWW-Authenticate", wwwAuthValue)
 			handlerutils.JSON(w, http.StatusUnauthorized, map[string]string{
@@ -41,7 +41,7 @@ func (p *TokenValidator) WithTokenValidation(next http.HandlerFunc) http.Handler
 		// Parse Authorization header
 		parts := strings.SplitN(authHeader, " ", 2)
 		if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" || parts[1] == "" {
-			resourceMetadataUrl := fmt.Sprintf("%s/.well-known/oauth-protected-resource/mcp", handlerutils.GetBaseURL(r))
+			resourceMetadataUrl := fmt.Sprintf("%s/.well-known/oauth-protected-resource", handlerutils.GetBaseURL(r))
 			wwwAuthValue := fmt.Sprintf(`Bearer error="invalid_token", error_description="Invalid Authorization header format, expected 'Bearer TOKEN'", resource_metadata="%s"`, resourceMetadataUrl)
 			w.Header().Set("WWW-Authenticate", wwwAuthValue)
 			handlerutils.JSON(w, http.StatusUnauthorized, map[string]string{
@@ -55,7 +55,7 @@ func (p *TokenValidator) WithTokenValidation(next http.HandlerFunc) http.Handler
 
 		tokenInfo, err := p.tokenManager.GetTokenInfo(token)
 		if err != nil {
-			resourceMetadataUrl := fmt.Sprintf("%s/.well-known/oauth-protected-resource/mcp", handlerutils.GetBaseURL(r))
+			resourceMetadataUrl := fmt.Sprintf("%s/.well-known/oauth-protected-resource", handlerutils.GetBaseURL(r))
 			wwwAuthValue := fmt.Sprintf(`Bearer error="invalid_token", error_description="Invalid or expired token", resource_metadata="%s"`, resourceMetadataUrl)
 			w.Header().Set("WWW-Authenticate", wwwAuthValue)
 			handlerutils.JSON(w, http.StatusUnauthorized, map[string]string{
@@ -69,7 +69,7 @@ func (p *TokenValidator) WithTokenValidation(next http.HandlerFunc) http.Handler
 		if tokenInfo.Props != nil {
 			decryptedProps, err := encryption.DecryptPropsIfNeeded(p.encryptionKey, tokenInfo.Props)
 			if err != nil {
-				resourceMetadataUrl := fmt.Sprintf("%s/.well-known/oauth-protected-resource/mcp", handlerutils.GetBaseURL(r))
+				resourceMetadataUrl := fmt.Sprintf("%s/.well-known/oauth-protected-resource", handlerutils.GetBaseURL(r))
 				wwwAuthValue := fmt.Sprintf(`Bearer error="invalid_token", error_description="Failed to decrypt token data", resource_metadata="%s"`, resourceMetadataUrl)
 				w.Header().Set("WWW-Authenticate", wwwAuthValue)
 				handlerutils.JSON(w, http.StatusUnauthorized, map[string]string{

pkg/proxy/proxy.go

@@ -48,8 +48,8 @@ type OAuthProxy struct {
 }
 
 // LoadConfigFromEnv loads configuration from environment variables
-func LoadConfigFromEnv() *types.Config {
-	return &types.Config{
+func LoadConfigFromEnv() (*types.Config, error) {
+	config := &types.Config{
 		DatabaseDSN:       os.Getenv("DATABASE_DSN"),
 		OAuthClientID:     os.Getenv("OAUTH_CLIENT_ID"),
 		OAuthClientSecret: os.Getenv("OAUTH_CLIENT_SECRET"),
@@ -58,6 +58,14 @@ func LoadConfigFromEnv() *types.Config {
 		EncryptionKey:     os.Getenv("ENCRYPTION_KEY"),
 		MCPServerURL:      os.Getenv("MCP_SERVER_URL"),
 	}
+
+	if u, err := url.Parse(config.MCPServerURL); err != nil || u.Scheme != "http" && u.Scheme != "https" {
+		return nil, fmt.Errorf("invalid MCP server URL: %w", err)
+	} else if u.Path != "" && u.Path != "/" || u.RawQuery != "" || u.Fragment != "" {
+		return nil, fmt.Errorf("MCP server URL must not contain a path, query, or fragment")
+	}
+
+	return config, nil
 }
 
 func NewOAuthProxy(config *types.Config) (*OAuthProxy, error) {
@@ -194,22 +202,21 @@ func (p *OAuthProxy) SetupRoutes(mux *http.ServeMux) {
 	revokeHandler := revoke.NewHandler(p.db)
 	tokenValidator := validate.NewTokenValidator(p.tokenManager, p.encryptionKey)
 
-	mux.HandleFunc("/health", p.withCORS(p.healthHandler))
+	mux.HandleFunc("GET /health", p.withCORS(p.healthHandler))
 
 	// OAuth endpoints
-	mux.HandleFunc("/authorize", p.withCORS(p.withRateLimit(authorizeHandler)))
-	mux.HandleFunc("/callback", p.withCORS(p.withRateLimit(callbackHandler)))
-	mux.HandleFunc("/token", p.withCORS(p.withRateLimit(tokenHandler)))
-	mux.HandleFunc("/revoke", p.withCORS(p.withRateLimit(revokeHandler)))
-	mux.HandleFunc("/register", p.withCORS(p.withRateLimit(register.NewHandler(p.db))))
+	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))))
 
 	// Metadata endpoints
-	mux.HandleFunc("/.well-known/oauth-authorization-server", p.withCORS(p.oauthMetadataHandler))
-	mux.HandleFunc("/.well-known/oauth-protected-resource/mcp", p.withCORS(p.protectedResourceMetadataHandler))
+	mux.HandleFunc("GET /.well-known/oauth-authorization-server", p.withCORS(p.oauthMetadataHandler))
+	mux.HandleFunc("GET /.well-known/oauth-protected-resource", p.withCORS(p.protectedResourceMetadataHandler))
 
-	// Protected resource endpoints
-	mux.HandleFunc("/mcp", p.withCORS(p.withRateLimit(tokenValidator.WithTokenValidation(p.mcpProxyHandler))))
-	mux.HandleFunc("/mcp/{path...}", p.withCORS(p.withRateLimit(tokenValidator.WithTokenValidation(p.mcpProxyHandler))))
+	// Protect everything else
+	mux.HandleFunc("/{path...}", p.withCORS(p.withRateLimit(tokenValidator.WithTokenValidation(p.mcpProxyHandler))))
 }
 
 // GetHandler returns an http.Handler for the OAuth proxy
@@ -289,9 +296,10 @@ func (p *OAuthProxy) oauthMetadataHandler(w http.ResponseWriter, r *http.Request
 }
 
 func (p *OAuthProxy) protectedResourceMetadataHandler(w http.ResponseWriter, r *http.Request) {
+	baseURL := handlerutils.GetBaseURL(r)
 	metadata := types.OAuthProtectedResourceMetadata{
-		Resource:              fmt.Sprintf("%s/mcp", handlerutils.GetBaseURL(r)),
-		AuthorizationServers:  []string{handlerutils.GetBaseURL(r)},
+		Resource:              baseURL,
+		AuthorizationServers:  []string{baseURL},
 		Scopes:                p.metadata.ScopesSupported,
 		ResourceName:          p.resourceName,
 		ResourceDocumentation: p.metadata.ServiceDocumentation,
@@ -387,15 +395,7 @@ func (p *OAuthProxy) mcpProxyHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Create target URL
-	var targetURL string
-	if path == "" {
-		// If no path is provided, use the MCP server URL directly
-		targetURL = p.GetMCPServerURL()
-	} else {
-		// If path is provided, append it to the MCP server URL
-		targetURL = p.GetMCPServerURL() + "/" + path
-	}
-
+	targetURL := p.GetMCPServerURL() + "/" + path
 	// Log the proxy request for debugging
 	log.Printf("Proxying request: %s %s -> %s", r.Method, r.URL.Path, targetURL)
 
@@ -404,12 +404,12 @@ func (p *OAuthProxy) mcpProxyHandler(w http.ResponseWriter, r *http.Request) {
 		Director: func(req *http.Request) {
 			req.Header.Del("Authorization")
 			req.Header.Set("X-Forwarded-Host", req.Host)
+			req.Header.Set("X-Forwarded-Proto", req.URL.Scheme)
 
 			newURL, _ := url.Parse(targetURL)
 			req.URL.Scheme = newURL.Scheme
 			req.URL.Host = newURL.Host
 			req.Host = newURL.Host
-			req.URL.Path = newURL.Path
 
 			// Add forwarded headers from token props
 			if tokenInfo.Props != nil {
@@ -427,6 +427,27 @@ func (p *OAuthProxy) mcpProxyHandler(w http.ResponseWriter, r *http.Request) {
 				}
 			}
 		},
+		ModifyResponse: func(resp *http.Response) error {
+			// Rewrite Location header to use proxy host instead of downstream server host
+			if location := resp.Header.Get("Location"); location != "" {
+				if locationURL, err := url.Parse(location); err == nil {
+					// Get the original request to extract proxy host
+					proxyHost := resp.Request.Header.Get("X-Forwarded-Host")
+					if proxyHost != "" {
+						// Parse downstream server URL to get scheme
+						downstreamURL, _ := url.Parse(p.GetMCPServerURL())
+
+						// Only rewrite if the location points to the downstream server
+						if locationURL.Host == downstreamURL.Host {
+							locationURL.Scheme = resp.Request.URL.Scheme
+							locationURL.Host = proxyHost
+							resp.Header.Set("Location", locationURL.String())
+						}
+					}
+				}
+			}
+			return nil
+		},
 		ErrorHandler: func(rw http.ResponseWriter, req *http.Request, err error) {
 			log.Printf("Proxy error: %v", err)
 			rw.WriteHeader(http.StatusBadGateway)