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

package mcp
 
import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"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, ok := s.dispatch[p.Name]
		s.mu.Unlock()
		if !ok {
			writeJSON(w, map[string]any{"jsonrpc": "2.0", "id": req.ID,
				"error": map[string]any{"code": -32602, "message": "unknown tool: " + p.Name}})
			return
		}
		// A failing tool is a result with isError set, not a protocol error:
		// the model reads the reason and can correct its next call.
		out, err := handler(r.Context(), sessionID, p.Input)
		text := ""
		if err != nil {
			text = err.Error()
		} else if b, mErr := json.Marshal(out); mErr != nil {
			http.Error(w, fmt.Sprintf("failed to encode tool output: %v", mErr), http.StatusInternalServerError)
			return
		} else {
			text = string(b)
		}
		result = map[string]any{
			"content": []map[string]any{{"type": "text", "text": text}},
			"isError": err != nil,
		}
	default:
		writeJSON(w, map[string]any{"jsonrpc": "2.0", "id": req.ID,
			"error": map[string]any{"code": -32601, "message": "method not found: " + req.Method}})
		return
	}
 
	writeJSON(w, map[string]any{"jsonrpc": "2.0", "id": req.ID, "result": result})
}
 
func writeJSON(w http.ResponseWriter, v any) {
	b, err := json.Marshal(v)
	if err != nil {
		http.Error(w, fmt.Sprintf("failed to encode response: %v", err), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	if _, err := w.Write(b); err != nil {
		log.Printf("mcp: write response: %v", err) // client went away; nothing left to send
	}
}

package mcp
 
import (
	"sort"
	"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 {
	terms := strings.Fields(strings.ToLower(query))
 
	type hit struct {
		i, score int
	}
	var hits []hit
	for i, entry := range idx.entries {
		score := 0
		for _, term := range terms {
			if strings.Contains(entry.searchText, term) {
				score++
			}
		}
		if score > 0 {
			hits = append(hits, hit{i, score})
		}
	}
 
	// Rank by descending relevance so the top-N are the best matches, not
	// just the first ones scanned. Without this sort, a weak early match
	// beats a strong later one and "top-N" is a lie.
	sort.SliceStable(hits, func(a, b int) bool { return hits[a].score > hits[b].score })
 
	if limit > len(hits) {
		limit = len(hits)
	}
	results := make([]map[string]any, 0, limit)
	for _, h := range hits[:limit] {
		entry := idx.entries[h.i]
		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, err := url.Parse(e.baseURL)
	if err != nil {
		return nil, fmt.Errorf("parse base URL: %w", err)
	}
	ref, err := url.Parse(path)
	if err != nil {
		return nil, fmt.Errorf("parse path: %w", err)
	}
	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 mcp
 
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
}

package mcp
 
import (
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"net/http"
)
 
// keys maps each API key to whether it is enabled; false means revoked.
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")
		valid := false
		for k, enabled := range keys {
			// Compare against every key, with no early exit, so response
			// time doesn't reveal which key matched.
			if subtle.ConstantTimeCompare([]byte(key), []byte(k)) == 1 && enabled {
				valid = true
			}
		}
		if !valid {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		// Session IDs reach metric labels, alerts and the audit log, so never
		// use the raw key. A truncated SHA-256 is stable but not reversible,
		// and prefixing it stops one client from landing in another's bucket.
		sum := sha256.Sum256([]byte(key))
		sessionID := hex.EncodeToString(sum[:8])
		if s := r.Header.Get("X-Session-ID"); s != "" {
			sessionID += ":" + s
		}
		r.Header.Set("X-Session-ID", sessionID)
		next.ServeHTTP(w, r)
	})
}

type SpecIndex struct {
    paths    map[string]*PathItem      // path template -> methods+ops
    baseURL  *url.URL
}
 
func LoadSpec(specPath, baseURL string) (*SpecIndex, error) {
    // The loader parses YAML or JSON and resolves $refs. gopkg.in/yaml.v3
    // decoding straight into openapi3.T leaves Paths empty and still validates.
    doc, err := openapi3.NewLoader().LoadFromFile(specPath)
    if err != nil {
        return nil, fmt.Errorf("load spec: %w", err)
    }
    if err := doc.Validate(context.Background()); err != nil {
        return nil, fmt.Errorf("invalid spec: %w", err)
    }
    if doc.Paths.Len() == 0 {
        return nil, errors.New("spec defines no paths")
    }
    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 !idx.matchesSpecPath(target.Path) {
        return nil, fmt.Errorf("path %q not in spec", target.Path)
    }
    return target, nil
}
 
// matchesSpecPath compares segment by segment, because spec keys are
// templates: /orders/123 must match /orders/{id}.
func (idx *SpecIndex) matchesSpecPath(p string) bool {
    got := strings.Split(p, "/")
    for tmpl := range idx.paths {
        want := strings.Split(tmpl, "/")
        if len(want) != len(got) {
            continue
        }
        match := true
        for i, seg := range want {
            param := strings.HasPrefix(seg, "{") && strings.HasSuffix(seg, "}")
            if (param && got[i] == "") || (!param && seg != got[i]) {
                match = false
                break
            }
        }
        if match {
            return true
        }
    }
    return false
}

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"`
    RawPath     string            `json:"raw_path,omitempty"` // model's input verbatim, kept for rejected paths
    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, error) {
    b, err := json.Marshal(e)
    if err != nil {
        return "", fmt.Errorf("marshal audit entry: %w", err)
    }
    return string(b), nil
}

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"},
    )
)