huggginface auth fix and fixed model list

app/auth.py

@@ -1,7 +1,10 @@
 from fastapi import HTTPException, Header, Depends
 from fastapi.security import APIKeyHeader
 from typing import Optional
-from config import API_KEY # Import API_KEY directly for use in local validation
+from config import API_KEY, HUGGINGFACE_API_KEY, HUGGINGFACE # Import API_KEY, HUGGINGFACE_API_KEY, HUGGINGFACE
+import os
+import json
+import base64
 
 # Function to validate API key (moved from config.py)
 def validate_api_key(api_key_to_validate: str) -> bool:
@@ -18,28 +21,89 @@ def validate_api_key(api_key_to_validate: str) -> bool:
 api_key_header = APIKeyHeader(name="Authorization", auto_error=False)
 
 # Dependency for API key validation
-async def get_api_key(authorization: Optional[str] = Header(None)):
-    if authorization is None:
-        raise HTTPException(
-            status_code=401,
-            detail="Missing API key. Please include 'Authorization: Bearer YOUR_API_KEY' header."
-        )
-    
-    # Check if the header starts with "Bearer "
-    if not authorization.startswith("Bearer "):
-        raise HTTPException(
-            status_code=401,
-            detail="Invalid API key format. Use 'Authorization: Bearer YOUR_API_KEY'"
-        )
-    
-    # Extract the API key
-    api_key = authorization.replace("Bearer ", "")
-    
-    # Validate the API key
-    if not validate_api_key(api_key): # Call local validate_api_key
-        raise HTTPException(
-            status_code=401,
-            detail="Invalid API key"
-        )
-    
-    return api_key
+async def get_api_key(
+    authorization: Optional[str] = Header(None),
+    x_ip_token: Optional[str] = Header(None, alias="x-ip-token")
+):
+    # Check if Hugging Face auth is enabled
+    if HUGGINGFACE:  # Use HUGGINGFACE from config
+        if x_ip_token is None:
+            raise HTTPException(
+                status_code=401, # Unauthorised - because x-ip-token is missing
+                detail="Missing x-ip-token header. This header is required for Hugging Face authentication."
+            )
+
+        try:
+            # Decode JWT payload
+            parts = x_ip_token.split('.')
+            if len(parts) < 2:
+                raise ValueError("Invalid JWT format: Not enough parts to extract payload.")
+            payload_encoded = parts[1]
+            # Add padding if necessary, as Python's base64.urlsafe_b64decode requires it
+            payload_encoded += '=' * (-len(payload_encoded) % 4)
+            decoded_payload_bytes = base64.urlsafe_b64decode(payload_encoded)
+            payload = json.loads(decoded_payload_bytes.decode('utf-8'))
+        except ValueError as ve:
+            # Log server-side for debugging, but return a generic client error
+            print(f"ValueError processing x-ip-token: {ve}")
+            raise HTTPException(status_code=400, detail=f"Invalid JWT format in x-ip-token: {str(ve)}")
+        except (json.JSONDecodeError, base64.binascii.Error, UnicodeDecodeError) as e:
+            print(f"Error decoding/parsing x-ip-token payload: {e}")
+            raise HTTPException(status_code=400, detail=f"Malformed x-ip-token payload: {str(e)}")
+        except Exception as e: # Catch any other unexpected errors during token processing
+            print(f"Unexpected error processing x-ip-token: {e}")
+            raise HTTPException(status_code=500, detail="Internal error processing x-ip-token.")
+
+        error_in_token = payload.get("error")
+
+        if error_in_token == "InvalidAccessToken":
+            raise HTTPException(
+                status_code=403,
+                detail="Access denied: x-ip-token indicates 'InvalidAccessToken'."
+            )
+        elif error_in_token is None:  # JSON 'null' is Python's None
+            # If error is null, auth is successful. Now check if HUGGINGFACE_API_KEY is configured.
+            if not HUGGINGFACE_API_KEY: # Check if the key is empty or not set.
+                print("Security configuration error: HUGGINGFACE_API_KEY is not configured, but HUGGINGFACE mode is active and x-ip-token is valid.")
+                raise HTTPException(
+                    status_code=500,
+                    detail="Service security configuration incomplete: HuggingFace API Key not set."
+                )
+            print(f"HuggingFace authentication successful via x-ip-token (error field was null).")
+            return HUGGINGFACE_API_KEY # Return the configured HUGGINGFACE_API_KEY
+        else:
+            # Any other non-null, non-"InvalidAccessToken" value in 'error' field
+            raise HTTPException(
+                status_code=403,
+                detail=f"Access denied: x-ip-token indicates an unhandled error: '{error_in_token}'."
+            )
+    else:
+        # Fallback to Bearer token authentication if HUGGINGFACE env var is not "true"
+        if authorization is None:
+            detail_message = "Missing API key. Please include 'Authorization: Bearer YOUR_API_KEY' header."
+            # Optionally, provide a hint if the HUGGINGFACE env var exists but is not "true"
+            if os.getenv("HUGGINGFACE") is not None: # Check for existence, not value
+                 detail_message += " (Note: HUGGINGFACE mode with x-ip-token is not currently active)."
+            raise HTTPException(
+                status_code=401,
+                detail=detail_message
+            )
+        
+        # Check if the header starts with "Bearer "
+        if not authorization.startswith("Bearer "):
+            raise HTTPException(
+                status_code=401,
+                detail="Invalid API key format. Use 'Authorization: Bearer YOUR_API_KEY'"
+            )
+        
+        # Extract the API key
+        api_key = authorization.replace("Bearer ", "")
+        
+        # Validate the API key
+        if not validate_api_key(api_key): # Call local validate_api_key
+            raise HTTPException(
+                status_code=401,
+                detail="Invalid API key"
+            )
+        
+        return api_key

app/config.py

@@ -6,6 +6,10 @@ DEFAULT_PASSWORD = "123456"
 # Get password from environment variable or use default
 API_KEY = os.environ.get("API_KEY", DEFAULT_PASSWORD)
 
+# HuggingFace Authentication Settings
+HUGGINGFACE = os.environ.get("HUGGINGFACE", "false").lower() == "true"
+HUGGINGFACE_API_KEY = os.environ.get("HUGGINGFACE_API_KEY", "") # Default to empty string, auth logic will verify if HF_MODE is true and this key is needed
+
 # Directory for service account credential files
 CREDENTIALS_DIR = os.environ.get("CREDENTIALS_DIR", "/app/credentials")
 

app/routes/models_api.py

@@ -25,6 +25,7 @@ async def list_models(fastapi_request: Request, api_key: str = Depends(get_api_k
     raw_express_models = await get_vertex_express_models()
     
     candidate_model_ids = set()
+    raw_vertex_models_set = set(raw_vertex_models)  # For checking origin during prefixing
 
     if has_express_key:
         candidate_model_ids.update(raw_express_models)
@@ -57,8 +58,12 @@ async def list_models(fastapi_request: Request, api_key: str = Depends(get_api_k
     for original_model_id in sorted(list(all_model_ids)):
         current_display_prefix = ""
         # Only add PAY_PREFIX if the model is not already an EXPRESS model (which has its own prefix)
-        if not original_model_id.startswith("[EXPRESS]") and \
-           has_sa_creds and not has_express_key and EXPERIMENTAL_MARKER not in original_model_id:
+        # Apply PAY_PREFIX if SA creds are present, it's a model from raw_vertex_models,
+        # it's not experimental, and not already an EXPRESS model.
+        if has_sa_creds and \
+           original_model_id in raw_vertex_models_set and \
+           EXPERIMENTAL_MARKER not in original_model_id and \
+           not original_model_id.startswith("[EXPRESS]"):
             current_display_prefix = PAY_PREFIX
         
         base_display_id = f"{current_display_prefix}{original_model_id}"