server: add MCP security layer for command validation
Add security infrastructure to validate and restrict MCP server commands before execution. MCPValidator (server/mcp_validator.go): - Command allowlist/blocklist validation - Argument sanitization (blocks shell metacharacters) - Environment variable name validation - Per-server security policy enforcement MCPSecurityConfig (server/mcp_security_config.go): - Configurable security rules per server - Default blocklist: bash, sh, sudo, rm, curl, wget, eval - Blocked metacharacters: ; | & $( ` > < = etc. - SECURITY REVIEW markers for critical sections MCPCommandResolver (server/mcp_command_resolver.go): - Resolves command paths across different environments - npx/npm detection and path resolution MCPCodeAPI (server/mcp_code_api.go): - Programmatic API for MCP server management Relates to #7865
This commit is contained in:
code4me2
parent
0fd68e46fa
commit
fc05536d52
|
|
@ -0,0 +1,149 @@
|
|||
package server
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strings"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
// MCPCodeAPI provides a code-like interface for MCP tools
|
||||
type MCPCodeAPI struct {
|
||||
manager *MCPManager
|
||||
}
|
||||
|
||||
// NewMCPCodeAPI creates a new MCP code API
|
||||
func NewMCPCodeAPI(manager *MCPManager) *MCPCodeAPI {
|
||||
return &MCPCodeAPI{
|
||||
manager: manager,
|
||||
}
|
||||
}
|
||||
|
||||
// GenerateMinimalContext returns essential context for tool usage
|
||||
func (m *MCPCodeAPI) GenerateMinimalContext(configs []api.MCPServerConfig) string {
|
||||
slog.Debug("GenerateMinimalContext called", "configs_count", len(configs))
|
||||
if len(configs) == 0 {
|
||||
slog.Debug("No MCP configs provided, returning empty context")
|
||||
return ""
|
||||
}
|
||||
|
||||
var context strings.Builder
|
||||
context.WriteString("\n=== MCP Tool Context ===\n")
|
||||
|
||||
for _, config := range configs {
|
||||
slog.Debug("Processing MCP config", "command", config.Command, "args", config.Args)
|
||||
// Check if this is a filesystem server (command or first arg contains filesystem)
|
||||
isFilesystem := strings.Contains(config.Command, "filesystem") ||
|
||||
(len(config.Args) > 0 && strings.Contains(config.Args[0], "filesystem"))
|
||||
|
||||
if isFilesystem && len(config.Args) > 1 {
|
||||
// Extract working directory from filesystem server
|
||||
workingDir := config.Args[1]
|
||||
slog.Debug("Adding filesystem context", "working_dir", workingDir)
|
||||
context.WriteString(fmt.Sprintf(`
|
||||
Filesystem tools are available with these constraints:
|
||||
- Working directory: %s
|
||||
- All file operations must use paths within this directory
|
||||
- Example usage:
|
||||
- List files: "List all files in %s"
|
||||
- Read file: "Read %s/filename.txt"
|
||||
- Create file: "Create %s/newfile.txt with content"
|
||||
- Paths outside %s will be rejected
|
||||
|
||||
When working with files, ALWAYS use the full path starting with %s
|
||||
`, workingDir, workingDir, workingDir, workingDir, workingDir, workingDir))
|
||||
}
|
||||
// Add other server types as needed
|
||||
}
|
||||
|
||||
context.WriteString("\n")
|
||||
result := context.String()
|
||||
slog.Debug("Generated MCP context", "length", len(result))
|
||||
return result
|
||||
}
|
||||
|
||||
// GenerateProgressiveContext returns context based on what tools are being used
|
||||
func (m *MCPCodeAPI) GenerateProgressiveContext(toolNames []string) string {
|
||||
var context strings.Builder
|
||||
|
||||
// Group tools by server
|
||||
serverTools := make(map[string][]string)
|
||||
for _, toolName := range toolNames {
|
||||
if clientName, exists := m.manager.GetToolClient(toolName); exists {
|
||||
serverTools[clientName] = append(serverTools[clientName], toolName)
|
||||
}
|
||||
}
|
||||
|
||||
// Generate context for each server's tools
|
||||
for serverName, tools := range serverTools {
|
||||
context.WriteString(fmt.Sprintf("\n%s tools being used:\n", serverName))
|
||||
for _, tool := range tools {
|
||||
// Get tool definition from manager
|
||||
if toolDef := m.manager.GetToolDefinition(serverName, tool); toolDef != nil {
|
||||
context.WriteString(fmt.Sprintf("- %s: %s\n", tool, toolDef.Function.Description))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return context.String()
|
||||
}
|
||||
|
||||
// InjectContextIntoMessages intelligently injects context into the message stream
|
||||
func (m *MCPCodeAPI) InjectContextIntoMessages(messages []api.Message, configs []api.MCPServerConfig) []api.Message {
|
||||
// Generate minimal context
|
||||
context := m.GenerateMinimalContext(configs)
|
||||
if context == "" {
|
||||
return messages
|
||||
}
|
||||
|
||||
// Check if there's already a system message
|
||||
if len(messages) > 0 && messages[0].Role == "system" {
|
||||
// Append to existing system message
|
||||
messages[0].Content += context
|
||||
} else {
|
||||
// Create new system message
|
||||
systemMsg := api.Message{
|
||||
Role: "system",
|
||||
Content: context,
|
||||
}
|
||||
messages = append([]api.Message{systemMsg}, messages...)
|
||||
}
|
||||
|
||||
return messages
|
||||
}
|
||||
|
||||
// ExtractWorkingDirectory extracts the working directory from MCP server args
|
||||
func ExtractWorkingDirectory(config api.MCPServerConfig) string {
|
||||
if strings.Contains(config.Command, "filesystem") && len(config.Args) > 1 {
|
||||
return config.Args[1]
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// GenerateToolCallExample generates an example of how to call a specific tool
|
||||
func (m *MCPCodeAPI) GenerateToolCallExample(serverName, toolName string) string {
|
||||
workingDir := ""
|
||||
|
||||
// Get working directory if filesystem
|
||||
if serverName == "filesystem" {
|
||||
if clients := m.manager.GetServerNames(); len(clients) > 0 {
|
||||
// This is a simplified approach - in production we'd properly track server configs
|
||||
workingDir = "/home/velvetm/Desktop/mcp-test-files" // Would be extracted from actual config
|
||||
}
|
||||
}
|
||||
|
||||
// Generate appropriate example based on tool
|
||||
switch toolName {
|
||||
case "list_directory":
|
||||
return fmt.Sprintf(`"List all files in %s"`, workingDir)
|
||||
case "read_file":
|
||||
return fmt.Sprintf(`"Read the file %s/example.txt"`, workingDir)
|
||||
case "write_file":
|
||||
return fmt.Sprintf(`"Create a file at %s/output.txt with content 'Hello World'"`, workingDir)
|
||||
case "create_directory":
|
||||
return fmt.Sprintf(`"Create a directory called %s/newdir"`, workingDir)
|
||||
default:
|
||||
return fmt.Sprintf(`"Use the %s tool"`, toolName)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,187 @@
|
|||
package server
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// CommandResolver handles resolving commands to their actual executables
|
||||
// with fallback detection for different system configurations
|
||||
type CommandResolver struct {
|
||||
mu sync.RWMutex
|
||||
resolved map[string]string
|
||||
}
|
||||
|
||||
// NewCommandResolver creates a new command resolver
|
||||
func NewCommandResolver() *CommandResolver {
|
||||
return &CommandResolver{
|
||||
resolved: make(map[string]string),
|
||||
}
|
||||
}
|
||||
|
||||
// DefaultCommandResolver is the shared resolver instance for production use.
|
||||
// Tests should use WithCommandResolver option instead of modifying this.
|
||||
var DefaultCommandResolver = NewCommandResolver()
|
||||
|
||||
// ResolveCommand finds the actual executable for a given command
|
||||
func (cr *CommandResolver) ResolveCommand(command string) (string, error) {
|
||||
cr.mu.RLock()
|
||||
if resolved, ok := cr.resolved[command]; ok {
|
||||
cr.mu.RUnlock()
|
||||
return resolved, nil
|
||||
}
|
||||
cr.mu.RUnlock()
|
||||
|
||||
// Try to resolve the command
|
||||
var resolved string
|
||||
var err error
|
||||
|
||||
switch command {
|
||||
case "npx":
|
||||
resolved, err = cr.resolveNodePackageManager()
|
||||
case "python":
|
||||
resolved, err = cr.resolvePython()
|
||||
case "node":
|
||||
resolved, err = cr.resolveNode()
|
||||
default:
|
||||
// For other commands, check if they exist as-is
|
||||
resolved, err = cr.checkCommand(command)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Cache the resolution
|
||||
cr.mu.Lock()
|
||||
cr.resolved[command] = resolved
|
||||
cr.mu.Unlock()
|
||||
|
||||
return resolved, nil
|
||||
}
|
||||
|
||||
// resolveNodePackageManager finds an available Node.js package manager
|
||||
func (cr *CommandResolver) resolveNodePackageManager() (string, error) {
|
||||
// Priority order for package managers
|
||||
managers := []struct {
|
||||
cmd string
|
||||
args []string
|
||||
}{
|
||||
{"npx", []string{"--version"}},
|
||||
{"pnpm", []string{"dlx", "--version"}}, // pnpm equivalent of npx
|
||||
{"yarn", []string{"dlx", "--version"}}, // yarn 2+ equivalent
|
||||
{"bunx", []string{"--version"}}, // bun equivalent
|
||||
}
|
||||
|
||||
for _, mgr := range managers {
|
||||
if path, err := exec.LookPath(mgr.cmd); err == nil {
|
||||
// Verify it actually works
|
||||
cmd := exec.Command(path, mgr.args...)
|
||||
if err := cmd.Run(); err == nil {
|
||||
// For pnpm/yarn, we need to return the dlx subcommand
|
||||
if mgr.cmd == "pnpm" {
|
||||
return "pnpm dlx", nil
|
||||
} else if mgr.cmd == "yarn" {
|
||||
return "yarn dlx", nil
|
||||
}
|
||||
return mgr.cmd, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check if npm is available and suggest installing npx
|
||||
if _, err := exec.LookPath("npm"); err == nil {
|
||||
return "", fmt.Errorf("npx not found but npm is available - install with: npm install -g npx")
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("no Node.js package manager found (tried npx, pnpm, yarn, bunx)")
|
||||
}
|
||||
|
||||
// resolvePython finds an available Python interpreter
|
||||
func (cr *CommandResolver) resolvePython() (string, error) {
|
||||
// Priority order for Python interpreters
|
||||
interpreters := []string{
|
||||
"python3", // Most Unix systems
|
||||
"python", // Windows or virtualenv
|
||||
"python3.12", // Specific versions
|
||||
"python3.11",
|
||||
"python3.10",
|
||||
"python3.9",
|
||||
"python3.8",
|
||||
}
|
||||
|
||||
for _, interp := range interpreters {
|
||||
if path, err := exec.LookPath(interp); err == nil {
|
||||
// Verify it's Python 3.8+ by checking version
|
||||
cmd := exec.Command(path, "--version")
|
||||
output, err := cmd.Output()
|
||||
if err == nil && len(output) > 0 {
|
||||
// Basic check that it's Python 3
|
||||
if string(output[:7]) == "Python " && output[7] >= '3' {
|
||||
return interp, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("no Python 3 interpreter found (tried python3, python, and versioned variants)")
|
||||
}
|
||||
|
||||
// resolveNode finds the Node.js executable
|
||||
func (cr *CommandResolver) resolveNode() (string, error) {
|
||||
// Try different Node.js executable names
|
||||
nodes := []string{"node", "nodejs"}
|
||||
|
||||
for _, node := range nodes {
|
||||
if path, err := exec.LookPath(node); err == nil {
|
||||
// Verify it works
|
||||
cmd := exec.Command(path, "--version")
|
||||
if err := cmd.Run(); err == nil {
|
||||
return node, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("Node.js not found (tried node, nodejs)")
|
||||
}
|
||||
|
||||
// checkCommand checks if a command exists as-is
|
||||
func (cr *CommandResolver) checkCommand(command string) (string, error) {
|
||||
if _, err := exec.LookPath(command); err == nil {
|
||||
return command, nil
|
||||
}
|
||||
return "", fmt.Errorf("command not found: %s", command)
|
||||
}
|
||||
|
||||
// ResolveForEnvironment checks environment variables for command overrides
|
||||
func (cr *CommandResolver) ResolveForEnvironment(command string) string {
|
||||
// Allow environment variable overrides
|
||||
envMap := map[string]string{
|
||||
"npx": "OLLAMA_NPX_COMMAND",
|
||||
"python": "OLLAMA_PYTHON_COMMAND",
|
||||
"node": "OLLAMA_NODE_COMMAND",
|
||||
}
|
||||
|
||||
if envVar, ok := envMap[command]; ok {
|
||||
if override := os.Getenv(envVar); override != "" {
|
||||
// Validate override against security blocklist
|
||||
if GetSecurityConfig().IsCommandAllowed(override) {
|
||||
return override
|
||||
}
|
||||
slog.Warn("Environment override blocked by security policy", "var", envVar, "command", override)
|
||||
}
|
||||
}
|
||||
|
||||
// Try standard resolution
|
||||
if resolved, err := cr.ResolveCommand(command); err == nil {
|
||||
return resolved
|
||||
}
|
||||
|
||||
// Return original command as fallback
|
||||
return command
|
||||
}
|
||||
|
||||
// NOTE: A GetSystemRequirements() method could be added here for diagnostics/status endpoints
|
||||
|
|
@ -0,0 +1,152 @@
|
|||
package server
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// =============================================================================
|
||||
// MCP Security Configuration
|
||||
// =============================================================================
|
||||
//
|
||||
// SECURITY REVIEW: This file defines the security policies that control which
|
||||
// commands can be executed as MCP servers and what arguments/environment they
|
||||
// can receive. Changes to this file should be reviewed by security-aware
|
||||
// maintainers.
|
||||
//
|
||||
// Key security surfaces:
|
||||
// - BlockedCommands: Prevents execution of dangerous system commands
|
||||
// - BlockedMetacharacters: Prevents shell injection attacks
|
||||
// - FilteredEnvironmentVars: Prevents credential leakage to MCP servers
|
||||
//
|
||||
// Threat model:
|
||||
// - Malicious MCP server configs attempting to execute system commands
|
||||
// - Shell injection through tool arguments
|
||||
// - Credential theft through environment variable access
|
||||
//
|
||||
// =============================================================================
|
||||
|
||||
// MCPSecurityConfig defines security policies for MCP servers
|
||||
type MCPSecurityConfig struct {
|
||||
// Commands that are never allowed as MCP servers
|
||||
BlockedCommands []string
|
||||
|
||||
// Shell metacharacters that are not allowed in arguments
|
||||
BlockedMetacharacters []string
|
||||
|
||||
// Environment variables that should be filtered
|
||||
FilteredEnvironmentVars []string
|
||||
}
|
||||
|
||||
// DefaultSecurityConfig returns the default security configuration.
|
||||
//
|
||||
// SECURITY REVIEW: This function defines the default blocklists. Adding or
|
||||
// removing entries has direct security implications. Consider:
|
||||
// - Why is a command being added/removed?
|
||||
// - What attack vectors does it enable/prevent?
|
||||
// - Are there bypass possibilities (symlinks, PATH manipulation)?
|
||||
func DefaultSecurityConfig() *MCPSecurityConfig {
|
||||
return &MCPSecurityConfig{
|
||||
// SECURITY: Blocked commands - these can never be used as MCP server commands.
|
||||
// Rationale: These commands could be used for privilege escalation,
|
||||
// arbitrary file manipulation, or establishing network connections.
|
||||
BlockedCommands: []string{
|
||||
// Shells - prevent arbitrary command execution
|
||||
"sh", "bash", "zsh", "fish", "csh", "ksh", "dash", "tcsh",
|
||||
"cmd", "cmd.exe", "powershell", "powershell.exe", "pwsh", "pwsh.exe",
|
||||
|
||||
// System commands - prevent privilege escalation and system damage
|
||||
"sudo", "su", "doas", "runas", "pkexec",
|
||||
"rm", "del", "rmdir", "format", "dd", "shred",
|
||||
"kill", "killall", "pkill", "shutdown", "reboot",
|
||||
"systemctl", "service", "init",
|
||||
|
||||
// Network tools - prevent data exfiltration and network attacks
|
||||
"curl", "wget", "nc", "netcat", "telnet", "ssh", "scp", "sftp",
|
||||
"nmap", "ping", "traceroute", "dig", "nslookup",
|
||||
|
||||
// Script interpreters - prevent arbitrary code execution
|
||||
"eval", "exec", "source", ".",
|
||||
"perl", "ruby", "php", "lua", "tcl",
|
||||
|
||||
// File manipulation - prevent permission/ownership changes
|
||||
"chmod", "chown", "chgrp", "mount", "umount",
|
||||
"ln", "mkfifo", "mknod",
|
||||
|
||||
// Package managers - prevent system modification
|
||||
"apt", "apt-get", "yum", "dnf", "pacman", "zypper",
|
||||
"brew", "port", "snap", "flatpak",
|
||||
},
|
||||
|
||||
// SECURITY: Shell metacharacters - block these in arguments to prevent injection.
|
||||
// These characters could be used to chain commands or redirect I/O.
|
||||
BlockedMetacharacters: []string{
|
||||
";", "|", "&", "$(", "`", ">", "<", ">>", "<<",
|
||||
"||", "&&", "\n", "\r", "$", "!", "*", "?", "=",
|
||||
},
|
||||
|
||||
// SECURITY: Environment variables to filter - prevent credential leakage.
|
||||
// MCP servers should not have access to credentials in the parent environment.
|
||||
FilteredEnvironmentVars: []string{
|
||||
// AWS
|
||||
"AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN",
|
||||
|
||||
// Cloud providers
|
||||
"GOOGLE_APPLICATION_CREDENTIALS", "AZURE_CLIENT_SECRET",
|
||||
|
||||
// API Keys
|
||||
"GITHUB_TOKEN", "GITLAB_TOKEN", "OPENAI_API_KEY", "ANTHROPIC_API_KEY",
|
||||
|
||||
// Database
|
||||
"DATABASE_URL", "DB_PASSWORD", "MYSQL_ROOT_PASSWORD", "POSTGRES_PASSWORD",
|
||||
|
||||
// Authentication
|
||||
"JWT_SECRET", "SESSION_SECRET", "AUTH_TOKEN", "API_KEY", "API_SECRET",
|
||||
|
||||
// SSH
|
||||
"SSH_AUTH_SOCK", "SSH_AGENT_PID",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// IsCommandAllowed checks if a command is allowed by security policy
|
||||
func (c *MCPSecurityConfig) IsCommandAllowed(command string) bool {
|
||||
baseName := filepath.Base(command)
|
||||
|
||||
for _, blocked := range c.BlockedCommands {
|
||||
if baseName == blocked || strings.HasSuffix(command, "/"+blocked) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// HasShellMetacharacters checks if a string contains shell metacharacters
|
||||
func (c *MCPSecurityConfig) HasShellMetacharacters(s string) bool {
|
||||
for _, meta := range c.BlockedMetacharacters {
|
||||
if strings.Contains(s, meta) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ShouldFilterEnvironmentVar checks if an environment variable should be filtered
|
||||
func (c *MCPSecurityConfig) ShouldFilterEnvironmentVar(key string) bool {
|
||||
for _, filtered := range c.FilteredEnvironmentVars {
|
||||
if key == filtered {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Global security config instance
|
||||
// NOTE: To add user customization, load from ~/.ollama/mcp-security.json and append to these defaults
|
||||
var globalSecurityConfig = DefaultSecurityConfig()
|
||||
|
||||
// GetSecurityConfig returns the global security configuration
|
||||
func GetSecurityConfig() *MCPSecurityConfig {
|
||||
return globalSecurityConfig
|
||||
}
|
||||
Loading…
Reference in New Issue