Fix null ref exception when websocket client retries. (#4589)

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

src/Misc/expressionFunc/hashFiles/package.json

@@ -41,10 +41,10 @@
     "@typescript-eslint/parser": "^8.65.0",
     "@vercel/ncc": "^0.38.3",
     "eslint": "^8.47.0",
-    "eslint-plugin-github": "^6.1.0",
+    "eslint-plugin-github": "^6.1.2",
     "eslint-plugin-prettier": "^5.0.0",
     "husky": "^9.1.7",
-    "lint-staged": "^16.4.0",
+    "lint-staged": "^17.2.0",
     "prettier": "^3.0.3",
     "typescript": "^6.0.3"
   }

src/Runner.Common/ResultsServer.cs

@@ -61,7 +61,7 @@ namespace GitHub.Runner.Common
             if (!string.IsNullOrEmpty(liveConsoleFeedUrl))
             {
                 _liveConsoleFeedUrl = liveConsoleFeedUrl;
-                InitializeWebsocketClient(liveConsoleFeedUrl, token, TimeSpan.Zero, retryConnection: true);
+                InitializeWebsocketClient(liveConsoleFeedUrl, TimeSpan.Zero, retryConnection: true);
             }
         }
 
@@ -164,9 +164,9 @@ namespace GitHub.Runner.Common
             return ValueTask.CompletedTask;
         }
 
-        private void InitializeWebsocketClient(string liveConsoleFeedUrl, string accessToken, TimeSpan delay, bool retryConnection = false)
+        private void InitializeWebsocketClient(string liveConsoleFeedUrl, TimeSpan delay, bool retryConnection = false)
         {
-            if (string.IsNullOrEmpty(accessToken))
+            if (string.IsNullOrEmpty(_token))
             {
                 Trace.Info($"No access token from server");
                 return;
@@ -179,12 +179,7 @@ namespace GitHub.Runner.Common
             }
 
             Trace.Info($"Creating websocket client ..." + liveConsoleFeedUrl);
-            this._websocketClient = new ClientWebSocket();
-            this._websocketClient.Options.SetRequestHeader("Authorization", $"Bearer {accessToken}");
-            var userAgentValues = new List<ProductInfoHeaderValue>();
-            userAgentValues.AddRange(UserAgentUtility.GetDefaultRestUserAgent());
-            userAgentValues.AddRange(HostContext.UserAgents);
-            this._websocketClient.Options.SetRequestHeader("User-Agent", string.Join(" ", userAgentValues.Select(x => x.ToString())));
+            this._websocketClient = CreateWebSocketClient();
 
             // during initialization, retry upto 3 times to setup connection
             this._websocketConnectTask = ConnectWebSocketClient(liveConsoleFeedUrl, delay, retryConnection);
@@ -201,8 +196,15 @@ namespace GitHub.Runner.Common
                 {
                     Trace.Info($"Attempting to start websocket client with delay {delay}.");
                     await Task.Delay(delay);
-                    using var connectTimeoutTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(30));
-                    await this._websocketClient.ConnectAsync(new Uri(feedStreamUrl), connectTimeoutTokenSource.Token);
+                    using (var connectTimeoutTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(30)))
+                    {
+                        if (this._websocketClient == null)
+                        {
+                            this._websocketClient = CreateWebSocketClient();
+                        }
+
+                        await this._websocketClient.ConnectAsync(new Uri(feedStreamUrl), connectTimeoutTokenSource.Token);
+                    }
                     Trace.Info($"Successfully started websocket client.");
                     connected = true;
                 }
@@ -211,6 +213,7 @@ namespace GitHub.Runner.Common
                     Trace.Info("Exception caught during websocket client connect, retry connection.");
                     Trace.Error(ex);
                     retries++;
+                    this._websocketClient?.Dispose();
                     this._websocketClient = null;
                     _lastConnectionFailure = DateTime.Now;
                 }
@@ -259,7 +262,7 @@ namespace GitHub.Runner.Common
                             Trace.Info($"Websocket is not open, let's attempt to connect back again with random backoff {delay} ms.");
                             Trace.Verbose(ex.ToString());
                             retries++;
-                            InitializeWebsocketClient(_liveConsoleFeedUrl, _token, delay);
+                            InitializeWebsocketClient(_liveConsoleFeedUrl, delay);
                         }
                     }
                 }
@@ -274,13 +277,24 @@ namespace GitHub.Runner.Common
                 if (_lastConnectionFailure.HasValue && DateTime.Now > _lastConnectionFailure.Value.AddMinutes(10))
                 {
                     // Some minutes passed since we retried last time, try connection again
-                    InitializeWebsocketClient(_liveConsoleFeedUrl, _token, TimeSpan.Zero);
+                    InitializeWebsocketClient(_liveConsoleFeedUrl, TimeSpan.Zero);
                 }
             }
 
             return delivered;
         }
 
+        private ClientWebSocket CreateWebSocketClient()
+        {
+            var client = new ClientWebSocket();
+            client.Options.SetRequestHeader("Authorization", $"Bearer {_token}");
+            var userAgentValues = new List<ProductInfoHeaderValue>();
+            userAgentValues.AddRange(UserAgentUtility.GetDefaultRestUserAgent());
+            userAgentValues.AddRange(HostContext.UserAgents);
+            client.Options.SetRequestHeader("User-Agent", string.Join(" ", userAgentValues.Select(x => x.ToString())));
+            return client;
+        }
+
         private void CloseWebSocket(WebSocketCloseStatus closeStatus, CancellationToken cancellationToken)
         {
             try

src/Runner.Common/VSockSecretNotifier.cs

@@ -0,0 +1,232 @@
+using System;
+using System.Net;
+using System.Net.Sockets;
+using System.Runtime.InteropServices;
+using System.Text;
+using System.Threading;
+using System.Threading.Channels;
+using System.Threading.Tasks;
+using GitHub.DistributedTask.Logging;
+using GitHub.Runner.Sdk;
+using Newtonsoft.Json;
+
+namespace GitHub.Runner.Common
+{
+
+    [ServiceLocator(Default = typeof(VSockSecretNotifier))]
+    public interface IVSockSecretNotifier : IRunnerService, IAsyncDisposable
+    {
+        bool TryStartNotifier();
+
+        void NotifyNewSecret(NewSecretEventArgs newSecret);
+    }
+
+    public sealed class VSockSecretNotifier : RunnerService, IVSockSecretNotifier
+    {
+        private Socket _vsock = null;
+
+        private CancellationTokenSource _cancellationTokenSource = null;
+
+        private Task _secretNotificationTask = null;
+
+        private Channel<byte[]> _channel = Channel.CreateUnbounded<byte[]>(new UnboundedChannelOptions() { SingleReader = true });
+
+        public bool TryStartNotifier()
+        {
+            if (_vsock != null)
+            {
+                Trace.Verbose("VSocket is already connected.");
+                return true;
+            }
+
+            // `GITHUB_ACTIONS_RUNNER_VSOCK_CID_PORT` is expected to be in the format "CID:PORT", e.g. "2:9999".
+            string vsockCidPort = Environment.GetEnvironmentVariable("GITHUB_ACTIONS_RUNNER_VSOCK_CID_PORT");
+            if (string.IsNullOrEmpty(vsockCidPort))
+            {
+                Trace.Verbose("VSocket CID/Port environment variable is not set.");
+                return false;
+            }
+
+            string[] parts = vsockCidPort.Split(':', 2, StringSplitOptions.TrimEntries | StringSplitOptions.RemoveEmptyEntries);
+            if (parts.Length != 2)
+            {
+                Trace.Verbose("VSocket CID/Port environment variable is not in the correct format.");
+                return false;
+            }
+
+            uint cid, port;
+            if (!uint.TryParse(parts[0], out cid) || !uint.TryParse(parts[1], out port))
+            {
+                Trace.Verbose("VSocket CID/Port environment variable contains invalid numbers.");
+                return false;
+            }
+
+            Trace.Info($"Attempting to start VSocket secret notifier with CID: {cid}, Port: {port}.");
+            try
+            {
+                SafeSocketHandle nativeSocket = NativeSocket((int)(AddressFamily)40, (int)SocketType.Stream, 0);
+                if (nativeSocket.IsInvalid)
+                {
+                    int error = Marshal.GetLastPInvokeError();
+                    nativeSocket.Dispose();
+                    throw new SocketException(error);
+                }
+
+                _vsock = new Socket(nativeSocket);
+                _vsock.Connect(new HostVsockEndPoint(cid, port));
+            }
+            catch (Exception ex)
+            {
+                Trace.Error($"Failed to create and connect VSocket: {ex}");
+                _vsock?.Dispose();
+                _vsock = null;
+                return false;
+            }
+
+            _cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(HostContext.RunnerShutdownToken);
+            _secretNotificationTask = ProcessSecretChannel();
+            Trace.Info($"VSocket secret notifier started successfully.");
+            return true;
+        }
+
+        public void NotifyNewSecret(NewSecretEventArgs newSecret)
+        {
+            if (_vsock == null)
+            {
+                Trace.Verbose("VSocket is not connected, skipping secret notification.");
+                return;
+            }
+
+            byte[] payloadBytes = Encoding.UTF8.GetBytes(StringUtil.ConvertToJson(new { RunnerSecrets = newSecret }, Formatting.None));
+            byte[] lengthPrefix = BitConverter.GetBytes(IPAddress.HostToNetworkOrder(payloadBytes.Length));
+            byte[] fullPayload = new byte[lengthPrefix.Length + payloadBytes.Length];
+            Buffer.BlockCopy(lengthPrefix, 0, fullPayload, 0, lengthPrefix.Length);
+            Buffer.BlockCopy(payloadBytes, 0, fullPayload, lengthPrefix.Length, payloadBytes.Length);
+
+            // we don't need to check return since unbounded channel will always accept the item.
+            _channel.Writer.TryWrite(fullPayload);
+        }
+
+        public async ValueTask DisposeAsync()
+        {
+            if (_vsock != null && _secretNotificationTask != null)
+            {
+                _cancellationTokenSource?.Cancel();
+                try
+                {
+                    await _secretNotificationTask;
+                }
+                catch (Exception ex)
+                {
+                    Trace.Error($"Secret notification task finished with error: {ex}");
+                }
+
+                _cancellationTokenSource?.Dispose();
+                _cancellationTokenSource = null;
+                _vsock?.Dispose();
+                _vsock = null;
+            }
+        }
+
+        private async Task ProcessSecretChannel()
+        {
+            try
+            {
+                while (!_cancellationTokenSource.Token.IsCancellationRequested &&
+                        await _channel.Reader.WaitToReadAsync(_cancellationTokenSource.Token))
+                {
+                    while (_channel.Reader.TryRead(out var payload))
+                    {
+                        try
+                        {
+                            // Socket.SendAsync on a stream socket may send fewer bytes than requested,
+                            // so keep sending until the entire payload has been written.
+                            int totalSent = 0;
+                            while (totalSent < payload.Length)
+                            {
+                                int bytesSent = await _vsock.SendAsync(payload.AsMemory(totalSent), SocketFlags.None, _cancellationTokenSource.Token);
+                                if (bytesSent == 0)
+                                {
+                                    throw new SocketException((int)SocketError.ConnectionReset);
+                                }
+
+                                totalSent += bytesSent;
+                            }
+                        }
+                        catch (OperationCanceledException)
+                        {
+                            Trace.Info("Secret notification task was canceled.");
+                        }
+                        catch (Exception ex)
+                        {
+                            Trace.Error($"Failed to notify new secret over VSocket: {ex}");
+                        }
+                    }
+                }
+            }
+            catch (OperationCanceledException)
+            {
+                Trace.Info("Secret notification task was canceled.");
+            }
+            catch (Exception ex)
+            {
+                Trace.Error($"Failed to process secret channel: {ex}");
+            }
+
+            _channel.Writer.TryComplete();
+        }
+
+        [DllImport("libc", SetLastError = true, EntryPoint = "socket")]
+        private static extern SafeSocketHandle NativeSocket(int domain, int type, int protocol);
+    }
+
+    internal sealed class HostVsockEndPoint : EndPoint
+    {
+        private const int SocketAddressSize = 16;
+        private readonly uint _cid;
+        private readonly uint _port;
+
+        public HostVsockEndPoint(uint cid, uint port)
+        {
+            _cid = cid;
+            _port = port;
+        }
+
+        public override AddressFamily AddressFamily => (AddressFamily)40;
+
+        public override SocketAddress Serialize()
+        {
+            SocketAddress socketAddress = new SocketAddress(AddressFamily.Unspecified, SocketAddressSize);
+            // sockaddr_vm layout: family(0-1), reserved1(2-3), port(4-7), cid(8-11)
+            ushort family = (ushort)AddressFamily;
+            socketAddress[0] = (byte)(family & 0xFF);
+            socketAddress[1] = (byte)((family >> 8) & 0xFF);
+            socketAddress[2] = 0;
+            socketAddress[3] = 0;
+            socketAddress[4] = (byte)(_port & 0xFF);
+            socketAddress[5] = (byte)((_port >> 8) & 0xFF);
+            socketAddress[6] = (byte)((_port >> 16) & 0xFF);
+            socketAddress[7] = (byte)((_port >> 24) & 0xFF);
+            socketAddress[8] = (byte)(_cid & 0xFF);
+            socketAddress[9] = (byte)((_cid >> 8) & 0xFF);
+            socketAddress[10] = (byte)((_cid >> 16) & 0xFF);
+            socketAddress[11] = (byte)((_cid >> 24) & 0xFF);
+            return socketAddress;
+        }
+
+        public override EndPoint Create(SocketAddress socketAddress)
+        {
+            uint port = (uint)socketAddress[4]
+                | ((uint)socketAddress[5] << 8)
+                | ((uint)socketAddress[6] << 16)
+                | ((uint)socketAddress[7] << 24);
+
+            uint cid = (uint)socketAddress[8]
+                | ((uint)socketAddress[9] << 8)
+                | ((uint)socketAddress[10] << 16)
+                | ((uint)socketAddress[11] << 24);
+
+            return new HostVsockEndPoint(cid, port);
+        }
+    }
+}

src/Runner.Worker/Worker.cs

@@ -1,15 +1,14 @@
-using GitHub.DistributedTask.WebApi;
-using Pipelines = GitHub.DistributedTask.Pipelines;
-using GitHub.Runner.Common.Util;
-using Newtonsoft.Json;
-using System;
+using System;
 using System.Collections.Generic;
+using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
-using GitHub.Services.WebApi;
+using GitHub.DistributedTask.WebApi;
 using GitHub.Runner.Common;
+using GitHub.Runner.Common.Util;
 using GitHub.Runner.Sdk;
-using System.Text;
+using Newtonsoft.Json;
+using Pipelines = GitHub.DistributedTask.Pipelines;
 
 namespace GitHub.Runner.Worker
 {
@@ -46,6 +45,7 @@ namespace GitHub.Runner.Worker
                 var jobRunner = HostContext.CreateService<IJobRunner>();
                 var terminal = HostContext.GetService<ITerminal>();
 
+                await using (var secretNotifier = HostContext.GetService<IVSockSecretNotifier>())
                 using (var channel = HostContext.CreateService<IProcessChannel>())
                 using (var jobRequestCancellationToken = CancellationTokenSource.CreateLinkedTokenSource(HostContext.RunnerShutdownToken))
                 using (var channelTokenSource = new CancellationTokenSource())
@@ -86,6 +86,14 @@ namespace GitHub.Runner.Worker
                     HostContext.WritePerfCounter($"WorkerJobMessageReceived_{jobMessage.RequestId.ToString()}");
 
                     // Initialize the secret masker and set the thread culture.
+                    if (Constants.Runner.Platform == Constants.OSPlatform.Linux &&
+                        secretNotifier.TryStartNotifier())
+                    {
+                        HostContext.SecretMasker.NewSecretAdded += (sender, e) =>
+                        {
+                            secretNotifier.NotifyNewSecret(e);
+                        };
+                    }
                     InitializeSecretMasker(jobMessage);
                     SetCulture(jobMessage);
 

src/Sdk/DTLogging/Logging/ISecretMasker.cs

@@ -1,5 +1,7 @@
 using System;
+using System.Collections.Generic;
 using System.ComponentModel;
+using System.Runtime.Serialization;
 
 namespace GitHub.DistributedTask.Logging
 {
@@ -11,5 +13,41 @@ namespace GitHub.DistributedTask.Logging
         void AddValueEncoder(ValueEncoder encoder);
         ISecretMasker Clone();
         String MaskSecrets(String input);
+        public event EventHandler<NewSecretEventArgs> NewSecretAdded;
+    }
+
+    public abstract class NewSecretEventArgs : EventArgs
+    {
+        public abstract String Type { get; }
+    }
+
+    [DataContract]
+    public sealed class NewRegexSecretEventArgs : NewSecretEventArgs
+    {
+        [DataMember]
+        public override String Type => "regex";
+
+        public NewRegexSecretEventArgs(String pattern)
+        {
+            Pattern = pattern;
+        }
+
+        [DataMember]
+        public String Pattern { get; private set; }
+    }
+
+    [DataContract]
+    public sealed class NewVariableSecretEventArgs : NewSecretEventArgs
+    {
+        [DataMember]
+        public override String Type => "variable";
+
+        public NewVariableSecretEventArgs(List<string> values)
+        {
+            Values.AddRange(values);
+        }
+
+        [DataMember]
+        public List<string> Values { get; private set; } = new List<string>();
     }
 }

src/Sdk/DTLogging/Logging/SecretMasker.cs

@@ -10,6 +10,8 @@ namespace GitHub.DistributedTask.Logging
     [EditorBrowsable(EditorBrowsableState.Never)]
     public sealed class SecretMasker : ISecretMasker, IDisposable
     {
+        public event EventHandler<NewSecretEventArgs> NewSecretAdded;
+
         public SecretMasker()
         {
             m_originalValueSecrets = new HashSet<ValueSecret>();
@@ -66,6 +68,8 @@ namespace GitHub.DistributedTask.Logging
                     m_lock.ExitWriteLock();
                 }
             }
+
+            NewSecretAdded?.Invoke(this, new NewRegexSecretEventArgs(pattern));
         }
 
         /// <summary>
@@ -133,6 +137,9 @@ namespace GitHub.DistributedTask.Logging
                     m_lock.ExitWriteLock();
                 }
             }
+
+            // valueSecrets contains all the values run through the encoders.
+            NewSecretAdded?.Invoke(this, new NewVariableSecretEventArgs(valueSecrets.Select(x => x.m_value).ToList()));
         }
 
         /// <summary>

src/Test/L0/Worker/WorkerL0.cs

@@ -16,11 +16,13 @@ namespace GitHub.Runner.Common.Tests.Worker
     {
         private Mock<IProcessChannel> _processChannel;
         private Mock<IJobRunner> _jobRunner;
+        private Mock<IVSockSecretNotifier> _vsockSecretNotifier;
 
         public WorkerL0()
         {
             _processChannel = new Mock<IProcessChannel>();
             _jobRunner = new Mock<IJobRunner>();
+            _vsockSecretNotifier = new Mock<IVSockSecretNotifier>();
         }
 
         private Pipelines.AgentJobRequestMessage CreateJobRequestMessage(string jobName)
@@ -88,6 +90,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 var worker = new GitHub.Runner.Worker.Worker();
                 hc.EnqueueInstance<IProcessChannel>(_processChannel.Object);
                 hc.EnqueueInstance<IJobRunner>(_jobRunner.Object);
+                hc.SetSingleton<IVSockSecretNotifier>(_vsockSecretNotifier.Object);
                 worker.Initialize(hc);
                 var jobMessage = CreateJobRequestMessage("job1");
                 var arWorkerMessages = new WorkerMessage[]
@@ -139,6 +142,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 var worker = new GitHub.Runner.Worker.Worker();
                 hc.EnqueueInstance<IProcessChannel>(_processChannel.Object);
                 hc.EnqueueInstance<IJobRunner>(_jobRunner.Object);
+                hc.SetSingleton<IVSockSecretNotifier>(_vsockSecretNotifier.Object);
                 worker.Initialize(hc);
                 var jobMessage = CreateJobRequestMessage("job1");
                 var cancelMessage = CreateJobCancelMessage(jobMessage.JobId);