From 23ebbaf58e3c35e3c9af915ec6e1c3cbc8ebd775 Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Wed, 24 Sep 2025 10:29:29 +0300 Subject: [PATCH] Add secrets management support to ToolHive MCP server Implement comprehensive secrets management functionality for the ToolHive MCP server: - Add list_secrets tool to list available secrets from ToolHive secrets store - Add set_secret tool to set secrets from file paths (file-based input only) - Enhance run_server tool with secrets parameter support - Add SecretMapping struct for structured secret name/target specification - Include comprehensive test coverage for all new functionality - Integrate with existing ToolHive secrets providers (encrypted, 1Password, etc.) The run_server tool now accepts a secrets array parameter allowing users to pass secrets to MCP servers when running them, matching the CLI --secret flag functionality but through the MCP interface. --- pkg/mcp/server/handler.go | 6 + pkg/mcp/server/handler_test.go | 89 +++++++++++-- pkg/mcp/server/list_secrets.go | 68 ++++++++++ pkg/mcp/server/list_secrets_test.go | 88 +++++++++++++ pkg/mcp/server/run_server.go | 36 +++++- pkg/mcp/server/server.go | 46 +++++++ pkg/mcp/server/set_secret.go | 118 ++++++++++++++++++ pkg/mcp/server/set_secret_test.go | 186 ++++++++++++++++++++++++++++ 8 files changed, 621 insertions(+), 16 deletions(-) create mode 100644 pkg/mcp/server/list_secrets.go create mode 100644 pkg/mcp/server/list_secrets_test.go create mode 100644 pkg/mcp/server/set_secret.go create mode 100644 pkg/mcp/server/set_secret_test.go diff --git a/pkg/mcp/server/handler.go b/pkg/mcp/server/handler.go index 7fb6ec576..14e46d82c 100644 --- a/pkg/mcp/server/handler.go +++ b/pkg/mcp/server/handler.go @@ -5,6 +5,7 @@ import ( "context" "fmt" + "github.com/stacklok/toolhive/pkg/config" "github.com/stacklok/toolhive/pkg/registry" "github.com/stacklok/toolhive/pkg/workloads" ) @@ -14,6 +15,7 @@ type Handler struct { ctx context.Context workloadManager workloads.Manager registryProvider registry.Provider + configProvider config.Provider } // NewHandler creates a new ToolHive handler @@ -30,9 +32,13 @@ func NewHandler(ctx context.Context) (*Handler, error) { return nil, fmt.Errorf("failed to get registry provider: %w", err) } + // Create config provider + configProvider := config.NewDefaultProvider() + return &Handler{ ctx: ctx, workloadManager: workloadManager, registryProvider: registryProvider, + configProvider: configProvider, }, nil } diff --git a/pkg/mcp/server/handler_test.go b/pkg/mcp/server/handler_test.go index 1fb49cb31..c6291b13e 100644 --- a/pkg/mcp/server/handler_test.go +++ b/pkg/mcp/server/handler_test.go @@ -31,6 +31,16 @@ func TestParseRunServerArgs(t *testing.T) { "KEY1": "value1", "KEY2": "value2", }, + "secrets": []interface{}{ + map[string]interface{}{ + "name": "github-token", + "target": "GITHUB_TOKEN", + }, + map[string]interface{}{ + "name": "api-key", + "target": "API_KEY", + }, + }, }, }, }, @@ -42,6 +52,10 @@ func TestParseRunServerArgs(t *testing.T) { "KEY1": "value1", "KEY2": "value2", }, + Secrets: []SecretMapping{ + {Name: "github-token", Target: "GITHUB_TOKEN"}, + {Name: "api-key", Target: "API_KEY"}, + }, }, wantErr: false, }, @@ -55,10 +69,11 @@ func TestParseRunServerArgs(t *testing.T) { }, }, expected: &runServerArgs{ - Server: "test-server", - Name: "test-server", // Should default to server name - Host: "127.0.0.1", // Should default to 127.0.0.1 - Env: nil, + Server: "test-server", + Name: "test-server", // Should default to server name + Host: "127.0.0.1", // Should default to 127.0.0.1 + Env: nil, + Secrets: nil, }, wantErr: false, }, @@ -73,10 +88,11 @@ func TestParseRunServerArgs(t *testing.T) { }, }, expected: &runServerArgs{ - Server: "my-server", - Name: "my-server", - Host: "127.0.0.1", - Env: nil, + Server: "my-server", + Name: "my-server", + Host: "127.0.0.1", + Env: nil, + Secrets: nil, }, wantErr: false, }, @@ -91,10 +107,11 @@ func TestParseRunServerArgs(t *testing.T) { }, }, expected: &runServerArgs{ - Server: "test-server", - Name: "test-server", - Host: "127.0.0.1", - Env: nil, + Server: "test-server", + Name: "test-server", + Host: "127.0.0.1", + Env: nil, + Secrets: nil, }, wantErr: false, }, @@ -306,3 +323,51 @@ func TestBuildServerConfig(t *testing.T) { }) } } + +func TestPrepareSecrets(t *testing.T) { + t.Parallel() + tests := []struct { + name string + secrets []SecretMapping + expected []string + }{ + { + name: "nil secrets", + secrets: nil, + expected: nil, + }, + { + name: "empty secrets", + secrets: []SecretMapping{}, + expected: nil, + }, + { + name: "single secret", + secrets: []SecretMapping{ + {Name: "github-token", Target: "GITHUB_TOKEN"}, + }, + expected: []string{"github-token,target=GITHUB_TOKEN"}, + }, + { + name: "multiple secrets", + secrets: []SecretMapping{ + {Name: "github-token", Target: "GITHUB_TOKEN"}, + {Name: "api-key", Target: "API_KEY"}, + {Name: "db-password", Target: "DATABASE_PASSWORD"}, + }, + expected: []string{ + "github-token,target=GITHUB_TOKEN", + "api-key,target=API_KEY", + "db-password,target=DATABASE_PASSWORD", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + result := prepareSecrets(tt.secrets) + assert.Equal(t, tt.expected, result) + }) + } +} diff --git a/pkg/mcp/server/list_secrets.go b/pkg/mcp/server/list_secrets.go new file mode 100644 index 000000000..46b6f6339 --- /dev/null +++ b/pkg/mcp/server/list_secrets.go @@ -0,0 +1,68 @@ +package server + +import ( + "context" + "fmt" + + "github.com/mark3labs/mcp-go/mcp" + + "github.com/stacklok/toolhive/pkg/secrets" +) + +// SecretInfo represents secret information returned by list +type SecretInfo struct { + Key string `json:"key"` + Description string `json:"description,omitempty"` +} + +// ListSecretsResponse represents the response from listing secrets +type ListSecretsResponse struct { + Secrets []SecretInfo `json:"secrets"` +} + +// ListSecrets lists all available secrets +func (h *Handler) ListSecrets(ctx context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + // Get the configuration to determine the secrets provider + cfg := h.configProvider.GetConfig() + + // Check if secrets setup has been completed + if !cfg.Secrets.SetupCompleted { + return mcp.NewToolResultError( + "Secrets provider not configured. Please run 'thv secret setup' to configure a secrets provider first"), nil + } + + // Get the provider type + providerType, err := cfg.Secrets.GetProviderType() + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("Failed to get secrets provider type: %v", err)), nil + } + + // Create the secrets provider + secretsProvider, err := secrets.CreateSecretProvider(providerType) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("Failed to create secrets provider: %v", err)), nil + } + + // List all secrets + secretDescriptions, err := secretsProvider.ListSecrets(ctx) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("Failed to list secrets: %v", err)), nil + } + + // Format results with structured data + var results []SecretInfo + for _, desc := range secretDescriptions { + info := SecretInfo{ + Key: desc.Key, + Description: desc.Description, + } + results = append(results, info) + } + + // Create structured response + response := ListSecretsResponse{ + Secrets: results, + } + + return mcp.NewToolResultStructuredOnly(response), nil +} diff --git a/pkg/mcp/server/list_secrets_test.go b/pkg/mcp/server/list_secrets_test.go new file mode 100644 index 000000000..b7722725c --- /dev/null +++ b/pkg/mcp/server/list_secrets_test.go @@ -0,0 +1,88 @@ +package server + +import ( + "context" + "testing" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/config" + configmocks "github.com/stacklok/toolhive/pkg/config/mocks" + registrymocks "github.com/stacklok/toolhive/pkg/registry/mocks" + workloadsmocks "github.com/stacklok/toolhive/pkg/workloads/mocks" +) + +func TestHandler_ListSecrets(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + t.Cleanup(func() { ctrl.Finish() }) + + tests := []struct { + name string + setupMocks func(*configmocks.MockProvider) + wantErr bool + checkResult func(*testing.T, *mcp.CallToolResult) + }{ + { + name: "secrets not setup", + setupMocks: func(configProvider *configmocks.MockProvider) { + // Mock config setup - not completed + cfg := &config.Config{ + Secrets: config.Secrets{ + SetupCompleted: false, + }, + } + configProvider.EXPECT().GetConfig().Return(cfg).AnyTimes() + }, + wantErr: false, + checkResult: func(t *testing.T, result *mcp.CallToolResult) { + t.Helper() + assert.NotNil(t, result) + assert.True(t, result.IsError) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // Create mocks + mockRegistry := registrymocks.NewMockProvider(ctrl) + mockWorkloadManager := workloadsmocks.NewMockManager(ctrl) + mockConfigProvider := configmocks.NewMockProvider(ctrl) + + // Setup mocks + if tt.setupMocks != nil { + tt.setupMocks(mockConfigProvider) + } + + handler := &Handler{ + ctx: context.Background(), + workloadManager: mockWorkloadManager, + registryProvider: mockRegistry, + configProvider: mockConfigProvider, + } + + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "list_secrets", + Arguments: map[string]interface{}{}, + }, + } + + result, err := handler.ListSecrets(context.Background(), request) + + if tt.wantErr { + assert.Error(t, err) + } else { + assert.NoError(t, err) + if tt.checkResult != nil { + tt.checkResult(t, result) + } + } + }) + } +} diff --git a/pkg/mcp/server/run_server.go b/pkg/mcp/server/run_server.go index b8ff182df..f00b4fcc6 100644 --- a/pkg/mcp/server/run_server.go +++ b/pkg/mcp/server/run_server.go @@ -14,12 +14,19 @@ import ( transporttypes "github.com/stacklok/toolhive/pkg/transport/types" ) +// SecretMapping represents a secret name and its target environment variable +type SecretMapping struct { + Name string `json:"name"` + Target string `json:"target"` +} + // runServerArgs holds the arguments for running a server type runServerArgs struct { - Server string `json:"server"` - Name string `json:"name,omitempty"` - Host string `json:"host,omitempty"` - Env map[string]string `json:"env,omitempty"` + Server string `json:"server"` + Name string `json:"name,omitempty"` + Host string `json:"host,omitempty"` + Env map[string]string `json:"env,omitempty"` + Secrets []SecretMapping `json:"secrets,omitempty"` } // RunServer runs an MCP server @@ -119,6 +126,12 @@ func buildServerConfig( // Prepare environment variables envVars := prepareEnvironmentVariables(imageMetadata, args.Env) + // Prepare secrets + secrets := prepareSecrets(args.Secrets) + if len(secrets) > 0 { + opts = append(opts, runner.WithSecrets(secrets)) + } + // Build the configuration envVarValidator := &runner.DetachedEnvVarValidator{} return runner.NewRunConfigBuilder(ctx, imageMetadata, envVars, envVarValidator, opts...) @@ -159,6 +172,21 @@ func prepareEnvironmentVariables(imageMetadata *registry.ImageMetadata, userEnv return envVarsMap } +// prepareSecrets converts SecretMapping array to the string format expected by the runner +func prepareSecrets(secretMappings []SecretMapping) []string { + if len(secretMappings) == 0 { + return nil + } + + secrets := make([]string, len(secretMappings)) + for i, mapping := range secretMappings { + // Convert to the format expected by runner: "secret_name,target=ENV_VAR_NAME" + secrets[i] = fmt.Sprintf("%s,target=%s", mapping.Name, mapping.Target) + } + + return secrets +} + // saveAndRunServer saves the configuration and runs the server func (h *Handler) saveAndRunServer(ctx context.Context, runConfig *runner.RunConfig, name string) error { // Save the run configuration state before starting diff --git a/pkg/mcp/server/server.go b/pkg/mcp/server/server.go index a1a08aef3..f17890624 100644 --- a/pkg/mcp/server/server.go +++ b/pkg/mcp/server/server.go @@ -136,6 +136,24 @@ func registerTools(mcpServer *server.MCPServer, handler *Handler) { "type": "string", }, }, + "secrets": map[string]interface{}{ + "type": "array", + "description": "Secrets to pass to the server as environment variables", + "items": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "name": map[string]interface{}{ + "type": "string", + "description": "Name of the secret in the ToolHive secrets store", + }, + "target": map[string]interface{}{ + "type": "string", + "description": "Target environment variable name in the server container", + }, + }, + "required": []string{"name", "target"}, + }, + }, }, Required: []string{"server"}, }, @@ -194,4 +212,32 @@ func registerTools(mcpServer *server.MCPServer, handler *Handler) { Required: []string{"name"}, }, }, handler.GetServerLogs) + + mcpServer.AddTool(mcp.Tool{ + Name: "list_secrets", + Description: "List all available secrets in the ToolHive secrets store", + InputSchema: mcp.ToolInputSchema{ + Type: "object", + Properties: map[string]interface{}{}, + }, + }, handler.ListSecrets) + + mcpServer.AddTool(mcp.Tool{ + Name: "set_secret", + Description: "Set a secret by reading its value from a file", + InputSchema: mcp.ToolInputSchema{ + Type: "object", + Properties: map[string]interface{}{ + "name": map[string]interface{}{ + "type": "string", + "description": "Name of the secret to set", + }, + "file_path": map[string]interface{}{ + "type": "string", + "description": "Path to the file containing the secret value", + }, + }, + Required: []string{"name", "file_path"}, + }, + }, handler.SetSecret) } diff --git a/pkg/mcp/server/set_secret.go b/pkg/mcp/server/set_secret.go new file mode 100644 index 000000000..6c66bf5a1 --- /dev/null +++ b/pkg/mcp/server/set_secret.go @@ -0,0 +1,118 @@ +package server + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/mark3labs/mcp-go/mcp" + + "github.com/stacklok/toolhive/pkg/secrets" +) + +// setSecretArgs holds the arguments for setting a secret +type setSecretArgs struct { + Name string `json:"name"` + FilePath string `json:"file_path"` +} + +// SetSecretResponse represents the response from setting a secret +type SetSecretResponse struct { + Status string `json:"status"` + Name string `json:"name"` +} + +// SetSecret sets a secret by reading its value from a file +func (h *Handler) SetSecret(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + // Parse arguments using BindArguments + args := &setSecretArgs{} + if err := request.BindArguments(args); err != nil { + return mcp.NewToolResultError(fmt.Sprintf("Failed to parse arguments: %v", err)), nil + } + + // Validate arguments + if args.Name == "" { + return mcp.NewToolResultError("Secret name cannot be empty"), nil + } + if args.FilePath == "" { + return mcp.NewToolResultError("File path cannot be empty"), nil + } + + // Clean and validate the file path + cleanPath := filepath.Clean(args.FilePath) + + // Check if file exists and is readable + fileInfo, err := os.Stat(cleanPath) + if err != nil { + if os.IsNotExist(err) { + return mcp.NewToolResultError(fmt.Sprintf("File does not exist: %s", cleanPath)), nil + } + return mcp.NewToolResultError(fmt.Sprintf("Cannot access file: %v", err)), nil + } + + // Check if it's a regular file (not a directory) + if !fileInfo.Mode().IsRegular() { + return mcp.NewToolResultError(fmt.Sprintf("Path is not a regular file: %s", cleanPath)), nil + } + + // Check file size (limit to 1MB for safety) + const maxFileSize = 1024 * 1024 // 1MB + if fileInfo.Size() > maxFileSize { + return mcp.NewToolResultError(fmt.Sprintf("File too large (max %d bytes): %d bytes", maxFileSize, fileInfo.Size())), nil + } + + // Read the file content + content, err := os.ReadFile(cleanPath) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("Failed to read file: %v", err)), nil + } + + // Trim whitespace from the content + secretValue := strings.TrimSpace(string(content)) + if secretValue == "" { + return mcp.NewToolResultError("File content is empty or contains only whitespace"), nil + } + + // Get the configuration to determine the secrets provider + cfg := h.configProvider.GetConfig() + + // Check if secrets setup has been completed + if !cfg.Secrets.SetupCompleted { + return mcp.NewToolResultError( + "Secrets provider not configured. Please run 'thv secret setup' to configure a secrets provider first"), nil + } + + // Get the provider type + providerType, err := cfg.Secrets.GetProviderType() + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("Failed to get secrets provider type: %v", err)), nil + } + + // Create the secrets provider + secretsProvider, err := secrets.CreateSecretProvider(providerType) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("Failed to create secrets provider: %v", err)), nil + } + + // Check if the provider supports writing + capabilities := secretsProvider.Capabilities() + if !capabilities.CanWrite { + return mcp.NewToolResultError(fmt.Sprintf( + "Secrets provider '%s' is read-only and does not support setting secrets", providerType)), nil + } + + // Set the secret + if err := secretsProvider.SetSecret(ctx, args.Name, secretValue); err != nil { + return mcp.NewToolResultError(fmt.Sprintf("Failed to set secret: %v", err)), nil + } + + // Create success response + response := SetSecretResponse{ + Status: "success", + Name: args.Name, + } + + return mcp.NewToolResultStructuredOnly(response), nil +} diff --git a/pkg/mcp/server/set_secret_test.go b/pkg/mcp/server/set_secret_test.go new file mode 100644 index 000000000..8f6bfcc96 --- /dev/null +++ b/pkg/mcp/server/set_secret_test.go @@ -0,0 +1,186 @@ +package server + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/config" + configmocks "github.com/stacklok/toolhive/pkg/config/mocks" + registrymocks "github.com/stacklok/toolhive/pkg/registry/mocks" + workloadsmocks "github.com/stacklok/toolhive/pkg/workloads/mocks" +) + +func TestHandler_SetSecret(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + t.Cleanup(func() { ctrl.Finish() }) + + // Create a temporary directory for test files + tempDir := t.TempDir() + + // Create test files + validFile := filepath.Join(tempDir, "valid_secret.txt") + err := os.WriteFile(validFile, []byte("my-secret-value\n"), 0600) + require.NoError(t, err) + + emptyFile := filepath.Join(tempDir, "empty_secret.txt") + err = os.WriteFile(emptyFile, []byte(" \n \n"), 0600) + require.NoError(t, err) + + largeFile := filepath.Join(tempDir, "large_secret.txt") + largeContent := make([]byte, 2*1024*1024) // 2MB + for i := range largeContent { + largeContent[i] = 'a' + } + err = os.WriteFile(largeFile, largeContent, 0600) + require.NoError(t, err) + + nonExistentFile := filepath.Join(tempDir, "nonexistent.txt") + + tests := []struct { + name string + args map[string]interface{} + setupMocks func(*configmocks.MockProvider) + wantErr bool + checkResult func(*testing.T, *mcp.CallToolResult) + }{ + { + name: "missing secret name", + args: map[string]interface{}{ + "file_path": validFile, + }, + setupMocks: func(_ *configmocks.MockProvider) {}, + wantErr: false, + checkResult: func(t *testing.T, result *mcp.CallToolResult) { + t.Helper() + assert.NotNil(t, result) + assert.True(t, result.IsError) + }, + }, + { + name: "missing file path", + args: map[string]interface{}{ + "name": "test-secret", + }, + setupMocks: func(_ *configmocks.MockProvider) {}, + wantErr: false, + checkResult: func(t *testing.T, result *mcp.CallToolResult) { + t.Helper() + assert.NotNil(t, result) + assert.True(t, result.IsError) + }, + }, + { + name: "file does not exist", + args: map[string]interface{}{ + "name": "test-secret", + "file_path": nonExistentFile, + }, + setupMocks: func(_ *configmocks.MockProvider) {}, + wantErr: false, + checkResult: func(t *testing.T, result *mcp.CallToolResult) { + t.Helper() + assert.NotNil(t, result) + assert.True(t, result.IsError) + }, + }, + { + name: "empty file content", + args: map[string]interface{}{ + "name": "test-secret", + "file_path": emptyFile, + }, + setupMocks: func(_ *configmocks.MockProvider) {}, + wantErr: false, + checkResult: func(t *testing.T, result *mcp.CallToolResult) { + t.Helper() + assert.NotNil(t, result) + assert.True(t, result.IsError) + }, + }, + { + name: "file too large", + args: map[string]interface{}{ + "name": "test-secret", + "file_path": largeFile, + }, + setupMocks: func(_ *configmocks.MockProvider) {}, + wantErr: false, + checkResult: func(t *testing.T, result *mcp.CallToolResult) { + t.Helper() + assert.NotNil(t, result) + assert.True(t, result.IsError) + }, + }, + { + name: "secrets not setup", + args: map[string]interface{}{ + "name": "test-secret", + "file_path": validFile, + }, + setupMocks: func(configProvider *configmocks.MockProvider) { + // Mock config setup - not completed + cfg := &config.Config{ + Secrets: config.Secrets{ + SetupCompleted: false, + }, + } + configProvider.EXPECT().GetConfig().Return(cfg).AnyTimes() + }, + wantErr: false, + checkResult: func(t *testing.T, result *mcp.CallToolResult) { + t.Helper() + assert.NotNil(t, result) + assert.True(t, result.IsError) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // Create mocks + mockRegistry := registrymocks.NewMockProvider(ctrl) + mockWorkloadManager := workloadsmocks.NewMockManager(ctrl) + mockConfigProvider := configmocks.NewMockProvider(ctrl) + + // Setup mocks + if tt.setupMocks != nil { + tt.setupMocks(mockConfigProvider) + } + + handler := &Handler{ + ctx: context.Background(), + workloadManager: mockWorkloadManager, + registryProvider: mockRegistry, + configProvider: mockConfigProvider, + } + + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "set_secret", + Arguments: tt.args, + }, + } + + result, err := handler.SetSecret(context.Background(), request) + + if tt.wantErr { + assert.Error(t, err) + } else { + assert.NoError(t, err) + if tt.checkResult != nil { + tt.checkResult(t, result) + } + } + }) + } +}