.Net: Updates to SessionsPythonPlugin (#11872)

### Motivation, Context and Description

This PR provides control over which domains requests can be sent to.
Additionally, it moves functionality that is common to all operations of
the plugin to one private method to remove duplication.

Contributes to:
#10070

Author: SergeyMenshykh

dotnet/src/Plugins/Plugins.Core/CodeInterpreter/SessionsPythonPlugin.cs

@@ -4,6 +4,7 @@
 using System.Collections.Generic;
 using System.ComponentModel;
 using System.IO;
+using System.Linq;
 using System.Net.Http;
 using System.Text;
 using System.Text.Json;
@@ -94,12 +95,9 @@ public async Task<string> ExecuteCodeAsync([Description("The valid Python code t
 
         var requestBody = new SessionsPythonCodeExecutionProperties(this._settings, code);
 
-        using var request = new HttpRequestMessage(HttpMethod.Post, $"{this._poolManagementEndpoint}/executions?identifier={this._settings.SessionId}&api-version={ApiVersion}")
-        {
-            Content = new StringContent(JsonSerializer.Serialize(requestBody), Encoding.UTF8, "application/json")
-        };
+        using var content = new StringContent(JsonSerializer.Serialize(requestBody), Encoding.UTF8, "application/json");
 
-        using var response = await httpClient.SendWithSuccessCheckAsync(request, CancellationToken.None).ConfigureAwait(false);
+        using var response = await this.SendAsync(httpClient, HttpMethod.Post, "executions", content).ConfigureAwait(false);
 
         var responseContent = JsonSerializer.Deserialize<JsonElement>(await response.Content.ReadAsStringAsync().ConfigureAwait(false));
 
@@ -139,15 +137,13 @@ public async Task<SessionsRemoteFileMetadata> UploadFileAsync(
         await this.AddHeadersAsync(httpClient).ConfigureAwait(false);
 
         using var fileContent = new ByteArrayContent(File.ReadAllBytes(localFilePath));
-        using var request = new HttpRequestMessage(HttpMethod.Post, $"{this._poolManagementEndpoint}files?identifier={this._settings.SessionId}&api-version={ApiVersion}")
+
+        using var multipartFormDataContent = new MultipartFormDataContent()
         {
-            Content = new MultipartFormDataContent
-            {
-                { fileContent, "file", remoteFileName },
-            }
+            { fileContent, "file", remoteFileName },
         };
 
-        using var response = await httpClient.SendWithSuccessCheckAsync(request, CancellationToken.None).ConfigureAwait(false);
+        using var response = await this.SendAsync(httpClient, HttpMethod.Post, "files", multipartFormDataContent).ConfigureAwait(false);
 
         var stringContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false);
 
@@ -172,9 +168,7 @@ public async Task<byte[]> DownloadFileAsync(
         using var httpClient = this._httpClientFactory.CreateClient();
         await this.AddHeadersAsync(httpClient).ConfigureAwait(false);
 
-        using var request = new HttpRequestMessage(HttpMethod.Get, $"{this._poolManagementEndpoint}/files/{Uri.EscapeDataString(remoteFileName)}/content?identifier={this._settings.SessionId}&api-version={ApiVersion}");
-
-        using var response = await httpClient.SendWithSuccessCheckAsync(request, CancellationToken.None).ConfigureAwait(false);
+        using var response = await this.SendAsync(httpClient, HttpMethod.Get, $"files/{Uri.EscapeDataString(remoteFileName)}/content").ConfigureAwait(false);
 
         var fileContent = await response.Content.ReadAsByteArrayAsync().ConfigureAwait(false);
 
@@ -205,9 +199,7 @@ public async Task<IReadOnlyList<SessionsRemoteFileMetadata>> ListFilesAsync()
         using var httpClient = this._httpClientFactory.CreateClient();
         await this.AddHeadersAsync(httpClient).ConfigureAwait(false);
 
-        using var request = new HttpRequestMessage(HttpMethod.Get, $"{this._poolManagementEndpoint}/files?identifier={this._settings.SessionId}&api-version={ApiVersion}");
-
-        using var response = await httpClient.SendWithSuccessCheckAsync(request, CancellationToken.None).ConfigureAwait(false);
+        using var response = await this.SendAsync(httpClient, HttpMethod.Get, "files").ConfigureAwait(false);
 
         var jsonElementResult = JsonSerializer.Deserialize<JsonElement>(await response.Content.ReadAsStringAsync().ConfigureAwait(false));
 
@@ -262,6 +254,36 @@ private async Task AddHeadersAsync(HttpClient httpClient)
         }
     }
 
+    /// <summary>
+    /// Sends an HTTP request to the specified path with the specified method and content.
+    /// </summary>
+    /// <param name="httpClient">The HTTP client to use.</param>
+    /// <param name="method">The HTTP method to use.</param>
+    /// <param name="path">The path to send the request to.</param>
+    /// <param name="httpContent">The content to send with the request.</param>
+    /// <returns>The HTTP response message.</returns>
+    private async Task<HttpResponseMessage> SendAsync(HttpClient httpClient, HttpMethod method, string path, HttpContent? httpContent = null)
+    {
+        // The query string is the same for all operations
+        var pathWithQueryString = $"{path}?identifier={this._settings.SessionId}&api-version={ApiVersion}";
+
+        var uri = new Uri(this._poolManagementEndpoint, pathWithQueryString);
+
+        // If a list of allowed domains has been provided, the host of the provided
+        // uri is checked to verify it is in the allowed domain list.
+        if (!this._settings.AllowedDomains?.Contains(uri.Host) ?? false)
+        {
+            throw new InvalidOperationException("Sending requests to the provided location is not allowed.");
+        }
+
+        using var request = new HttpRequestMessage(method, uri)
+        {
+            Content = httpContent,
+        };
+
+        return await httpClient.SendWithSuccessCheckAsync(request, CancellationToken.None).ConfigureAwait(false);
+    }
+
 #if NET
     [GeneratedRegex(@"^(\s|`)*(?i:python)?\s*", RegexOptions.ExplicitCapture)]
     private static partial Regex RemoveLeadingWhitespaceBackticksPython();

dotnet/src/Plugins/Plugins.Core/CodeInterpreter/SessionsPythonSettings.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
+using System.Collections.Generic;
 using System.ComponentModel;
 using System.Text.Json.Serialization;
 
@@ -23,6 +24,11 @@ public class SessionsPythonSettings
     [JsonIgnore]
     public Uri Endpoint { get; set; }
 
+    /// <summary>
+    /// List of allowed domains to download from.
+    /// </summary>
+    public IEnumerable<string>? AllowedDomains { get; set; }
+
     /// <summary>
     /// The session identifier.
     /// </summary>

dotnet/src/Plugins/Plugins.UnitTests/Core/SessionsPythonPluginTests.cs

@@ -291,6 +291,47 @@ public async Task ItShouldDownloadFileSavingInDiskAsync()
         Assert.Equal(responseContent, await File.ReadAllBytesAsync(downloadDiskPath));
     }
 
+    /// <summary>
+    /// Test the allowed domains for the endpoint.
+    /// </summary>
+    /// <remarks>
+    /// Considering that the functionality which verifies endpoints against the allowed domains is located in one private method,
+    /// and the method is reused for all operations of the plugin, we test it only for one operation (ListFilesAsync).
+    /// </remarks>
+    [Theory]
+    [InlineData("fake-test-host.io", "https://fake-test-host.io/subscriptions/123/rg/456/sps/test-pool", true)]
+    [InlineData("prod.fake-test-host.io", "https://prod.fake-test-host.io/subscriptions/123/rg/456/sps/test-pool", true)]
+    [InlineData("www.fake-test-host.io", "https://www.fake-test-host.io/subscriptions/123/rg/456/sps/test-pool", true)]
+    [InlineData("www.prod.fake-test-host.io", "https://www.prod.fake-test-host.io/subscriptions/123/rg/456/sps/test-pool", true)]
+    [InlineData("fake-test-host.io", "https://fake-test-host-1.io/subscriptions/123/rg/456/sps/test-pool", false)]
+    [InlineData("fake-test-host.io", "https://www.fake-test-host.io/subscriptions/123/rg/456/sps/test-pool", false)]
+    [InlineData("www.fake-test-host.io", "https://fake-test-host.io/subscriptions/123/rg/456/sps/test-pool", false)]
+    public async Task ItShouldRespectAllowedDomainsAsync(string allowedDomain, string actualEndpoint, bool isAllowed)
+    {
+        // Arrange
+        this._defaultSettings.AllowedDomains = [allowedDomain];
+        this._defaultSettings.Endpoint = new Uri(actualEndpoint);
+
+        this._messageHandlerStub.ResponseToReturn = new HttpResponseMessage(HttpStatusCode.OK)
+        {
+            Content = new StringContent(File.ReadAllText(ListFilesTestDataFilePath)),
+        };
+
+        var sut = new SessionsPythonPlugin(this._defaultSettings, this._httpClientFactory);
+
+        // Act
+#pragma warning disable CA1031 // Do not catch general exception types
+        try
+        {
+            await sut.ListFilesAsync();
+        }
+        catch when (!isAllowed)
+        {
+            // Ignore exception if the endpoint is not allowed since we expect it
+        }
+#pragma warning restore CA1031 // Do not catch general exception types
+    }
+
     public void Dispose()
     {
         this._httpClient.Dispose();