diff --git a/src/ApiService/ApiService/Functions/NodeAddSshKey.cs b/src/ApiService/ApiService/Functions/NodeAddSshKey.cs new file mode 100644 index 000000000..ae8ac9c29 --- /dev/null +++ b/src/ApiService/ApiService/Functions/NodeAddSshKey.cs @@ -0,0 +1,57 @@ +using System.Net; +using Microsoft.Azure.Functions.Worker; +using Microsoft.Azure.Functions.Worker.Http; + +namespace Microsoft.OneFuzz.Service.Functions; + +public class NodeAddSshKey { + + private readonly ILogTracer _log; + private readonly IEndpointAuthorization _auth; + private readonly IOnefuzzContext _context; + + public NodeAddSshKey(ILogTracer log, IEndpointAuthorization auth, IOnefuzzContext context) { + _log = log; + _auth = auth; + _context = context; + } + + private async Async.Task Post(HttpRequestData req) { + var request = await RequestHandling.ParseRequest(req); + if (!request.IsOk) { + return await _context.RequestHandling.NotOk( + req, + request.ErrorV, + "NodeAddSshKey"); + } + + var node = await _context.NodeOperations.GetByMachineId(request.OkV.MachineId); + + if (node == null) { + return await _context.RequestHandling.NotOk( + req, + new Error(ErrorCode.UNABLE_TO_FIND, new[] { "unable to find node" }), + $"{request.OkV.MachineId}"); + } + + var result = await _context.NodeOperations.AddSshPublicKey(node, request.OkV.PublicKey); + if (!result.IsOk) { + return await _context.RequestHandling.NotOk(req, result.ErrorV, "NodeAddSshKey"); + } + + var response = req.CreateResponse(HttpStatusCode.OK); + await response.WriteAsJsonAsync(new BoolResult(true)); + return response; + + + } + + [Function("node_add_ssh_key")] + public Async.Task Run([HttpTrigger(AuthorizationLevel.Anonymous, "POST", Route = "node/add_ssh_key")] HttpRequestData req) { + return _auth.CallIfUser(req, r => r.Method switch { + "POST" => Post(r), + _ => throw new InvalidOperationException("Unsupported HTTP method"), + }); + } + +} diff --git a/src/ApiService/ApiService/OneFuzzTypes/Requests.cs b/src/ApiService/ApiService/OneFuzzTypes/Requests.cs index ffa94975d..3751a8da8 100644 --- a/src/ApiService/ApiService/OneFuzzTypes/Requests.cs +++ b/src/ApiService/ApiService/OneFuzzTypes/Requests.cs @@ -137,3 +137,5 @@ public record JobSearch( List? TaskState = null, bool? WithTasks = null ); + +public record NodeAddSshKeyPost(Guid MachineId, string PublicKey); diff --git a/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs b/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs index 3ecb2ff8f..b6bdcadfc 100644 --- a/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs +++ b/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs @@ -45,6 +45,8 @@ public interface INodeOperations : IStatefulOrm { static readonly TimeSpan NODE_REIMAGE_TIME = TimeSpan.FromDays(6.0); Async.Task StopTask(Guid task_id); + + Async.Task> AddSshPublicKey(Node node, string publicKey); } @@ -421,6 +423,23 @@ public class NodeOperations : StatefulOrm, INod } + public async Task> AddSshPublicKey(Node node, string publicKey) { + if (publicKey == null) { + throw new ArgumentNullException(nameof(publicKey)); + } + + if (node.ScalesetId == null) { + return OneFuzzResult.Error(new Error(ErrorCode.INVALID_REQUEST, + new[] { "only able to add ssh keys to scaleset nodes" })); + } + + var key = publicKey.EndsWith('\n') ? publicKey : $"{publicKey}\n"; + + await SendMessage(node, new NodeCommand { AddSshKey = new NodeCommandAddSshKey(key) }); + + return OneFuzzResult.Ok(true); + } + /// returns True on stopping the node and False if this doesn't stop the node private async Task StopIfComplete(Node node, bool done = false) { var nodeTaskIds = await _context.NodeTasksOperations.GetByMachineId(node.MachineId).Select(nt => nt.TaskId).ToArrayAsync(); diff --git a/src/ApiService/ApiService/onefuzzlib/TaskOperations.cs b/src/ApiService/ApiService/onefuzzlib/TaskOperations.cs index 421f741dc..00710649f 100644 --- a/src/ApiService/ApiService/onefuzzlib/TaskOperations.cs +++ b/src/ApiService/ApiService/onefuzzlib/TaskOperations.cs @@ -110,7 +110,7 @@ public class TaskOperations : StatefulOrm, ITas } private async Async.Task MarkDependantsFailed(Task task, List? taskInJob = null) { - taskInJob ??= await SearchByPartitionKeys(new[] { task.JobId.ToString() }).ToListAsync(); + taskInJob ??= await SearchByPartitionKeys(new[] { $"{task.JobId}" }).ToListAsync(); foreach (var t in taskInJob) { if (t.Config.PrereqTasks != null) {