improvement: handle azure workload identity authentication

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,16 @@ import { ProviderAPIConfig } from '../types';
 import {
   getAccessTokenFromEntraId,
   getAzureManagedIdentityToken,
+  getAzureWorkloadIdentityToken,
 } from './utils';
+import { env, getRuntimeKey } from 'hono/adapter';
 
 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 +41,40 @@ 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;
+
+      const runtime = getRuntimeKey();
+      if (
+        authorityHost &&
+        tenantId &&
+        clientId &&
+        federatedTokenFile &&
+        runtime === 'node'
+      ) {
+        const fs = await import('fs');
+        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;