From 2b66cbe69937b1ee29541c877b59ccea9fe7f304 Mon Sep 17 00:00:00 2001
From: Pavel Iakovenko <paveliak@users.noreply.github.com>
Date: Thu, 31 Aug 2023 17:35:50 -0400
Subject: [PATCH] Delegating handler for Http redirects (#2814)

* Delegating handler for Http redirects

* Renamed varible name
---
 src/Runner.Common/RedirectMessageHandler.cs | 73 +++++++++++++++++++++
 src/Runner.Worker/JobRunner.cs              |  9 ++-
 2 files changed, 81 insertions(+), 1 deletion(-)
 create mode 100644 src/Runner.Common/RedirectMessageHandler.cs

diff --git a/src/Runner.Common/RedirectMessageHandler.cs b/src/Runner.Common/RedirectMessageHandler.cs
new file mode 100644
index 000000000..81d216c01
--- /dev/null
+++ b/src/Runner.Common/RedirectMessageHandler.cs
@@ -0,0 +1,73 @@
+using System;
+using System.ComponentModel;
+using System.Net;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+using GitHub.Runner.Sdk;
+using GitHub.Services.Common;
+
+namespace GitHub.Runner.Common
+{
+    /// <summary>
+    /// Handles redirects for Http requests
+    /// </summary>
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    public class RedirectMessageHandler : DelegatingHandler
+    {
+        public RedirectMessageHandler(ITraceWriter trace)
+        {
+            Trace = trace;
+        }
+
+        protected override async Task<HttpResponseMessage> SendAsync(
+            HttpRequestMessage request,
+            CancellationToken cancellationToken)
+        {
+            HttpResponseMessage response = await base.SendAsync(request, cancellationToken).ConfigureAwait(false);
+
+            if (response != null &&
+                IsRedirect(response.StatusCode) &&
+                response.Headers.Location != null)
+            {
+                Trace.Info($"Redirecting to '{response.Headers.Location}'.");
+
+                request = await CloneAsync(request, response.Headers.Location).ConfigureAwait(false);
+
+                response.Dispose();
+
+                response = await base.SendAsync(request, cancellationToken).ConfigureAwait(false);
+            }
+
+            return response;
+        }
+
+        private static bool IsRedirect(HttpStatusCode statusCode)
+        {
+            return (int)statusCode >= 300 && (int)statusCode < 400;
+        }
+
+        private static async Task<HttpRequestMessage> CloneAsync(HttpRequestMessage request, Uri requestUri)
+        {
+            var clone = new HttpRequestMessage(request.Method, requestUri)
+            {
+                Version = request.Version
+            };
+
+            request.Headers.ForEach(header => clone.Headers.TryAddWithoutValidation(header.Key, header.Value));
+
+            request.Options.ForEach(option => clone.Options.Set(new HttpRequestOptionsKey<object>(option.Key), option.Value));
+
+            if (request.Content != null)
+            {
+                clone.Content = new ByteArrayContent(await request.Content.ReadAsByteArrayAsync().ConfigureAwait(false));
+
+                request.Content.Headers.ForEach(header => clone.Content.Headers.TryAddWithoutValidation(header.Key, header.Value));
+            }
+
+            return clone;
+        }
+
+        private readonly ITraceWriter Trace;
+    }
+}
diff --git a/src/Runner.Worker/JobRunner.cs b/src/Runner.Worker/JobRunner.cs
index 5e50fd680..88ac3fc68 100644
--- a/src/Runner.Worker/JobRunner.cs
+++ b/src/Runner.Worker/JobRunner.cs
@@ -84,7 +84,14 @@ namespace GitHub.Runner.Worker
                 Trace.Info($"Creating job server with URL: {jobServerUrl}");
                 // jobServerQueue is the throttling reporter.
                 _jobServerQueue = HostContext.GetService<IJobServerQueue>();
-                VssConnection jobConnection = VssUtil.CreateConnection(jobServerUrl, jobServerCredential, new DelegatingHandler[] { new ThrottlingReportHandler(_jobServerQueue) });
+                var delegatingHandlers = new List<DelegatingHandler>() { new ThrottlingReportHandler(_jobServerQueue) };
+                message.Variables.TryGetValue("Actions.EnableHttpRedirects", out VariableValue enableHttpRedirects);
+                if (StringUtil.ConvertToBoolean(enableHttpRedirects?.Value) &&
+                    !StringUtil.ConvertToBoolean(Environment.GetEnvironmentVariable("GITHUB_ACTIONS_RUNNER_NO_HTTP_REDIRECTS")))
+                {
+                    delegatingHandlers.Add(new RedirectMessageHandler(Trace));
+                }
+                VssConnection jobConnection = VssUtil.CreateConnection(jobServerUrl, jobServerCredential, delegatingHandlers);
                 await jobServer.ConnectAsync(jobConnection);
 
                 _jobServerQueue.Start(message);