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

430 lines
12 KiB
Go

// 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")
}