Building an MCP Server in Go with Code Mode: From 1.17M Tokens to 1,000

package mcp
 
import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"sync"
)
 
type Request struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      any             `json:"id"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params,omitempty"`
}
 
type Tool struct {
	Name        string         `json:"name"`
	Description string         `json:"description"`
	InputSchema map[string]any `json:"inputSchema"`
}
 
type ToolHandler func(ctx context.Context, sessionID string, input json.RawMessage) (any, error)
 
type Server struct {
	tools    []Tool
	dispatch map[string]ToolHandler
	mu       sync.Mutex
}
 
func NewServer() *Server {
	return &Server{dispatch: make(map[string]ToolHandler)}
}
 
func (s *Server) Register(tool Tool, handler ToolHandler) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.tools = append(s.tools, tool)
	s.dispatch[tool.Name] = handler
}
 
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	body, err := io.ReadAll(io.LimitReader(r.Body, 1<<20))
	if err != nil {
		http.Error(w, fmt.Sprintf("failed to read request body: %v", err), http.StatusBadRequest)
		return
	}
	var req Request
	if err := json.Unmarshal(body, &req); err != nil {
		http.Error(w, fmt.Sprintf("failed to unmarshal request: %v", err), http.StatusBadRequest)
		return
	}
 
	sessionID := r.Header.Get("X-Session-ID")
	var result any
 
	switch req.Method {
	case "tools/list":
		s.mu.Lock()
		result = map[string]any{"tools": s.tools}
		s.mu.Unlock()
	case "tools/call":
		var p struct {
			Name  string          `json:"name"`
			Input json.RawMessage `json:"arguments"`
		}
		if err := json.Unmarshal(req.Params, &p); err != nil {
			http.Error(w, fmt.Sprintf("failed to unmarshal params: %v", err), http.StatusBadRequest)
			return
		}
		s.mu.Lock()
		handler := s.dispatch[p.Name]
		s.mu.Unlock()
		if handler != nil {
			var handlerErr error
			result, handlerErr = handler(r.Context(), sessionID, p.Input)
			if handlerErr != nil {
				http.Error(w, fmt.Sprintf("tool call failed: %v", handlerErr), http.StatusInternalServerError)
				return
			}
		}
	}
 
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]any{
		"jsonrpc": "2.0",
		"id":      req.ID,
		"result":  result,
	})
}

package mcp
 
import "strings"
 
type SpecIndex struct {
	entries []struct {
		Path, Method, Summary, searchText string
		Tags                              []string
		Responses                         any
	}
}
 
func BuildIndex(spec map[string]any) *SpecIndex {
	idx := &SpecIndex{}
	paths, _ := spec["paths"].(map[string]any)
	for path, pathItem := range paths {
		methods, _ := pathItem.(map[string]any)
		for method, operation := range methods {
			op, _ := operation.(map[string]any)
			summary, _ := op["summary"].(string)
			searchText := strings.ToLower(path + " " + method + " " + summary)
			idx.entries = append(idx.entries, struct {
				Path, Method, Summary, searchText string
				Tags                              []string
				Responses                         any
			}{
				Path:       path,
				Method:     strings.ToUpper(method),
				Summary:    summary,
				searchText: searchText,
				Responses:  op["responses"],
			})
		}
	}
	return idx
}
 
func (idx *SpecIndex) Search(query string, limit int) []map[string]any {
	var results []map[string]any
	terms := strings.Fields(strings.ToLower(query))
	
	for _, entry := range idx.entries {
		score := 0
		for _, term := range terms {
			if strings.Contains(entry.searchText, term) {
				score++
			}
		}
		if score > 0 && len(results) < limit {
			results = append(results, map[string]any{
				"path":     entry.Path,
				"method":   entry.Method,
				"summary":  entry.Summary,
				"responses": entry.Responses,
			})
		}
	}
	return results
}

package mcp
 
import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
)
 
type Executor struct {
	baseURL string
	authHeader string
	allowedHeaders map[string]bool
	client *http.Client
}
 
func (e *Executor) Execute(ctx context.Context, method, path string, body any, modelHeaders map[string]string) (map[string]any, error) {
	// Validate path is relative (no scheme/host)
	if strings.Contains(path, "://") {
		return nil, fmt.Errorf("path must be relative")
	}
 
	// Resolve path against baseURL
	base, _ := url.Parse(e.baseURL)
	ref, _ := url.Parse(path)
	target := base.ResolveReference(ref)
	
	// Ensure target stays within baseURL (prevent /v1evil from matching /v1)
	basePrefix := strings.TrimRight(e.baseURL, "/") + "/"
	if !strings.HasPrefix(target.String(), basePrefix) && target.String() != strings.TrimRight(e.baseURL, "/") {
		return nil, fmt.Errorf("path outside base URL")
	}
 
	// Build request with injected auth
	var bodyReader io.Reader
	if body != nil {
		b, err := json.Marshal(body)
		if err != nil {
			return nil, fmt.Errorf("marshal body: %w", err)
		}
		bodyReader = bytes.NewReader(b)
	}
	req, err := http.NewRequestWithContext(ctx, strings.ToUpper(method), target.String(), bodyReader)
	if err != nil {
		return nil, fmt.Errorf("create request: %w", err)
	}
	req.Header.Set("Authorization", e.authHeader)
	req.Header.Set("Content-Type", "application/json")
 
	// Copy only whitelisted headers from the model. Anything not in
	// allowedHeaders is dropped, and Authorization is never overridable —
	// it was injected server-side above and the model cannot replace it.
	for name, value := range modelHeaders {
		canonical := http.CanonicalHeaderKey(name)
		if canonical == "Authorization" || !e.allowedHeaders[canonical] {
			continue
		}
		req.Header.Set(canonical, value)
	}
 
	resp, err := e.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("execute request: %w", err)
	}
	defer resp.Body.Close()
 
	respBody, err := io.ReadAll(io.LimitReader(resp.Body, 512*1024))
	if err != nil {
		return nil, fmt.Errorf("read response body: %w", err)
	}
	var result any
	if err := json.Unmarshal(respBody, &result); err != nil {
		return nil, fmt.Errorf("unmarshal response: %w", err)
	}
 
	return map[string]any{
		"status_code": resp.StatusCode,
		"body":        result,
	}, nil
}

package main
 
import (
	"sync"
	"time"
)
 
type SessionLimiter struct {
	mu       sync.Mutex
	sessions map[string]*bucket
	limit    int
	window   time.Duration
}
 
type bucket struct {
	tokens  int
	resetAt time.Time
}
 
func (l *SessionLimiter) Allow(sessionID string) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	
	now := time.Now()
	b, ok := l.sessions[sessionID]
	if !ok || now.After(b.resetAt) {
		l.sessions[sessionID] = &bucket{tokens: l.limit - 1, resetAt: now.Add(l.window)}
		return true
	}
	if b.tokens <= 0 {
		return false
	}
	b.tokens--
	return true
}

func AuthMiddleware(keys map[string]bool, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		key := r.Header.Get("X-API-Key")
		if !keys[key] {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		sessionID := r.Header.Get("X-Session-ID")
		if sessionID == "" {
			sessionID = key
		}
		r.Header.Set("X-Session-ID", sessionID)
		next.ServeHTTP(w, r)
	})
}

type SpecIndex struct {
    paths    map[string]*PathItem      // path -> methods+ops
    baseURL  *url.URL
}
 
func LoadSpec(specPath, baseURL string) (*SpecIndex, error) {
    raw, err := os.ReadFile(specPath)
    if err != nil { return nil, fmt.Errorf("read spec: %w", err) }
 
    var doc openapi3.T
    if err := yaml.Unmarshal(raw, &doc); err != nil {
        return nil, fmt.Errorf("parse spec: %w", err)
    }
    if err := doc.Validate(context.Background()); err != nil {
        return nil, fmt.Errorf("invalid spec: %w", err)
    }
    base, err := url.Parse(baseURL)
    if err != nil || base.Host == "" {
        return nil, fmt.Errorf("invalid baseURL %q", baseURL)
    }
 
    idx := &SpecIndex{paths: make(map[string]*PathItem), baseURL: base}
    for path, item := range doc.Paths.Map() {
        idx.paths[path] = newPathItem(item)
    }
    return idx, nil
}

var ErrPathEscape = errors.New("path escapes spec baseURL; refusing")
 
// resolveAndValidate parses the model's path argument against baseURL.
// Refuses any path that resolves to a different scheme or host, even
// after URL normalisation. The defence is "host MUST equal baseURL.Host
// AND scheme MUST equal baseURL.Scheme", checked AFTER url.Parse, so
// schemes like file:// or relative-protocol //evil.com are caught.
func (idx *SpecIndex) resolveAndValidate(rawPath string) (*url.URL, error) {
    if strings.Contains(rawPath, "://") || strings.HasPrefix(rawPath, "//") {
        return nil, ErrPathEscape
    }
    target := idx.baseURL.ResolveReference(&url.URL{Path: rawPath})
    if target.Host != idx.baseURL.Host || target.Scheme != idx.baseURL.Scheme {
        return nil, ErrPathEscape
    }
    if _, ok := idx.paths[target.Path]; !ok {
        return nil, fmt.Errorf("path %q not in spec", target.Path)
    }
    return target, nil
}

type sessionBucket struct {
    limiter  *rate.Limiter
    lastSeen atomic.Int64 // UnixNano of the most recent Allow call
}
 
type SessionLimiter struct {
    buckets sync.Map   // sessionID -> *sessionBucket
    rate    rate.Limit // ops per second per session
    burst   int
}
 
func (l *SessionLimiter) Allow(sessionID string) bool {
    v, _ := l.buckets.LoadOrStore(sessionID, &sessionBucket{
        limiter: rate.NewLimiter(l.rate, l.burst),
    })
    b := v.(*sessionBucket)
    b.lastSeen.Store(time.Now().UnixNano())
    return b.limiter.Allow()
}
 
// Reaper: drop buckets that haven't been touched in maxIdle.
// Run from a 5-minute ticker so dropped sessions don't pin memory.
func (l *SessionLimiter) Reap(maxIdle time.Duration) {
    cutoff := time.Now().Add(-maxIdle).UnixNano()
    l.buckets.Range(func(k, v any) bool {
        b := v.(*sessionBucket)
        // Only reap fully-replenished, idle buckets so an in-flight
        // session can't be dropped mid-burst.
        if b.limiter.Tokens() >= float64(l.burst) && b.lastSeen.Load() < cutoff {
            l.buckets.Delete(k)
        }
        return true
    })
}

type AuditEntry struct {
    OccurredAt  time.Time         `json:"occurred_at"`
    SessionID   string            `json:"session_id"`
    Method      string            `json:"method"`
    Path        string            `json:"path"`
    Status      int               `json:"status"`
    LatencyMs   int64             `json:"latency_ms"`
    BodyHash    string            `json:"body_hash"`         // sha256, for replay forensics
    ResponseHash string           `json:"response_hash"`
}
 
func (e AuditEntry) MarshalForLog() string {
    b, _ := json.Marshal(e); return string(b)
}

groups:
  - name: mcp-server-agent-misbehaviour
    interval: 30s
    rules:
      - alert: MCPExecuteErrorBurst
        expr: |
          sum by (session_id) (rate(mcp_execute_errors_total[2m]))
            > 0.5
        for: 3m
        labels: { severity: warning }
        annotations:
          summary: "Session {{ $labels.session_id }} bursting execute errors"
          runbook: "https://runbooks/mcp/execute-error-burst"
 
      - alert: MCPPathEscapeAttempt
        expr: increase(mcp_path_escape_total[5m]) > 0
        labels: { severity: critical }
        annotations:
          summary: "SSRF or path-escape attempt detected"
          description: "Inspect audit log for raw_path; review caller identity"
 
      - alert: MCPRateLimiterSaturated
        expr: |
          sum(rate(mcp_rate_limited_total[5m]))
            / sum(rate(mcp_execute_calls_total[5m])) > 0.1
        for: 10m
        labels: { severity: warning }
        annotations:
          summary: "More than 10% of execute calls are being throttled"
 
      - alert: MCPResponseHashCollision
        expr: |
          sum by (session_id) (rate(mcp_response_hash_repeats_total[10m]))
            > 5
        for: 5m
        labels: { severity: warning }
        annotations:
          summary: "Session repeating identical response hashes — possible exfil loop"

var (
    executeCalls = prometheus.NewCounterVec(
        prometheus.CounterOpts{Name: "mcp_execute_calls_total"},
        []string{"method", "path_prefix", "status"},
    )
    executeErrors = prometheus.NewCounterVec(
        prometheus.CounterOpts{Name: "mcp_execute_errors_total"},
        []string{"session_id", "reason"},
    )
    pathEscape = prometheus.NewCounter(
        prometheus.CounterOpts{Name: "mcp_path_escape_total"},
    )
    executeLatency = prometheus.NewHistogramVec(
        prometheus.HistogramOpts{
            Name:    "mcp_execute_latency_seconds",
            Buckets: prometheus.ExponentialBuckets(0.005, 2, 12),
        },
        []string{"path_prefix"},
    )
)