Pass Token Reliably

.gitignore

@@ -1,3 +1,4 @@
+client.sh
 TOKEN
 WriteToken
 diffvg/

rest_api.py

@@ -3,6 +3,7 @@ from pydantic import BaseModel
 from concurrent.futures import ThreadPoolExecutor
 import uuid, time
 import wai_service
+from typing import Optional
 
 app = FastAPI()
 
@@ -14,6 +15,7 @@ class InferenceRequest(BaseModel):
     optimized_letter: str
     font: str = "KaushanScript-Regular"
     seed: int = 0
+    token: Optional[str] = None
 
     class Config:
         extra = "allow" 
@@ -21,6 +23,9 @@ class InferenceRequest(BaseModel):
 @app.post("/generate")
 def enqueue(req: InferenceRequest):
     """Start a job and return its UUID immediately."""
+    if req.token is None:
+        raise HTTPException(422, "field 'token' is required")
+    payload = req.dict(exclude_none=True)
     job_id = str(uuid.uuid4())
     fut = EXEC.submit(wai_service.handler, dict(req.__dict__))
     JOBS[job_id] = (time.time(), fut)

wai_service.py

@@ -40,7 +40,8 @@ def _sanitize(cfg):
 
 
 def handler(payload: dict) -> str:
-
+    logging.info("handler received job (word=%s, has_token=%s)",
+                 payload.get("word"), "token" in payload)
     init_dir = os.path.join("code", "data", "init")
     if not os.path.isdir(init_dir):
         os.makedirs(init_dir, exist_ok=True)