diff --git a/src/ApiService/ApiService/Functions/AgentEvents.cs b/src/ApiService/ApiService/Functions/AgentEvents.cs index 23a3417940..5f68b8ce59 100644 --- a/src/ApiService/ApiService/Functions/AgentEvents.cs +++ b/src/ApiService/ApiService/Functions/AgentEvents.cs @@ -71,6 +71,10 @@ private async Async.Task Post(HttpRequestData req) { } 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 diff --git a/src/ApiService/ApiService/Functions/Pool.cs b/src/ApiService/ApiService/Functions/Pool.cs index 497a8fd2c8..d6d8170561 100644 --- a/src/ApiService/ApiService/Functions/Pool.cs +++ b/src/ApiService/ApiService/Functions/Pool.cs @@ -17,11 +17,12 @@ public Pool(ILogTracer log, IEndpointAuthorization auth, IOnefuzzContext context } [Function("Pool")] - public Async.Task Run([HttpTrigger(AuthorizationLevel.Anonymous, "GET", "POST", "DELETE")] HttpRequestData req) + public Async.Task 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 @@ private async Task Post(HttpRequestData req) { return await RequestHandling.Ok(req, await Populate(PoolToPoolResponse(newPool), true)); } + + private async Task Patch(HttpRequestData req) { + var request = await RequestHandling.ParseRequest(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 Get(HttpRequestData req) { var request = await RequestHandling.ParseRequest(req); if (!request.IsOk) { diff --git a/src/ApiService/ApiService/OneFuzzTypes/Requests.cs b/src/ApiService/ApiService/OneFuzzTypes/Requests.cs index a673b5f7ea..a66ebbf2c7 100644 --- a/src/ApiService/ApiService/OneFuzzTypes/Requests.cs +++ b/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; diff --git a/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs b/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs index 173f1552a5..a1e3cf156c 100644 --- a/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs +++ b/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs @@ -195,12 +195,12 @@ public async Task CanProcessNewWork(Node node) { 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 bool IsOutdated(Node node) { } 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; } diff --git a/src/cli/onefuzz/api.py b/src/cli/onefuzz/api.py index 8fe5d66e61..ce38724212 100644 --- a/src/cli/onefuzz/api.py +++ b/src/cli/onefuzz/api.py @@ -1297,6 +1297,24 @@ def create( ), ) + 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""" diff --git a/src/pytypes/onefuzztypes/models.py b/src/pytypes/onefuzztypes/models.py index 57094dffff..9f9cb9fbea 100644 --- a/src/pytypes/onefuzztypes/models.py +++ b/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]] diff --git a/src/pytypes/onefuzztypes/requests.py b/src/pytypes/onefuzztypes/requests.py index 5d7547f955..ea107ea8cf 100644 --- a/src/pytypes/onefuzztypes/requests.py +++ b/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]