16. Tool Security

Function calling expands the attack surface of LLM applications. A compromised or poorly designed tool can expose filesystem access, network capabilities, or sensitive data to malicious queries. Security must be built into the tool architecture from the start.

Input Validation

Every tool must validate inputs before execution:

import re
from typing import Any
from pydantic import BaseModel, validator, constr

class PathInput(BaseModel):
    path: constr(max_length=255, pattern=r'^[a-zA-Z0-9/_.\-]+$')  # Alphanumeric and safe chars only
    
    @validator('path')
    def prevent_path_traversal(cls, v):
        # Block common traversal attempts
        dangerous_patterns = ['../', '..\\', '%2e%2e', '~']
        for pattern in dangerous_patterns:
            if pattern.lower() in v.lower():
                raise ValueError(f"Path traversal attempt detected: {pattern}")
        return v

class ReadFileTool:
    def _validate_path(self, path: str) -> PathInput:
        try:
            return PathInput(path=path)
        except Exception as e:
            raise ValueError(f"Invalid path: {e}")
    
    def _run(self, path: str, max_size: int = 1024 * 1024) -> dict:
        validated = self._validate_path(path)
        
        full_path = f"/allowed/base/dir/{validated.path}"
        
        # Additional size check
        if os.path.getsize(full_path) > max_size:
            raise ValueError(f"File exceeds maximum size: {max_size}")
        
        with open(full_path, 'r') as f:
            return {"content": f.read(), "path": validated.path}

Sandboxed Execution

Tools should execute with minimal privileges:

import os
import resource

class SandboxedExecutor:
    """Execute tools with resource limits and restricted permissions."""
    
    def __init__(self, base_dir: str, max_memory_mb: int = 512):
        self.base_dir = os.path.abspath(base_dir)
        self.max_memory_bytes = max_memory_mb * 1024 * 1024
    
    def execute(self, func: callable, *args, **kwargs):
        # Set resource limits before execution
        resource.setrlimit(resource.RLIMIT_AS, (self.max_memory_bytes, self.max_memory_bytes))
        resource.setrlimit(resource.RLIMIT_CPU, (30, 30))  # 30 seconds CPU time
        
        # Chroot-like restriction using os.chdir
        original_cwd = os.getcwd()
        try:
            os.chdir(self.base_dir)
            
            # Run with restricted environment
            result = func(*args, **kwargs)
            
            # Verify we stayed within allowed directory
            if not self._is_safe_path(os.getcwd()):
                raise SecurityError("Process escaped base directory")
            
            return result
        finally:
            os.chdir(original_cwd)
    
    def _is_safe_path(self, path: str) -> bool:
        real_path = os.path.realpath(path)
        return real_path.startswith(self.base_dir)

Tool Allowlist

Only explicitly allowed tools should be available:

from typing import Callable
from dataclasses import dataclass

@dataclass
class ToolPermission:
    name: str
    allowed: bool
    max_calls_per_minute: int = 60
    requires_audit: bool = False

class SecureToolRegistry:
    def __init__(self):
        self._permissions: dict[str, ToolPermission] = {}
        self._implementations: dict[str, Callable] = {}
    
    def register(self, name: str, implementation: Callable, permission: ToolPermission):
        self._permissions[name] = permission
        self._implementations[name] = implementation
    
    def execute(self, tool_name: str, params: dict) -> Any:
        if tool_name not in self._permissions:
            raise SecurityError(f"Tool '{tool_name}' not found in registry")
        
        permission = self._permissions[tool_name]
        
        if not permission.allowed:
            raise SecurityError(f"Tool '{tool_name}' is not permitted")
        
        if tool_name not in self._implementations:
            raise SecurityError(f"Tool '{tool_name}' has no implementation")
        
        # Audit logging for restricted tools
        if permission.requires_audit:
            self._audit_call(tool_name, params)
        
        return self._implementations[tool_name](**params)
    
    def _audit_call(self, tool_name: str, params: dict):
        import structlog
        logger = structlog.get_logger()
        logger.warning(
            "restricted_tool_call",
            tool=tool_name,
            params=self._redact_sensitive(params)
        )
    
    def _redact_sensitive(self, params: dict) -> dict:
        sensitive_keys = ['password', 'token', 'api_key', 'secret']
        redacted = params.copy()
        for key in sensitive_keys:
            if key in redacted:
                redacted[key] = '***REDACTED***'
        return redacted

SQL Injection Prevention

Database tools require special attention:

import sqlite3

class SafeDatabaseTool:
    def __init__(self, db_path: str):
        self.db_path = db_path
        self._validate_connection()
    
    def _validate_connection(self):
        # Ensure database is in allowed location
        real_path = os.path.realpath(self.db_path)
        if not real_path.startswith("/allowed/databases/"):
            raise SecurityError("Database not in allowed directory")
    
    def execute_query(self, query: str, params: tuple):
        # Whitelist validation
        allowed_keywords = ['SELECT', 'FROM', 'WHERE', 'AND', 'OR', 'LIMIT', 'ORDER BY']
        
        upper_query = query.upper()
        for keyword in allowed_keywords:
            if keyword in upper_query:
                # Verify it's used correctly
                pattern = rf'\b{keyword}\b'
                if not re.search(pattern, upper_query):
                    raise SecurityError(f"Invalid {keyword} usage")
        
        # Parameterized query prevents injection
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        try:
            cursor.execute(query, params)
            return cursor.fetchall()
        finally:
            conn.close()