Scoped notification pause (#3579)

* Laying groundwork * Almost there * CLI updates * Remove unused code * cleanup * Broken test * fmt * . * 🥹 * forgot a file * Move from PUT to PATCH
microsoft · Oct 30, 2023 · d50fd48 · d50fd48
1 parent 72d775f
commit d50fd48
Show file tree

Hide file tree

Showing 10 changed files with 100 additions and 14 deletions.
diff --git a/src/ApiService/ApiService/Functions/Containers.cs b/src/ApiService/ApiService/Functions/Containers.cs
@@ -16,11 +16,12 @@ public ContainersFunction(ILogger<ContainersFunction> logger, IOnefuzzContext co
 
     [Function("Containers")]
     [Authorize(Allow.User)]
-    public Async.Task<HttpResponseData> Run([HttpTrigger(AuthorizationLevel.Anonymous, "GET", "POST", "DELETE")] HttpRequestData req)
+    public Async.Task<HttpResponseData> Run([HttpTrigger(AuthorizationLevel.Anonymous, "GET", "POST", "PATCH", "DELETE")] HttpRequestData req)
         => req.Method switch {
             "GET" => Get(req),
             "POST" => Post(req),
             "DELETE" => Delete(req),
+            "PATCH" => Patch(req),
             _ => throw new NotSupportedException(),
         };
 
@@ -108,4 +109,21 @@ private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
                 SasUrl: sas,
                 Metadata: post.Metadata));
     }
+
+    private async Async.Task<HttpResponseData> Patch(HttpRequestData req) {
+        var request = await RequestHandling.ParseRequest<ContainerUpdate>(req);
+        if (!request.IsOk) {
+            return await _context.RequestHandling.NotOk(req, request.ErrorV, context: "container update");
+        }
+
+        var toUpdate = request.OkV;
+        _logger.LogInformation("updating {ContainerName}", toUpdate.Name);
+        var updated = await _context.Containers.CreateOrUpdateContainerTag(toUpdate.Name, StorageType.Corpus, toUpdate.Metadata.ToDictionary(x => x.Key, x => x.Value));
+
+        if (!updated.IsOk) {
+            return await _context.RequestHandling.NotOk(req, updated.ErrorV, "container update");
+        }
+
+        return await RequestHandling.Ok(req, new ContainerInfoBase(toUpdate.Name, toUpdate.Metadata));
+    }
 }
diff --git a/src/ApiService/ApiService/Functions/QueueFileChanges.cs b/src/ApiService/ApiService/Functions/QueueFileChanges.cs
@@ -2,6 +2,8 @@
 using System.Text.Json.Nodes;
 using System.Threading.Tasks;
 using Azure.Core;
+using Azure.Storage.Blobs;
+using Azure.Storage.Blobs.Models;
 using Microsoft.Azure.Functions.Worker;
 using Microsoft.Extensions.Logging;
 using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
@@ -60,7 +62,11 @@ public async Async.Task Run(
         try {
             var result = await FileAdded(storageAccount, fileChangeEvent);
             if (!result.IsOk) {
-                await RequeueMessage(msg, result.ErrorV.Code == ErrorCode.ADO_WORKITEM_PROCESSING_DISABLED ? TimeSpan.FromDays(1) : null);
+                if (result.ErrorV.Code == ErrorCode.ADO_WORKITEM_PROCESSING_DISABLED) {
+                    await RequeueMessage(msg, TimeSpan.FromDays(1), incrementDequeueCount: false);
+                } else {
+                    await RequeueMessage(msg);
+                }
             }
         } catch (Exception e) {
             _log.LogError(e, "File Added failed");
@@ -83,21 +89,26 @@ private async Async.Task<OneFuzzResultVoid> FileAdded(ResourceIdentifier storage
 
         _log.LogInformation("file added : {Container} - {Path}", container.String, path);
 
+        var account = await _storage.GetBlobServiceClientForAccount(storageAccount);
+        var containerClient = account.GetBlobContainerClient(container.String);
+        var containerProps = await containerClient.GetPropertiesAsync();
+
+        if (_context.NotificationOperations.ShouldPauseNotificationsForContainer(containerProps.Value.Metadata)) {
+            return Error.Create(ErrorCode.ADO_WORKITEM_PROCESSING_DISABLED, $"container {container} has a metadata tag set to pause notifications processing");
+        }
+
         var (_, result) = await (
-            ApplyRetentionPolicy(storageAccount, container, path),
+            ApplyRetentionPolicy(containerClient, containerProps, path),
             _notificationOperations.NewFiles(container, path));
 
         return result;
     }
 
-    private async Async.Task<bool> ApplyRetentionPolicy(ResourceIdentifier storageAccount, Container container, string path) {
+    private async Async.Task<bool> ApplyRetentionPolicy(BlobContainerClient containerClient, BlobContainerProperties containerProps, string path) {
         if (await _context.FeatureManagerSnapshot.IsEnabledAsync(FeatureFlagConstants.EnableContainerRetentionPolicies)) {
             // default retention period can be applied to the container
             // if one exists, we will set the expiry date on the newly-created blob, if it doesn't already have one
-            var account = await _storage.GetBlobServiceClientForAccount(storageAccount);
-            var containerClient = account.GetBlobContainerClient(container.String);
-            var containerProps = await containerClient.GetPropertiesAsync();
-            var retentionPeriod = RetentionPolicyUtils.GetContainerRetentionPeriodFromMetadata(containerProps.Value.Metadata);
+            var retentionPeriod = RetentionPolicyUtils.GetContainerRetentionPeriodFromMetadata(containerProps.Metadata);
             if (!retentionPeriod.IsOk) {
                 _log.LogError("invalid retention period: {Error}", retentionPeriod.ErrorV);
             } else if (retentionPeriod.OkV is TimeSpan period) {
@@ -116,7 +127,7 @@ private async Async.Task<bool> ApplyRetentionPolicy(ResourceIdentifier storageAc
         return false;
     }
 
-    private async Async.Task RequeueMessage(string msg, TimeSpan? visibilityTimeout = null) {
+    private async Async.Task RequeueMessage(string msg, TimeSpan? visibilityTimeout = null, bool incrementDequeueCount = true) {
         var json = JsonNode.Parse(msg);
 
         // Messages that are 'manually' requeued by us as opposed to being requeued by the azure functions runtime
@@ -135,7 +146,9 @@ await _context.Queue.QueueObject(
                 StorageType.Config)
                 .IgnoreResult();
         } else {
-            json!["data"]!["customDequeueCount"] = newCustomDequeueCount + 1;
+            if (incrementDequeueCount) {
+                json!["data"]!["customDequeueCount"] = newCustomDequeueCount + 1;
+            }
             await _context.Queue.QueueObject(
                 QueueFileChangesQueueName,
                 json,

diff --git a/src/ApiService/ApiService/OneFuzzTypes/Enums.cs b/src/ApiService/ApiService/OneFuzzTypes/Enums.cs
@@ -53,6 +53,9 @@ public enum ErrorCode {
     INVALID_RETENTION_PERIOD = 497,
     INVALID_CLI_VERSION = 498,
     TRANSIENT_NOTIFICATION_FAILURE = 499,
+
+    FAILED_CONTAINER_PROPERTIES_ACCESS = 500,
+    FAILED_SAVING_CONTAINER_METADATA = 501,
     // NB: if you update this enum, also update enums.py
 }
 

diff --git a/src/ApiService/ApiService/OneFuzzTypes/Requests.cs b/src/ApiService/ApiService/OneFuzzTypes/Requests.cs
@@ -128,6 +128,11 @@ public record ContainerDelete(
     IDictionary<string, string>? Metadata = null
 ) : BaseRequest;
 
+public record ContainerUpdate(
+    [property: Required] Container Name,
+    [property: Required] IDictionary<string, string> Metadata
+) : BaseRequest;
+
 public record NotificationCreate(
     [property: Required] Container Container,
     [property: Required] bool ReplaceExisting,

diff --git a/src/ApiService/ApiService/onefuzzlib/Containers.cs b/src/ApiService/ApiService/onefuzzlib/Containers.cs
@@ -1,5 +1,6 @@
 using System.IO;
 using System.IO.Compression;
+using System.Text.Json;
 using System.Threading.Tasks;
 using ApiService.OneFuzzLib.Orm;
 using Azure;
@@ -41,6 +42,8 @@ public interface IContainers {
     public Async.Task<OneFuzzResultVoid> DownloadAsZip(Container container, StorageType storageType, Stream stream, string? prefix = null);
 
     public Async.Task DeleteAllExpiredBlobs();
+
+    public Async.Task<OneFuzzResultVoid> CreateOrUpdateContainerTag(Container container, StorageType storageType, Dictionary<string, string> tags);
 }
 
 public class Containers : Orm<ContainerInformation>, IContainers {
@@ -448,4 +451,29 @@ private async Async.Task DeleteExpiredBlobsForAccount(ResourceIdentifier storage
             }
         }
     }
+
+    public async Task<OneFuzzResultVoid> CreateOrUpdateContainerTag(Container container, StorageType storageType, Dictionary<string, string> tags) {
+        var client = await FindContainer(container, storageType);
+        if (client is null || !await client.ExistsAsync()) {
+            return Error.Create(ErrorCode.INVALID_CONTAINER, $"Could not find container {container} in {storageType}");
+        }
+
+        var metadataRequest = await client.GetPropertiesAsync();
+        if (metadataRequest is null || metadataRequest.GetRawResponse().IsError) {
+            return Error.Create(ErrorCode.FAILED_CONTAINER_PROPERTIES_ACCESS, $"Could not access container properties for container: {container} in {storageType}");
+        }
+
+        var metadata = metadataRequest.Value.Metadata ?? new Dictionary<string, string>();
+
+        foreach (var kvp in tags) {
+            metadata[kvp.Key] = kvp.Value;
+        }
+
+        var saveMetadataRequest = await client.SetMetadataAsync(metadata);
+        if (saveMetadataRequest is null || saveMetadataRequest.GetRawResponse().IsError) {
+            return Error.Create(ErrorCode.FAILED_SAVING_CONTAINER_METADATA, $"Could not save metadata to container: {container} in {storageType}. Metadata: {JsonSerializer.Serialize(metadata)}");
+        }
+
+        return OneFuzzResultVoid.Ok;
+    }
 }
diff --git a/src/ApiService/ApiService/onefuzzlib/NotificationOperations.cs b/src/ApiService/ApiService/onefuzzlib/NotificationOperations.cs
@@ -11,12 +11,11 @@ public interface INotificationOperations : IOrm<Notification> {
     Async.Task<OneFuzzResult<Notification>> Create(Container container, NotificationTemplate config, bool replaceExisting);
     Async.Task<Notification?> GetNotification(Guid notifificationId);
 
-    System.Threading.Tasks.Task<OneFuzzResultVoid> TriggerNotification(Container container,
-        Notification notification, IReport? reportOrRegression);
+    System.Threading.Tasks.Task<OneFuzzResultVoid> TriggerNotification(Container container, Notification notification, IReport? reportOrRegression);
+    bool ShouldPauseNotificationsForContainer(IDictionary<string, string> containerMetadata);
 }
 
 public class NotificationOperations : Orm<Notification>, INotificationOperations {
-
     public NotificationOperations(ILogger<NotificationOperations> log, IOnefuzzContext context)
         : base(log, context) {
 
@@ -190,4 +189,7 @@ private async Async.Task<NotificationTemplate> HideSecrets(NotificationTemplate
     public async Async.Task<Notification?> GetNotification(Guid notifificationId) {
         return await SearchByPartitionKeys(new[] { notifificationId.ToString() }).SingleOrDefaultAsync();
     }
+
+    private const string PAUSE_NOTIFICATIONS_TAG = "pauseNotifications";
+    public bool ShouldPauseNotificationsForContainer(IDictionary<string, string> containerMetadata) => containerMetadata.ContainsKey(PAUSE_NOTIFICATIONS_TAG) && containerMetadata[PAUSE_NOTIFICATIONS_TAG] == "true";
 }
diff --git a/src/ApiService/ApiService/onefuzzlib/Reports.cs b/src/ApiService/ApiService/onefuzzlib/Reports.cs
@@ -67,7 +67,6 @@ public Reports(ILogger<Reports> log, IContainers containers) {
     }
 
     private static T? TryDeserialize<T>(string content) where T : class {
-
         try {
             return JsonSerializer.Deserialize<T>(content, EntityConverter.GetJsonSerializerOptions());
         } catch (JsonException) {

diff --git a/src/cli/onefuzz/api.py b/src/cli/onefuzz/api.py
@@ -486,6 +486,17 @@ def delete(self, name: str) -> responses.BoolResult:
             "DELETE", responses.BoolResult, data=requests.ContainerDelete(name=name)
         )
 
+    def update(
+        self, name: str, metadata: Dict[str, str]
+    ) -> responses.ContainerInfoBase:
+        """Update a container's metadata"""
+        self.logger.debug("update container: %s", name)
+        return self._req_model(
+            "PATCH",
+            responses.ContainerInfoBase,
+            data=requests.ContainerUpdate(name=name, metadata=metadata),
+        )
+
     def list(self) -> List[responses.ContainerInfoBase]:
         """Get a list of containers"""
         self.logger.debug("list containers")

diff --git a/src/pytypes/onefuzztypes/enums.py b/src/pytypes/onefuzztypes/enums.py
@@ -307,6 +307,8 @@ class ErrorCode(Enum):
     INVALID_RETENTION_PERIOD = 497
     INVALID_CLI_VERSION = 498
     TRANSIENT_NOTIFICATION_FAILURE = 499
+    FAILED_CONTAINER_PROPERTIES_ACCESS = 500
+    FAILED_SAVING_CONTAINER_METADATA = 501
     # NB: if you update this enum, also update Enums.cs
 
 

diff --git a/src/pytypes/onefuzztypes/requests.py b/src/pytypes/onefuzztypes/requests.py
@@ -220,6 +220,11 @@ class ContainerDelete(BaseRequest):
     metadata: Optional[Dict[str, str]]
 
 
+class ContainerUpdate(BaseRequest):
+    name: Container
+    metadata: Dict[str, str]
+
+
 class ReproGet(BaseRequest):
     vm_id: Optional[UUID]