From 6fcae3042a6c1344c569976274ebc22f00e349ee Mon Sep 17 00:00:00 2001
From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com>
Date: Sat, 4 May 2024 00:38:01 +0100
Subject: [PATCH] feat(engine): Suppress stack traces using catch-all handler
 (#126)

---
 tracecat/api/app.py    | 13 +++++++++++--
 tracecat/runner/app.py |  9 +++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/tracecat/api/app.py b/tracecat/api/app.py
index d9d94656..c0c09f8d 100644
--- a/tracecat/api/app.py
+++ b/tracecat/api/app.py
@@ -5,10 +5,10 @@
 import polars as pl
 from aio_pika import Channel
 from aio_pika.pool import Pool
-from fastapi import Depends, FastAPI, HTTPException, status
+from fastapi import Depends, FastAPI, HTTPException, Request, status
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.params import Body
-from fastapi.responses import StreamingResponse
+from fastapi.responses import ORJSONResponse, StreamingResponse
 from sqlalchemy import Engine, or_
 from sqlalchemy.exc import NoResultFound
 from sqlmodel import Session, select
@@ -132,6 +132,15 @@ async def lifespan(app: FastAPI):
 )
 
 
+# Catch-all exception handler to prevent stack traces from leaking
+@app.exception_handler(Exception)
+async def custom_exception_handler(request: Request, exc: Exception):
+    return ORJSONResponse(
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        content={"message": "An unexpected error occurred. Please try again later."},
+    )
+
+
 @app.get("/")
 def root() -> dict[str, str]:
     return {"message": "Hello world. I am the API."}
diff --git a/tracecat/runner/app.py b/tracecat/runner/app.py
index 0df7bd80..d1089880 100644
--- a/tracecat/runner/app.py
+++ b/tracecat/runner/app.py
@@ -173,6 +173,15 @@ async def valid_workflow(workflow_id: str) -> str:
     return workflow_id
 
 
+# Catch-all exception handler to prevent stack traces from leaking
+@app.exception_handler(Exception)
+async def custom_exception_handler(request: Request, exc: Exception):
+    return ORJSONResponse(
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        content={"message": "An unexpected error occurred. Please try again later."},
+    )
+
+
 # Endpoints
 @app.get("/")
 def root() -> dict[str, str]: