migrate add_node_ssh_key (#2193)

* migrate add_node_ssh_key

* build fix

* fix

* format

* fix return value and function name

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<HttpResponseData> Post(HttpRequestData req) {
+        var request = await RequestHandling.ParseRequest<NodeAddSshKeyPost>(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<HttpResponseData> 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"),
+        });
+    }
+
+}

src/ApiService/ApiService/OneFuzzTypes/Requests.cs

@@ -137,3 +137,5 @@ public record JobSearch(
     List<TaskState>? TaskState = null,
     bool? WithTasks = null
 );
+
+public record NodeAddSshKeyPost(Guid MachineId, string PublicKey);

src/ApiService/ApiService/onefuzzlib/NodeOperations.cs

@@ -45,6 +45,8 @@ public interface INodeOperations : IStatefulOrm<Node, NodeState> {
     static readonly TimeSpan NODE_REIMAGE_TIME = TimeSpan.FromDays(6.0);
 
     Async.Task StopTask(Guid task_id);
+
+    Async.Task<OneFuzzResult<bool>> AddSshPublicKey(Node node, string publicKey);
 }
 
 
@@ -421,6 +423,23 @@ public class NodeOperations : StatefulOrm<Node, NodeState, NodeOperations>, INod
 
     }
 
+    public async Task<OneFuzzResult<bool>> AddSshPublicKey(Node node, string publicKey) {
+        if (publicKey == null) {
+            throw new ArgumentNullException(nameof(publicKey));
+        }
+
+        if (node.ScalesetId == null) {
+            return OneFuzzResult<bool>.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<bool>(true);
+    }
+
     /// returns True on stopping the node and False if this doesn't stop the node
     private async Task<bool> StopIfComplete(Node node, bool done = false) {
         var nodeTaskIds = await _context.NodeTasksOperations.GetByMachineId(node.MachineId).Select(nt => nt.TaskId).ToArrayAsync();

src/ApiService/ApiService/onefuzzlib/TaskOperations.cs

@@ -110,7 +110,7 @@ public class TaskOperations : StatefulOrm<Task, TaskState, TaskOperations>, ITas
     }
 
     private async Async.Task MarkDependantsFailed(Task task, List<Task>? 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) {