From a1bf8401d7dac2d22f89612c06708eaef39bed09 Mon Sep 17 00:00:00 2001
From: eric sciple <ericsciple@users.noreply.github.com>
Date: Thu, 9 Jun 2022 13:07:44 -0500
Subject: [PATCH] Handle message from broker (#1934)

---
 src/Runner.Listener/Runner.cs              | 29 ++++++++++++++++++++++
 src/Runner.Listener/RunnerJobRequestRef.cs | 13 ++++++++++
 2 files changed, 42 insertions(+)
 create mode 100644 src/Runner.Listener/RunnerJobRequestRef.cs

diff --git a/src/Runner.Listener/Runner.cs b/src/Runner.Listener/Runner.cs
index 1b145be5d..8455ebef2 100644
--- a/src/Runner.Listener/Runner.cs
+++ b/src/Runner.Listener/Runner.cs
@@ -458,6 +458,35 @@ namespace GitHub.Runner.Listener
                                     }
                                 }
                             }
+                            // Broker flow
+                            else if (string.Equals(message.MessageType, JobRequestMessageTypes.RunnerJobRequest, StringComparison.OrdinalIgnoreCase))
+                            {
+                                if (autoUpdateInProgress || runOnceJobReceived)
+                                {
+                                    skipMessageDeletion = true;
+                                    Trace.Info($"Skip message deletion for job request message '{message.MessageId}'.");
+                                }
+                                else
+                                {
+                                    var messageRef = StringUtil.ConvertFromJson<RunnerJobRequestRef>(message.Body);
+
+                                    // Create connection
+                                    var credMgr = HostContext.GetService<ICredentialManager>();
+                                    var creds = credMgr.LoadCredentials();
+
+                                    // todo: add retries https://github.com/github/actions-broker/issues/49
+                                    var runServer = HostContext.CreateService<IRunServer>();
+                                    await runServer.ConnectAsync(new Uri(settings.ServerUrl), creds);
+                                    var jobMessage = await runServer.GetJobMessageAsync(messageRef.RunnerRequestId);
+
+                                    jobDispatcher.Run(jobMessage, runOnce);
+                                    if (runOnce)
+                                    {
+                                        Trace.Info("One time used runner received job message.");
+                                        runOnceJobReceived = true;
+                                    }
+                                }
+                            }
                             else if (string.Equals(message.MessageType, JobCancelMessage.MessageType, StringComparison.OrdinalIgnoreCase))
                             {
                                 var cancelJobMessage = JsonUtility.FromString<JobCancelMessage>(message.Body);
diff --git a/src/Runner.Listener/RunnerJobRequestRef.cs b/src/Runner.Listener/RunnerJobRequestRef.cs
new file mode 100644
index 000000000..fad71d121
--- /dev/null
+++ b/src/Runner.Listener/RunnerJobRequestRef.cs
@@ -0,0 +1,13 @@
+using System.Runtime.Serialization;
+
+namespace GitHub.Runner.Listener
+{
+    [DataContract]
+    public sealed class RunnerJobRequestRef
+    {
+        [DataMember(Name = "id")]
+        public string Id { get; set; }
+        [DataMember(Name = "runner_request_id")]
+        public string RunnerRequestId { get; set; }
+    }
+}
\ No newline at end of file