feat: Dual Authentication Mechanism API + NUC

.env.ci

@@ -7,8 +7,8 @@
 HF_TOKEN="Hugging Face API Token"
 
 ENVIRONMENT = "mainnet"
-NILAI_GUNICORN_WORKERS = 10
-AUTH_STRATEGY = "nuc"
+NILAI_GUNICORN_WORKERS = 2
+AUTH_STRATEGY = "api_key"
 
 # The domain name of the server
 # - It must be written as "localhost" or "test.nilai.nillion"

.github/workflows/ci.yml

@@ -63,7 +63,7 @@ jobs:
           github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
           runners-per-machine: 1
           number-of-machines: 1
-          ec2-image-id: ami-0ac221d824dd88706
+          ec2-image-id: ami-0174a246556e8750b
           ec2-instance-type: g4dn.xlarge
           subnet-id: subnet-0ec4c353621eabae2
           security-group-id: sg-03ee5c56e1f467aa0
@@ -94,7 +94,7 @@ jobs:
           cache-dependency-glob: "**/pyproject.toml"
       - name: Install dependencies
         run: |
-          apt-get update && apt-get install curl git pkg-config automake file -y
+          apt-get update && apt-get install curl git pkg-config automake file python3.12-dev -y
           uv sync
 
       - name: Build vllm
@@ -103,7 +103,7 @@ jobs:
       - name: Build attestation
         run: docker build -t nillion/nilai-attestation:latest -f docker/attestation.Dockerfile .
 
-      - name: Build nilal API
+      - name: Build nilai API
         run: docker build -t nillion/nilai-api:latest -f docker/api.Dockerfile --target nilai --platform linux/amd64 .
 
       - name: Create .env
@@ -124,12 +124,22 @@ jobs:
       - name: Wait for services to be healthy
         run: bash scripts/wait_for_ci_services.sh
 
-      - name: Run E2E tests
+      - name: Run E2E tests for NUC
         run: |
           set -e
           export ENVIRONMENT=ci
           uv run pytest -v tests/e2e
 
+      - name: Run E2E tests for API Key
+        run: |
+          set -e
+          # Create a user with a rate limit of 1000 requests per minute, hour, and day
+          export AUTH_TOKEN=$(docker exec nilai-api uv run src/nilai_api/commands/add_user.py --name test1 --ratelimit-minute 1000 --ratelimit-hour 1000 --ratelimit-day 1000 | jq ".apikey" -r)
+          export ENVIRONMENT=ci
+          # Set the environment variable for the API key
+          export AUTH_STRATEGY=api_key
+          uv run pytest -v tests/e2e
+
       - name: Stop Services
         run: |
           docker-compose -f docker-compose.yml \

caddy/Caddyfile

@@ -2,7 +2,7 @@
 	tls {
 		protocols tls1.2 tls1.3
 	}
- }
+}
 
 {$NILAI_SERVER_DOMAIN} {
 	import ssl_config
@@ -12,12 +12,12 @@
 		reverse_proxy grafana:3000
 	}
 
-	handle_path /grafana {
-		uri strip_prefix /grafana
-		reverse_proxy grafana:3000
+	handle_path /nuc/* {
+		uri strip_prefix /nuc
+		reverse_proxy nuc-api:8080
 	}
 
 	handle {
 		reverse_proxy api:8080
 	}
- }
+}

docker-compose.dev.yml

@@ -6,6 +6,17 @@ services:
     volumes:
       - ./nilai-api/:/app/nilai-api/
       - ./packages/:/app/packages/
+      - ./nilai-auth/nuc-helpers/:/app/nilai-auth/nuc-helpers/
+    networks:
+      - nilauth
+  nuc-api:
+    platform: linux/amd64 # for macOS to force running on Rosetta 2
+    ports:
+      - "8088:8080"
+    volumes:
+      - ./nilai-api/:/app/nilai-api/
+      - ./packages/:/app/packages/
+      - ./nilai-auth/nuc-helpers/:/app/nilai-auth/nuc-helpers/
     networks:
       - nilauth
   attestation:
@@ -20,6 +31,9 @@ services:
   postgres:
     ports:
       - "5432:5432"
+  nuc-postgres:
+    ports:
+      - "5433:5432"
   grafana:
     ports:
       - "3000:3000"

docker-compose.yml

@@ -40,6 +40,25 @@ services:
       retries: 5
       start_period: 10s
       timeout: 10s
+
+  nuc-postgres:
+    image: postgres:16
+    container_name: nuc-postgres
+    restart: always
+    env_file:
+      - .env
+    environment:
+      - POSTGRES_HOST=nuc-postgres
+    networks:
+      - frontend_net
+    volumes:
+      - nuc_postgres_data:/var/lib/postgresql/data
+    healthcheck:
+      test: ["CMD", "sh", "-c", "pg_isready -U ${POSTGRES_USER} -d ${POSTGRES_DB} -h localhost"]
+      interval: 30s
+      retries: 5
+      start_period: 10s
+      timeout: 10s
   prometheus:
     container_name: prometheus
     image: prom/prometheus:v3.1.0
@@ -128,6 +147,35 @@ services:
       retries: 3
       start_period: 15s
       timeout: 10s
+  nuc-api:
+    container_name: nilai-nuc-api
+    image: nillion/nilai-api:latest
+    privileged: true
+    volumes:
+      - /dev/sev-guest:/dev/sev-guest # for AMD SEV
+    depends_on:
+      etcd:
+        condition: service_healthy
+      nuc-postgres:
+        condition: service_healthy
+      api:
+        condition: service_healthy
+    restart: unless-stopped
+    networks:
+      - frontend_net
+      - backend_net
+      - proxy_net
+    env_file:
+      - .env
+    environment:
+      - AUTH_STRATEGY=nuc # Overwrite the default strategy
+      - POSTGRES_HOST=nuc-postgres
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:8080/v1/health"]
+      interval: 30s
+      retries: 3
+      start_period: 15s
+      timeout: 10s
   attestation:
     image: nillion/nilai-attestation:latest
     restart: unless-stopped
@@ -171,3 +219,4 @@ networks:
 
 volumes:
   postgres_data:
+  nuc_postgres_data:

nilai-api/alembic/env.py

@@ -9,8 +9,13 @@
 
 from alembic import context
 from nilai_api.db import Base
+from nilai_api.db.users import UserModel
+from nilai_api.db.logs import QueryLog
 import nilai_api.config as nilai_config
 
+# If we don't use the models, they remain unused, and the migration fails
+# This is a workaround to ensure the models are loaded
+_, _ = UserModel, QueryLog
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
 config = context.config

nilai-api/alembic/versions/b9642f45db1d_fix_changed_to_timestamps_with_timezone.py

@@ -0,0 +1,75 @@
+"""fix: changed to timestamps with timezone
+
+Revision ID: b9642f45db1d
+Revises: ca76e3ebe6ee
+Create Date: 2025-05-13 09:47:30.506632
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision: str = "b9642f45db1d"
+down_revision: Union[str, None] = "ca76e3ebe6ee"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column(
+        "query_logs",
+        "query_timestamp",
+        existing_type=postgresql.TIMESTAMP(),
+        type_=sa.DateTime(timezone=True),
+        existing_nullable=False,
+        existing_server_default=sa.text("now()"),
+    )
+    op.alter_column(
+        "users",
+        "signup_date",
+        existing_type=postgresql.TIMESTAMP(),
+        type_=sa.DateTime(timezone=True),
+        existing_nullable=False,
+        existing_server_default=sa.text("now()"),
+    )
+    op.alter_column(
+        "users",
+        "last_activity",
+        existing_type=postgresql.TIMESTAMP(),
+        type_=sa.DateTime(timezone=True),
+        existing_nullable=True,
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column(
+        "users",
+        "last_activity",
+        existing_type=sa.DateTime(timezone=True),
+        type_=postgresql.TIMESTAMP(),
+        existing_nullable=True,
+    )
+    op.alter_column(
+        "users",
+        "signup_date",
+        existing_type=sa.DateTime(timezone=True),
+        type_=postgresql.TIMESTAMP(),
+        existing_nullable=False,
+        existing_server_default=sa.text("now()"),
+    )
+    op.alter_column(
+        "query_logs",
+        "query_timestamp",
+        existing_type=sa.DateTime(timezone=True),
+        type_=postgresql.TIMESTAMP(),
+        existing_nullable=False,
+        existing_server_default=sa.text("now()"),
+    )
+    # ### end Alembic commands ###

nilai-api/pyproject.toml

@@ -32,6 +32,7 @@ dependencies = [
     "web3>=7.8.0",
     "click>=8.1.8",
     "nuc",
+    "nuc-helpers",
 ]
 
 
@@ -42,3 +43,4 @@ build-backend = "hatchling.build"
 [tool.uv.sources]
 nilai-common = { workspace = true }
 nuc = { git = "https://github.com/NillionNetwork/nuc-py.git" }
+nuc-helpers = { workspace = true }

nilai-api/src/nilai_api/app.py

@@ -3,7 +3,7 @@
 
 from prometheus_fastapi_instrumentator import Instrumentator
 from fastapi import Depends, FastAPI
-from nilai_api.auth import get_user
+from nilai_api.auth import get_auth_info
 from nilai_api.rate_limiting import setup_redis_conn
 from nilai_api.routers import private, public
 from nilai_api import config
@@ -86,7 +86,7 @@ async def lifespan(app: FastAPI):
 
 
 app.include_router(public.router)
-app.include_router(private.router, dependencies=[Depends(get_user)])
+app.include_router(private.router, dependencies=[Depends(get_auth_info)])
 
 origins = [
     "https://docs.nillion.com",  # TODO: When users want to connect from browser

nilai-api/src/nilai_api/auth/__init__.py

@@ -1,49 +1,53 @@
-from fastapi import HTTPException, Security, status
+from fastapi import Security
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 
 from logging import getLogger
 
 from nilai_api import config
-from nilai_api.auth.jwt import validate_jwt
-from nilai_api.db.users import UserManager, UserModel
-from nilai_api.auth.strategies import STRATEGIES
+from nilai_api.db.users import UserManager
+from nilai_api.auth.strategies import AuthenticationStrategy
 
 from nuc.validate import ValidationException
+from nuc_helpers.usage import UsageLimitError
+
+from nilai_api.auth.common import (
+    AuthenticationInfo,
+    AuthenticationError,
+    TokenRateLimit,
+    TokenRateLimits,
+)
 
 logger = getLogger(__name__)
 bearer_scheme = HTTPBearer()
 
 
-class AuthenticationError(HTTPException):
-    def __init__(self, detail: str):
-        super().__init__(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=detail,
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-
-
-async def get_user(
+async def get_auth_info(
     credentials: HTTPAuthorizationCredentials = Security(bearer_scheme),
-) -> UserModel:
+) -> AuthenticationInfo:
     try:
-        if config.AUTH_STRATEGY not in STRATEGIES:
-            logger.error(f"Invalid auth strategy: {config.AUTH_STRATEGY}")
-            raise AuthenticationError("Server misconfiguration: invalid auth strategy")
-
-        user = await STRATEGIES[config.AUTH_STRATEGY](credentials.credentials)
-        if not user:
-            raise AuthenticationError("Missing or invalid API key")
-        await UserManager.update_last_activity(userid=user.userid)
-        return user
+        strategy_name: str = config.AUTH_STRATEGY.upper()
+
+        try:
+            strategy = AuthenticationStrategy[strategy_name]
+        except KeyError:  # If the strategy is not found, we raise an error
+            logger.error(f"Invalid auth strategy: {strategy_name}")
+            raise AuthenticationError(
+                f"Server misconfiguration: invalid auth strategy: {strategy_name}"
+            )
+
+        auth_info = await strategy(credentials.credentials)
+        await UserManager.update_last_activity(userid=auth_info.user.userid)
+        return auth_info
     except AuthenticationError as e:
         raise e
     except ValueError as e:
         raise AuthenticationError(detail="Authentication failed: " + str(e))
     except ValidationException as e:
         raise AuthenticationError(detail="NUC validation failed: " + str(e))
+    except UsageLimitError as e:
+        raise AuthenticationError(detail="Usage limit error: " + str(e))
     except Exception as e:
         raise AuthenticationError(detail="Unexpected authentication error: " + str(e))
 
 
-__all__ = ["get_user", "validate_jwt"]
+__all__ = ["get_auth_info", "AuthenticationInfo", "TokenRateLimits", "TokenRateLimit"]