[ENH] CIP-2: Auth Providers Proposal

chromadb/api/fastapi.py

@@ -50,6 +50,8 @@ def __init__(self, system: System):
         self._session = requests.Session()
         if self._header is not None:
             self._session.headers.update(self._header)
+        if system.auth_provider is not None:
+            system.auth_provider.authenticate(self._session)
 
     @override
     def heartbeat(self) -> int:

chromadb/auth/README.md

@@ -0,0 +1,46 @@
+# Chroma Server Middlewares
+
+## Basic Auth
+
+This is very rudimentary security middleware that checks Authorization headers for basic auth credentials.
+
+The basic auth user and pass are configured on the server side using `CHROMA_SERVER_AUTH_PROVIDER_CONFIG`. Make sure to
+also define the auth provider `CHROMA_SERVER_AUTH_PROVIDER="chromadb.auth.BasicAuthServerProvider"`.
+
+### Usage
+
+Start the server:
+
+```bash
+CHROMA_SERVER_AUTH_PROVIDER="chromadb.BasicAuthServerProvider" \
+CHROMA_SERVER_AUTH_PROVIDER_CONFIG='{"username":"admin","password":"admin"}' \
+ALLOW_RESET=1 \
+IS_PERSISTENT=1 \
+uvicorn chromadb.app:app --workers 1 --host 0.0.0.0 --port 8000  --proxy-headers --log-config log_config.yml --reload
+```
+
+Test request with `curl`:
+
+```bash
+curl http://localhost:8000/api/v1 -v -H "Authorization: Basic YWRtaW46YWRtaW4="
+```
+
+Test with client side auth provider:
+
+```python
+import chromadb
+from chromadb import Settings
+
+client = chromadb.HttpClient(settings=Settings(chroma_client_auth_provider="chromadb.auth.BasicAuthClientProvider",
+                                               chroma_client_auth_provider_config={"username": "admin", "password": "admin"}))
+client.heartbeat()
+```
+
+Test with Http Client and basic auth header:
+
+```python
+import chromadb
+
+client = chromadb.HttpClient(host="localhost", port="8000", headers={"Authorization": "Basic YWRtaW46YWRtaW4="})
+client.heartbeat()
+```

chromadb/auth/__init__.py

@@ -0,0 +1,4 @@
+from chromadb.config import BasicAuthClientProvider, BasicAuthServerProvider
+
+# Re-export types from chromadb
+__all__ = ["BasicAuthClientProvider", "BasicAuthServerProvider"]

chromadb/config.py

@@ -1,12 +1,22 @@
-from pydantic import BaseSettings
-from typing import Optional, List, Any, Dict, TypeVar, Set, cast, Iterable, Type
-from typing_extensions import Literal
-from abc import ABC
+import base64
 import importlib
+import inspect
 import logging
-from overrides import EnforceOverrides, override
+import os
+from abc import ABC, abstractmethod
 from graphlib import TopologicalSorter
-import inspect
+from typing import Optional, List, Any, Dict, Set, Iterable, Union
+from typing import Type, TypeVar, cast
+
+import requests
+from overrides import override
+from overrides import overrides, EnforceOverrides
+from pydantic import BaseSettings, SecretStr
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.requests import Request
+from starlette.responses import Response, JSONResponse
+from starlette.types import ASGIApp
+from typing_extensions import Literal
 
 # The thin client will have a flag to control which implementations to use
 is_thin_client = False
@@ -15,10 +25,8 @@
 except ImportError:
     is_thin_client = False
 
-
 logger = logging.getLogger(__name__)
 
-
 LEGACY_ERROR = """\033[91mYou are using a deprecated configuration of Chroma.
 
 \033[94mIf you do not have data you wish to migrate, you only need to change how you construct
@@ -59,7 +67,25 @@
 }
 
 
-class Settings(BaseSettings):
+class ClientAuthProvider(ABC, EnforceOverrides):
+    def __init__(self, settings: "Settings") -> None:
+        self._settings = settings
+
+    @abstractmethod
+    def authenticate(self, session: requests.Session) -> None:
+        pass
+
+
+class ServerAuthProvider(ABC, EnforceOverrides):
+    def __init__(self, settings: "Settings") -> None:
+        self._settings = settings
+
+    @abstractmethod
+    def authenticate(self, request: Request) -> Union[Response, None]:
+        pass
+
+
+class Settings(BaseSettings):  # type: ignore
     environment: str = ""
 
     # Legacy config has to be kept around because pydantic will error on nonexisting keys
@@ -88,7 +114,13 @@ class Settings(BaseSettings):
     chroma_server_ssl_enabled: Optional[bool] = False
     chroma_server_grpc_port: Optional[str] = None
     chroma_server_cors_allow_origins: List[str] = []  # eg ["http://localhost:3000"]
+    # eg ["chromadb.api.fastapi.middlewares.auth.AuthMiddleware"]
+    chroma_server_middlewares: List[str] = []
 
+    chroma_server_auth_provider: Optional[str] = None
+    chroma_server_auth_provider_config: Optional[Union[str, Dict[str, Any]]] = None
+    chroma_client_auth_provider: Optional[str] = None
+    chroma_client_auth_provider_config: Optional[Union[str, Dict[str, Any]]] = None
     anonymized_telemetry: bool = True
 
     allow_reset: bool = False
@@ -106,7 +138,7 @@ def require(self, key: str) -> Any:
     def __getitem__(self, key: str) -> Any:
         val = getattr(self, key)
         # Error on legacy config values
-        if val in _legacy_config_values:
+        if isinstance(val, str) and val in _legacy_config_values:
             raise ValueError(LEGACY_ERROR)
         return val
 
@@ -158,7 +190,7 @@ def reset_state(self) -> None:
 
 class System(Component):
     settings: Settings
-
+    auth_provider: Optional["ClientAuthProvider"]
     _instances: Dict[Type[Component], Component]
 
     def __init__(self, settings: Settings):
@@ -169,7 +201,18 @@ def __init__(self, settings: Settings):
                     "Chroma is running in http-only client mode, and can only be run with 'chromadb.api.fastapi.FastAPI' as the chroma_api_impl. \
             see https://docs.trychroma.com/usage-guide?lang=py#using-the-python-http-only-client for more information."
                 )
-
+        if (
+            settings.chroma_client_auth_provider is not None
+            and settings.chroma_client_auth_provider.strip() != ""
+        ):
+            logger.debug(
+                f"Client Auth Provider: {settings.chroma_client_auth_provider}"
+            )
+            self.auth_provider = get_class(
+                settings.chroma_client_auth_provider, ClientAuthProvider
+            )(settings)
+        else:
+            self.auth_provider = None
         # Validate settings don't contain any legacy config values
         for key in _legacy_config_keys:
             if settings[key] is not None:
@@ -225,11 +268,141 @@ def stop(self) -> None:
     def reset_state(self) -> None:
         """Reset the state of this system and all constituents in reverse dependency order"""
         if not self.settings.allow_reset:
-            raise ValueError("Resetting is not allowed by this configuration (to enable it, set `allow_reset` to `True` in your Settings() or include `ALLOW_RESET=TRUE` in your environment variables)")
+            raise ValueError(
+                "Resetting is not allowed by this configuration (to enable it, set `allow_reset` to `True` in your Settings() or include `ALLOW_RESET=TRUE` in your environment variables)"
+            )
         for component in reversed(list(self.components())):
             component.reset_state()
 
 
+class BasicAuthClientProvider(ClientAuthProvider):
+    _basic_auth_token: SecretStr
+
+    def __init__(self, settings: "Settings") -> None:
+        super().__init__(settings)
+        self._settings = settings
+        if os.environ.get("CHROMA_CLIENT_AUTH_BASIC_USERNAME") and os.environ.get(
+            "CHROMA_CLIENT_AUTH_BASIC_PASSWORD"
+        ):
+            self._basic_auth_token = _create_token(
+                os.environ.get("CHROMA_CLIENT_AUTH_BASIC_USERNAME", ""),
+                os.environ.get("CHROMA_CLIENT_AUTH_BASIC_PASSWORD", ""),
+            )
+        elif isinstance(
+            self._settings.chroma_client_auth_provider_config, str
+        ) and os.path.exists(self._settings.chroma_client_auth_provider_config):
+            with open(self._settings.chroma_client_auth_provider_config) as f:
+                # read first line of file which should be user:password
+                _auth_data = f.readline().strip().split(":")
+                # validate auth data
+                if len(_auth_data) != 2:
+                    raise ValueError("Invalid auth data")
+                self._basic_auth_token = _create_token(_auth_data[0], _auth_data[1])
+        elif self._settings.chroma_client_auth_provider_config and isinstance(
+            self._settings.chroma_client_auth_provider_config, dict
+        ):
+            self._basic_auth_token = _create_token(
+                self._settings.chroma_client_auth_provider_config["username"],
+                self._settings.chroma_client_auth_provider_config["password"],
+            )
+        else:
+            raise ValueError("Basic auth credentials not found")
+
+    @overrides
+    def authenticate(self, session: requests.Session) -> None:
+        session.headers.update(
+            {"Authorization": f"Basic {self._basic_auth_token.get_secret_value()}"}
+        )
+
+
+class ChromaAuthMiddleware(BaseHTTPMiddleware):  # type: ignore
+    def __init__(self, app: ASGIApp, settings: "Settings") -> None:
+        super().__init__(app)
+        self._settings = settings
+        self._settings.require("chroma_server_auth_provider")
+        if settings.chroma_server_auth_provider:
+            _cls = get_class(settings.chroma_server_auth_provider, ServerAuthProvider)
+            logger.debug(f"Server Auth Provider: {_cls}")
+            self._auth_provider = _cls(settings)
+
+    async def dispatch(
+        self, request: Request, call_next: RequestResponseEndpoint
+    ) -> Response:
+        response = self._auth_provider.authenticate(request)
+        if response is not None:
+            return response
+        return await call_next(request)
+
+
+def _create_token(username: str, password: str) -> SecretStr:
+    return SecretStr(
+        base64.b64encode(f"{username}:{password}".encode("utf-8")).decode("utf-8")
+    )
+
+
+class BasicAuthServerProvider(ServerAuthProvider):
+    _basic_auth_token: SecretStr
+    _ignore_auth_paths: List[str] = ["/api/v1", "/api/v1/heartbeat", "/api/v1/version"]
+
+    def __init__(self, settings: "Settings") -> None:
+        super().__init__(settings)
+        self._settings = settings
+        self._basic_auth_token = SecretStr("")
+        if os.environ.get("CHROMA_SERVER_AUTH_BASIC_USERNAME") and os.environ.get(
+            "CHROMA_SERVER_AUTH_BASIC_PASSWORD"
+        ):
+            self._basic_auth_token = _create_token(
+                os.environ.get("CHROMA_SERVER_AUTH_BASIC_USERNAME", ""),
+                os.environ.get("CHROMA_SERVER_AUTH_BASIC_PASSWORD", ""),
+            )
+            self._ignore_auth_paths = os.environ.get(
+                "CHROMA_SERVER_AUTH_IGNORE_PATHS", ",".join(self._ignore_auth_paths)
+            ).split(",")
+        elif isinstance(
+            self._settings.chroma_server_auth_provider_config, str
+        ) and os.path.exists(self._settings.chroma_server_auth_provider_config):
+            with open(self._settings.chroma_server_auth_provider_config) as f:
+                # read first line of file which should be user:password
+                _auth_data = f.readline().strip().split(":")
+                # validate auth data
+                if len(_auth_data) != 2:
+                    raise ValueError("Invalid auth data")
+                self._basic_auth_token = _create_token(_auth_data[0], _auth_data[1])
+            self._ignore_auth_paths = os.environ.get(
+                "CHROMA_SERVER_AUTH_IGNORE_PATHS", ",".join(self._ignore_auth_paths)
+            ).split(",")
+        elif self._settings.chroma_server_auth_provider_config and isinstance(
+            self._settings.chroma_server_auth_provider_config, dict
+        ):
+            # encode the username and password base64
+            self._basic_auth_token = _create_token(
+                self._settings.chroma_server_auth_provider_config["username"],
+                self._settings.chroma_server_auth_provider_config["password"],
+            )
+            if "ignore_auth_paths" in self._settings.chroma_server_auth_provider_config:
+                self._ignore_auth_paths = (
+                    self._settings.chroma_server_auth_provider_config[
+                        "ignore_auth_paths"
+                    ]
+                )
+        else:
+            raise ValueError("Basic auth credentials not found")
+
+    @overrides
+    def authenticate(self, request: Request) -> Union[Response, None]:
+        auth_header = request.headers.get("Authorization", "").split()
+        # Check if the header exists and the token is correct
+        if request.url.path in self._ignore_auth_paths:
+            logger.debug(f"Skipping auth for path {request.url.path}")
+            return None
+        if (
+            len(auth_header) != 2
+            or auth_header[1] != self._basic_auth_token.get_secret_value()
+        ):
+            return JSONResponse({"error": "Unauthorized"}, status_code=401)
+        return None
+
+
 C = TypeVar("C")
 
 

chromadb/server/fastapi/__init__.py

@@ -8,7 +8,6 @@
 from fastapi import HTTPException, status
 from uuid import UUID
 
-
 import chromadb
 from chromadb.api.models.Collection import Collection
 from chromadb.api.types import GetResult, QueryResult
@@ -110,6 +109,13 @@ def __init__(self, settings: Settings):
             allow_origins=settings.chroma_server_cors_allow_origins,
             allow_methods=["*"],
         )
+        if (
+            settings.chroma_server_auth_provider is not None
+            and settings.chroma_server_auth_provider.strip() != ""
+        ):
+            self._app.add_middleware(
+                chromadb.config.ChromaAuthMiddleware, settings=settings
+            )
 
         self.router = ChromaAPIRouter()