Fix token comparison

app.py

@@ -1,3 +1,4 @@
+import hmac
 import os
 from typing import Union
 
@@ -13,6 +14,7 @@ app = FastAPI(
 )
 
 api_key = APIKeyQuery(name="token", auto_error=False)
+client_token: str = os.getenv("CLIENT_TOKEN", "")
 
 model = Model(
     os.getenv("MODEL_REPO_ID", ""),
@@ -26,7 +28,7 @@ async def validate_token(
 ):
     if token is None:
         raise HTTPException(status.HTTP_401_UNAUTHORIZED, "No token provided")
-    if token != os.getenv("CLIENT_TOKEN"):
+    if not hmac.compare_digest(token, client_token):
         raise HTTPException(status.HTTP_401_UNAUTHORIZED, "Invalid token")
     return token
 
@@ -42,6 +44,8 @@ async def calculate_embedding(
 ):
     try:
         image_content = await image.read()
+        if isinstance(image_content, str):
+            image_content = image_content.encode()
         pred = model.predict(model.preprocess(image_content))
         return EmbeddingResponse(embedding=pred[0].tolist())
     except Exception as e:
@@ -62,7 +66,11 @@ async def calculate_similarity_score(
 ):
     try:
         image1_content = await image1.read()
+        if isinstance(image1_content, str):
+            image1_content = image1_content.encode()
         image2_content = await image2.read()
+        if isinstance(image2_content, str):
+            image2_content = image2_content.encode()
         pred = model.predict(
             np.vstack(
                 [model.preprocess(image1_content), model.preprocess(image2_content)]

requirements.txt

@@ -8,4 +8,5 @@ python-multipart
 mediapipe
 pandas
 pytest
-python-dotenv
+python-dotenv
+bcrypt