From 7232981d95402e03f2476eeb552c16e14ed21d25 Mon Sep 17 00:00:00 2001
From: Cheick Keita <chkeita@microsoft.com>
Date: Thu, 27 Oct 2022 10:26:34 -0700
Subject: [PATCH] Use records to unpack the request parameters in
 AgentRegistration (#2570)

---
 .../ApiService/Functions/AgentRegistration.cs | 55 ++++++++++++-------
 .../ApiService/OneFuzzTypes/Requests.cs       | 16 ++++++
 2 files changed, 51 insertions(+), 20 deletions(-)

diff --git a/src/ApiService/ApiService/Functions/AgentRegistration.cs b/src/ApiService/ApiService/Functions/AgentRegistration.cs
index 19dd42e2e6..00ccd2925f 100644
--- a/src/ApiService/ApiService/Functions/AgentRegistration.cs
+++ b/src/ApiService/ApiService/Functions/AgentRegistration.cs
@@ -1,5 +1,4 @@
-using System.Web;
-using Azure.Storage.Sas;
+using Azure.Storage.Sas;
 using Microsoft.Azure.Functions.Worker;
 using Microsoft.Azure.Functions.Worker.Http;
 
@@ -31,9 +30,14 @@ public Async.Task<HttpResponseData> Run(
             });
 
     private async Async.Task<HttpResponseData> Get(HttpRequestData req) {
-        var uri = HttpUtility.ParseQueryString(req.Url.Query);
-        var rawMachineId = uri["machine_id"];
-        if (rawMachineId is null || !Guid.TryParse(rawMachineId, out var machineId)) {
+        var request = await RequestHandling.ParseUri<AgentRegistrationGet>(req);
+        if (!request.IsOk) {
+            return await _context.RequestHandling.NotOk(req, request.ErrorV, "agent registration");
+        }
+
+        var machineId = request.OkV.MachineId;
+
+        if (machineId == Guid.Empty) {
             return await _context.RequestHandling.NotOk(
                 req,
                 new Error(
@@ -83,9 +87,19 @@ private async Async.Task<AgentRegistrationResponse> CreateRegistrationResponse(S
     }
 
     private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
-        var uri = HttpUtility.ParseQueryString(req.Url.Query);
-        var rawMachineId = uri["machine_id"];
-        if (rawMachineId is null || !Guid.TryParse(rawMachineId, out var machineId)) {
+        var request = await RequestHandling.ParseUri<AgentRegistrationPost>(req);
+        if (!request.IsOk) {
+            return await _context.RequestHandling.NotOk(req, request.ErrorV, "agent registration");
+        }
+
+        var machineId = request.OkV.MachineId;
+        var poolName = request.OkV.PoolName;
+        var scalesetId = request.OkV.ScalesetId;
+        var version = request.OkV.Version;
+        var os = request.OkV.Os;
+        var machineName = request.OkV.MachineName;
+
+        if (machineId == Guid.Empty) {
             return await _context.RequestHandling.NotOk(
                 req,
                 new Error(
@@ -94,8 +108,7 @@ private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
                 "agent registration");
         }
 
-        var rawPoolName = uri["pool_name"];
-        if (rawPoolName is null || !PoolName.TryParse(rawPoolName, out var poolName)) {
+        if (poolName is null) {
             return await _context.RequestHandling.NotOk(
                 req,
                 new Error(
@@ -104,16 +117,8 @@ private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
                 "agent registration");
         }
 
-        var rawScalesetId = uri["scaleset_id"];
-        var scalesetId = rawScalesetId is null ? (Guid?)null : Guid.Parse(rawScalesetId);
+        var instanceId = machineName is not null ? InstanceIds.InstanceIdFromMachineName(machineName) : null;
 
-        var version = uri["version"] ?? "1.0.0";
-
-        var machineName = uri["machine_name"];
-        string? instanceId = null;
-        if (machineName is not null) {
-            instanceId = InstanceIds.InstanceIdFromMachineName(machineName);
-        }
 
         _log.Info($"registration request: {machineId:Tag:MachineId} {poolName:Tag:PoolName} {scalesetId:Tag:ScalesetId} {version:Tag:Version}");
         var poolResult = await _context.PoolOperations.GetByName(poolName);
@@ -133,13 +138,23 @@ private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
             await _context.NodeOperations.Delete(existingNode);
         }
 
+        if (os != null && pool.Os != os) {
+            return await _context.RequestHandling.NotOk(
+                req,
+                new Error(
+                    Code: ErrorCode.INVALID_REQUEST,
+                    Errors: new[] { $"OS mismatch: pool '{poolName}' is configured for '{pool.Os}', but agent is running '{os}'" }),
+                "agent registration");
+        }
+
         var node = new Service.Node(
             PoolName: poolName,
             PoolId: pool.PoolId,
             MachineId: machineId,
             ScalesetId: scalesetId,
             InstanceId: instanceId,
-            Version: version);
+            Version: version
+            );
 
         var r = await _context.NodeOperations.Replace(node);
         if (!r.IsOk) {
diff --git a/src/ApiService/ApiService/OneFuzzTypes/Requests.cs b/src/ApiService/ApiService/OneFuzzTypes/Requests.cs
index 63488a4922..2ca98016a4 100644
--- a/src/ApiService/ApiService/OneFuzzTypes/Requests.cs
+++ b/src/ApiService/ApiService/OneFuzzTypes/Requests.cs
@@ -290,3 +290,19 @@ public record WebhookUpdate(
 public record InstanceConfigUpdate(
     [property: Required] InstanceConfig config
 ) : BaseRequest;
+
+
+public record AgentRegistrationGet(
+    [property: Required] Guid MachineId
+) : BaseRequest;
+
+
+public record AgentRegistrationPost(
+    [property: Required] PoolName PoolName,
+    Guid? ScalesetId,
+    [property: Required] Guid MachineId,
+    Os? Os,
+    string? MachineName,
+    [property: Required] string Version = "1.0.0"
+) : BaseRequest;
+