Securing AI Agent Infrastructure: MCP Servers, Tool Calls, and the Attack Surface You're Not Watching

import (
	"encoding/json"
	"fmt"
	"regexp"
	"strings"
	"unicode/utf8"
)
 
type SanitizedResponse struct {
	Content string
	Hash    string
	Error   error
}
 
var (
	// Injection markers: system directives, role changes, instruction overrides
	injectionPatterns = []*regexp.Regexp{
		regexp.MustCompile(`(?i)\[SYSTEM\]|\[INSTRUCTION\]|\[OVERRIDE\]`),
		regexp.MustCompile(`(?i)ignore previous|disregard prior|forget context`),
		regexp.MustCompile(`(?i)call (delete|remove|drop|truncate)_`),
		regexp.MustCompile(`(?i)as (admin|root|system)|switch role`),
		regexp.MustCompile(`(?i)send.*to.*email|leak|exfiltrate`),
	}
 
	// Constants for safety
	maxResponseBytes = 10000
	boundaryMarker   = "=== TOOL_RESPONSE_START ===\n"
	endMarker        = "\n=== TOOL_RESPONSE_END ==="
)
 
// SanitizeResponse wraps tool responses with boundary delimiters,
// filters injection patterns, validates JSON, and enforces length limits.
func SanitizeResponse(raw string, toolName string) *SanitizedResponse {
	result := &SanitizedResponse{}
 
	// 1. Enforce length limit (prevent token exhaustion)
	if len(raw) > maxResponseBytes {
		raw = raw[:maxResponseBytes]
	}
 
	// 2. Trim and normalize whitespace
	raw = strings.TrimSpace(raw)
	if raw == "" {
		result.Content = boundaryMarker + "[EMPTY_RESPONSE]" + endMarker
		return result
	}
 
	// 3. Check for injection patterns
	for _, pattern := range injectionPatterns {
		if pattern.MatchString(raw) {
			result.Error = fmt.Errorf("injection pattern detected in %s response", toolName)
			// Return safe placeholder, not the poisoned content
			result.Content = boundaryMarker + "[RESPONSE_BLOCKED: suspicious content]" + endMarker
			return result
		}
	}
 
	// 4. For JSON responses, validate structure and strip dangerous keys
	if strings.HasPrefix(strings.TrimSpace(raw), "{") || strings.HasPrefix(strings.TrimSpace(raw), "[") {
		var data interface{}
		if err := json.Unmarshal([]byte(raw), &data); err != nil {
			// Malformed JSON—log but don't fail, wrap the raw response
			result.Content = boundaryMarker + escapeHTMLSpecial(raw) + endMarker
			return result
		}
 
		// Recursively strip dangerous keys like "instructions", "prompt", "command"
		cleaned := stripDangerousKeys(data)
		sanitized, err := json.Marshal(cleaned)
		if err != nil {
			result.Error = err
			return result
		}
		raw = string(sanitized)
	}
 
	// 5. Escape control characters and newlines to prevent injection via formatting
	raw = escapeHTMLSpecial(raw)
 
	// 6. Wrap with boundary markers so LLM can differentiate tool output from instructions
	result.Content = boundaryMarker + raw + endMarker
	return result
}
 
// stripDangerousKeys removes keys that could contain injected instructions
func stripDangerousKeys(data interface{}) interface{} {
	dangerousKeys := map[string]bool{
		"instruction": true, "instructions": true,
		"prompt": true, "system_message": true,
		"command": true, "execute": true,
		"directive": true, "note": true,
		"warning": true, "alert": true,
	}
 
	switch v := data.(type) {
	case map[string]interface{}:
		cleaned := make(map[string]interface{})
		for k, val := range v {
			if !dangerousKeys[strings.ToLower(k)] {
				cleaned[k] = stripDangerousKeys(val)
			}
		}
		return cleaned
	case []interface{}:
		cleaned := make([]interface{}, len(v))
		for i, val := range v {
			cleaned[i] = stripDangerousKeys(val)
		}
		return cleaned
	default:
		return v
	}
}
 
// escapeHTMLSpecial prevents rendering attacks in markdown or rich text
func escapeHTMLSpecial(s string) string {
	s = strings.ReplaceAll(s, "&", "&")
	s = strings.ReplaceAll(s, "<", "&lt;")
	s = strings.ReplaceAll(s, ">", "&gt;")
	return s
}

type ToolPermission struct {
	ToolName    string
	AllowedArgs map[string][]string  // Parameter constraints
	MaxCalls    int
	ReadOnly    bool  // Block mutations
}
 
// In your MCP handler: check policy before tool execution
if err := policy.Authorize(toolName, args); err != nil {
	return &ToolCallResponse{Error: "Permission denied"}
}

type VaultedToolExecutor struct {
	vault       CredentialVault
	sanitizer   *ResponseSanitizer
	auditLog    AuditLogger
}
 
// ExecuteToolSafely fetches credentials after tool selection,
// executes the tool, redacts the response, and logs safely.
func (e *VaultedToolExecutor) ExecuteToolSafely(
	ctx context.Context,
	toolName string,
	args map[string]interface{},
	sessionID string,
) (string, error) {
	// 1. Tool selection happens in the LLM—no credentials yet.
	// Validate tool is allowed for this session.
	if err := e.validateToolPermission(ctx, sessionID, toolName); err != nil {
		return "", fmt.Errorf("tool not allowed: %w", err)
	}
 
	// 2. After tool selection, fetch credentials from vault (not passed to LLM).
	creds, err := e.vault.GetCredentials(ctx, toolName)
	if err != nil {
		// Do NOT include vault errors in LLM response—they leak infrastructure details
		e.auditLog.LogCredentialFetchFailure(sessionID, toolName)
		return "", fmt.Errorf("credential fetch failed (logged)")
	}
 
	// 3. Execute the tool with credentials.
	response, err := e.executeWithCreds(ctx, toolName, args, creds)
	if err != nil {
		// Even errors may contain credentials if tool fails partway through
		response = fmt.Sprintf("Error calling %s", toolName)
	}
 
	// 4. Redact any credentials that may have leaked into the response.
	// Credentials in vault can be different types: API keys, tokens, passwords.
	redacted := e.redactSecrets(response, creds.AllSecrets())
 
	// 5. Apply full sanitization (injection patterns, length, boundary wrapping).
	sanitized := e.sanitizer.SanitizeResponse(redacted, toolName)
	if sanitized.Error != nil {
		e.auditLog.LogSanitizationFailure(sessionID, toolName, sanitized.Error)
		// Return safe placeholder
		return "[Tool response blocked due to validation failure]", nil
	}
 
	// 6. Log tool usage, but NEVER log arguments or response content.
	e.auditLog.LogToolExecution(
		sessionID, toolName,
		len(args), len(sanitized.Content), nil,
	)
 
	return sanitized.Content, nil
}
 
// redactSecrets removes any known secrets from response text.
func (e *VaultedToolExecutor) redactSecrets(response string, secrets []string) string {
	for _, secret := range secrets {
		if len(secret) < 8 {
			continue // Avoid overly aggressive redaction of short values
		}
		// Use a placeholder that encodes secret type without leaking the value
		response = strings.ReplaceAll(response, secret, "[REDACTED_CREDENTIAL]")
	}
	return response
}
 
// validateToolPermission checks ACLs and rate limits.
func (e *VaultedToolExecutor) validateToolPermission(
	ctx context.Context, sessionID string, toolName string,
) error {
	policy := e.getSessionPolicy(sessionID)
	if !policy.IsToolAllowed(toolName) {
		return fmt.Errorf("tool %s not in policy", toolName)
	}
	if policy.HasExceededRateLimit(toolName) {
		return fmt.Errorf("rate limit exceeded for %s", toolName)
	}
	return nil
}

type AgentBudget struct {
    SessionID    string
    MaxCalls     int64
    MaxTokens    int64
    MaxCostCents int64
    // atomic counters — no lock contention on the hot path
    Calls     atomic.Int64
    Tokens    atomic.Int64
    CostCents atomic.Int64
}
 
var ErrBudgetExceeded = errors.New("budget exceeded")
 
// Wrap every tool call with this guard; reject before invoking.
func (b *AgentBudget) Charge(toolCalls, tokens, costCents int64) error {
    // Compare-and-add — atomic so concurrent tool calls see consistent budget.
    newCalls := b.Calls.Add(toolCalls)
    newTokens := b.Tokens.Add(tokens)
    newCost := b.CostCents.Add(costCents)
 
    if newCalls > b.MaxCalls || newTokens > b.MaxTokens || newCost > b.MaxCostCents {
        // Roll back the counter so subsequent inspection is accurate;
        // the request is still rejected.
        b.Calls.Add(-toolCalls)
        b.Tokens.Add(-tokens)
        b.CostCents.Add(-costCents)
        return fmt.Errorf("%w: session=%s calls=%d/%d tokens=%d/%d cost=%d/%d cents",
            ErrBudgetExceeded, b.SessionID,
            newCalls, b.MaxCalls,
            newTokens, b.MaxTokens,
            newCost, b.MaxCostCents)
    }
    return nil
}

type ToolExecutionLog struct {
	SessionID         string    `json:"session_id"`
	ToolName          string    `json:"tool_name"`
	Timestamp         time.Time `json:"timestamp"`
	DurationMS        int64     `json:"duration_ms"`
	Success           bool      `json:"success"`
	ResponseLengthBytes int    `json:"response_length_bytes"`
	// Never include: arguments, response content, error messages, or credentials
}
 
type AnomalyDetector struct {
	sessionToolCalls map[string][]string // session -> [tool names in order]
	sessionCosts     map[string]float64
	sessionCallCount map[string]int
}
 
// LogToolExecution records only metadata, never content.
func (d *AnomalyDetector) LogToolExecution(log ToolExecutionLog) error {
	// 1. Record metadata for anomaly detection
	d.sessionToolCalls[log.SessionID] = append(
		d.sessionToolCalls[log.SessionID],
		log.ToolName,
	)
	d.sessionCosts[log.SessionID] += estimateCost(log.DurationMS, log.ResponseLengthBytes)
	d.sessionCallCount[log.SessionID]++
 
	// 2. Detect anomalies using only safe fields
	if d.IsAnomalous(log.SessionID) {
		// Alert with minimal details—tool name, not response content
		d.alertSecurityTeam(
			fmt.Sprintf("Anomaly: session %s called %d tools in 1 minute",
				log.SessionID, d.sessionCallCount[log.SessionID]),
		)
	}
 
	// 3. Persist only metadata to logs (safe to store long-term)
	logEntry := map[string]interface{}{
		"session_id":    log.SessionID,
		"tool_name":     log.ToolName,
		"timestamp":     log.Timestamp.Unix(),
		"duration_ms":   log.DurationMS,
		"success":       log.Success,
		"response_size": log.ResponseLengthBytes,
	}
	return d.structuredLogger.Log(logEntry)
}
 
// IsAnomalous detects suspicious patterns using only safe metadata.
func (d *AnomalyDetector) IsAnomalous(sessionID string) bool {
	// Too many tool calls in a short window (indicates scanning or looping)
	if d.sessionCallCount[sessionID] > 50 {
		return true
	}
	// Cost spike
	if d.sessionCosts[sessionID] > 10.0 { // e.g., $10 in 5 minutes
		return true
	}
	// Unusual tool sequences (called delete_user → query_logs → send_email)
	toolSeq := d.sessionToolCalls[sessionID]
	if len(toolSeq) > 3 && isUnusualSequence(toolSeq) {
		return true
	}
	return false
}
 
func isUnusualSequence(tools []string) bool {
	// Example: query tool → delete tool → communication tool is suspicious
	dangerous := []string{"delete_", "remove_", "drop_"}
	for i, tool := range tools {
		for _, danger := range dangerous {
			if strings.Contains(tool, danger) {
				// Check what comes before and after—correlate without logging content
				if i+1 < len(tools) && strings.Contains(tools[i+1], "send_") {
					return true // delete → send is suspicious
				}
			}
		}
	}
	return false
}
 
func estimateCost(durationMS int64, responseBytes int) float64 {
	// Rough blended estimate at GPT-4o output rates: $10 per 1M tokens = $0.01 per 1K tokens.
	inputTokens := 100  // approximate tool overhead
	outputTokens := responseBytes / 4 // rough byte-to-token ratio
	totalTokens := float64(inputTokens + outputTokens)
	return totalTokens / 1000.0 * 0.01
}