porting the proxy state machine

microsoft · Aug 23, 2022 · 0a60a47 · 0a60a47
1 parent f92d207
commit 0a60a47
Show file tree

Hide file tree

Showing 5 changed files with 177 additions and 15 deletions.
diff --git a/src/ApiService/ApiService/OneFuzzTypes/Model.cs b/src/ApiService/ApiService/OneFuzzTypes/Model.cs
@@ -328,7 +328,7 @@ public record InstanceConfig
     [DefaultValue(InitMethod.DefaultConstructor)] NetworkConfig NetworkConfig,
     [DefaultValue(InitMethod.DefaultConstructor)] NetworkSecurityGroupConfig ProxyNsgConfig,
     AzureVmExtensionConfig? Extensions,
-    string? ProxyVmSku,
+    string ProxyVmSku,
     bool AllowPoolManagement = true,
     IDictionary<Endpoint, ApiAccessRule>? ApiAccessRules = null,
     IDictionary<PrincipalId, GroupId[]>? GroupMembership = null,

diff --git a/src/ApiService/ApiService/onefuzzlib/Extension.cs b/src/ApiService/ApiService/onefuzzlib/Extension.cs
@@ -11,6 +11,7 @@ public interface IExtensions {
     Async.Task<IList<VirtualMachineScaleSetExtensionData>> FuzzExtensions(Pool pool, Scaleset scaleset);
 
     Async.Task<Dictionary<string, VirtualMachineExtensionData>> ReproExtensions(AzureLocation region, Os reproOs, Guid reproId, ReproConfig reproConfig, Container? setupContainer);
+    Task<IList<VMExtensionWrapper>> ProxyManagerExtensions(string region, Guid proxyId);
 }
 
 public class Extensions : IExtensions {
@@ -449,4 +450,19 @@ await _context.Containers.GetFileSasUrl(
         return extensionsDict;
     }
 
+    public async Task<IList<VMExtensionWrapper>> ProxyManagerExtensions(string region, Guid proxyId) {
+        var config = await _context.Containers.GetFileSasUrl(new Container("proxy-config"),
+            $"{region}/{proxyId}/config.json", StorageType.Config, BlobSasPermissions.Read);
+
+        var proxyManager = await _context.Containers.GetFileSasUrl(new Container("tools"),
+            $"linux/onefuzz-proxy-manager", StorageType.Config, BlobSasPermissions.Read);
+
+
+        var baseExtension =
+            await AgentConfig(region, Os.Linux, AgentMode.Proxy, new List<Uri> { config, proxyManager }, true);
+
+        var extensions = await GenericExtensions(region, Os.Linux);
+        extensions.Add(baseExtension);
+        return extensions;
+    }
 }
diff --git a/src/ApiService/ApiService/onefuzzlib/ProxyOperations.cs b/src/ApiService/ApiService/onefuzzlib/ProxyOperations.cs
@@ -1,5 +1,7 @@
 using System.Threading.Tasks;
 using ApiService.OneFuzzLib.Orm;
+using Azure.ResourceManager.Compute;
+using Azure.ResourceManager.Compute.Models;
 using Azure.Storage.Sas;
 using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
 
@@ -8,7 +10,7 @@ namespace Microsoft.OneFuzz.Service;
 public interface IProxyOperations : IStatefulOrm<Proxy, VmState> {
     Task<Proxy?> GetByProxyId(Guid proxyId);
 
-    Async.Task SetState(Proxy proxy, VmState state);
+    Async.Task<Proxy> SetState(Proxy proxy, VmState state);
     bool IsAlive(Proxy proxy);
     Async.Task SaveProxyConfig(Proxy proxy);
     bool IsOutdated(Proxy proxy);
@@ -110,14 +112,15 @@ public async Async.Task SaveProxyConfig(Proxy proxy) {
     }
 
 
-    public async Async.Task SetState(Proxy proxy, VmState state) {
+    public async Async.Task<Proxy> SetState(Proxy proxy, VmState state) {
         if (proxy.State == state) {
-            return;
+            return proxy;
         }
 
-        await Replace(proxy with { State = state });
-
+        var newProxy = proxy with { State = state };
+        await Replace(newProxy);
         await _context.Events.SendEvent(new EventProxyStateUpdated(proxy.Region, proxy.ProxyId, proxy.State));
+        return newProxy;
     }
 
 
@@ -133,4 +136,138 @@ public async Async.Task<List<Forward>> GetForwards(Proxy proxy) {
         }
         return forwards;
     }
+
+    public async Async.Task<Proxy> Init(Proxy proxy) {
+        var config = await _context.ConfigOperations.Fetch();
+        var vm = GetVm(proxy, config);
+        var vmData = await _context.VmOperations.GetVm(vm.Name);
+
+        if (vmData != null) {
+            if (vmData.ProvisioningState == "Failed") {
+                return await SetProvisionFailed(proxy, vmData);
+            } else {
+                await SaveProxyConfig(proxy);
+                return await SetState(proxy, VmState.ExtensionsLaunch);
+            }
+        } else {
+            var nsg = new Nsg(proxy.Region, proxy.Region);
+            var result = await _context.NsgOperations.Create(nsg);
+            if (!result.IsOk) {
+                return await SetFailed(proxy, result.ErrorV);
+            }
+
+            var nsgConfig = config.ProxyNsgConfig;
+            var result2 = await _context.NsgOperations.SetAllowedSources(nsg, nsgConfig);
+
+            if (!result2.IsOk) {
+                return await SetFailed(proxy, result2.ErrorV);
+            }
+
+            var result3 = await _context.VmOperations.Create(vm with { Nsg = nsg });
+
+            if (!result3.IsOk) {
+                return await SetFailed(proxy, result3.ErrorV);
+            }
+            return proxy;
+        }
+    }
+
+    private async System.Threading.Tasks.Task<Proxy> SetProvisionFailed(Proxy proxy, VirtualMachineData vmData) {
+        var errors = GetErrors(proxy, vmData).ToArray();
+        await SetFailed(proxy, new Error(ErrorCode.PROXY_FAILED, errors));
+        return proxy;
+    }
+
+    private async Task<Proxy> SetFailed(Proxy proxy, Error error) {
+        if (proxy.Error != null) {
+            return proxy;
+        }
+
+        _logTracer.Error($"vm failed: {proxy.Region} -{error}");
+        await _context.Events.SendEvent(new EventProxyFailed(proxy.Region, proxy.ProxyId, error));
+        return await SetState(proxy with { Error = error }, VmState.Stopping);
+    }
+
+
+    private static IEnumerable<string> GetErrors(Proxy proxy, VirtualMachineData vmData) {
+        var instanceView = vmData.InstanceView;
+        yield return "provisioning failed";
+        if (instanceView is null) {
+            yield break;
+        }
+
+        foreach (var status in instanceView.Statuses) {
+            if (status.Level == StatusLevelTypes.Error) {
+                yield return $"code:{status.Code} status:{status.DisplayStatus} message:{status.Message}";
+            }
+        }
+    }
+
+    public static Vm GetVm(Proxy proxy, InstanceConfig config) {
+        var tags = config.VmssTags;
+        const string PROXY_IMAGE = "Canonical:UbuntuServer:18.04-LTS:latest";
+        return new Vm(
+            // name should be less than 40 chars otherwise it gets truncated by azure
+            Name: $"proxy-{proxy.ProxyId:N}",
+            Region: proxy.Region,
+            Sku: config.ProxyVmSku,
+            Image: PROXY_IMAGE,
+            Auth: proxy.Auth,
+            Tags: config.VmssTags,
+            Nsg: null
+        );
+    }
+
+    public async Task<Proxy> ExtensionsLaunch(Proxy proxy) {
+        var config = await _context.ConfigOperations.Fetch();
+        var vm = GetVm(proxy, config);
+        var vmData = await _context.VmOperations.GetVm(vm.Name);
+
+        if (vmData == null) {
+            return await SetFailed(proxy, new Error(ErrorCode.PROXY_FAILED, new[] { "azure not able to find vm" }));
+        }
+
+        if (vmData.ProvisioningState == "Failed") {
+            return await SetProvisionFailed(proxy, vmData);
+        }
+
+        var ip = await _context.IpOperations.GetPublicIp(vmData.NetworkProfile.NetworkInterfaces[0].Id);
+        if (ip == null) {
+            return proxy;
+        }
+
+        var newProxy = proxy with { Ip = ip };
+
+        var extensions = await _context.Extensions.ProxyManagerExtensions(newProxy.Region, newProxy.ProxyId);
+        var result = await _context.VmOperations.AddExtensions(vm,
+            extensions
+                .Select(e => e.GetAsVirtualMachineExtension())
+                .ToDictionary(x => x.Item1, x => x.Item2));
+
+        if (!result.IsOk) {
+            return await SetFailed(newProxy, result.ErrorV);
+        }
+
+        return await SetState(newProxy, VmState.Running);
+    }
+
+    public async Task<Proxy> Stopping(Proxy proxy) {
+        var config = await _context.ConfigOperations.Fetch();
+        var vm = GetVm(proxy, config);
+        if (!await _context.VmOperations.IsDeleted(vm)) {
+            _logTracer.Error($"stopping proxy: {proxy.Region}");
+            await _context.VmOperations.Delete(vm);
+            return proxy;
+        }
+
+        return await Stopped(proxy);
+    }
+
+    private async Task<Proxy> Stopped(Proxy proxy) {
+        var stoppedVm = await SetState(proxy, VmState.Stopped);
+        _logTracer.Info($"removing proxy: {proxy.Region}");
+        await _context.Events.SendEvent(new EventProxyDeleted(proxy.Region, proxy.ProxyId));
+        await Delete(proxy);
+        return stoppedVm;
+    }
 }
diff --git a/src/ApiService/ApiService/onefuzzlib/ReproOperations.cs b/src/ApiService/ApiService/onefuzzlib/ReproOperations.cs
@@ -18,7 +18,7 @@ public interface IReproOperations : IStatefulOrm<Repro, VmState> {
 
     public Async.Task<Repro> Stopped(Repro repro);
 
-    public Async.Task<Repro> SetFailed(Repro repro, VirtualMachineResource vmData);
+    public Async.Task<Repro> SetFailed(Repro repro, VirtualMachineData vmData);
 
     public Async.Task<Repro> SetError(Repro repro, Error result);
 
@@ -120,7 +120,7 @@ public async Async.Task<Repro> Init(Repro repro) {
         var vm = await GetVm(repro, config);
         var vmData = await _context.VmOperations.GetVm(vm.Name);
         if (vmData != null) {
-            if (vmData.Data.ProvisioningState == "Failed") {
+            if (vmData.ProvisioningState == "Failed") {
                 return await _context.ReproOperations.SetFailed(repro, vmData);
             } else {
                 var scriptResult = await BuildReproScript(repro);
@@ -167,13 +167,13 @@ public async Async.Task<Repro> ExtensionsLaunch(Repro repro) {
             );
         }
 
-        if (vmData.Data.ProvisioningState == "Failed") {
+        if (vmData.ProvisioningState == "Failed") {
             return await _context.ReproOperations.SetFailed(repro, vmData);
         }
 
         if (string.IsNullOrEmpty(repro.Ip)) {
             repro = repro with {
-                Ip = await _context.IpOperations.GetPublicIp(vmData.Data.NetworkProfile.NetworkInterfaces.First().Id)
+                Ip = await _context.IpOperations.GetPublicIp(vmData.NetworkProfile.NetworkInterfaces.First().Id)
             };
         }
 
@@ -196,8 +196,8 @@ await _context.ReproOperations.GetSetupContainer(repro)
         return repro;
     }
 
-    public async Async.Task<Repro> SetFailed(Repro repro, VirtualMachineResource vmData) {
-        var errors = (await vmData.InstanceViewAsync()).Value.Statuses
+    public async Async.Task<Repro> SetFailed(Repro repro, VirtualMachineData vmData) {
+        var errors = vmData.InstanceView.Statuses
             .Where(status => status.Level.HasValue && string.Equals(status.Level?.ToString(), "error", StringComparison.OrdinalIgnoreCase))
             .Select(status => $"{status.Code} {status.DisplayStatus} {status.Message}")
             .ToArray();

diff --git a/src/ApiService/ApiService/onefuzzlib/VmOperations.cs b/src/ApiService/ApiService/onefuzzlib/VmOperations.cs
@@ -11,7 +11,7 @@ public interface IVmOperations {
 
     Async.Task<bool> HasComponents(string name);
 
-    Async.Task<VirtualMachineResource?> GetVm(string name);
+    Task<VirtualMachineData?> GetVm(string name);
 
     Async.Task<bool> Delete(Vm vm);
 
@@ -64,14 +64,23 @@ public async Async.Task<bool> HasComponents(string name) {
         return false;
     }
 
-    public async Async.Task<VirtualMachineResource?> GetVm(string name) {
+    public async Task<VirtualMachineData?> GetVm(string name) {
         // _logTracer.Debug($"getting vm: {name}");
         try {
-            return await _context.Creds.GetResourceGroupResource().GetVirtualMachineAsync(name);
+            var result = await _context.Creds.GetResourceGroupResource().GetVirtualMachineAsync(name, InstanceViewTypes.InstanceView);
+            if (result == null) {
+                return null;
+            }
+            if (result.Value.HasData) {
+                return result.Value.Data;
+            }
+
         } catch (RequestFailedException) {
             // _logTracer.Debug($"vm does not exist {ex});
             return null;
         }
+
+        return null;
     }
 
     public async Async.Task<bool> Delete(Vm vm) {