Skip to content
This repository was archived by the owner on Nov 1, 2023. It is now read-only.

Commit d50fd48

Browse files
authored
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
1 parent 72d775f commit d50fd48

File tree

10 files changed

+100
-14
lines changed

10 files changed

+100
-14
lines changed

src/ApiService/ApiService/Functions/Containers.cs

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,12 @@ public ContainersFunction(ILogger<ContainersFunction> logger, IOnefuzzContext co
1616

1717
[Function("Containers")]
1818
[Authorize(Allow.User)]
19-
public Async.Task<HttpResponseData> Run([HttpTrigger(AuthorizationLevel.Anonymous, "GET", "POST", "DELETE")] HttpRequestData req)
19+
public Async.Task<HttpResponseData> Run([HttpTrigger(AuthorizationLevel.Anonymous, "GET", "POST", "PATCH", "DELETE")] HttpRequestData req)
2020
=> req.Method switch {
2121
"GET" => Get(req),
2222
"POST" => Post(req),
2323
"DELETE" => Delete(req),
24+
"PATCH" => Patch(req),
2425
_ => throw new NotSupportedException(),
2526
};
2627

@@ -108,4 +109,21 @@ private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
108109
SasUrl: sas,
109110
Metadata: post.Metadata));
110111
}
112+
113+
private async Async.Task<HttpResponseData> Patch(HttpRequestData req) {
114+
var request = await RequestHandling.ParseRequest<ContainerUpdate>(req);
115+
if (!request.IsOk) {
116+
return await _context.RequestHandling.NotOk(req, request.ErrorV, context: "container update");
117+
}
118+
119+
var toUpdate = request.OkV;
120+
_logger.LogInformation("updating {ContainerName}", toUpdate.Name);
121+
var updated = await _context.Containers.CreateOrUpdateContainerTag(toUpdate.Name, StorageType.Corpus, toUpdate.Metadata.ToDictionary(x => x.Key, x => x.Value));
122+
123+
if (!updated.IsOk) {
124+
return await _context.RequestHandling.NotOk(req, updated.ErrorV, "container update");
125+
}
126+
127+
return await RequestHandling.Ok(req, new ContainerInfoBase(toUpdate.Name, toUpdate.Metadata));
128+
}
111129
}

src/ApiService/ApiService/Functions/QueueFileChanges.cs

Lines changed: 22 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
using System.Text.Json.Nodes;
33
using System.Threading.Tasks;
44
using Azure.Core;
5+
using Azure.Storage.Blobs;
6+
using Azure.Storage.Blobs.Models;
57
using Microsoft.Azure.Functions.Worker;
68
using Microsoft.Extensions.Logging;
79
using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
@@ -60,7 +62,11 @@ public async Async.Task Run(
6062
try {
6163
var result = await FileAdded(storageAccount, fileChangeEvent);
6264
if (!result.IsOk) {
63-
await RequeueMessage(msg, result.ErrorV.Code == ErrorCode.ADO_WORKITEM_PROCESSING_DISABLED ? TimeSpan.FromDays(1) : null);
65+
if (result.ErrorV.Code == ErrorCode.ADO_WORKITEM_PROCESSING_DISABLED) {
66+
await RequeueMessage(msg, TimeSpan.FromDays(1), incrementDequeueCount: false);
67+
} else {
68+
await RequeueMessage(msg);
69+
}
6470
}
6571
} catch (Exception e) {
6672
_log.LogError(e, "File Added failed");
@@ -83,21 +89,26 @@ private async Async.Task<OneFuzzResultVoid> FileAdded(ResourceIdentifier storage
8389

8490
_log.LogInformation("file added : {Container} - {Path}", container.String, path);
8591

92+
var account = await _storage.GetBlobServiceClientForAccount(storageAccount);
93+
var containerClient = account.GetBlobContainerClient(container.String);
94+
var containerProps = await containerClient.GetPropertiesAsync();
95+
96+
if (_context.NotificationOperations.ShouldPauseNotificationsForContainer(containerProps.Value.Metadata)) {
97+
return Error.Create(ErrorCode.ADO_WORKITEM_PROCESSING_DISABLED, $"container {container} has a metadata tag set to pause notifications processing");
98+
}
99+
86100
var (_, result) = await (
87-
ApplyRetentionPolicy(storageAccount, container, path),
101+
ApplyRetentionPolicy(containerClient, containerProps, path),
88102
_notificationOperations.NewFiles(container, path));
89103

90104
return result;
91105
}
92106

93-
private async Async.Task<bool> ApplyRetentionPolicy(ResourceIdentifier storageAccount, Container container, string path) {
107+
private async Async.Task<bool> ApplyRetentionPolicy(BlobContainerClient containerClient, BlobContainerProperties containerProps, string path) {
94108
if (await _context.FeatureManagerSnapshot.IsEnabledAsync(FeatureFlagConstants.EnableContainerRetentionPolicies)) {
95109
// default retention period can be applied to the container
96110
// if one exists, we will set the expiry date on the newly-created blob, if it doesn't already have one
97-
var account = await _storage.GetBlobServiceClientForAccount(storageAccount);
98-
var containerClient = account.GetBlobContainerClient(container.String);
99-
var containerProps = await containerClient.GetPropertiesAsync();
100-
var retentionPeriod = RetentionPolicyUtils.GetContainerRetentionPeriodFromMetadata(containerProps.Value.Metadata);
111+
var retentionPeriod = RetentionPolicyUtils.GetContainerRetentionPeriodFromMetadata(containerProps.Metadata);
101112
if (!retentionPeriod.IsOk) {
102113
_log.LogError("invalid retention period: {Error}", retentionPeriod.ErrorV);
103114
} else if (retentionPeriod.OkV is TimeSpan period) {
@@ -116,7 +127,7 @@ private async Async.Task<bool> ApplyRetentionPolicy(ResourceIdentifier storageAc
116127
return false;
117128
}
118129

119-
private async Async.Task RequeueMessage(string msg, TimeSpan? visibilityTimeout = null) {
130+
private async Async.Task RequeueMessage(string msg, TimeSpan? visibilityTimeout = null, bool incrementDequeueCount = true) {
120131
var json = JsonNode.Parse(msg);
121132

122133
// 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(
135146
StorageType.Config)
136147
.IgnoreResult();
137148
} else {
138-
json!["data"]!["customDequeueCount"] = newCustomDequeueCount + 1;
149+
if (incrementDequeueCount) {
150+
json!["data"]!["customDequeueCount"] = newCustomDequeueCount + 1;
151+
}
139152
await _context.Queue.QueueObject(
140153
QueueFileChangesQueueName,
141154
json,

src/ApiService/ApiService/OneFuzzTypes/Enums.cs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,9 @@ public enum ErrorCode {
5353
INVALID_RETENTION_PERIOD = 497,
5454
INVALID_CLI_VERSION = 498,
5555
TRANSIENT_NOTIFICATION_FAILURE = 499,
56+
57+
FAILED_CONTAINER_PROPERTIES_ACCESS = 500,
58+
FAILED_SAVING_CONTAINER_METADATA = 501,
5659
// NB: if you update this enum, also update enums.py
5760
}
5861

src/ApiService/ApiService/OneFuzzTypes/Requests.cs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,11 @@ public record ContainerDelete(
128128
IDictionary<string, string>? Metadata = null
129129
) : BaseRequest;
130130

131+
public record ContainerUpdate(
132+
[property: Required] Container Name,
133+
[property: Required] IDictionary<string, string> Metadata
134+
) : BaseRequest;
135+
131136
public record NotificationCreate(
132137
[property: Required] Container Container,
133138
[property: Required] bool ReplaceExisting,

src/ApiService/ApiService/onefuzzlib/Containers.cs

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using System.IO;
22
using System.IO.Compression;
3+
using System.Text.Json;
34
using System.Threading.Tasks;
45
using ApiService.OneFuzzLib.Orm;
56
using Azure;
@@ -41,6 +42,8 @@ public interface IContainers {
4142
public Async.Task<OneFuzzResultVoid> DownloadAsZip(Container container, StorageType storageType, Stream stream, string? prefix = null);
4243

4344
public Async.Task DeleteAllExpiredBlobs();
45+
46+
public Async.Task<OneFuzzResultVoid> CreateOrUpdateContainerTag(Container container, StorageType storageType, Dictionary<string, string> tags);
4447
}
4548

4649
public class Containers : Orm<ContainerInformation>, IContainers {
@@ -448,4 +451,29 @@ private async Async.Task DeleteExpiredBlobsForAccount(ResourceIdentifier storage
448451
}
449452
}
450453
}
454+
455+
public async Task<OneFuzzResultVoid> CreateOrUpdateContainerTag(Container container, StorageType storageType, Dictionary<string, string> tags) {
456+
var client = await FindContainer(container, storageType);
457+
if (client is null || !await client.ExistsAsync()) {
458+
return Error.Create(ErrorCode.INVALID_CONTAINER, $"Could not find container {container} in {storageType}");
459+
}
460+
461+
var metadataRequest = await client.GetPropertiesAsync();
462+
if (metadataRequest is null || metadataRequest.GetRawResponse().IsError) {
463+
return Error.Create(ErrorCode.FAILED_CONTAINER_PROPERTIES_ACCESS, $"Could not access container properties for container: {container} in {storageType}");
464+
}
465+
466+
var metadata = metadataRequest.Value.Metadata ?? new Dictionary<string, string>();
467+
468+
foreach (var kvp in tags) {
469+
metadata[kvp.Key] = kvp.Value;
470+
}
471+
472+
var saveMetadataRequest = await client.SetMetadataAsync(metadata);
473+
if (saveMetadataRequest is null || saveMetadataRequest.GetRawResponse().IsError) {
474+
return Error.Create(ErrorCode.FAILED_SAVING_CONTAINER_METADATA, $"Could not save metadata to container: {container} in {storageType}. Metadata: {JsonSerializer.Serialize(metadata)}");
475+
}
476+
477+
return OneFuzzResultVoid.Ok;
478+
}
451479
}

src/ApiService/ApiService/onefuzzlib/NotificationOperations.cs

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,11 @@ public interface INotificationOperations : IOrm<Notification> {
1111
Async.Task<OneFuzzResult<Notification>> Create(Container container, NotificationTemplate config, bool replaceExisting);
1212
Async.Task<Notification?> GetNotification(Guid notifificationId);
1313

14-
System.Threading.Tasks.Task<OneFuzzResultVoid> TriggerNotification(Container container,
15-
Notification notification, IReport? reportOrRegression);
14+
System.Threading.Tasks.Task<OneFuzzResultVoid> TriggerNotification(Container container, Notification notification, IReport? reportOrRegression);
15+
bool ShouldPauseNotificationsForContainer(IDictionary<string, string> containerMetadata);
1616
}
1717

1818
public class NotificationOperations : Orm<Notification>, INotificationOperations {
19-
2019
public NotificationOperations(ILogger<NotificationOperations> log, IOnefuzzContext context)
2120
: base(log, context) {
2221

@@ -190,4 +189,7 @@ private async Async.Task<NotificationTemplate> HideSecrets(NotificationTemplate
190189
public async Async.Task<Notification?> GetNotification(Guid notifificationId) {
191190
return await SearchByPartitionKeys(new[] { notifificationId.ToString() }).SingleOrDefaultAsync();
192191
}
192+
193+
private const string PAUSE_NOTIFICATIONS_TAG = "pauseNotifications";
194+
public bool ShouldPauseNotificationsForContainer(IDictionary<string, string> containerMetadata) => containerMetadata.ContainsKey(PAUSE_NOTIFICATIONS_TAG) && containerMetadata[PAUSE_NOTIFICATIONS_TAG] == "true";
193195
}

src/ApiService/ApiService/onefuzzlib/Reports.cs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,6 @@ public Reports(ILogger<Reports> log, IContainers containers) {
6767
}
6868

6969
private static T? TryDeserialize<T>(string content) where T : class {
70-
7170
try {
7271
return JsonSerializer.Deserialize<T>(content, EntityConverter.GetJsonSerializerOptions());
7372
} catch (JsonException) {

src/cli/onefuzz/api.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -486,6 +486,17 @@ def delete(self, name: str) -> responses.BoolResult:
486486
"DELETE", responses.BoolResult, data=requests.ContainerDelete(name=name)
487487
)
488488

489+
def update(
490+
self, name: str, metadata: Dict[str, str]
491+
) -> responses.ContainerInfoBase:
492+
"""Update a container's metadata"""
493+
self.logger.debug("update container: %s", name)
494+
return self._req_model(
495+
"PATCH",
496+
responses.ContainerInfoBase,
497+
data=requests.ContainerUpdate(name=name, metadata=metadata),
498+
)
499+
489500
def list(self) -> List[responses.ContainerInfoBase]:
490501
"""Get a list of containers"""
491502
self.logger.debug("list containers")

src/pytypes/onefuzztypes/enums.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -307,6 +307,8 @@ class ErrorCode(Enum):
307307
INVALID_RETENTION_PERIOD = 497
308308
INVALID_CLI_VERSION = 498
309309
TRANSIENT_NOTIFICATION_FAILURE = 499
310+
FAILED_CONTAINER_PROPERTIES_ACCESS = 500
311+
FAILED_SAVING_CONTAINER_METADATA = 501
310312
# NB: if you update this enum, also update Enums.cs
311313

312314

src/pytypes/onefuzztypes/requests.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,11 @@ class ContainerDelete(BaseRequest):
220220
metadata: Optional[Dict[str, str]]
221221

222222

223+
class ContainerUpdate(BaseRequest):
224+
name: Container
225+
metadata: Dict[str, str]
226+
227+
223228
class ReproGet(BaseRequest):
224229
vm_id: Optional[UUID]
225230

0 commit comments

Comments
 (0)