feat: add MS Defender option

Author: ofirschwartz1

README.md

@@ -114,6 +114,20 @@ There are multiple ways to run this sample: locally using Ollama or Azure OpenAI
 
 See the [cost estimation](./docs/cost.md) details for running this sample on Azure.
 
+#### (Optional) Enable additional user context to Microsoft Defender for Cloud
+In case you have Microsoft Defender for Cloud protection on your Azure OpenAI resource and you want to have additional user context on the alerts, run this command:
+
+```bash
+azd env set MS_DEFENDER_ENABLED true
+```
+
+To customize the application name of the user context, run this command:
+```bash
+azd env set APPLICATION_NAME <your application name>
+```
+
+For more details, refer to the [Microsoft Defender for Cloud documentation](https://learn.microsoft.com/azure/defender-for-cloud/gain-end-user-context-ai).
+
 #### Deploy the sample
 
 1. Open a terminal and navigate to the root of the project.

packages/api/src/functions/chat-post.ts

@@ -10,6 +10,7 @@ import {
 import { AzureOpenAI, OpenAI } from 'openai';
 import 'dotenv/config';
 import { type ChatCompletionChunk } from 'openai/resources';
+import { getMsDefenderUserJson, type UserSecurityContext } from './security/ms-defender-utils.js';
 
 const azureOpenAiScope = 'https://cognitiveservices.azure.com/.default';
 const systemPrompt = `Assistant helps the user with cooking questions. Be brief in your answers. Answer only plain text, DO NOT use Markdown.
@@ -61,12 +62,19 @@ export async function postChat(
       throw new Error('No OpenAI API key or Azure OpenAI deployment provided');
     }
 
+    let userSecurityContext: UserSecurityContext | undefined;
+    if (process.env.MS_DEFENDER_ENABLED) {
+      userSecurityContext = getMsDefenderUserJson(request);
+    }
+
     if (stream) {
+      // @ts-expect-error user_security_context field is unsupported via openai client
       const responseStream = await openai.chat.completions.create({
         messages: [{ role: 'system', content: systemPrompt }, ...messages],
         temperature: 0.7,
         model,
         stream: true,
+        user_security_context: userSecurityContext,
       });
       const jsonStream = Readable.from(createJsonStream(responseStream));
 
@@ -83,6 +91,8 @@ export async function postChat(
       messages: [{ role: 'system', content: systemPrompt }, ...messages],
       temperature: 0.7,
       model,
+      // @ts-expect-error user_security_context field is unsupported via openai client
+      user_security_context: userSecurityContext,
     });
 
     return {

packages/api/src/functions/security/ms-defender-utils.ts

@@ -0,0 +1,97 @@
+import process from 'node:process';
+import { type HttpRequest } from '@azure/functions';
+
+/**
+ * Generates the user security context which contains several parameters that describe the AI application itself, and the end user that interacts with the AI application.
+ * These fields assist your security operations teams to investigate and mitigate security incidents by providing a comprehensive approach to protecting your AI applications.
+ * [Learn more](https://learn.microsoft.com/azure/defender-for-cloud/gain-end-user-context-ai) about protecting AI applications using Microsoft Defender for Cloud.
+ * @param request - The HTTP request
+ * @returns A json string which represents the user context
+ */
+export function getMsDefenderUserJson(request: HttpRequest): UserSecurityContext {
+  const sourceIp = getSourceIp(request);
+  const authenticatedUserDetails = getAuthenticatedUserDetails(request);
+
+  const userSecurityContext = {
+    end_user_tenant_id: authenticatedUserDetails.get('tenantId'),
+    end_user_id: authenticatedUserDetails.get('userId'),
+    source_ip: sourceIp,
+    application_name: process.env.APPLICATION_NAME,
+  } as UserSecurityContext;
+
+  return userSecurityContext;
+}
+
+/**
+ * Extracts user authentication details from the 'X-Ms-Client-Principal' header.
+ * This is based on [Azure Static Web App documentation](https://learn.microsoft.com/en-us/azure/static-web-apps/user-information)
+ * @param request - The HTTP request
+ * @returns A dictionary containing authentication details of the user
+ */
+function getAuthenticatedUserDetails(request: HttpRequest): Map<string, string> {
+  const authenticatedUserDetails = new Map<string, string>();
+  const principalHeader = request.headers.get('X-Ms-Client-Principal');
+  if (principalHeader === null) {
+    return authenticatedUserDetails;
+  }
+
+  const principal = parsePrincipal(principalHeader);
+  if (principal === null) {
+    return authenticatedUserDetails;
+  }
+
+  const tenantId = process.env.AZURE_TENANT_ID;
+  if (principal!.identityProvider === 'aad') {
+    // TODO: add only when userId represents actual IDP user id
+    // authenticatedUserDetails.set('userId', principal['userId']);
+    authenticatedUserDetails.set('tenantId', tenantId!);
+  }
+
+  return authenticatedUserDetails;
+}
+
+function parsePrincipal(principal: string | undefined): Principal | undefined {
+  if (principal === null) {
+    return undefined;
+  }
+
+  try {
+    return JSON.parse(Buffer.from(principal!, 'base64').toString('utf8')) as Principal;
+  } catch {
+    return undefined;
+  }
+}
+
+function getSourceIp(request: HttpRequest) {
+  const xForwardFor = request.headers.get('X-Forwarded-For');
+  if (xForwardFor === null) {
+    return null;
+  }
+
+  const ip = xForwardFor.split(',')[0];
+  const colonIndex = ip.lastIndexOf(':');
+
+  // Case of ipv4
+  if (colonIndex !== -1 && ip.indexOf(':') === colonIndex) {
+    return ip.slice(0, Math.max(0, colonIndex));
+  }
+
+  // Case of ipv6
+  if (ip.startsWith('[') && ip.includes(']:')) {
+    return ip.slice(0, Math.max(0, ip.indexOf(']:') + 1));
+  }
+
+  return ip;
+}
+
+type Principal = {
+  identityProvider: string;
+  userId: string;
+};
+
+export type UserSecurityContext = {
+  application_name: string;
+  end_user_id: string;
+  end_user_tenant_id: string;
+  source_ip: string;
+};