refactor: migrate MCP SDK to main repo (#3516)

plugins/wasm-go/pkg/mcp/server/proxy_auth_test.go

@@ -0,0 +1,429 @@
+// Copyright (c) 2022 Alibaba Group Holding Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package server
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestApiKeyAuthentication tests API key authentication forwarding
+func TestApiKeyAuthentication(t *testing.T) {
+	server := NewMcpProxyServer("auth-test")
+
+	// Configure security scheme
+	scheme := SecurityScheme{
+		ID:                "ApiKeyAuth",
+		Type:              "apiKey",
+		In:                "header",
+		Name:              "X-API-Key",
+		DefaultCredential: "default-api-key",
+	}
+
+	server.AddSecurityScheme(scheme)
+
+	// Set server fields directly
+	server.SetMcpServerURL("http://secure-backend.example.com/mcp")
+	server.SetTimeout(5000)
+
+	// Create tool with client-to-gateway and gateway-to-backend security
+	toolConfig := McpProxyToolConfig{
+		Name:        "secure_tool",
+		Description: "Tool requiring authentication",
+		Security: SecurityRequirement{
+			ID:          "ApiKeyAuth", // Client-to-gateway authentication
+			Passthrough: true,         // Extract client credential for backend use
+		},
+		Args: []ToolArg{
+			{
+				Name:        "data",
+				Description: "Data parameter",
+				Type:        "string",
+				Required:    true,
+			},
+		},
+		OutputSchema: map[string]any{
+			"type": "object",
+			"properties": map[string]any{
+				"result": map[string]any{
+					"type":        "string",
+					"description": "The result of the operation",
+				},
+			},
+		},
+		RequestTemplate: RequestTemplate{
+			Security: SecurityRequirement{
+				ID: "ApiKeyAuth", // Gateway-to-backend authentication (same scheme for simplicity)
+			},
+		},
+	}
+
+	err := server.AddProxyTool(toolConfig)
+	require.NoError(t, err)
+
+	tool, exists := server.GetMCPTools()["secure_tool"]
+	require.True(t, exists)
+
+	params := map[string]interface{}{
+		"data": "test data",
+	}
+	paramsBytes, err := json.Marshal(params)
+	require.NoError(t, err)
+
+	toolInstance := tool.Create(paramsBytes)
+	require.NotNil(t, toolInstance)
+
+	// Authentication is now handled automatically during tool calls
+	// The actual authentication flow is tested in integration tests
+}
+
+// TestBearerAuthentication tests Bearer token authentication
+func TestBearerAuthentication(t *testing.T) {
+	server := NewMcpProxyServer("bearer-auth-test")
+
+	// Configure Bearer security scheme
+	scheme := SecurityScheme{
+		ID:     "BearerAuth",
+		Type:   "http",
+		Scheme: "bearer",
+	}
+
+	server.AddSecurityScheme(scheme)
+
+	// Set server fields directly
+	server.SetMcpServerURL("https://secure-backend.example.com/mcp")
+	server.SetTimeout(8000)
+
+	// Create tool with Bearer authentication
+	// Create tool using only gateway-to-backend authentication (no client auth required)
+	toolConfig := McpProxyToolConfig{
+		Name:        "bearer_tool",
+		Description: "Tool with Bearer authentication to backend only",
+		Args: []ToolArg{
+			{
+				Name:        "query",
+				Description: "Query parameter",
+				Type:        "string",
+				Required:    true,
+			},
+		},
+		RequestTemplate: RequestTemplate{
+			Security: SecurityRequirement{
+				ID: "BearerAuth", // Only gateway-to-backend authentication
+			},
+		},
+	}
+
+	err := server.AddProxyTool(toolConfig)
+	require.NoError(t, err)
+
+	tool, exists := server.GetMCPTools()["bearer_tool"]
+	require.True(t, exists)
+
+	params := map[string]interface{}{
+		"query": "test query",
+	}
+	paramsBytes, err := json.Marshal(params)
+	require.NoError(t, err)
+
+	toolInstance := tool.Create(paramsBytes)
+	require.NotNil(t, toolInstance)
+
+	// Authentication is now handled automatically during tool calls
+	// The actual authentication flow is tested in integration tests
+
+	// Test backward compatibility: this tool uses RequestTemplate.Security (legacy way)
+	// which should still work
+}
+
+// TestBasicAuthentication tests Basic authentication
+func TestBasicAuthentication(t *testing.T) {
+	server := NewMcpProxyServer("basic-auth-test")
+
+	// Configure Basic security scheme
+	scheme := SecurityScheme{
+		ID:     "BasicAuth",
+		Type:   "http",
+		Scheme: "basic",
+	}
+
+	server.AddSecurityScheme(scheme)
+
+	// Test tool call with Basic authentication
+	toolConfig := McpProxyToolConfig{
+		Name:        "basic_tool",
+		Description: "Tool with Basic authentication",
+		Args: []ToolArg{
+			{
+				Name:        "resource",
+				Description: "Resource identifier",
+				Type:        "string",
+				Required:    true,
+			},
+		},
+		RequestTemplate: RequestTemplate{
+			Security: SecurityRequirement{
+				ID: "BasicAuth",
+			},
+		},
+	}
+
+	err := server.AddProxyTool(toolConfig)
+	require.NoError(t, err)
+
+	tool, exists := server.GetMCPTools()["basic_tool"]
+	require.True(t, exists)
+
+	params := map[string]interface{}{
+		"resource": "test-resource",
+	}
+	paramsBytes, err := json.Marshal(params)
+	require.NoError(t, err)
+
+	toolInstance := tool.Create(paramsBytes)
+	require.NotNil(t, toolInstance)
+
+	// Authentication is now handled automatically during tool calls
+	// The actual authentication flow is tested in integration tests
+
+	// Test OutputSchema functionality (only for tools that have it configured)
+	if toolWithOutputSchema, ok := tool.(ToolWithOutputSchema); ok {
+		outputSchema := toolWithOutputSchema.OutputSchema()
+		if outputSchema != nil {
+			// Only validate if outputSchema is configured
+			assert.Equal(t, "object", outputSchema["type"])
+			properties, hasProperties := outputSchema["properties"].(map[string]any)
+			require.True(t, hasProperties)
+			resultSchema, hasResult := properties["result"].(map[string]any)
+			require.True(t, hasResult)
+			assert.Equal(t, "string", resultSchema["type"])
+			assert.Equal(t, "The result of the operation", resultSchema["description"])
+		}
+	}
+}
+
+// TestMultipleSecuritySchemes tests multiple security schemes in one server
+func TestMultipleSecuritySchemes(t *testing.T) {
+	server := NewMcpProxyServer("multi-auth-test")
+
+	// Add multiple security schemes
+	schemes := []SecurityScheme{
+		{
+			ID:   "ApiKeyAuth",
+			Type: "apiKey",
+			In:   "header",
+			Name: "X-API-Key",
+		},
+		{
+			ID:     "BearerAuth",
+			Type:   "http",
+			Scheme: "bearer",
+		},
+	}
+
+	for _, scheme := range schemes {
+		server.AddSecurityScheme(scheme)
+	}
+
+	// Test that both schemes are available
+	for _, scheme := range schemes {
+		retrievedScheme, exists := server.GetSecurityScheme(scheme.ID)
+		assert.True(t, exists)
+		assert.Equal(t, scheme.ID, retrievedScheme.ID)
+		assert.Equal(t, scheme.Type, retrievedScheme.Type)
+	}
+}
+
+// ProxyAuthContext, RequestTemplate, SecurityConfig and authentication methods
+// are now implemented in proxy_server.go
+
+// TestToolsListAuthentication tests authentication configuration for tools/list requests
+func TestToolsListAuthentication(t *testing.T) {
+	server := NewMcpProxyServer("test-server")
+
+	// Add a security scheme for global authentication
+	scheme := SecurityScheme{
+		ID:                "GlobalAuth",
+		Type:              "apiKey",
+		In:                "header",
+		Name:              "X-API-Key",
+		DefaultCredential: "default-global-key",
+	}
+	server.AddSecurityScheme(scheme)
+
+	// Test that we can retrieve the security scheme
+	retrievedScheme, exists := server.GetSecurityScheme("GlobalAuth")
+	assert.True(t, exists)
+	assert.Equal(t, "GlobalAuth", retrievedScheme.ID)
+	assert.Equal(t, "apiKey", retrievedScheme.Type)
+	assert.Equal(t, "header", retrievedScheme.In)
+	assert.Equal(t, "X-API-Key", retrievedScheme.Name)
+
+	// Test setting default security directly on server
+	defaultDownstreamSecurity := SecurityRequirement{
+		ID:          "GlobalAuth",
+		Passthrough: true,
+	}
+	defaultUpstreamSecurity := SecurityRequirement{
+		ID: "GlobalAuth",
+	}
+
+	server.SetDefaultDownstreamSecurity(defaultDownstreamSecurity)
+	server.SetDefaultUpstreamSecurity(defaultUpstreamSecurity)
+
+	// Verify default security settings
+	retrievedDownstream := server.GetDefaultDownstreamSecurity()
+	assert.Equal(t, "GlobalAuth", retrievedDownstream.ID)
+	assert.True(t, retrievedDownstream.Passthrough)
+
+	retrievedUpstream := server.GetDefaultUpstreamSecurity()
+	assert.Equal(t, "GlobalAuth", retrievedUpstream.ID)
+
+	t.Logf("Tools/list authentication configuration test completed successfully")
+}
+
+// TestDefaultSecurityFallback tests the fallback mechanism from tool-level to default security
+func TestDefaultSecurityFallback(t *testing.T) {
+	server := NewMcpProxyServer("test-server")
+
+	// Add security schemes
+	defaultScheme := SecurityScheme{
+		ID:                "DefaultAuth",
+		Type:              "apiKey",
+		In:                "header",
+		Name:              "X-Default-Key",
+		DefaultCredential: "default-key",
+	}
+	toolScheme := SecurityScheme{
+		ID:                "ToolAuth",
+		Type:              "apiKey",
+		In:                "header",
+		Name:              "X-Tool-Key",
+		DefaultCredential: "tool-key",
+	}
+	server.AddSecurityScheme(defaultScheme)
+	server.AddSecurityScheme(toolScheme)
+
+	// Test tool configuration with tool-level security (should use tool-level, not default)
+	toolConfigWithSecurity := McpProxyToolConfig{
+		Name:        "secure_tool",
+		Description: "Tool with its own security",
+		Security: SecurityRequirement{
+			ID:          "ToolAuth",
+			Passthrough: true,
+		},
+		RequestTemplate: RequestTemplate{
+			Security: SecurityRequirement{
+				ID: "ToolAuth",
+			},
+		},
+	}
+
+	// Test tool configuration without tool-level security (should fallback to default)
+	toolConfigWithoutSecurity := McpProxyToolConfig{
+		Name:        "fallback_tool",
+		Description: "Tool that falls back to default security",
+		// No Security field configured, should use default
+		RequestTemplate: RequestTemplate{
+			// No Security field configured, should use default
+		},
+	}
+
+	// Set default security directly on server
+	server.SetDefaultDownstreamSecurity(SecurityRequirement{
+		ID:          "DefaultAuth",
+		Passthrough: false,
+	})
+	server.SetDefaultUpstreamSecurity(SecurityRequirement{
+		ID: "DefaultAuth",
+	})
+
+	// Set server configuration directly
+	server.SetMcpServerURL("http://backend.example.com")
+	server.SetTimeout(5000)
+
+	// Add tools to server
+	err := server.AddProxyTool(toolConfigWithSecurity)
+	assert.NoError(t, err)
+	err = server.AddProxyTool(toolConfigWithoutSecurity)
+	assert.NoError(t, err)
+
+	// Verify tools were added
+	tools := server.GetMCPTools()
+	assert.Contains(t, tools, "secure_tool")
+	assert.Contains(t, tools, "fallback_tool")
+
+	t.Logf("Default security fallback test completed successfully")
+}
+
+// TestURLModificationInAuthentication tests that authentication can modify the URL (e.g., adding query parameters)
+func TestURLModificationInAuthentication(t *testing.T) {
+	server := NewMcpProxyServer("test-server")
+
+	// Add a security scheme that adds parameters to query (apiKey in query)
+	scheme := SecurityScheme{
+		ID:                "QueryApiKey",
+		Type:              "apiKey",
+		In:                "query",
+		Name:              "api_key",
+		DefaultCredential: "test-key-123",
+	}
+	server.AddSecurityScheme(scheme)
+
+	// Verify the security scheme was added correctly
+	retrievedScheme, exists := server.GetSecurityScheme("QueryApiKey")
+	assert.True(t, exists)
+	assert.Equal(t, "apiKey", retrievedScheme.Type)
+	assert.Equal(t, "query", retrievedScheme.In)
+	assert.Equal(t, "api_key", retrievedScheme.Name)
+
+	t.Logf("URL modification authentication configuration test completed successfully")
+}
+
+// TestProxyServerFields tests the server-level field setting and getting
+func TestProxyServerFields(t *testing.T) {
+	server := NewMcpProxyServer("test-server")
+
+	// Test mcpServerURL
+	testURL := "http://mcp.example.com:8080/mcp"
+	server.SetMcpServerURL(testURL)
+	assert.Equal(t, testURL, server.GetMcpServerURL())
+
+	// Test timeout
+	testTimeout := 10000
+	server.SetTimeout(testTimeout)
+	assert.Equal(t, testTimeout, server.GetTimeout())
+
+	// Test default security settings
+	downstreamSec := SecurityRequirement{
+		ID:          "test-downstream",
+		Passthrough: true,
+	}
+	upstreamSec := SecurityRequirement{
+		ID: "test-upstream",
+	}
+
+	server.SetDefaultDownstreamSecurity(downstreamSec)
+	server.SetDefaultUpstreamSecurity(upstreamSec)
+
+	assert.Equal(t, "test-downstream", server.GetDefaultDownstreamSecurity().ID)
+	assert.True(t, server.GetDefaultDownstreamSecurity().Passthrough)
+	assert.Equal(t, "test-upstream", server.GetDefaultUpstreamSecurity().ID)
+
+	t.Logf("Proxy server fields test completed successfully")
+}