Add New Endpoint to update the pool authentication (#3059)

* check that the node in managed before sending the stop message

* added endpoint to update the pool

* Update src/ApiService/ApiService/Functions/Pool.cs

Co-authored-by: Teo Voinea <58236992+tevoinea@users.noreply.github.com>

---------

Co-authored-by: Teo Voinea <58236992+tevoinea@users.noreply.github.com>

src/ApiService/ApiService/Functions/AgentEvents.cs

@@ -71,6 +71,10 @@ public class AgentEvents {
         }
 
         if (ev.State == NodeState.Free) {
+            if (!node.Managed) {
+                return null;
+            }
+
             if (node.ReimageRequested || node.DeleteRequested) {
                 _log.Info($"stopping free node with reset flags: {machineId:Tag:MachineId}");
                 // discard result: node not used after this point

src/ApiService/ApiService/Functions/Pool.cs

@@ -17,11 +17,12 @@ public class Pool {
     }
 
     [Function("Pool")]
-    public Async.Task<HttpResponseData> Run([HttpTrigger(AuthorizationLevel.Anonymous, "GET", "POST", "DELETE")] HttpRequestData req)
+    public Async.Task<HttpResponseData> Run([HttpTrigger(AuthorizationLevel.Anonymous, "GET", "POST", "DELETE", "PATCH")] HttpRequestData req)
         => _auth.CallIfUser(req, r => r.Method switch {
             "GET" => Get(r),
             "POST" => Post(r),
             "DELETE" => Delete(r),
+            "PATCH" => Patch(r),
             var m => throw new InvalidOperationException("Unsupported HTTP method {m}"),
         });
 
@@ -71,6 +72,42 @@ public class Pool {
         return await RequestHandling.Ok(req, await Populate(PoolToPoolResponse(newPool), true));
     }
 
+
+    private async Task<HttpResponseData> Patch(HttpRequestData req) {
+        var request = await RequestHandling.ParseRequest<PoolUpdate>(req);
+        if (!request.IsOk) {
+            return await _context.RequestHandling.NotOk(req, request.ErrorV, "PoolUpdate");
+        }
+
+        var answer = await _auth.CheckRequireAdmins(req);
+        if (!answer.IsOk) {
+            return await _context.RequestHandling.NotOk(req, answer.ErrorV, "PoolUpdate");
+        }
+
+        var update = request.OkV;
+        var pool = await _context.PoolOperations.GetByName(update.Name);
+        if (!pool.IsOk) {
+            return await _context.RequestHandling.NotOk(
+                req,
+                new Error(
+                    Code: ErrorCode.INVALID_REQUEST,
+                    Errors: new string[] { "pool with that name does not exist" }),
+                "PoolUpdate");
+        }
+
+        var updated = pool.OkV with { ObjectId = update.ObjectId };
+        var updatePool = await _context.PoolOperations.Update(updated);
+        if (updatePool.IsOk) {
+            return await RequestHandling.Ok(req, await Populate(PoolToPoolResponse(updated), true));
+        } else {
+            return await _context.RequestHandling.NotOk(req, new Error(
+                Code: ErrorCode.INVALID_REQUEST,
+                Errors: new string[] { updatePool.ErrorV.Reason }), "PoolUpdate");
+        }
+
+
+    }
+
     private async Task<HttpResponseData> Get(HttpRequestData req) {
         var request = await RequestHandling.ParseRequest<PoolSearch>(req);
         if (!request.IsOk) {

src/ApiService/ApiService/OneFuzzTypes/Requests.cs

@@ -270,6 +270,13 @@ public record PoolCreate(
     [property: Required] Os Os,
     [property: Required] Architecture Arch,
     [property: Required] bool Managed,
+    Guid? ObjectId = null,
+    bool Update = false
+) : BaseRequest;
+
+
+public record PoolUpdate(
+    [property: Required] PoolName Name,
     Guid? ObjectId = null
 ) : BaseRequest;
 

src/ApiService/ApiService/onefuzzlib/NodeOperations.cs

@@ -195,12 +195,12 @@ public class NodeOperations : StatefulOrm<Node, NodeState, NodeOperations>, INod
             return CanProcessNewWorkResponse.NotAllowed("node is set to be deleted");
         }
 
-        if (node.ReimageRequested) {
+        if (node.ReimageRequested && node.Managed) {
             _ = await Stop(node, done: true);
             return CanProcessNewWorkResponse.NotAllowed("node is set to be reimaged");
         }
 
-        if (await CouldShrinkScaleset(node)) {
+        if (await CouldShrinkScaleset(node) && node.Managed) {
             _ = await SetHalt(node);
             return CanProcessNewWorkResponse.NotAllowed("node is scheduled to shrink");
         }
@@ -488,7 +488,8 @@ public class NodeOperations : StatefulOrm<Node, NodeState, NodeOperations>, INod
     }
 
     public bool IsTooOld(Node node) {
-        return node.ScalesetId != null
+        return node.Managed
+            && node.ScalesetId != null
             && node.InitializedAt != null
             && node.InitializedAt < DateTime.UtcNow - INodeOperations.NODE_REIMAGE_TIME;
     }

src/cli/onefuzz/api.py

@@ -1297,6 +1297,24 @@ class Pool(Endpoint):
             ),
         )
 
+    def update(
+        self,
+        name: str,
+        object_id: Optional[UUID] = None,
+    ) -> models.Pool:
+        """
+        Update a worker pool
+
+        :param str name: Name of the worker-pool
+        """
+        self.logger.debug("create worker pool")
+
+        return self._req_model(
+            "PATCH",
+            models.Pool,
+            data=requests.PoolUpdate(name=name, object_id=object_id),
+        )
+
     def get_config(self, pool_name: primitives.PoolName) -> models.AgentConfig:
         """Get the agent configuration for the pool"""
 

src/pytypes/onefuzztypes/models.py

@@ -653,6 +653,7 @@ class Pool(BaseModel):
     # intended to be used to pass the information to the CLI when the CLI asks
     # for information about what work is in the queue for the pool.
     work_queue: Optional[List[WorkSetSummary]]
+    object_id: Optional[UUID]
 
     # explicitly excluded from Tables
     scaleset_summary: Optional[List[ScalesetSummary]]

src/pytypes/onefuzztypes/requests.py

@@ -105,6 +105,11 @@ class PoolCreate(BaseRequest):
     autoscale: Optional[AutoScaleConfig]
 
 
+class PoolUpdate(BaseRequest):
+    name: PoolName
+    object_id: Optional[UUID]
+
+
 class PoolSearch(BaseRequest):
     pool_id: Optional[UUID]
     name: Optional[PoolName]