Extract anvilmcp to standalone repo
anvilmcp now at github.com/lneely/anvillm-mcp Sandbox package kept (used by tmux backend).
This commit is contained in:
Levi Neely
parent
8e50b84707
commit
0da6f22551
Binary file not shown.
|
|
@ -1,387 +0,0 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"9fans.net/go/plan9/client"
|
||||
"anvillm/internal/sandbox"
|
||||
)
|
||||
|
||||
// readTool reads a tool from the 9P tools directory.
|
||||
func readTool(name string) (string, error) {
|
||||
// Prevent path traversal
|
||||
if strings.Contains(name, "/") || strings.Contains(name, "..") {
|
||||
return "", fmt.Errorf("invalid tool name")
|
||||
}
|
||||
|
||||
ns := fmt.Sprintf("/tmp/ns.%s.:0", os.Getenv("USER"))
|
||||
fsys, err := client.Mount("unix", filepath.Join(ns, "anvillm"))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to mount 9P: %v", err)
|
||||
}
|
||||
defer fsys.Close()
|
||||
|
||||
fid, err := fsys.Open("/tools/"+name, 0)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("tool not found: %s", name)
|
||||
}
|
||||
defer fid.Close()
|
||||
|
||||
var buf []byte
|
||||
tmp := make([]byte, 8192)
|
||||
for {
|
||||
n, err := fid.Read(tmp)
|
||||
if n > 0 {
|
||||
buf = append(buf, tmp[:n]...)
|
||||
}
|
||||
if err != nil || n < len(tmp) {
|
||||
break
|
||||
}
|
||||
}
|
||||
return string(buf), nil
|
||||
}
|
||||
|
||||
// PipeStep is one stage in a tool pipeline.
|
||||
// Exactly one of Tool or Code must be set.
|
||||
type PipeStep struct {
|
||||
Tool string // named tool read from 9P (trusted)
|
||||
Code string // inline bash code (untrusted, validated)
|
||||
Args []string
|
||||
}
|
||||
|
||||
// buildPipeline constructs a single bash pipeline string from the given steps.
|
||||
// Each step is wrapped in a subshell: ( set -- args; <code> ) | ...
|
||||
// Tool steps are trusted (sourced from 9P); inline code steps are validated
|
||||
// individually here so the combined string is always returned as trusted.
|
||||
func buildPipeline(steps []PipeStep) (string, bool, error) {
|
||||
if len(steps) == 0 {
|
||||
return "", false, fmt.Errorf("pipe requires at least one step")
|
||||
}
|
||||
parts := make([]string, 0, len(steps))
|
||||
for _, step := range steps {
|
||||
var code string
|
||||
if step.Tool != "" {
|
||||
var err error
|
||||
code, err = readTool(step.Tool)
|
||||
if err != nil {
|
||||
return "", false, fmt.Errorf("pipe step %q: %v", step.Tool, err)
|
||||
}
|
||||
} else if step.Code != "" {
|
||||
if err := validateCode(step.Code); err != nil {
|
||||
return "", false, fmt.Errorf("pipe step code: %v", err)
|
||||
}
|
||||
code = step.Code
|
||||
} else {
|
||||
return "", false, fmt.Errorf("each pipe step requires either 'tool' or 'code'")
|
||||
}
|
||||
if len(step.Args) > 0 {
|
||||
var escaped []string
|
||||
for _, arg := range step.Args {
|
||||
escaped = append(escaped, "'"+strings.ReplaceAll(arg, "'", "'\\''")+"'")
|
||||
}
|
||||
parts = append(parts, fmt.Sprintf("( set -- %s\n%s )", strings.Join(escaped, " "), code))
|
||||
} else {
|
||||
parts = append(parts, fmt.Sprintf("(\n%s\n)", code))
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, " |\n"), true, nil
|
||||
}
|
||||
|
||||
var dangerousPatterns = []*regexp.Regexp{
|
||||
regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s*/(home|var|usr|etc|boot|root|bin|sbin|lib|opt|srv)?`), // rm -rf, rm -r -f on sensitive paths
|
||||
regexp.MustCompile(`rm\s+(-[a-z]*f[a-z]*\s+)*-[a-z]*r[a-z]*\s*/(home|var|usr|etc|boot|root|bin|sbin|lib|opt|srv)?`), // rm -fr, rm -f -r on sensitive paths
|
||||
regexp.MustCompile(`rm\s+.*--recursive.*--force`), // rm --recursive --force
|
||||
regexp.MustCompile(`rm\s+.*--force.*--recursive`), // rm --force --recursive
|
||||
regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s+\.\.?(/|$)`), // rm -rf ./ or rm -rf ../
|
||||
regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s+~`), // rm -rf ~ (home dir)
|
||||
regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s+\*`), // rm -rf * (glob expansion)
|
||||
regexp.MustCompile(`:\s*\(\s*\)\s*\{\s*:\s*\|\s*:\s*&`), // fork bomb
|
||||
regexp.MustCompile(`\bmkfs\b`), // filesystem format
|
||||
regexp.MustCompile(`\bdd\b.*\bif=/dev/`), // dd from device
|
||||
regexp.MustCompile(`>\s*/dev/sd`), // write to block device
|
||||
regexp.MustCompile(`\beval\s+".*\$`), // eval with variable expansion
|
||||
regexp.MustCompile(`\b(sudo|su)\s`), // privilege escalation
|
||||
regexp.MustCompile(`/etc/(shadow|sudoers)`), // sensitive files (not passwd)
|
||||
}
|
||||
|
||||
var whitespacePattern = regexp.MustCompile(`\s+`)
|
||||
|
||||
// Rate limiting for validation failures
|
||||
var (
|
||||
rateLimitMu sync.Mutex
|
||||
validationFailures int
|
||||
lastFailure time.Time
|
||||
blockedUntil time.Time
|
||||
)
|
||||
|
||||
const (
|
||||
maxFailures = 5
|
||||
blockDuration = 30 * time.Second
|
||||
failureWindow = 60 * time.Second
|
||||
)
|
||||
|
||||
func checkRateLimit() error {
|
||||
rateLimitMu.Lock()
|
||||
defer rateLimitMu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
if now.Before(blockedUntil) {
|
||||
remaining := blockedUntil.Sub(now).Round(time.Second)
|
||||
return fmt.Errorf("rate limited: too many validation failures, blocked for %v", remaining)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func recordValidationFailure() {
|
||||
rateLimitMu.Lock()
|
||||
defer rateLimitMu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
// Reset counter if outside failure window
|
||||
if now.Sub(lastFailure) > failureWindow {
|
||||
validationFailures = 0
|
||||
}
|
||||
|
||||
validationFailures++
|
||||
lastFailure = now
|
||||
|
||||
if validationFailures >= maxFailures {
|
||||
blockedUntil = now.Add(blockDuration)
|
||||
validationFailures = 0
|
||||
logSecurityEvent(SecurityEvent{
|
||||
Timestamp: now,
|
||||
EventType: "rate_limit_triggered",
|
||||
Details: fmt.Sprintf("blocked for %v after %d failures", blockDuration, maxFailures),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func isPermissionError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
msg := strings.ToLower(err.Error())
|
||||
return strings.Contains(msg, "permission denied") ||
|
||||
strings.Contains(msg, "no such file or directory")
|
||||
}
|
||||
|
||||
func validateCode(code string) error {
|
||||
if err := checkRateLimit(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Normalize: collapse whitespace, lowercase for pattern matching
|
||||
normalized := strings.ToLower(code)
|
||||
normalized = whitespacePattern.ReplaceAllString(normalized, " ")
|
||||
|
||||
for _, pattern := range dangerousPatterns {
|
||||
if pattern.MatchString(normalized) {
|
||||
recordValidationFailure()
|
||||
logSecurityEvent(SecurityEvent{
|
||||
Timestamp: time.Now(),
|
||||
EventType: "validation_failure",
|
||||
Details: fmt.Sprintf("dangerous pattern: %s", pattern.String()),
|
||||
})
|
||||
return fmt.Errorf("dangerous pattern detected")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// loadLayeredConfig loads sandbox config using the layered approach:
|
||||
// global.yaml -> backend (anvilmcp) -> sandbox/<name>.yaml
|
||||
func loadLayeredConfig(name string) (*sandbox.Config, error) {
|
||||
// Load global.yaml as base
|
||||
baseCfg, err := sandbox.Load()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to load global config: %w", err)
|
||||
}
|
||||
|
||||
// Convert base config to layered format
|
||||
baseLayer := sandbox.LayeredConfig{
|
||||
Filesystem: baseCfg.Filesystem,
|
||||
Network: baseCfg.Network,
|
||||
Env: baseCfg.Env,
|
||||
}
|
||||
layers := []sandbox.LayeredConfig{baseLayer}
|
||||
|
||||
// Load sandbox layer
|
||||
if name == "" {
|
||||
name = "anvilmcp"
|
||||
}
|
||||
sbxLayer, err := sandbox.LoadSandbox(name)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to load sandbox %q: %w", name, err)
|
||||
}
|
||||
layers = append(layers, sbxLayer)
|
||||
|
||||
// Merge layers
|
||||
general := sandbox.GeneralConfig{
|
||||
BestEffort: baseCfg.General.BestEffort,
|
||||
LogLevel: baseCfg.General.LogLevel,
|
||||
}
|
||||
advanced := sandbox.AdvancedConfig{
|
||||
LDD: baseCfg.Advanced.LDD,
|
||||
AddExec: baseCfg.Advanced.AddExec,
|
||||
}
|
||||
|
||||
return sandbox.Merge(general, advanced, layers...), nil
|
||||
}
|
||||
|
||||
func executeCode(code, language string, timeout int, sandboxName string, trusted bool) (string, error) {
|
||||
start := time.Now()
|
||||
|
||||
if timeout <= 0 {
|
||||
timeout = 30
|
||||
}
|
||||
|
||||
if !trusted {
|
||||
if err := validateCode(code); err != nil {
|
||||
logExecution(ExecutionLog{
|
||||
Timestamp: start,
|
||||
CodeHash: hashCode(code),
|
||||
Language: language,
|
||||
Duration: time.Since(start),
|
||||
Success: false,
|
||||
OutputSize: 0,
|
||||
Error: err.Error(),
|
||||
})
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
// Load layered sandbox config
|
||||
cfg, err := loadLayeredConfig(sandboxName)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
homeDir, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to get home dir: %v", err)
|
||||
}
|
||||
|
||||
// Use current working directory
|
||||
workDir, _ := os.Getwd()
|
||||
|
||||
// For anvilmcp sandbox, use temp workspace
|
||||
var cleanupWorkDir bool
|
||||
if sandboxName == "" || sandboxName == "anvilmcp" {
|
||||
workspaceBase := filepath.Join(homeDir, ".cache", "anvillm", "exec")
|
||||
if err := os.MkdirAll(workspaceBase, 0700); err != nil {
|
||||
return "", fmt.Errorf("failed to create workspace base: %v", err)
|
||||
}
|
||||
workDir, err = os.MkdirTemp(workspaceBase, "anvilmcp-*")
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to create workspace: %v", err)
|
||||
}
|
||||
cleanupWorkDir = true
|
||||
}
|
||||
|
||||
if cleanupWorkDir {
|
||||
defer os.RemoveAll(workDir)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
|
||||
defer cancel()
|
||||
|
||||
var cmd *exec.Cmd
|
||||
switch language {
|
||||
case "bash", "":
|
||||
wrapped := sandbox.WrapCommand(cfg, []string{"bash", "-c", code}, workDir)
|
||||
cmd = exec.CommandContext(ctx, wrapped[0], wrapped[1:]...)
|
||||
cmd.Dir = workDir
|
||||
// Add new language cases here:
|
||||
// case "python":
|
||||
// wrapped := sandbox.WrapCommand(cfg, []string{"python3", "-c", code}, workDir)
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported language: %s (supported: bash)", language)
|
||||
}
|
||||
|
||||
// Limit output size (10MB raw, 8KB returned)
|
||||
const maxToolOutputSize = 8000
|
||||
var outputBuf bytes.Buffer
|
||||
lw := &limitedWriter{w: &outputBuf, limit: 10 * 1024 * 1024}
|
||||
cmd.Stdout = lw
|
||||
cmd.Stderr = lw
|
||||
|
||||
err = cmd.Run()
|
||||
output := outputBuf.Bytes()
|
||||
|
||||
if lw.truncated {
|
||||
output = append(output, []byte("\n[output truncated at 10MB]")...)
|
||||
}
|
||||
|
||||
duration := time.Since(start)
|
||||
|
||||
execLog := ExecutionLog{
|
||||
Timestamp: start,
|
||||
CodeHash: hashCode(code),
|
||||
Language: language,
|
||||
Duration: duration,
|
||||
Success: err == nil && ctx.Err() != context.DeadlineExceeded,
|
||||
OutputSize: len(output),
|
||||
}
|
||||
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
execLog.Error = fmt.Sprintf("execution timeout after %d seconds", timeout)
|
||||
logExecution(execLog)
|
||||
logSecurityEvent(SecurityEvent{
|
||||
Timestamp: start,
|
||||
EventType: "timeout",
|
||||
Language: language,
|
||||
Details: fmt.Sprintf("timeout after %d seconds", timeout),
|
||||
})
|
||||
return "", fmt.Errorf("execution timeout after %d seconds", timeout)
|
||||
}
|
||||
if err != nil {
|
||||
execLog.Error = err.Error()
|
||||
logExecution(execLog)
|
||||
return string(output), fmt.Errorf("execution failed: %v\nOutput: %s", err, string(output))
|
||||
}
|
||||
|
||||
logExecution(execLog)
|
||||
result := string(output)
|
||||
if len(result) > maxToolOutputSize {
|
||||
result = result[:maxToolOutputSize] + "\n... (output truncated)"
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
type limitedWriter struct {
|
||||
w io.Writer
|
||||
written int
|
||||
limit int
|
||||
truncated bool
|
||||
}
|
||||
|
||||
func (lw *limitedWriter) Write(p []byte) (n int, err error) {
|
||||
if lw.written >= lw.limit {
|
||||
lw.truncated = true
|
||||
return len(p), nil // Discard, report all bytes consumed
|
||||
}
|
||||
|
||||
remaining := lw.limit - lw.written
|
||||
toWrite := p
|
||||
if len(p) > remaining {
|
||||
toWrite = p[:remaining]
|
||||
lw.truncated = true
|
||||
}
|
||||
|
||||
written, err := lw.w.Write(toWrite)
|
||||
lw.written += written
|
||||
if err != nil {
|
||||
return written, err
|
||||
}
|
||||
return len(p), nil // Report all input bytes as consumed
|
||||
}
|
||||
|
|
@ -1,188 +0,0 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestValidateCode(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
code string
|
||||
wantErr bool
|
||||
}{
|
||||
{"safe code", "echo hello", false},
|
||||
{"rm -rf /", "rm -rf /", true},
|
||||
{"rm -rf / with spaces", "rm -rf /", true},
|
||||
{"rm -rf / with tabs", "rm\t-rf\t/", true},
|
||||
{"rm -rf/ no space", "rm -rf/", true},
|
||||
{"rm -fr /", "rm -fr /", true},
|
||||
{"rm -fr/ no space", "rm -fr/", true},
|
||||
{"rm -rf /home", "rm -rf /home", true},
|
||||
{"rm -rf /var", "rm -rf /var", true},
|
||||
{"rm -rf /etc", "rm -rf /etc", true},
|
||||
{"rm -rf /usr", "rm -rf /usr", true},
|
||||
{"fork bomb", ":(){ :|:& };:", true},
|
||||
{"fork bomb with spaces", ": ( ) { : | : & } ; :", true},
|
||||
{"mkfs", "mkfs.ext4 /dev/sda", true},
|
||||
{"dd from device", "dd if=/dev/zero of=/dev/sda", true},
|
||||
{"chmod 777", "chmod 777 /etc/passwd", true},
|
||||
{"curl http", "curl http://example.com", true},
|
||||
{"curl https", "curl https://example.com", true},
|
||||
{"curl ftp", "curl ftp://example.com/file", true},
|
||||
{"curl file", "curl file:///etc/passwd", true},
|
||||
{"wget http", "wget http://example.com/file", true},
|
||||
{"wget https", "wget https://example.com/file", true},
|
||||
{"wget ftp", "wget ftp://example.com/file", true},
|
||||
{"device write", "echo data > /dev/sda", true},
|
||||
{"exec", "exec(malicious)", true},
|
||||
{"eval", "eval(dangerous)", true},
|
||||
{"sudo", "sudo rm file", true},
|
||||
{"su", "su - root", true},
|
||||
{"etc passwd", "cat /etc/passwd", true},
|
||||
{"etc shadow", "cat /etc/shadow", true},
|
||||
{"safe 9p", "9p read anvillm/inbox/user", false},
|
||||
{"safe jq", "echo '{}' | jq .field", false},
|
||||
{"safe rm", "rm file.txt", false},
|
||||
{"rm -rf ./", "rm -rf ./", true},
|
||||
{"rm -rf ../", "rm -rf ../", true},
|
||||
{"rm -rf .", "rm -rf .", true},
|
||||
{"rm -r -f /", "rm -r -f /", true},
|
||||
{"rm -f -r /", "rm -f -r /", true},
|
||||
{"rm --recursive --force /", "rm --recursive --force /", true},
|
||||
{"rm --force --recursive /", "rm --force --recursive /", true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := validateCode(tt.code)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("validateCode() error = %v, wantErr %v", err, tt.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDangerousPatternsCompile(t *testing.T) {
|
||||
// Verify all patterns are valid regexes (they compile at init, but this documents the expectation)
|
||||
if len(dangerousPatterns) == 0 {
|
||||
t.Error("dangerousPatterns should not be empty")
|
||||
}
|
||||
for i, p := range dangerousPatterns {
|
||||
if p == nil {
|
||||
t.Errorf("dangerousPatterns[%d] is nil", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimitedWriter(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
limit int
|
||||
writes []string
|
||||
wantTotal int
|
||||
wantTruncated bool
|
||||
}{
|
||||
{
|
||||
name: "under limit",
|
||||
limit: 100,
|
||||
writes: []string{"hello", " ", "world"},
|
||||
wantTotal: 11,
|
||||
wantTruncated: false,
|
||||
},
|
||||
{
|
||||
name: "at limit",
|
||||
limit: 5,
|
||||
writes: []string{"hello"},
|
||||
wantTotal: 5,
|
||||
wantTruncated: false,
|
||||
},
|
||||
{
|
||||
name: "over limit",
|
||||
limit: 5,
|
||||
writes: []string{"hello", "world"},
|
||||
wantTotal: 5,
|
||||
wantTruncated: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
var buf strings.Builder
|
||||
lw := &limitedWriter{w: &buf, limit: tt.limit}
|
||||
|
||||
for _, write := range tt.writes {
|
||||
lw.Write([]byte(write))
|
||||
}
|
||||
|
||||
if lw.truncated != tt.wantTruncated {
|
||||
t.Errorf("limitedWriter truncated = %v, want %v", lw.truncated, tt.wantTruncated)
|
||||
}
|
||||
if buf.Len() != tt.wantTotal {
|
||||
t.Errorf("limitedWriter wrote %d bytes, want %d", buf.Len(), tt.wantTotal)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func resetRateLimitState() {
|
||||
rateLimitMu.Lock()
|
||||
defer rateLimitMu.Unlock()
|
||||
validationFailures = 0
|
||||
lastFailure = time.Time{}
|
||||
blockedUntil = time.Time{}
|
||||
}
|
||||
|
||||
func TestRateLimitCounterResetAfterWindow(t *testing.T) {
|
||||
resetRateLimitState()
|
||||
defer resetRateLimitState()
|
||||
|
||||
// Record failures but not enough to trigger block
|
||||
for i := 0; i < maxFailures-1; i++ {
|
||||
recordValidationFailure()
|
||||
}
|
||||
|
||||
// Simulate time passing beyond failureWindow
|
||||
rateLimitMu.Lock()
|
||||
lastFailure = time.Now().Add(-failureWindow - time.Second)
|
||||
rateLimitMu.Unlock()
|
||||
|
||||
// Next failure should reset counter, not trigger block
|
||||
recordValidationFailure()
|
||||
|
||||
if err := checkRateLimit(); err != nil {
|
||||
t.Errorf("expected no rate limit after window reset, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimitBlocksAfterMaxFailures(t *testing.T) {
|
||||
resetRateLimitState()
|
||||
defer resetRateLimitState()
|
||||
|
||||
for i := 0; i < maxFailures; i++ {
|
||||
recordValidationFailure()
|
||||
}
|
||||
|
||||
if err := checkRateLimit(); err == nil {
|
||||
t.Error("expected rate limit error after maxFailures")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimitUnblocksAfterDuration(t *testing.T) {
|
||||
resetRateLimitState()
|
||||
defer resetRateLimitState()
|
||||
|
||||
for i := 0; i < maxFailures; i++ {
|
||||
recordValidationFailure()
|
||||
}
|
||||
|
||||
// Simulate block duration passing
|
||||
rateLimitMu.Lock()
|
||||
blockedUntil = time.Now().Add(-time.Second)
|
||||
rateLimitMu.Unlock()
|
||||
|
||||
if err := checkRateLimit(); err != nil {
|
||||
t.Errorf("expected no rate limit after block duration, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
|
@ -1,285 +0,0 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
var (
|
||||
executionSemaphore = make(chan struct{}, 3) // Max 3 concurrent executions
|
||||
)
|
||||
|
||||
type MCPRequest struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID interface{} `json:"id"`
|
||||
Method string `json:"method"`
|
||||
Params json.RawMessage `json:"params,omitempty"`
|
||||
}
|
||||
|
||||
type MCPResponse struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID interface{} `json:"id"`
|
||||
Result interface{} `json:"result,omitempty"`
|
||||
Error *MCPError `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type MCPError struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
type Tool struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
InputSchema InputSchema `json:"inputSchema"`
|
||||
}
|
||||
|
||||
type InputSchema struct {
|
||||
Type string `json:"type"`
|
||||
Properties map[string]Property `json:"properties"`
|
||||
Required []string `json:"required,omitempty"`
|
||||
}
|
||||
|
||||
type Property struct {
|
||||
Type string `json:"type"`
|
||||
Description string `json:"description"`
|
||||
Items *Items `json:"items,omitempty"`
|
||||
Enum []string `json:"enum,omitempty"`
|
||||
}
|
||||
|
||||
type Items struct {
|
||||
Type string `json:"type"`
|
||||
}
|
||||
|
||||
func main() {
|
||||
fmt.Fprintln(os.Stderr, "[anvilmcp] Starting MCP server")
|
||||
scanner := bufio.NewScanner(os.Stdin)
|
||||
for scanner.Scan() {
|
||||
line := scanner.Bytes()
|
||||
fmt.Fprintf(os.Stderr, "[anvilmcp] Received: %s\n", string(line))
|
||||
var req MCPRequest
|
||||
if err := json.Unmarshal(line, &req); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "[anvilmcp] Parse error: %v\n", err)
|
||||
sendError(nil, -32700, "Parse error")
|
||||
continue
|
||||
}
|
||||
|
||||
fmt.Fprintf(os.Stderr, "[anvilmcp] Method: %s\n", req.Method)
|
||||
switch req.Method {
|
||||
case "initialize":
|
||||
sendResponse(req.ID, map[string]interface{}{
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": map[string]interface{}{
|
||||
"tools": map[string]bool{},
|
||||
},
|
||||
"serverInfo": map[string]string{
|
||||
"name": "anvilmcp",
|
||||
"version": "0.1.0",
|
||||
},
|
||||
})
|
||||
case "notifications/initialized":
|
||||
// Notification - no response needed
|
||||
fmt.Fprintln(os.Stderr, "[anvilmcp] Initialized notification received")
|
||||
case "tools/list":
|
||||
sendResponse(req.ID, map[string]interface{}{
|
||||
"tools": []Tool{
|
||||
{
|
||||
Name: "execute_code",
|
||||
Description: "Execute bash via one of three mutually exclusive modes (pipe wins if present):\n1. tool+args — run a named library tool: {tool: 'list_sessions.sh', args: ['foo']}\n2. code+args — run inline bash: {code: 'echo hello'}\n3. pipe — Unix pipeline of stages, stdout→stdin: {pipe: [{tool: 'list_sessions.sh'}, {code: 'grep active'}, {tool: 'format_table.sh', args: ['--compact']}]}\nEach pipe stage is {tool?, args?, code?}. All-tool pipes are trusted (skip validation); any inline code stage triggers validation.",
|
||||
InputSchema: InputSchema{
|
||||
Type: "object",
|
||||
Properties: map[string]Property{
|
||||
"tool": {Type: "string", Description: "Mode 1: named tool from anvillm/tools/ (e.g. 'list_sessions.sh')"},
|
||||
"args": {Type: "array", Description: "Positional args ($1 $2 …) for the tool or code", Items: &Items{Type: "string"}},
|
||||
"code": {Type: "string", Description: "Mode 2: inline bash to execute directly"},
|
||||
"language": {Type: "string", Description: "Programming language", Enum: []string{"bash"}},
|
||||
"timeout": {Type: "integer", Description: "Timeout in seconds (default: 30). Recommended: 600 for builds, 120 for network ops, 1800 for remote builds"},
|
||||
"sandbox": {Type: "string", Description: "Sandbox config name (default: anvilmcp)"},
|
||||
"pipe": {Type: "array", Description: "Mode 3: ordered pipeline stages, each {tool?, args?, code?}. Stages run as a single bash -c; stdout of each feeds stdin of the next.", Items: &Items{Type: "object"}},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
case "tools/call":
|
||||
handleToolCall(req)
|
||||
default:
|
||||
sendError(req.ID, -32601, "Method not found")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func handleToolCall(req MCPRequest) {
|
||||
var params struct {
|
||||
Name string `json:"name"`
|
||||
Arguments map[string]interface{} `json:"arguments"`
|
||||
}
|
||||
if err := json.Unmarshal(req.Params, ¶ms); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "[anvilmcp] Invalid params: %v\n", err)
|
||||
sendError(req.ID, -32602, "Invalid params")
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Fprintf(os.Stderr, "[anvilmcp] Tool call: %s with args: %v\n", params.Name, params.Arguments)
|
||||
switch params.Name {
|
||||
case "execute_code":
|
||||
tool, _ := params.Arguments["tool"].(string)
|
||||
var toolArgs []string
|
||||
if argsRaw, ok := params.Arguments["args"].([]interface{}); ok {
|
||||
for _, a := range argsRaw {
|
||||
if s, ok := a.(string); ok {
|
||||
toolArgs = append(toolArgs, s)
|
||||
}
|
||||
}
|
||||
} else if argsStr, ok := params.Arguments["args"].(string); ok {
|
||||
// LLM passed args as a JSON-encoded string instead of an array
|
||||
var parsed []string
|
||||
if err := json.Unmarshal([]byte(argsStr), &parsed); err == nil {
|
||||
toolArgs = parsed
|
||||
}
|
||||
}
|
||||
code, _ := params.Arguments["code"].(string)
|
||||
language, _ := params.Arguments["language"].(string)
|
||||
|
||||
// Parse pipe steps
|
||||
var pipeSteps []PipeStep
|
||||
if pipeRaw, ok := params.Arguments["pipe"].([]interface{}); ok {
|
||||
for _, stepRaw := range pipeRaw {
|
||||
stepMap, ok := stepRaw.(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
step := PipeStep{}
|
||||
step.Tool, _ = stepMap["tool"].(string)
|
||||
step.Code, _ = stepMap["code"].(string)
|
||||
if argsRaw, ok := stepMap["args"].([]interface{}); ok {
|
||||
for _, a := range argsRaw {
|
||||
if s, ok := a.(string); ok {
|
||||
step.Args = append(step.Args, s)
|
||||
}
|
||||
}
|
||||
} else if argsStr, ok := stepMap["args"].(string); ok {
|
||||
var parsed []string
|
||||
if err := json.Unmarshal([]byte(argsStr), &parsed); err == nil {
|
||||
step.Args = parsed
|
||||
}
|
||||
}
|
||||
pipeSteps = append(pipeSteps, step)
|
||||
}
|
||||
}
|
||||
|
||||
// Resolve execution mode: pipe > tool > inline code
|
||||
trusted := false
|
||||
if len(pipeSteps) > 0 {
|
||||
var err error
|
||||
code, trusted, err = buildPipeline(pipeSteps)
|
||||
if err != nil {
|
||||
sendError(req.ID, -32000, err.Error())
|
||||
return
|
||||
}
|
||||
} else if tool != "" {
|
||||
toolCode, err := readTool(tool)
|
||||
if err != nil {
|
||||
sendError(req.ID, -32000, fmt.Sprintf("failed to read tool %s: %v", tool, err))
|
||||
return
|
||||
}
|
||||
code = toolCode
|
||||
trusted = true
|
||||
// For bash with args, wrap script to receive positional params
|
||||
if len(toolArgs) > 0 {
|
||||
// Escape args for bash using single quotes (prevents all expansion)
|
||||
var escaped []string
|
||||
for _, arg := range toolArgs {
|
||||
// Replace ' with '\'' (end quote, escaped quote, start quote)
|
||||
escaped = append(escaped, "'"+strings.ReplaceAll(arg, "'", "'\\''")+"'")
|
||||
}
|
||||
code = fmt.Sprintf("set -- %s\n%s", strings.Join(escaped, " "), code)
|
||||
}
|
||||
}
|
||||
|
||||
if code == "" {
|
||||
sendError(req.ID, -32602, "either 'tool', 'code', or 'pipe' is required")
|
||||
return
|
||||
}
|
||||
|
||||
if language == "" {
|
||||
language = "bash"
|
||||
}
|
||||
timeout := 30
|
||||
if t, ok := params.Arguments["timeout"].(float64); ok {
|
||||
timeout = int(t)
|
||||
}
|
||||
sandbox, _ := params.Arguments["sandbox"].(string)
|
||||
sandboxExplicit := sandbox != ""
|
||||
if sandbox == "" {
|
||||
sandbox = "anvilmcp"
|
||||
}
|
||||
|
||||
// Acquire execution slot
|
||||
executionSemaphore <- struct{}{}
|
||||
defer func() { <-executionSemaphore }()
|
||||
fmt.Fprintf(os.Stderr, "[anvilmcp] Executing %s code (timeout: %ds, sandbox: %s, trusted: %v)\n", language, timeout, sandbox, trusted)
|
||||
result, err := executeCode(code, language, timeout, sandbox, trusted)
|
||||
|
||||
// Fallback to default sandbox on permission errors (only if sandbox wasn't explicit)
|
||||
if err != nil && !sandboxExplicit && isPermissionError(err) {
|
||||
fmt.Fprintf(os.Stderr, "[anvilmcp] Permission error in %s sandbox, falling back to default\n", sandbox)
|
||||
result, err = executeCode(code, language, timeout, "default", trusted)
|
||||
}
|
||||
|
||||
// Log token comparison
|
||||
codeTokens := estimateTokens(code)
|
||||
outputTokens := estimateTokens(result)
|
||||
reduction := 0.0
|
||||
if codeTokens > 0 {
|
||||
reduction = (1.0 - float64(outputTokens)/float64(codeTokens)) * 100
|
||||
}
|
||||
logTokens(TokenLog{
|
||||
Timestamp: time.Now(),
|
||||
Method: "execute_code",
|
||||
DirectTokens: codeTokens,
|
||||
CodeExecTokens: outputTokens,
|
||||
Reduction: reduction,
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "[anvilmcp] Error: %v\n", err)
|
||||
sendError(req.ID, -32000, err.Error())
|
||||
return
|
||||
}
|
||||
fmt.Fprintf(os.Stderr, "[anvilmcp] Execution complete: %d bytes\n", len(result))
|
||||
sendResponse(req.ID, map[string]interface{}{
|
||||
"content": []map[string]string{
|
||||
{"type": "text", "text": result},
|
||||
},
|
||||
})
|
||||
default:
|
||||
sendError(req.ID, -32601, "Tool not found")
|
||||
}
|
||||
}
|
||||
|
||||
func sendResponse(id interface{}, result interface{}) {
|
||||
resp := MCPResponse{
|
||||
JSONRPC: "2.0",
|
||||
ID: id,
|
||||
Result: result,
|
||||
}
|
||||
data, _ := json.Marshal(resp)
|
||||
fmt.Println(string(data))
|
||||
os.Stdout.Sync()
|
||||
}
|
||||
|
||||
func sendError(id interface{}, code int, message string) {
|
||||
resp := MCPResponse{
|
||||
JSONRPC: "2.0",
|
||||
ID: id,
|
||||
Error: &MCPError{Code: code, Message: message},
|
||||
}
|
||||
data, _ := json.Marshal(resp)
|
||||
fmt.Println(string(data))
|
||||
os.Stdout.Sync()
|
||||
}
|
||||
|
|
@ -1,129 +0,0 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
)
|
||||
|
||||
type ExecutionLog struct {
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
CodeHash string `json:"code_hash"`
|
||||
Language string `json:"language"`
|
||||
Duration time.Duration `json:"duration"`
|
||||
Success bool `json:"success"`
|
||||
OutputSize int `json:"output_size"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type TokenLog struct {
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Method string `json:"method"`
|
||||
DirectTokens int `json:"direct_tokens"`
|
||||
CodeExecTokens int `json:"code_exec_tokens"`
|
||||
Reduction float64 `json:"reduction_percent"`
|
||||
}
|
||||
|
||||
type SecurityEvent struct {
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
EventType string `json:"event_type"`
|
||||
Language string `json:"language"`
|
||||
Details string `json:"details"`
|
||||
}
|
||||
|
||||
func hashCode(code string) string {
|
||||
h := sha256.Sum256([]byte(code))
|
||||
return hex.EncodeToString(h[:])
|
||||
}
|
||||
|
||||
func logExecution(log ExecutionLog) error {
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
logDir := filepath.Join(home, ".local", "state", "anvilmcp")
|
||||
if err := os.MkdirAll(logDir, 0755); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
logFile := filepath.Join(logDir, "executions.jsonl")
|
||||
// 0644: world-readable, contains only hashes and metadata (no sensitive data)
|
||||
f, err := os.OpenFile(logFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
data, err := json.Marshal(log)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
_, err = fmt.Fprintf(f, "%s\n", data)
|
||||
return err
|
||||
}
|
||||
|
||||
func logTokens(log TokenLog) error {
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
logDir := filepath.Join(home, ".local", "state", "anvilmcp")
|
||||
if err := os.MkdirAll(logDir, 0755); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
logFile := filepath.Join(logDir, "tokens.jsonl")
|
||||
// 0644: world-readable, contains only aggregate metrics (no sensitive data)
|
||||
f, err := os.OpenFile(logFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
data, err := json.Marshal(log)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
_, err = fmt.Fprintf(f, "%s\n", data)
|
||||
return err
|
||||
}
|
||||
|
||||
func logSecurityEvent(event SecurityEvent) error {
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
logDir := filepath.Join(home, ".local", "state", "anvilmcp")
|
||||
if err := os.MkdirAll(logDir, 0755); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
logFile := filepath.Join(logDir, "security.jsonl")
|
||||
// 0600: owner-only, contains security events (rate limits, validation failures, timeouts)
|
||||
f, err := os.OpenFile(logFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
data, err := json.Marshal(event)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
_, err = fmt.Fprintf(f, "%s\n", data)
|
||||
return err
|
||||
}
|
||||
|
||||
func estimateTokens(text string) int {
|
||||
return len(text) / 4
|
||||
}
|
||||
Reference in New Issue