Skip to content
This repository has been archived by the owner on Nov 1, 2023. It is now read-only.

Add New Endpoint to update the pool authentication #3059

Merged
merged 4 commits into from
Apr 24, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/ApiService/ApiService/Functions/AgentEvents.cs
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,10 @@ private async Async.Task<HttpResponseData> 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
Expand Down
39 changes: 38 additions & 1 deletion src/ApiService/ApiService/Functions/Pool.cs
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,12 @@ public Pool(ILogTracer log, IEndpointAuthorization auth, IOnefuzzContext context
}

[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}"),
});

Expand Down Expand Up @@ -71,6 +72,42 @@ private async Task<HttpResponseData> Post(HttpRequestData req) {
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) {
Expand Down
7 changes: 7 additions & 0 deletions src/ApiService/ApiService/OneFuzzTypes/Requests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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
chkeita marked this conversation as resolved.
Show resolved Hide resolved
) : BaseRequest;


public record PoolUpdate(
[property: Required] PoolName Name,
Guid? ObjectId = null
) : BaseRequest;

Expand Down
7 changes: 4 additions & 3 deletions src/ApiService/ApiService/onefuzzlib/NodeOperations.cs
Original file line number Diff line number Diff line change
Expand Up @@ -195,12 +195,12 @@ public async Task<CanProcessNewWorkResponse> 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");
}
Expand Down Expand Up @@ -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;
}
Expand Down
18 changes: 18 additions & 0 deletions src/cli/onefuzz/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"""

Expand Down
1 change: 1 addition & 0 deletions src/pytypes/onefuzztypes/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]]
Expand Down
5 changes: 5 additions & 0 deletions src/pytypes/onefuzztypes/requests.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down