From 9b9324e5ed9f6d1c7937c6e5f0591eac793e4d9d Mon Sep 17 00:00:00 2001
From: alexmg <alex.meyergleaves@gmail.com>
Date: Mon, 18 Nov 2024 16:04:25 +1000
Subject: [PATCH] Add TrustedCertificateThumbprints option to AzureOpenAIConfig

---
 .../AzureOpenAI/AzureOpenAIConfig.cs          |  7 +++
 .../Internals/AzureOpenAIClientBuilder.cs     | 51 ++++++++++++++++++-
 service/Service/appsettings.json              | 10 +++-
 3 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/extensions/AzureOpenAI/AzureOpenAI/AzureOpenAIConfig.cs b/extensions/AzureOpenAI/AzureOpenAI/AzureOpenAIConfig.cs
index 993b65983..a50938d59 100644
--- a/extensions/AzureOpenAI/AzureOpenAI/AzureOpenAIConfig.cs
+++ b/extensions/AzureOpenAI/AzureOpenAI/AzureOpenAIConfig.cs
@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
+using System.Collections.Generic;
 using Azure.Core;
 
 #pragma warning disable IDE0130 // reduce number of "using" statements
@@ -83,6 +84,12 @@ public enum APITypes
     /// </summary>
     public int MaxRetries { get; set; } = 10;
 
+    /// <summary>
+    /// Thumbprints of certificates that should be trusted for HTTPS requests when SSL policy errors are detected.
+    /// This should only be used for local development when using a proxy to call the OpenAI endpoints.
+    /// </summary>
+    public HashSet<string> TrustedCertificateThumbprints { get; set; } = [];
+
     /// <summary>
     /// Set credentials manually from code
     /// </summary>
diff --git a/extensions/AzureOpenAI/AzureOpenAI/Internals/AzureOpenAIClientBuilder.cs b/extensions/AzureOpenAI/AzureOpenAI/Internals/AzureOpenAIClientBuilder.cs
index 9ccfdce8d..9a7932140 100644
--- a/extensions/AzureOpenAI/AzureOpenAI/Internals/AzureOpenAIClientBuilder.cs
+++ b/extensions/AzureOpenAI/AzureOpenAI/Internals/AzureOpenAIClientBuilder.cs
@@ -1,9 +1,10 @@
-// Copyright (c) Microsoft. All rights reserved.
+// Copyright (c) Microsoft. All rights reserved.
 
 using System;
 using System.ClientModel;
 using System.ClientModel.Primitives;
 using System.Net.Http;
+using System.Net.Security;
 using Azure.AI.OpenAI;
 using Azure.Identity;
 using Microsoft.Extensions.Logging;
@@ -32,6 +33,14 @@ internal static AzureOpenAIClient Build(
         // See https://github.com/Azure/azure-sdk-for-net/issues/46109
         options.AddPolicy(new SingleAuthorizationHeaderPolicy(), PipelinePosition.PerTry);
 
+        if (httpClient is null && config.TrustedCertificateThumbprints.Count > 0)
+        {
+#pragma warning disable CA2000
+            // False Positive: https://github.com/dotnet/roslyn-analyzers/issues/4636
+            httpClient = BuildHttpClientWithCustomCertificateValidation(config);
+#pragma warning restore CA2000
+        }
+
         if (httpClient is not null)
         {
             options.Transport = new HttpClientPipelineTransport(httpClient);
@@ -57,6 +66,46 @@ internal static AzureOpenAIClient Build(
                 throw new ConfigurationException($"Azure OpenAI: authentication type '{config.Auth:G}' is not supported");
         }
     }
+
+    private static HttpClient BuildHttpClientWithCustomCertificateValidation(AzureOpenAIConfig config)
+    {
+#pragma warning disable CA2000
+        // False Positive: https://github.com/dotnet/roslyn-analyzers/issues/4636
+        var handler = new HttpClientHandler();
+#pragma warning restore CA2000
+
+        handler.ClientCertificateOptions = ClientCertificateOption.Manual;
+        handler.ServerCertificateCustomValidationCallback =
+            (_, cert, _, policyErrors) =>
+            {
+                if (policyErrors == SslPolicyErrors.None)
+                {
+                    // Return true if there are no policy errors.
+                    return true;
+                }
+
+                // Attempt to get the thumbprint of the remote certificate.
+                string? remoteCertThumbprint = cert?.GetCertHashString();
+                if (remoteCertThumbprint is null)
+                {
+                    return false;
+                }
+
+                // Check if the thumbprint matches any of the trusted thumbprints.
+                foreach (string trustedThumbprint in config.TrustedCertificateThumbprints)
+                {
+                    if (string.Equals(remoteCertThumbprint, trustedThumbprint, StringComparison.OrdinalIgnoreCase))
+                    {
+                        return true;
+                    }
+                }
+
+                // Reject if no thumbprint matches.
+                return false;
+            };
+
+        return new HttpClient(handler);
+    }
 }
 
 // Use only for local debugging - Usage:
diff --git a/service/Service/appsettings.json b/service/Service/appsettings.json
index 55564bada..a9940b9b5 100644
--- a/service/Service/appsettings.json
+++ b/service/Service/appsettings.json
@@ -339,7 +339,10 @@
         // See https://learn.microsoft.com/azure/ai-services/openai/reference#embeddings
         "MaxEmbeddingBatchSize": 1,
         // How many times to retry in case of throttling.
-        "MaxRetries": 10
+        "MaxRetries": 10,
+        // Thumbprints of certificates that should be trusted for HTTPS requests when SSL policy errors are detected.
+        // This should only be used for local development when using a proxy to call the OpenAI endpoints.
+        "TrustedCertificateThumbprints": []
       },
       "AzureOpenAIText": {
         // "ApiKey" or "AzureIdentity"
@@ -355,7 +358,10 @@
         // "ChatCompletion" or "TextCompletion"
         "APIType": "ChatCompletion",
         // How many times to retry in case of throttling.
-        "MaxRetries": 10
+        "MaxRetries": 10,
+        // Thumbprints of certificates that should be trusted for HTTPS requests when SSL policy errors are detected.
+        // This should only be used for local development when using a proxy to call the OpenAI endpoints.
+        "TrustedCertificateThumbprints": []
       },
       "AzureQueues": {
         // "ConnectionString" or "AzureIdentity". For other options see <AzureQueueConfig>.