fix some bugs (#2406)

* - fix queries in timer retention

- do not discard proxy record after proxy state is processed, since that record needs to persist

* addressing comments

Co-authored-by: stas <statis@microsoft.com>

src/ApiService/ApiService/Functions/TimerProxy.cs

@@ -20,7 +20,9 @@ public async Async.Task Run([TimerTrigger("00:00:30")] TimerInfo myTimer) {
 
         var proxies = await proxyOperations.QueryAsync().ToListAsync();
 
-        foreach (var proxy in proxies) {
+        foreach (var p in proxies) {
+            var proxy = p;
+
             if (VmStateHelper.Available.Contains(proxy.State)) {
                 // Note, outdated checked at the start, but set at the end of this loop.
                 // As this function is called via a timer, this works around a user
@@ -40,11 +42,14 @@ public async Async.Task Run([TimerTrigger("00:00:30")] TimerInfo myTimer) {
 
             if (VmStateHelper.NeedsWork.Contains(proxy.State)) {
                 _logger.Info($"scaleset-proxy: update state. proxy:{proxy.Region} state:{proxy.State}");
-                await proxyOperations.ProcessStateUpdate(proxy);
+                proxy = await proxyOperations.ProcessStateUpdate(proxy);
             }
 
-            if (proxy.State != VmState.Stopped && proxyOperations.IsOutdated(proxy)) {
-                await proxyOperations.Replace(proxy with { Outdated = true });
+            if (proxy is not null && proxy.State != VmState.Stopped && proxyOperations.IsOutdated(proxy)) {
+                var r = await proxyOperations.Replace(proxy with { Outdated = true });
+                if (!r.IsOk) {
+                    _logger.Error($"Failed to replace proxy recordy for proxy {proxy.ProxyId} due to {r.ErrorV}");
+                }
             }
         }
 
@@ -54,8 +59,8 @@ public async Async.Task Run([TimerTrigger("00:00:30")] TimerInfo myTimer) {
         foreach (var region in regions) {
             var allOutdated = proxies.Where(x => x.Region == region).All(p => p.Outdated);
             if (allOutdated) {
-                await proxyOperations.GetOrCreate(region);
-                _logger.Info($"Creating new proxy in region {region}");
+                var proxy = await proxyOperations.GetOrCreate(region);
+                _logger.Info($"Creating new proxy with id {proxy.ProxyId} in region {region}");
             }
 
             // this is required in order to support upgrade from non-nsg to

src/ApiService/ApiService/Functions/TimerRetention.cs

@@ -1,5 +1,5 @@
-using Microsoft.Azure.Functions.Worker;
-
+using ApiService.OneFuzzLib.Orm;
+using Microsoft.Azure.Functions.Worker;
 namespace Microsoft.OneFuzz.Service.Functions;
 
 public class TimerRetention {
@@ -33,8 +33,8 @@ public async Async.Task Run([TimerTrigger("20:00:00")] TimerInfo t) {
         var timeRetainedOlder = now - RETENTION_POLICY;
         var timeRetainedNewer = now + SEARCH_EXTENT;
 
-        var timeFilter = $"Timestamp lt datetime'{timeRetainedOlder.ToString("o")}' and Timestamp gt datetime'{timeRetainedNewer.ToString("o")}'";
-        var timeFilterNewer = $"Timestamp gt datetime '{timeRetainedOlder.ToString("o")}'";
+        var timeFilter = Query.TimeRange(timeRetainedOlder, timeRetainedNewer);
+        var timeFilterNewer = Query.NewerThan(timeRetainedOlder);
 
         // Collecting 'still relevant' task containers.
         // NOTE: This must be done before potentially modifying tasks otherwise
@@ -52,7 +52,6 @@ from container in task.Config.Containers
             }
         }
 
-
         await foreach (var notification in _notificaitonOps.QueryAsync(timeFilter)) {
             _log.Verbose($"checking expired notification for removal: {notification.NotificationId}");
             var container = notification.Container;
@@ -66,7 +65,7 @@ from container in task.Config.Containers
             }
         }
 
-        await foreach (var job in _jobOps.QueryAsync($"{timeFilter} and state eq '{JobState.Enabled}'")) {
+        await foreach (var job in _jobOps.QueryAsync(Query.And(timeFilter, Query.EqualEnum("state", JobState.Enabled)))) {
             if (job.UserInfo is not null && job.UserInfo.Upn is not null) {
                 _log.Info($"removing PII from job {job.JobId}");
                 var userInfo = job.UserInfo with { Upn = null };
@@ -78,7 +77,7 @@ from container in task.Config.Containers
             }
         }
 
-        await foreach (var task in _taskOps.QueryAsync($"{timeFilter} and state eq '{TaskState.Stopped}'")) {
+        await foreach (var task in _taskOps.QueryAsync(Query.And(timeFilter, Query.EqualEnum("state", TaskState.Stopped)))) {
             if (task.UserInfo is not null && task.UserInfo.Upn is not null) {
                 _log.Info($"removing PII from task {task.TaskId}");
                 var userInfo = task.UserInfo with { Upn = null };

src/ApiService/ApiService/UserCredentials.cs

@@ -82,7 +82,7 @@ from t in token.Claims
 
                     return OneFuzzResult<UserInfo>.Ok(new(applicationId, objectId, upn));
                 } else {
-                    _log.Error($"issuer not from allowed tenant: {token.Issuer} - {allowedTenants}");
+                    _log.Error($"issuer not from allowed tenant. issuer: {token.Issuer} - tenants: {allowedTenants.OkV}");
                     return OneFuzzResult<UserInfo>.Error(ErrorCode.INVALID_REQUEST, new[] { "unauthorized AAD issuer" });
                 }
             } else {

src/ApiService/ApiService/onefuzzlib/ProxyOperations.cs

@@ -15,7 +15,7 @@ public interface IProxyOperations : IStatefulOrm<Proxy, VmState> {
     bool IsAlive(Proxy proxy);
     Async.Task SaveProxyConfig(Proxy proxy);
     bool IsOutdated(Proxy proxy);
-    Async.Task<Proxy?> GetOrCreate(Region region);
+    Async.Task<Proxy> GetOrCreate(Region region);
     Task<bool> IsUsed(Proxy proxy);
 
     // state transitions:
@@ -42,7 +42,7 @@ public ProxyOperations(ILogTracer log, IOnefuzzContext context)
         return await data.FirstOrDefaultAsync();
     }
 
-    public async Async.Task<Proxy?> GetOrCreate(Region region) {
+    public async Async.Task<Proxy> GetOrCreate(Region region) {
         var proxyList = QueryAsync(filter: TableClient.CreateQueryFilter($"region eq {region.String} and outdated eq false"));
 
         await foreach (var proxy in proxyList) {

src/ApiService/ApiService/onefuzzlib/VmOperations.cs

@@ -218,8 +218,10 @@ await vm.Value.GetVirtualMachineExtensions().CreateOrUpdateAsync(
                 extensionName,
                 extension
             );
-        } catch (RequestFailedException ex) when (ex.Status == 409 && ex.Message.Contains("VM is marked for deletion")) {
-            _logTracer.Info(ex.Message);
+        } catch (RequestFailedException ex) when
+            (ex.Status == 409 &&
+                (ex.Message.Contains("VM is marked for deletion") || ex.Message.Contains("The request failed due to conflict with a concurrent request."))) {
+            _logTracer.Info($"Tried to create extension but failed due to {ex.Message}");
         }
         return;
     }

src/ApiService/ApiService/onefuzzlib/orm/Queries.cs

@@ -46,12 +46,25 @@ public static string EqualAnyEnum<T>(string property, IEnumerable<T> enums) wher
             return EqualAny(property, convertedEnums);
         }
 
+        public static string EqualEnum<T>(string property, T e) where T : Enum {
+            string convertedEnum = JsonSerializer.Serialize(e, EntityConverter.GetJsonSerializerOptions()).Trim('"');
+            return $"{property} eq '{EscapeString(convertedEnum)}'";
+        }
+
         public static string TimeRange(DateTimeOffset min, DateTimeOffset max) {
             // NB: this uses the auto-populated Timestamp property, and will result in a table scan
             // TODO: should this be inclusive at the endpoints?
             return TableClient.CreateQueryFilter($"Timestamp lt {max} and Timestamp gt {min}");
         }
 
+        public static string NewerThan(DateTimeOffset t) {
+            return TableClient.CreateQueryFilter($"Timestamp gt {t}");
+        }
+
+        public static string OlderThan(DateTimeOffset t) {
+            return TableClient.CreateQueryFilter($"Timestamp lt {t}");
+        }
+
         public static string StartsWith(string property, string prefix) {
             var upperBound = prefix[..(prefix.Length - 1)] + (char)(prefix.Last() + 1);
             // property name should not be escaped, but strings should be:

src/ApiService/FunctionalTests/1f-api/ApiBase.cs

@@ -176,10 +176,9 @@ public ApiBase(Uri endpoint, string relativeUri, Microsoft.OneFuzz.Service.Reque
         }
     }
 
-    public async Task<JsonElement> Get(JsonObject root) {
+    public async Task<JsonElement> Get(JsonObject root, string? subPath = null) {
         var body = root.ToJsonString();
-        var r = await _request.Get(_endpoint, body);
-        var ss = await r.Content.ReadAsStringAsync();
+        var r = await _request.Get(subPath is null ? _endpoint : new Uri($"{_endpoint}{subPath}"), body);
         return (await JsonDocument.ParseAsync(r.Content.ReadAsStream())).RootElement;
     }