From d50fd48a8b36cd0f902d79457b8dcd3d56832c8b Mon Sep 17 00:00:00 2001 From: Teo Voinea <58236992+tevoinea@users.noreply.github.com> Date: Mon, 30 Oct 2023 13:21:33 -0400 Subject: [PATCH] Scoped notification pause (#3579) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Laying groundwork * Almost there * CLI updates * Remove unused code * cleanup * Broken test * fmt * . * 🥹 * forgot a file * Move from PUT to PATCH --- .../ApiService/Functions/Containers.cs | 20 +++++++++++- .../ApiService/Functions/QueueFileChanges.cs | 31 +++++++++++++------ .../ApiService/OneFuzzTypes/Enums.cs | 3 ++ .../ApiService/OneFuzzTypes/Requests.cs | 5 +++ .../ApiService/onefuzzlib/Containers.cs | 28 +++++++++++++++++ .../onefuzzlib/NotificationOperations.cs | 8 +++-- .../ApiService/onefuzzlib/Reports.cs | 1 - src/cli/onefuzz/api.py | 11 +++++++ src/pytypes/onefuzztypes/enums.py | 2 ++ src/pytypes/onefuzztypes/requests.py | 5 +++ 10 files changed, 100 insertions(+), 14 deletions(-) diff --git a/src/ApiService/ApiService/Functions/Containers.cs b/src/ApiService/ApiService/Functions/Containers.cs index 96554c880e..5178f5be0f 100644 --- a/src/ApiService/ApiService/Functions/Containers.cs +++ b/src/ApiService/ApiService/Functions/Containers.cs @@ -16,11 +16,12 @@ public ContainersFunction(ILogger logger, IOnefuzzContext co [Function("Containers")] [Authorize(Allow.User)] - public Async.Task Run([HttpTrigger(AuthorizationLevel.Anonymous, "GET", "POST", "DELETE")] HttpRequestData req) + public Async.Task 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 Post(HttpRequestData req) { SasUrl: sas, Metadata: post.Metadata)); } + + private async Async.Task Patch(HttpRequestData req) { + var request = await RequestHandling.ParseRequest(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 index 8ef77bd2a5..f2aa08b306 100644 --- 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 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 ApplyRetentionPolicy(ResourceIdentifier storageAccount, Container container, string path) { + private async Async.Task 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 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 index 5a8b22527d..c0e3c68eba 100644 --- 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 index db63499d30..f3a4c32965 100644 --- a/src/ApiService/ApiService/OneFuzzTypes/Requests.cs +++ b/src/ApiService/ApiService/OneFuzzTypes/Requests.cs @@ -128,6 +128,11 @@ public record ContainerDelete( IDictionary? Metadata = null ) : BaseRequest; +public record ContainerUpdate( + [property: Required] Container Name, + [property: Required] IDictionary 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 index f7bb3086a0..e004815abc 100644 --- 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 DownloadAsZip(Container container, StorageType storageType, Stream stream, string? prefix = null); public Async.Task DeleteAllExpiredBlobs(); + + public Async.Task CreateOrUpdateContainerTag(Container container, StorageType storageType, Dictionary tags); } public class Containers : Orm, IContainers { @@ -448,4 +451,29 @@ private async Async.Task DeleteExpiredBlobsForAccount(ResourceIdentifier storage } } } + + public async Task CreateOrUpdateContainerTag(Container container, StorageType storageType, Dictionary 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(); + + 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 index 5744d6d2f5..67369989e1 100644 --- a/src/ApiService/ApiService/onefuzzlib/NotificationOperations.cs +++ b/src/ApiService/ApiService/onefuzzlib/NotificationOperations.cs @@ -11,12 +11,11 @@ public interface INotificationOperations : IOrm { Async.Task> Create(Container container, NotificationTemplate config, bool replaceExisting); Async.Task GetNotification(Guid notifificationId); - System.Threading.Tasks.Task TriggerNotification(Container container, - Notification notification, IReport? reportOrRegression); + System.Threading.Tasks.Task TriggerNotification(Container container, Notification notification, IReport? reportOrRegression); + bool ShouldPauseNotificationsForContainer(IDictionary containerMetadata); } public class NotificationOperations : Orm, INotificationOperations { - public NotificationOperations(ILogger log, IOnefuzzContext context) : base(log, context) { @@ -190,4 +189,7 @@ private async Async.Task HideSecrets(NotificationTemplate public async Async.Task GetNotification(Guid notifificationId) { return await SearchByPartitionKeys(new[] { notifificationId.ToString() }).SingleOrDefaultAsync(); } + + private const string PAUSE_NOTIFICATIONS_TAG = "pauseNotifications"; + public bool ShouldPauseNotificationsForContainer(IDictionary 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 index fdda7259e9..ac2e1029b2 100644 --- a/src/ApiService/ApiService/onefuzzlib/Reports.cs +++ b/src/ApiService/ApiService/onefuzzlib/Reports.cs @@ -67,7 +67,6 @@ public Reports(ILogger log, IContainers containers) { } private static T? TryDeserialize(string content) where T : class { - try { return JsonSerializer.Deserialize(content, EntityConverter.GetJsonSerializerOptions()); } catch (JsonException) { diff --git a/src/cli/onefuzz/api.py b/src/cli/onefuzz/api.py index 64cad8c368..4f4e152484 100644 --- 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 index 446193f1d2..14315ddb5d 100644 --- 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 index d284fb416d..df9fb3e1f3 100644 --- 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]