feat: add optional swarm API key auth

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 All notable changes to the **LLM Swarm** project will be documented in this file.
 
+## [Unreleased]
+### Security
+- Added optional `SWARM_API_KEY` bearer authentication for tracker requests and node-to-node layer processing.
+
 ## [0.1.0] - 2026-04-28
 ### Added
 - **Initial Prototype:** Core P2P Pooled Compute architecture implemented in Python.

README.md

@@ -47,6 +47,7 @@ Point your node to the Leader's Public Tracker. The node will automatically prob
 **Example (Qwen3.5-27B POC):**
 ```bash
 export TRACKER_URL="https://remedy-unwatched-styling.ngrok-free.dev"
+export SWARM_API_KEY="ask-the-swarm-leader-for-this-shared-key"  # Optional, but recommended for public swarms
 export NODE_ID="Volunteer_Node_$(hostname)"
 export LAYER_START=5
 export LAYER_END=10
@@ -65,6 +66,11 @@ If you are the Swarm Leader (hosting the Tracker and initial layers), use the au
 
 This script starts the Tracker and your Entry Node (hosting Layers 0-4).
 
+For public swarms, set `SWARM_API_KEY` to the same shared secret on the tracker
+and every node before launching. When set, tracker control-plane requests and
+node-to-node layer processing require `Authorization: Bearer <SWARM_API_KEY>`;
+leaving it unset preserves unauthenticated local development behavior.
+
 ### Port Forwarding Details
 If you are hosting from home (e.g., behind an Orbi or Eero router), ensure your ports are reachable:
 
@@ -104,6 +110,11 @@ The easiest way to run LLM Swarm is using Docker.
 ```bash
 docker-compose up --build
 ```
+To enable shared-key authentication in Docker, export `SWARM_API_KEY` first:
+```bash
+export SWARM_API_KEY="replace-with-a-long-random-secret"
+docker-compose up --build
+```
 This command starts:
 - A **Tracker** on port `12345`
 - An **Entry Node** on port `9000` (serving layers 0-10)
@@ -163,7 +174,7 @@ The Entry Node will receive the request, orchestrate the inference across the gl
 
 LLM Swarm is an experimental prototype. We are looking for contributors to help with the following:
 
-- [ ] **Security Hardening:** Implement Swarm-wide API Keys for node-to-node authentication.
+- [x] **Security Hardening:** Implement Swarm-wide API Keys for tracker and node-to-node authentication via `SWARM_API_KEY`.
 - [ ] **Encrypted Communication:** Move from raw HTTP to `libp2p` with Noise/TLS encryption.
 - [ ] **Tensor Validation:** Implement checksums and basic verification to prevent malicious nodes from poisoning the inference.
 - [ ] **Compression:** Implement tensor quantization/compression for faster transmission over slow internet connections.

auth_utils.py

@@ -0,0 +1,37 @@
+import os
+from typing import Dict
+
+from fastapi import Header, HTTPException, status
+
+
+def configured_api_key() -> str:
+    """Return the optional shared API key used for swarm requests."""
+    return os.getenv("SWARM_API_KEY", "").strip()
+
+
+def auth_headers() -> Dict[str, str]:
+    """Return Authorization headers for outgoing swarm requests when configured."""
+    api_key = configured_api_key()
+    if not api_key:
+        return {}
+    return {"Authorization": f"Bearer {api_key}"}
+
+
+async def require_swarm_api_key(authorization: str = Header(default="")) -> None:
+    """Require a bearer token only when SWARM_API_KEY is configured.
+
+    Leaving SWARM_API_KEY unset preserves the prototype's unauthenticated local
+    development behavior. Setting it enables a shared secret for tracker and
+    node-to-node HTTP calls.
+    """
+    api_key = configured_api_key()
+    if not api_key:
+        return
+
+    expected = f"Bearer {api_key}"
+    if authorization != expected:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Missing or invalid swarm API key",
+            headers={"WWW-Authenticate": "Bearer"},
+        )

docker-compose.yml

@@ -5,6 +5,8 @@ services:
       dockerfile: Dockerfile.tracker
     ports:
       - "12345:12345"
+    environment:
+      - SWARM_API_KEY=${SWARM_API_KEY:-}
     networks:
       - swarm-net
 
@@ -20,6 +22,7 @@ services:
       - LAYER_START=0
       - LAYER_END=4
       - TRACKER_URL=http://tracker:12345
+      - SWARM_API_KEY=${SWARM_API_KEY:-}
       - MODEL_PATH=/app/models/Qwen_Qwen3.5-27B-Q4_K_M.gguf
     volumes:
       - ./Qwen_Qwen3.5-27B-Q4_K_M.gguf:/app/models/Qwen_Qwen3.5-27B-Q4_K_M.gguf
@@ -40,6 +43,7 @@ services:
       - LAYER_START=11
       - LAYER_END=20
       - TRACKER_URL=http://tracker:12345
+      - SWARM_API_KEY=${SWARM_API_KEY:-}
       - MODEL_PATH=/app/models/Qwen_Qwen3.5-27B-Q4_K_M.gguf
     volumes:
       - ./Qwen_Qwen3.5-27B-Q4_K_M.gguf:/app/models/Qwen_Qwen3.5-27B-Q4_K_M.gguf

swarm_node.py

@@ -3,14 +3,15 @@
 import httpx
 import asyncio
 import numpy as np
-from fastapi import FastAPI, Request
+from fastapi import Depends, FastAPI, Request
 from pydantic import BaseModel
 from typing import List, Optional
 import os
 import json
 import threading
 import time
 from llama_cpp import Llama
+from auth_utils import auth_headers, require_swarm_api_key
 
 app = FastAPI()
 
@@ -41,7 +42,7 @@ def heartbeat_task():
     """Periodically send heartbeat to the tracker."""
     while True:
         try:
-            httpx.post(f"{config.tracker_url}/heartbeat?node_id={config.node_id}")
+            httpx.post(f"{config.tracker_url}/heartbeat?node_id={config.node_id}", headers=auth_headers())
         except Exception:
             pass
         time.sleep(20)
@@ -83,7 +84,7 @@ async def probe_network():
     print(f"\n--- 🛰️  Swarm Network Probe ---")
     try:
         async with httpx.AsyncClient() as client:
-            resp = await client.get(f"{config.tracker_url}/coverage")
+            resp = await client.get(f"{config.tracker_url}/coverage", headers=auth_headers())
             if resp.status_code == 200:
                 data = resp.json()
                 total = data["total_layers"]
@@ -146,15 +147,15 @@ async def startup_event():
     # We try to register immediately so the tracker knows we are "Coming Soon"
     try:
         async with httpx.AsyncClient() as client:
-            await client.post(f"{config.tracker_url}/register", json=registration_data)
+            await client.post(f"{config.tracker_url}/register", json=registration_data, headers=auth_headers())
             print(f"[{config.node_id}] 📡 Registered with tracker at {config.tracker_url}")
     except Exception as e:
         print(f"[{config.node_id}] ❌ Failed to register with tracker: {e}")
 
     # 3. Start Heartbeat Thread
     threading.Thread(target=heartbeat_task, daemon=True).start()
 
-@app.post("/process_layers")
+@app.post("/process_layers", dependencies=[Depends(require_swarm_api_key)])
 async def process_layers(state: LayerState):
     print(f"[{config.node_id}] Swarm processing layers {config.layers[0]}-{config.layers[1]}")
 
@@ -173,7 +174,7 @@ async def process_layers(state: LayerState):
     if next_peer_url:
         print(f"[{config.node_id}] Forwarding to next peer: {next_peer_url}")
         async with httpx.AsyncClient() as client:
-            await client.post(f"{next_peer_url}/process_layers", json=new_state.model_dump())
+            await client.post(f"{next_peer_url}/process_layers", json=new_state.model_dump(), headers=auth_headers())
         return {"status": "forwarded", "next": next_peer_url}
     else:
         print(f"[{config.node_id}] No more peers. Pipeline complete.")
@@ -183,7 +184,7 @@ async def find_next_peer(model_id: str, target_layer: int) -> Optional[str]:
     """Query tracker for the next peer in the chain."""
     try:
         async with httpx.AsyncClient() as client:
-            resp = await client.get(f"{config.tracker_url}/find_peer", params={"model_id": model_id, "layer": target_layer})
+            resp = await client.get(f"{config.tracker_url}/find_peer", params={"model_id": model_id, "layer": target_layer}, headers=auth_headers())
             if resp.status_code == 200:
                 return resp.json()["url"]
     except Exception:
@@ -218,7 +219,7 @@ async def generate(prompt: str):
     next_peer_url = await find_next_peer(config.model_id, config.layers[1] + 1)
     if next_peer_url:
         async with httpx.AsyncClient() as client:
-            resp = await client.post(f"{next_peer_url}/process_layers", json=state.model_dump())
+            resp = await client.post(f"{next_peer_url}/process_layers", json=state.model_dump(), headers=auth_headers())
             return resp.json()
 
     return {"error": "No peers found to complete the chain"}

tests/test_api_key_auth.py

@@ -0,0 +1,83 @@
+import importlib
+from pathlib import Path
+import sys
+import types
+
+from fastapi.testclient import TestClient
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+
+def reload_module(name):
+    sys.modules.pop(name, None)
+    return importlib.import_module(name)
+
+
+def test_tracker_rejects_register_without_swarm_api_key(monkeypatch):
+    monkeypatch.setenv("SWARM_API_KEY", "test-secret")
+    tracker = reload_module("tracker")
+    tracker.peers_db.clear()
+
+    client = TestClient(tracker.app)
+    response = client.post(
+        "/register",
+        json={
+            "node_id": "node-a",
+            "url": "http://node-a:9000",
+            "model_id": "swarm-mesh-v1",
+            "layer_start": 0,
+            "layer_end": 4,
+        },
+    )
+
+    assert response.status_code == 401
+    assert tracker.peers_db == {}
+
+
+def test_tracker_accepts_register_with_swarm_api_key(monkeypatch):
+    monkeypatch.setenv("SWARM_API_KEY", "test-secret")
+    tracker = reload_module("tracker")
+    tracker.peers_db.clear()
+
+    client = TestClient(tracker.app)
+    response = client.post(
+        "/register",
+        headers={"Authorization": "Bearer test-secret"},
+        json={
+            "node_id": "node-a",
+            "url": "http://node-a:9000",
+            "model_id": "swarm-mesh-v1",
+            "layer_start": 0,
+            "layer_end": 4,
+        },
+    )
+
+    assert response.status_code == 200
+    assert "node-a" in tracker.peers_db
+
+
+def test_swarm_node_rejects_layer_processing_without_swarm_api_key(monkeypatch):
+    monkeypatch.setenv("SWARM_API_KEY", "test-secret")
+    fake_llama_cpp = types.ModuleType("llama_cpp")
+    fake_llama_cpp.Llama = object
+    monkeypatch.setitem(sys.modules, "llama_cpp", fake_llama_cpp)
+    swarm_node = reload_module("swarm_node")
+    swarm_node.config = swarm_node.NodeConfig(port=9000)
+
+    client = TestClient(swarm_node.app)
+    response = client.post(
+        "/process_layers",
+        json={
+            "model_id": "swarm-mesh-v1",
+            "layer_start": 0,
+            "layer_end": 4,
+            "tensor_data": [1.0, 2.0],
+            "shape": [1, 2],
+            "prompt_tokens": [1],
+            "current_token_index": 0,
+        },
+    )
+
+    assert response.status_code == 401

tracker.py

@@ -1,9 +1,10 @@
 import uvicorn
-from fastapi import FastAPI, HTTPException
+from fastapi import Depends, FastAPI, HTTPException
 from pydantic import BaseModel
 from typing import List, Dict, Optional
 import time
 import threading
+from auth_utils import require_swarm_api_key
 
 app = FastAPI()
 
@@ -33,25 +34,25 @@ def cleanup_stale_peers():
 def start_cleanup_thread():
     threading.Thread(target=cleanup_stale_peers, daemon=True).start()
 
-@app.post("/register")
+@app.post("/register", dependencies=[Depends(require_swarm_api_key)])
 async def register_peer(peer: PeerInfo):
     peer.last_seen = time.time()
     peers_db[peer.node_id] = peer
     print(f"[Tracker] Registered/Updated peer: {peer.node_id} ({peer.url}) layers {peer.layer_start}-{peer.layer_end}")
     return {"status": "registered"}
 
-@app.post("/heartbeat")
+@app.post("/heartbeat", dependencies=[Depends(require_swarm_api_key)])
 async def heartbeat(node_id: str):
     if node_id in peers_db:
         peers_db[node_id].last_seen = time.time()
         return {"status": "ok"}
     raise HTTPException(status_code=404, detail="Peer not found")
 
-@app.get("/peers")
+@app.get("/peers", dependencies=[Depends(require_swarm_api_key)])
 async def list_peers():
     return list(peers_db.values())
 
-@app.get("/coverage")
+@app.get("/coverage", dependencies=[Depends(require_swarm_api_key)])
 async def get_coverage():
     """Return a map of layer coverage to help nodes identify gaps."""
     coverage = {}
@@ -73,7 +74,7 @@ async def get_coverage():
         "active_peers": len(peers_db)
     }
 
-@app.get("/find_peer")
+@app.get("/find_peer", dependencies=[Depends(require_swarm_api_key)])
 async def find_peer(model_id: str, layer: int):
     # Find a peer that hosts the requested layer
     for peer in peers_db.values():