[ENH]: JS Client Static Token support (chroma-core#1114)

Refs: chroma-core#1083

## Description of changes

*Summarize the changes made by this PR.*
 - New functionality
- JS Client now supports Authorization, and X-Chroma-Token auths
supported
         - Tests and integration tests updated

## Test plan
*How are these changes tested?*

- [x] Tests pass locally `yarn test` for js

## Documentation Changes
TBD

bin/integration-test

@@ -9,15 +9,38 @@ function cleanup {
   rm server.htpasswd .chroma_env
 }
 
-function setup_basic_auth {
-  # Generate htpasswd file
-  docker run --rm --entrypoint htpasswd httpd:2 -Bbn admin admin > server.htpasswd
-  # Create .chroma_env file
-  cat <<EOF > .chroma_env
+function setup_auth {
+  local auth_type="$1"
+  case "$auth_type" in
+    basic)
+      docker run --rm --entrypoint htpasswd httpd:2 -Bbn admin admin > server.htpasswd
+      cat <<EOF > .chroma_env
 CHROMA_SERVER_AUTH_CREDENTIALS_FILE="/chroma/server.htpasswd"
-CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER='chromadb.auth.providers.HtpasswdFileServerAuthCredentialsProvider'
-CHROMA_SERVER_AUTH_PROVIDER='chromadb.auth.basic.BasicAuthServerProvider'
+CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER="chromadb.auth.providers.HtpasswdFileServerAuthCredentialsProvider"
+CHROMA_SERVER_AUTH_PROVIDER="chromadb.auth.basic.BasicAuthServerProvider"
 EOF
+      ;;
+    token)
+      cat <<EOF > .chroma_env
+CHROMA_SERVER_AUTH_CREDENTIALS="test-token"
+CHROMA_SERVER_AUTH_TOKEN_TRANSPORT_HEADER="AUTHORIZATION"
+CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER="chromadb.auth.token.TokenConfigServerAuthCredentialsProvider"
+CHROMA_SERVER_AUTH_PROVIDER="chromadb.auth.token.TokenAuthServerProvider"
+EOF
+      ;;
+    xtoken)
+          cat <<EOF > .chroma_env
+CHROMA_SERVER_AUTH_CREDENTIALS="test-token"
+CHROMA_SERVER_AUTH_TOKEN_TRANSPORT_HEADER="X_CHROMA_TOKEN"
+CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER="chromadb.auth.token.TokenConfigServerAuthCredentialsProvider"
+CHROMA_SERVER_AUTH_PROVIDER="chromadb.auth.token.TokenAuthServerProvider"
+EOF
+    ;;
+    *)
+      echo "Unknown auth type: $auth_type"
+      exit 1
+      ;;
+  esac
 }
 
 trap cleanup EXIT
@@ -28,19 +51,21 @@ export CHROMA_INTEGRATION_TEST_ONLY=1
 export CHROMA_API_IMPL=chromadb.api.fastapi.FastAPI
 export CHROMA_SERVER_HOST=localhost
 export CHROMA_SERVER_HTTP_PORT=8000
-
-echo testing: python -m pytest "$@"
-python -m pytest "$@"
+#
+#echo testing: python -m pytest "$@"
+#python -m pytest "$@"
 
 cd clients/js
 yarn
 yarn test:run
 docker compose down
 cd ../..
-echo "Testing auth"
-setup_basic_auth #this is specific to the auth type, later on we'll have other auth types
-cd clients/js
-# Start docker compose - this should be auth agnostic
-docker compose --env-file ../../.chroma_env -f ../../docker-compose.test-auth.yml up --build -d
-yarn test:run-auth
-cd ../..
+for auth_type in basic token xtoken; do
+  echo "Testing $auth_type auth"
+  setup_auth "$auth_type"
+  cd clients/js
+  docker compose --env-file ../../.chroma_env -f ../../docker-compose.test-auth.yml up --build -d
+  yarn test:run-auth-"$auth_type"
+  cd ../..
+  docker compose down
+done

clients/js/package.json

@@ -33,15 +33,22 @@
     "testnoauth": "run-s db:clean db:run test:runfull db:clean",
     "testauth": "run-s db:cleanauth db:run-auth test:runfull-authonly db:cleanauth",
     "test:set-port": "cross-env URL=localhost:8001",
-    "test:run": "jest --runInBand --testPathIgnorePatterns=test/auth.basic.test.ts",
-    "test:run-auth": "jest --runInBand --testPathPattern=test/auth.basic.test.ts",
-    "test:runfull": "PORT=8001 jest --runInBand --testPathIgnorePatterns=test/auth.basic.test.ts",
-    "test:runfull-authonly": "PORT=8001 jest --runInBand --testPathPattern=test/auth.basic.test.ts",
+    "test:run": "jest --runInBand --testPathIgnorePatterns=test/auth.*.test.ts",
+    "test:run-auth-basic": "jest --runInBand --testPathPattern=test/auth.basic.test.ts",
+    "test:run-auth-token": "jest --runInBand --testPathPattern=test/auth.token.test.ts",
+    "test:run-auth-xtoken": "XTOKEN_TEST=true jest --runInBand --testPathPattern=test/auth.token.test.ts",
+    "test:runfull": "PORT=8001 jest --runInBand --testPathIgnorePatterns=test/auth.*.test.ts",
+    "test:runfull-authonly": "run-s db:run-auth-basic test:runfull-authonly-basic db:clean db:run-auth-token test:runfull-authonly-token db:clean db:run-auth-xtoken test:runfull-authonly-xtoken db:clean",
+    "test:runfull-authonly-basic": "PORT=8001 jest --runInBand --testPathPattern=test/auth.basic.test.ts",
+    "test:runfull-authonly-token": "PORT=8001 jest --runInBand --testPathPattern=test/auth.token.test.ts",
+    "test:runfull-authonly-xtoken": "PORT=8001 XTOKEN_TEST=true jest --runInBand --testPathPattern=test/auth.token.test.ts",
     "test:update": "run-s db:clean db:run && jest --runInBand --updateSnapshot && run-s db:clean",
     "db:clean": "cd ../.. && CHROMA_PORT=8001 docker-compose -f docker-compose.test.yml down --volumes",
     "db:cleanauth": "cd ../.. && CHROMA_PORT=8001 docker-compose -f docker-compose.test-auth.yml down --volumes",
     "db:run": "cd ../.. && CHROMA_PORT=8001 docker-compose -f docker-compose.test.yml up --detach && sleep 5",
-    "db:run-auth": "cd ../.. && CHROMA_PORT=8001 docker-compose -f docker-compose.test-auth.yml up --detach && sleep 5",
+    "db:run-auth-basic": "cd ../.. && docker run --rm --entrypoint htpasswd httpd:2 -Bbn admin admin > server.htpasswd && echo \"CHROMA_SERVER_AUTH_CREDENTIALS_FILE=/chroma/server.htpasswd\\nCHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER=chromadb.auth.providers.HtpasswdFileServerAuthCredentialsProvider\\nCHROMA_SERVER_AUTH_PROVIDER=chromadb.auth.basic.BasicAuthServerProvider\\nCHROMA_PORT=8001\" > .chroma_env && docker-compose -f docker-compose.test-auth.yml --env-file ./.chroma_env up --detach && sleep 5",
+    "db:run-auth-token": "cd ../.. && echo \"CHROMA_SERVER_AUTH_CREDENTIALS=test-token\nCHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER=chromadb.auth.token.TokenConfigServerAuthCredentialsProvider\nCHROMA_SERVER_AUTH_PROVIDER=chromadb.auth.token.TokenAuthServerProvider\\nCHROMA_PORT=8001\" > .chroma_env && docker-compose -f docker-compose.test-auth.yml --env-file ./.chroma_env up --detach && sleep 5",
+    "db:run-auth-xtoken": "cd ../.. && echo \"CHROMA_SERVER_AUTH_TOKEN_TRANSPORT_HEADER=X_CHROMA_TOKEN\nCHROMA_SERVER_AUTH_CREDENTIALS=test-token\nCHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER=chromadb.auth.token.TokenConfigServerAuthCredentialsProvider\nCHROMA_SERVER_AUTH_PROVIDER=chromadb.auth.token.TokenAuthServerProvider\\nCHROMA_PORT=8001\" > .chroma_env && docker-compose -f docker-compose.test-auth.yml --env-file ./.chroma_env up --detach && sleep 5",
     "clean": "rimraf dist",
     "build": "run-s clean build:*",
     "build:main": "tsc -p tsconfig.json",

clients/js/src/auth.ts

@@ -118,7 +118,10 @@ class BasicAuthClientAuthProvider implements ClientAuthProvider {
      * @throws {Error} If neither credentials provider or text credentials are supplied.
      */
 
-    constructor(options: { textCredentials: any; credentialsProvider: ClientAuthCredentialsProvider<any> | undefined }) {
+    constructor(options: {
+        textCredentials: any;
+        credentialsProvider: ClientAuthCredentialsProvider<any> | undefined
+    }) {
         if (!options.credentialsProvider && !options.textCredentials) {
             throw new Error("Either credentials provider or text credentials must be supplied.");
         }
@@ -130,6 +133,85 @@ class BasicAuthClientAuthProvider implements ClientAuthProvider {
     }
 }
 
+class TokenAuthCredentials implements AbstractCredentials<SecretStr> {
+    private readonly credentials: SecretStr;
+
+    constructor(_creds: string) {
+        this.credentials = new SecretStr(_creds)
+    }
+
+    getCredentials(): SecretStr {
+        return this.credentials;
+    }
+}
+
+export class TokenCredentialsProvider implements ClientAuthCredentialsProvider<TokenAuthCredentials> {
+    private readonly credentials: TokenAuthCredentials;
+
+    constructor(_creds: string | undefined) {
+        if (_creds === undefined && !process.env.CHROMA_CLIENT_AUTH_CREDENTIALS) throw new Error("Credentials must be supplied via environment variable (CHROMA_CLIENT_AUTH_CREDENTIALS) or passed in as configuration.");
+        this.credentials = new TokenAuthCredentials((_creds ?? process.env.CHROMA_CLIENT_AUTH_CREDENTIALS) as string);
+    }
+
+    getCredentials(): TokenAuthCredentials {
+        return this.credentials;
+    }
+}
+
+export class TokenClientAuthProvider implements ClientAuthProvider {
+    private readonly credentialsProvider: ClientAuthCredentialsProvider<any>;
+    private readonly providerOptions: { headerType: TokenHeaderType };
+
+    constructor(options: {
+        textCredentials: any;
+        credentialsProvider: ClientAuthCredentialsProvider<any> | undefined,
+        providerOptions?: { headerType: TokenHeaderType }
+    }) {
+        if (!options.credentialsProvider && !options.textCredentials) {
+            throw new Error("Either credentials provider or text credentials must be supplied.");
+        }
+        if (options.providerOptions === undefined || !options.providerOptions.hasOwnProperty("headerType")) {
+            this.providerOptions = {headerType: "AUTHORIZATION"};
+        } else {
+            this.providerOptions = {headerType: options.providerOptions.headerType};
+        }
+        this.credentialsProvider = options.credentialsProvider || new TokenCredentialsProvider(options.textCredentials);
+    }
+
+    authenticate(): ClientAuthResponse {
+        return new TokenClientAuthResponse(this.credentialsProvider.getCredentials(), this.providerOptions.headerType);
+    }
+
+}
+
+
+type TokenHeaderType = 'AUTHORIZATION' | 'X_CHROMA_TOKEN';
+
+const TokenHeader: Record<TokenHeaderType, (value: string) => { key: string; value: string; }> = {
+    AUTHORIZATION: (value: string) => ({key: "Authorization", value: `Bearer ${value}`}),
+    X_CHROMA_TOKEN: (value: string) => ({key: "X-Chroma-Token", value: value})
+}
+
+class TokenClientAuthResponse implements ClientAuthResponse {
+    constructor(private readonly credentials: TokenAuthCredentials, private readonly headerType: TokenHeaderType = 'AUTHORIZATION') {
+    }
+
+    getAuthInfo(): { key: string; value: string } {
+        if (this.headerType === 'AUTHORIZATION') {
+            return TokenHeader.AUTHORIZATION(this.credentials.getCredentials().getSecret());
+        } else if (this.headerType === 'X_CHROMA_TOKEN') {
+            return TokenHeader.X_CHROMA_TOKEN(this.credentials.getCredentials().getSecret());
+        } else {
+            throw new Error("Invalid header type: " + this.headerType + ". Valid types are: " + Object.keys(TokenHeader).join(", "));
+        }
+    }
+
+    getAuthInfoType(): AuthInfoType {
+        return AuthInfoType.HEADER;
+    }
+}
+
+
 export class IsomorphicFetchClientAuthProtocolAdapter implements ClientAuthProtocolAdapter<RequestInit> {
     authProvider: ClientAuthProvider | undefined;
     wrapperApi: DefaultApi | undefined;
@@ -144,7 +226,17 @@ export class IsomorphicFetchClientAuthProtocolAdapter implements ClientAuthProto
 
         switch (authConfiguration.provider) {
             case "basic":
-                this.authProvider = new BasicAuthClientAuthProvider({textCredentials: authConfiguration.credentials, credentialsProvider: authConfiguration.credentialsProvider});
+                this.authProvider = new BasicAuthClientAuthProvider({
+                    textCredentials: authConfiguration.credentials,
+                    credentialsProvider: authConfiguration.credentialsProvider
+                });
+                break;
+            case "token":
+                this.authProvider = new TokenClientAuthProvider({
+                    textCredentials: authConfiguration.credentials,
+                    credentialsProvider: authConfiguration.credentialsProvider,
+                    providerOptions: authConfiguration.providerOptions
+                });
                 break;
             default:
                 this.authProvider = undefined;
@@ -225,4 +317,5 @@ export type AuthOptions = {
     credentialsProvider?: ClientAuthCredentialsProvider<any> | undefined,
     configProvider?: ClientAuthConfigurationProvider<any> | undefined,
     credentials?: any | undefined,
+    providerOptions?: any | undefined
 }

clients/js/test/auth.basic.test.ts

@@ -1,6 +1,6 @@
 import {expect, test} from "@jest/globals";
 import {ChromaClient} from "../src/ChromaClient";
-import chroma from "./initClientWithAuth";
+import {chromaBasic} from "./initClientWithAuth";
 import chromaNoAuth from "./initClient";
 
 test("it should get the version without auth needed", async () => {
@@ -22,12 +22,12 @@ test("it should raise error when non authenticated", async () => {
 });
 
 test('it should list collections', async () => {
-    await chroma.reset()
-    let collections = await chroma.listCollections()
+    await chromaBasic.reset()
+    let collections = await chromaBasic.listCollections()
     expect(collections).toBeDefined()
     expect(collections).toBeInstanceOf(Array)
     expect(collections.length).toBe(0)
-    const collection = await chroma.createCollection({name: "test"});
-    collections = await chroma.listCollections()
+    await chromaBasic.createCollection({name: "test"});
+    collections = await chromaBasic.listCollections()
     expect(collections.length).toBe(1)
 })

clients/js/test/auth.token.test.ts

@@ -0,0 +1,59 @@
+import {expect, test} from "@jest/globals";
+import {ChromaClient} from "../src/ChromaClient";
+import {chromaTokenDefault, chromaTokenBearer, chromaTokenXToken} from "./initClientWithAuth";
+import chromaNoAuth from "./initClient";
+
+test("it should get the version without auth needed", async () => {
+    const version = await chromaNoAuth.version();
+    expect(version).toBeDefined();
+    expect(version).toMatch(/^[0-9]+\.[0-9]+\.[0-9]+$/);
+});
+
+test("it should get the heartbeat without auth needed", async () => {
+    const heartbeat = await chromaNoAuth.heartbeat();
+    expect(heartbeat).toBeDefined();
+    expect(heartbeat).toBeGreaterThan(0);
+});
+
+test("it should raise error when non authenticated", async () => {
+    await expect(chromaNoAuth.listCollections()).rejects.toMatchObject({
+        status: 401
+    });
+});
+
+if (!process.env.XTOKEN_TEST) {
+    test('it should list collections with default token config', async () => {
+        await chromaTokenDefault.reset()
+        let collections = await chromaTokenDefault.listCollections()
+        expect(collections).toBeDefined()
+        expect(collections).toBeInstanceOf(Array)
+        expect(collections.length).toBe(0)
+        const collection = await chromaTokenDefault.createCollection({name: "test"});
+        collections = await chromaTokenDefault.listCollections()
+        expect(collections.length).toBe(1)
+    })
+
+    test('it should list collections with explicit bearer token config', async () => {
+        await chromaTokenBearer.reset()
+        let collections = await chromaTokenBearer.listCollections()
+        expect(collections).toBeDefined()
+        expect(collections).toBeInstanceOf(Array)
+        expect(collections.length).toBe(0)
+        const collection = await chromaTokenBearer.createCollection({name: "test"});
+        collections = await chromaTokenBearer.listCollections()
+        expect(collections.length).toBe(1)
+    })
+} else {
+
+    test('it should list collections with explicit x-token token config', async () => {
+        await chromaTokenXToken.reset()
+        let collections = await chromaTokenXToken.listCollections()
+        expect(collections).toBeDefined()
+        expect(collections).toBeInstanceOf(Array)
+        expect(collections.length).toBe(0)
+        const collection = await chromaTokenXToken.createCollection({name: "test"});
+        collections = await chromaTokenXToken.listCollections()
+        expect(collections.length).toBe(1)
+    })
+
+}

clients/js/test/initClientWithAuth.ts

@@ -2,6 +2,13 @@ import {ChromaClient} from "../src/ChromaClient";
 
 const PORT = process.env.PORT || "8000";
 const URL = "http://localhost:" + PORT;
-const chroma = new ChromaClient({path: URL, auth: {provider: "basic", credentials: "admin:admin"}});
-
-export default chroma;
+export const chromaBasic = new ChromaClient({path: URL, auth: {provider: "basic", credentials: "admin:admin"}});
+export const chromaTokenDefault = new ChromaClient({path: URL, auth: {provider: "token", credentials: "test-token"}});
+export const chromaTokenBearer = new ChromaClient({
+    path: URL,
+    auth: {provider: "token", credentials: "test-token", providerOptions: {headerType: "AUTHORIZATION"}}
+});
+export const chromaTokenXToken = new ChromaClient({
+    path: URL,
+    auth: {provider: "token", credentials: "test-token", providerOptions: {headerType: "X_CHROMA_TOKEN"}}
+});

docker-compose.test-auth.yml

@@ -17,9 +17,11 @@ services:
       - ANONYMIZED_TELEMETRY=False
       - ALLOW_RESET=True
       - IS_PERSISTENT=TRUE
-      - CHROMA_SERVER_AUTH_CREDENTIALS_FILE=/chroma/server.htpasswd
-      - CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER=chromadb.auth.providers.HtpasswdFileServerAuthCredentialsProvider
-      - CHROMA_SERVER_AUTH_PROVIDER=chromadb.auth.basic.BasicAuthServerProvider
+      - CHROMA_SERVER_AUTH_CREDENTIALS_FILE=${CHROMA_SERVER_AUTH_CREDENTIALS_FILE}
+      - CHROMA_SERVER_AUTH_CREDENTIALS=${CHROMA_SERVER_AUTH_CREDENTIALS}
+      - CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER=${CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER}
+      - CHROMA_SERVER_AUTH_PROVIDER=${CHROMA_SERVER_AUTH_PROVIDER}
+      - CHROMA_SERVER_AUTH_TOKEN_TRANSPORT_HEADER=${CHROMA_SERVER_AUTH_TOKEN_TRANSPORT_HEADER}
     ports:
       - ${CHROMA_PORT}:8000
     networks: