improvement: handle azure workload identity authentication

So far the Azure OpenAI integration was handling authentication using Client ID / Client Secret and Managed identity using the IMDS endpoint which is deprecated in favor of Workload Identity (using the public OAuth2 endpoint of Entra ID).

This changeset aims at handling this new authentication type.

Note that this requires reading environment variables set by the Azure runtime onto the virtual machine / pod using a workload identity. It also needs to read a file on disk (containing an assertion to use to exchange against a JWT).

Author: Adrien Bestel

src/handlers/handlerUtils.ts

@@ -914,6 +914,8 @@ export function constructConfigFromRequestHeaders(
     azureAuthMode: requestHeaders[`x-${POWERED_BY}-azure-auth-mode`],
     azureManagedClientId:
       requestHeaders[`x-${POWERED_BY}-azure-managed-client-id`],
+    azureWorkloadClientId:
+      requestHeaders[`x-${POWERED_BY}-azure-workload-client-id`],
     azureEntraClientId: requestHeaders[`x-${POWERED_BY}-azure-entra-client-id`],
     azureEntraClientSecret:
       requestHeaders[`x-${POWERED_BY}-azure-entra-client-secret`],

src/middlewares/requestValidator/schema/config.ts

@@ -4,6 +4,7 @@ import {
   VALID_PROVIDERS,
   GOOGLE_VERTEX_AI,
   TRITON,
+  AZURE_OPEN_AI,
 } from '../../../globals';
 
 export const configSchema: any = z
@@ -107,6 +108,7 @@ export const configSchema: any = z
     openai_organization: z.string().optional(),
     // AzureOpenAI specific
     azure_model_name: z.string().optional(),
+    azure_auth_mode: z.string().optional(),
     strict_open_ai_compliance: z.boolean().optional(),
   })
   .refine(
@@ -123,6 +125,8 @@ export const configSchema: any = z
         (value.vertex_service_account_json || value.vertex_project_id);
       const hasAWSDetails =
         value.aws_access_key_id && value.aws_secret_access_key;
+      const hasAzureAuth =
+        value.provider == AZURE_OPEN_AI && value.azure_auth_mode;
 
       return (
         hasProviderApiKey ||
@@ -137,7 +141,8 @@ export const configSchema: any = z
         value.after_request_hooks ||
         value.before_request_hooks ||
         value.input_guardrails ||
-        value.output_guardrails
+        value.output_guardrails ||
+        hasAzureAuth
       );
     },
     {

src/providers/azure-openai/api.ts

@@ -2,14 +2,17 @@ import { ProviderAPIConfig } from '../types';
 import {
   getAccessTokenFromEntraId,
   getAzureManagedIdentityToken,
+  getAzureWorkloadIdentityToken,
 } from './utils';
+import { env } from 'hono/adapter';
+import fs from 'fs';
 
 const AzureOpenAIAPIConfig: ProviderAPIConfig = {
   getBaseURL: ({ providerOptions }) => {
     const { resourceName } = providerOptions;
     return `https://${resourceName}.openai.azure.com/openai`;
   },
-  headers: async ({ providerOptions, fn }) => {
+  headers: async ({ c, providerOptions, fn }) => {
     const { apiKey, azureAuthMode } = providerOptions;
 
     if (azureAuthMode === 'entra') {
@@ -39,6 +42,32 @@ const AzureOpenAIAPIConfig: ProviderAPIConfig = {
         Authorization: `Bearer ${accessToken}`,
       };
     }
+    if (azureAuthMode === 'workload') {
+      const { azureWorkloadClientId } = providerOptions;
+
+      const authorityHost = env(c).AZURE_AUTHORITY_HOST;
+      const tenantId = env(c).AZURE_TENANT_ID;
+      const clientId = azureWorkloadClientId || env(c).AZURE_CLIENT_ID;
+      const federatedTokenFile = env(c).AZURE_FEDERATED_TOKEN_FILE;
+
+      if (authorityHost && tenantId && clientId && federatedTokenFile) {
+        const federatedToken = fs.readFileSync(federatedTokenFile, 'utf8');
+
+        if (federatedToken) {
+          const scope = 'https://cognitiveservices.azure.com/.default';
+          const accessToken = await getAzureWorkloadIdentityToken(
+            authorityHost,
+            tenantId,
+            clientId,
+            federatedToken,
+            scope
+          );
+          return {
+            Authorization: `Bearer ${accessToken}`,
+          };
+        }
+      }
+    }
     const headersObj: Record<string, string> = {
       'api-key': `${apiKey}`,
     };

src/providers/azure-openai/utils.ts

@@ -63,6 +63,44 @@ export async function getAzureManagedIdentityToken(
   }
 }
 
+export async function getAzureWorkloadIdentityToken(
+  authorityHost: string,
+  tenantId: string,
+  clientId: string,
+  federatedToken: string,
+  scope = 'https://cognitiveservices.azure.com/.default'
+) {
+  try {
+    const url = `${authorityHost}/${tenantId}/oauth2/v2.0/token`;
+    const params = new URLSearchParams({
+      client_id: clientId,
+      client_assertion: federatedToken,
+      client_assertion_type:
+        'urn:ietf:params:oauth:client-assertion-type:jwt-bearer',
+      scope: scope,
+      grant_type: 'client_credentials',
+    });
+
+    const response = await fetch(url, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/x-www-form-urlencoded',
+      },
+      body: params,
+    });
+
+    if (!response.ok) {
+      const errorMessage = await response.text();
+      console.log({ message: `Error from Entra ${errorMessage}` });
+      return undefined;
+    }
+    const data: { access_token: string } = await response.json();
+    return data.access_token;
+  } catch (error) {
+    console.log(error);
+  }
+}
+
 export const AzureOpenAIFinetuneResponseTransform = (
   response: Response | ErrorResponse,
   responseStatus: number

src/types/requestBody.ts

@@ -61,6 +61,7 @@ export interface Options {
   adAuth?: string;
   azureAuthMode?: string;
   azureManagedClientId?: string;
+  azureWorkloadClientId?: string;
   azureEntraClientId?: string;
   azureEntraClientSecret?: string;
   azureEntraTenantId?: string;