DEV Community

Saivedant Hava
Saivedant Hava

Posted on

I security-audited my own AI gateway and added WASM plugin support. Here's what I found.

I've been building AegisFlow, an open-source AI gateway in Go. It sits between your apps and LLM
  providers (OpenAI, Anthropic, Ollama, etc.) and handles routing, security, rate limiting, and
  observability.

  Yesterday I sat down and did a proper security audit of the whole thing. Found more issues than I'd
  like to admit.

  The security stuff

  Timing attacks on API key validation. The tenant key lookup was using plain string comparison. An
  attacker could measure response times to progressively guess keys character by character. Switched
  to SHA-256 hashing both sides and comparing with subtle.ConstantTimeCompare. Also iterates all
  tenants on every check so there's no early-exit timing leak.

  inputHash := sha256.Sum256([]byte(apiKey))
  var match *TenantConfig
  for i := range c.Tenants {
      for _, key := range c.Tenants[i].APIKeys {
          keyHash := sha256.Sum256([]byte(key))
          if subtle.ConstantTimeCompare(inputHash[:], keyHash[:]) == 1 {
              match = &c.Tenants[i]
          }
      }
  }

  Admin panel was open by default. If you didn't set admin.token in your config, every admin endpoint
  was unauthenticated. Usage stats, provider health, tenant config, policy rules  all exposed. An
  attacker could read your exact jailbreak keywords and craft prompts around them. Now it blocks all
  data endpoints when no token is configured and logs a warning at startup.

  Rate limiter was failing open. When the limiter errored (Redis down, memory issue), it silently let
  requests through. Flipped it to return 503 instead.

  Other fixes: 10MB request body limit (was unlimited  one curl command could OOM the process), SSE
  injection in streaming responses (policy violation messages were string-concatenated into the event
  stream instead of JSON-serialized), and response cache was keyed by model+messages but not tenant
  ID, so tenant A could get tenant B's cached response.

  Jailbreak detection got real

  The keyword filter had 3 patterns. That's not a filter, that's a suggestion. Expanded to 25 covering
   the common techniques.

  But the bigger fix: added NFKC Unicode normalization. Before this, someone could write "іgnore
  previous instructions" with a Cyrillic і and walk right past the filter. Or add extra spaces:
  "ignore  previous  instructions". The normalizeText function now handles both:

  func normalizeText(s string) string {
      s = norm.NFKC.String(s)
      s = strings.Map(unicode.ToLower, s)
      s = multiSpaceRe.ReplaceAllString(s, " ")
      return strings.TrimSpace(s)
  }

  WASM plugin support

  This is the part I'm most excited about. You can now write custom policy filters in any language
  that compiles to WebAssembly and load them at runtime through config.

  policies:
    input:
      - name: "custom-toxicity"
        type: "wasm"
        action: "block"
        path: "plugins/toxicity.wasm"
        timeout: 100ms
        on_error: "block"

  The plugin ABI is four exported functions:

  ┌───────────────────────────────────────────────────┬───────────────────────────────────────────┐
                        Export                                      What it does                
  ├───────────────────────────────────────────────────┼───────────────────────────────────────────┤
   alloc(size)  ptr                                  Plugin allocates memory for host to write 
                                                       inputs                                   
  ├───────────────────────────────────────────────────┼───────────────────────────────────────────┤
   check(content_ptr, content_len, meta_ptr,          Main check. Returns 0 (allow) or 1        
   meta_len)  i32                                    (violation)                               
  ├───────────────────────────────────────────────────┼───────────────────────────────────────────┤
   get_result_ptr()  ptr                             Pointer to result JSON after violation    
  ├───────────────────────────────────────────────────┼───────────────────────────────────────────┤
   get_result_len()  i32                             Length of result JSON                     
  └───────────────────────────────────────────────────┴───────────────────────────────────────────┘

  The host passes request metadata as JSON (tenant ID, model, provider, phase) so plugins can make
  context-aware decisions. Config-level action always overrides whatever the plugin returns  a plugin
   can't escalate its own permissions.

  Runtime is wazero, pure Go, zero CGO. Each plugin gets its own timeout and on_error behavior (block
  or allow). If a plugin crashes or returns garbage JSON, the gateway handles it gracefully based on
  your config.

  Wrote a Go example plugin that blocks messages containing "forbidden":

  //go:wasmexport check
  func check(contentPtr *byte, contentLen uint32, metaPtr *byte, metaLen uint32) int32 {
      content := unsafe.String(contentPtr, contentLen)
      if strings.Contains(strings.ToLower(content), "forbidden") {
          result, _ := json.Marshal(map[string]string{
              "action":  "block",
              "message": "content contains forbidden word",
          })
          resultBuf = result
          return 1
      }
      return 0
  }

  Build it with GOOS=wasip1 GOARCH=wasm go build -o plugin.wasm . and point your config at it. No
  gateway recompilation.

  Live request feed

  Added a fifth page to the admin dashboard. Shows every request flowing through the gateway in real
  time  latency, status codes, cache hits, policy violations. Auto-refreshes every 2 seconds. Nothing
   fancy, but useful when you're debugging or demoing.

  What's next

  Phase 3: Kubernetes operator with CRDs, multi-region routing, and A/B testing for models. The goal
  is making this work across clusters, not just on a single node.

  Repo: github.com/saivedant169/AegisFlow

  57 tests, all passing with race detector. MIT licensed. If you're proxying LLM traffic and want
  something self-hosted, take a look.
Enter fullscreen mode Exit fullscreen mode

Top comments (0)