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

Commit

Permalink
Rework use of storage clients (#2302)
Browse files Browse the repository at this point in the history 
- Reuse Storage clients by caching them against account name.
  - This follows recommendations of Storage documentation: https://devblogs.microsoft.com/azure-sdk/lifetime-management-and-thread-safety-guarantees-of-azure-sdk-net-clients/
- Use strongly-typed `ResourceIdentifier` when passing around Account IDs, to prevent mixing them up with Account Names. 
- Simplify & centralize SAS generation in one place.
  • Loading branch information
@Porges
Porges authored Sep 6, 2022
1 parent 318e568 commit 84b2cc9
Show file tree
Hide file tree
Showing 14 changed files with 265 additions and 217 deletions.
3 changes: 2 additions & 1 deletion src/ApiService/ApiService/Functions/QueueFileChanges.cs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
using System.Text.Json;
using Azure.Core;
using Microsoft.Azure.Functions.Worker;
using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;

Expand Down Expand Up @@ -39,7 +40,7 @@ public async Async.Task Run(

const string topic = "topic";
if (!fileChangeEvent.RootElement.TryGetProperty(topic, out var topicElement)
|| !_storage.CorpusAccounts().Contains(topicElement.GetString())) {
|| !_storage.CorpusAccounts().Contains(new ResourceIdentifier(topicElement.GetString()!))) {
return;
}

Expand Down
21 changes: 17 additions & 4 deletions src/ApiService/ApiService/ServiceConfiguration.cs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
using System.Reflection;
using Azure.Core;

namespace Microsoft.OneFuzz.Service;

Expand All @@ -25,8 +26,8 @@ public interface IServiceConfig {
public string? DiagnosticsAzureBlobRetentionDays { get; }

public string? MultiTenantDomain { get; }
public string? OneFuzzDataStorage { get; }
public string? OneFuzzFuncStorage { get; }
public ResourceIdentifier? OneFuzzDataStorage { get; }
public ResourceIdentifier? OneFuzzFuncStorage { get; }
public string? OneFuzzInstance { get; }
public string? OneFuzzInstanceName { get; }
public string? OneFuzzKeyvault { get; }
Expand Down Expand Up @@ -92,8 +93,20 @@ public ServiceConfiguration() {

public string? MultiTenantDomain { get => GetEnv("MULTI_TENANT_DOMAIN"); }

public string? OneFuzzDataStorage { get => GetEnv("ONEFUZZ_DATA_STORAGE"); }
public string? OneFuzzFuncStorage { get => GetEnv("ONEFUZZ_FUNC_STORAGE"); }
public ResourceIdentifier? OneFuzzDataStorage {
get {
var env = GetEnv("ONEFUZZ_DATA_STORAGE");
return env is null ? null : new ResourceIdentifier(env);
}
}

public ResourceIdentifier? OneFuzzFuncStorage {
get {
var env = GetEnv("ONEFUZZ_FUNC_STORAGE");
return env is null ? null : new ResourceIdentifier(env);
}
}

public string? OneFuzzInstance { get => GetEnv("ONEFUZZ_INSTANCE"); }
public string? OneFuzzInstanceName { get => GetEnv("ONEFUZZ_INSTANCE_NAME"); }
public string? OneFuzzKeyvault { get => GetEnv("ONEFUZZ_KEYVAULT"); }
Expand Down
88 changes: 27 additions & 61 deletions src/ApiService/ApiService/onefuzzlib/Containers.cs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
using System.Threading;
using System.Threading.Tasks;
using Azure;
using Azure.ResourceManager;
using Azure.Storage;
using Azure.Storage.Blobs;
using Azure.Storage.Blobs.Models;
using Azure.Storage.Sas;
Expand Down Expand Up @@ -34,19 +32,15 @@ public interface IContainers {
}

public class Containers : IContainers {
private ILogTracer _log;
private IStorage _storage;
private ICreds _creds;
private ArmClient _armClient;
private readonly ILogTracer _log;
private readonly IStorage _storage;
private readonly IServiceConfig _config;

static TimeSpan CONTAINER_SAS_DEFAULT_DURATION = TimeSpan.FromDays(30);

public Containers(ILogTracer log, IStorage storage, ICreds creds, IServiceConfig config) {
public Containers(ILogTracer log, IStorage storage, IServiceConfig config) {
_log = log;
_storage = storage;
_creds = creds;
_armClient = creds.ArmClient;
_config = config;

_getInstanceId = new Lazy<Async.Task<Guid>>(async () => {
Expand All @@ -64,7 +58,7 @@ public Containers(ILogTracer log, IStorage storage, ICreds creds, IServiceConfig
if (client is null)
return null;

return new Uri($"{_storage.GetBlobEndpoint(client.AccountName)}{container}/{name}");
return client.GetBlobClient(name).Uri;
}

public async Async.Task<BinaryData?> GetBlob(Container container, string name, StorageType storageType) {
Expand Down Expand Up @@ -131,33 +125,22 @@ private static readonly BlobContainerSasPermissions _containerCreatePermissions

var containerName = _config.OneFuzzStoragePrefix + container.ContainerName;

var containers =
_storage.GetAccounts(storageType)
.Reverse()
.Select(async account => (await _storage.GetBlobServiceClientForAccount(account)).GetBlobContainerClient(containerName));

foreach (var c in containers) {
var client = await c;
if ((await client.ExistsAsync()).Value) {
return client;
foreach (var account in _storage.GetAccounts(storageType).Reverse()) {
var accountClient = await _storage.GetBlobServiceClientForAccount(account);
var containerClient = accountClient.GetBlobContainerClient(containerName);
if (await containerClient.ExistsAsync()) {
return containerClient;
}
}

return null;
}

public async Async.Task<Uri> GetFileSasUrl(Container container, string name, StorageType storageType, BlobSasPermissions permissions, TimeSpan? duration = null) {
var client = await FindContainer(container, storageType) ?? throw new Exception($"unable to find container: {container.ContainerName} - {storageType}");

var (startTime, endTime) = SasTimeWindow(duration ?? TimeSpan.FromDays(30));

var sasBuilder = new BlobSasBuilder(permissions, endTime) {
StartsOn = startTime,
BlobContainerName = _config.OneFuzzStoragePrefix + container.ContainerName,
BlobName = name
};

var sasUrl = client.GetBlobClient(name).GenerateSasUri(sasBuilder);
return sasUrl;
var blobClient = client.GetBlobClient(name);
var timeWindow = SasTimeWindow(duration ?? TimeSpan.FromDays(30));
return _storage.GenerateBlobSasUri(permissions, blobClient, timeWindow);
}

public static (DateTimeOffset, DateTimeOffset) SasTimeWindow(TimeSpan timeSpan) {
Expand All @@ -184,47 +167,34 @@ public async Async.Task SaveBlob(Container container, string name, string data,
public virtual Async.Task<Guid> GetInstanceId() => _getInstanceId.Value;
private readonly Lazy<Async.Task<Guid>> _getInstanceId;

public static Uri? GetContainerSasUrlService(
public Uri GetContainerSasUrlService(
BlobContainerClient client,
BlobContainerSasPermissions permissions,
bool tag = false,
TimeSpan? timeSpan = null) {
var (start, expiry) = SasTimeWindow(timeSpan ?? TimeSpan.FromDays(30.0));
var sasBuilder = new BlobSasBuilder(permissions, expiry) { StartsOn = start };
return client.GenerateSasUri(sasBuilder);
var timeWindow = SasTimeWindow(timeSpan ?? TimeSpan.FromDays(30.0));
return _storage.GenerateBlobContainerSasUri(permissions, client, timeWindow);
}

public async Async.Task<Uri> AddContainerSasUrl(Uri uri, TimeSpan? duration = null) {
if (uri.Query.Contains("sig")) {
return uri;
}

var (startTime, endTime) = SasTimeWindow(duration ?? CONTAINER_SAS_DEFAULT_DURATION);
var blobUriBuilder = new BlobUriBuilder(uri);
var accountKey = await _storage.GetStorageAccountNameKeyByName(blobUriBuilder.AccountName);
var sasBuilder = new BlobSasBuilder(
BlobContainerSasPermissions.Read | BlobContainerSasPermissions.Write | BlobContainerSasPermissions.Delete | BlobContainerSasPermissions.List,
endTime) {
BlobContainerName = blobUriBuilder.BlobContainerName,
StartsOn = startTime
};

var sas = sasBuilder.ToSasQueryParameters(new StorageSharedKeyCredential(blobUriBuilder.AccountName, accountKey)).ToString();
return new UriBuilder(uri) {
Query = sas
}.Uri;
var serviceClient = await _storage.GetBlobServiceClientForAccountName(blobUriBuilder.AccountName);
var containerClient = serviceClient.GetBlobContainerClient(blobUriBuilder.BlobContainerName);

var permissions = BlobContainerSasPermissions.Read | BlobContainerSasPermissions.Write | BlobContainerSasPermissions.Delete | BlobContainerSasPermissions.List;

var timeWindow = SasTimeWindow(duration ?? CONTAINER_SAS_DEFAULT_DURATION);

return _storage.GenerateBlobContainerSasUri(permissions, containerClient, timeWindow);
}

public async Async.Task<Uri> GetContainerSasUrl(Container container, StorageType storageType, BlobContainerSasPermissions permissions, TimeSpan? duration = null) {
public async Task<Uri> GetContainerSasUrl(Container container, StorageType storageType, BlobContainerSasPermissions permissions, TimeSpan? duration = null) {
var client = await FindContainer(container, storageType) ?? throw new Exception($"unable to find container: {container.ContainerName} - {storageType}");
var (startTime, endTime) = SasTimeWindow(duration ?? CONTAINER_SAS_DEFAULT_DURATION);
var sasBuilder = new BlobSasBuilder(permissions, endTime) {
StartsOn = startTime,
BlobContainerName = _config.OneFuzzStoragePrefix + container.ContainerName,
};

var sasUrl = client.GenerateSasUri(sasBuilder);
return sasUrl;
var timeWindow = SasTimeWindow(duration ?? CONTAINER_SAS_DEFAULT_DURATION);
return _storage.GenerateBlobContainerSasUri(permissions, client, timeWindow);
}

public async Async.Task<bool> BlobExists(Container container, string name, StorageType storageType) {
Expand All @@ -237,10 +207,6 @@ public async Task<Dictionary<string, IDictionary<string, string>>> GetContainers
IEnumerable<IEnumerable<KeyValuePair<string, IDictionary<string, string>>>> data =
await Async.Task.WhenAll(accounts.Select(async acc => {
var service = await _storage.GetBlobServiceClientForAccount(acc);
if (service is null) {
throw new InvalidOperationException($"unable to get blob service for account {acc}");
}

return await service.GetBlobContainersAsync(BlobContainerTraits.Metadata).Select(container =>
KeyValuePair.Create(container.Name, container.Properties.Metadata)).ToListAsync();
}));
Expand Down
8 changes: 4 additions & 4 deletions src/ApiService/ApiService/onefuzzlib/Creds.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,15 +65,15 @@ public DefaultAzureCredential GetIdentity() {
public string GetSubscription() {
var storageResourceId = _config.OneFuzzDataStorage
?? throw new System.Exception("Data storage env var is not present");
var storageResource = new ResourceIdentifier(storageResourceId);
return storageResource.SubscriptionId!;
return storageResourceId.SubscriptionId
?? throw new Exception("OneFuzzDataStorage did not have subscription ID");
}

public string GetBaseResourceGroup() {
var storageResourceId = _config.OneFuzzDataStorage
?? throw new System.Exception("Data storage env var is not present");
var storageResource = new ResourceIdentifier(storageResourceId);
return storageResource.ResourceGroupName!;
return storageResourceId.ResourceGroupName
?? throw new Exception("OneFuzzDataStorage did not have resource group name");
}

public ResourceIdentifier GetResourceGroupResourceIdentifier() {
Expand Down
31 changes: 12 additions & 19 deletions src/ApiService/ApiService/onefuzzlib/Queue.cs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
using System.Text.Json;
using System.Threading.Tasks;
using Azure.Core;
using Azure.Storage;
using Azure.Storage.Queues;
using Azure.Storage.Sas;
using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
Expand All @@ -10,7 +9,7 @@ namespace Microsoft.OneFuzz.Service;
public interface IQueue {
Async.Task SendMessage(string name, string message, StorageType storageType, TimeSpan? visibilityTimeout = null, TimeSpan? timeToLive = null);
Async.Task<bool> QueueObject<T>(string name, T obj, StorageType storageType, TimeSpan? visibilityTimeout = null, TimeSpan? timeToLive = null);
Async.Task<Uri> GetQueueSas(string name, StorageType storageType, QueueSasPermissions permissions, TimeSpan? duration = null);
Task<Uri> GetQueueSas(string name, StorageType storageType, QueueSasPermissions permissions, TimeSpan? duration = null);
ResourceIdentifier GetResourceId(string queueName, StorageType storageType);
Task<IList<T>> PeekQueue<T>(string name, StorageType storageType);
Async.Task<bool> RemoveFirstMessage(string name, StorageType storageType);
Expand Down Expand Up @@ -42,22 +41,14 @@ public async Async.Task SendMessage(string name, string message, StorageType sto
}
}

public async Task<QueueClient> GetQueueClient(string name, StorageType storageType) {
var client = await GetQueueClientService(storageType);
return client.GetQueueClient(name);
}
public async Task<QueueClient> GetQueueClient(string name, StorageType storageType)
=> (await GetQueueClientService(storageType)).GetQueueClient(name);

public async Task<QueueServiceClient> GetQueueClientService(StorageType storageType) {
var accountId = _storage.GetPrimaryAccount(storageType);
_log.Verbose($"getting blob container (account_id: {accountId})");
var (name, key) = await _storage.GetStorageAccountNameAndKey(accountId);
var endpoint = _storage.GetQueueEndpoint(name);
var options = new QueueClientOptions { MessageEncoding = QueueMessageEncoding.Base64 };
return new QueueServiceClient(endpoint, new StorageSharedKeyCredential(name, key), options);
}
public Task<QueueServiceClient> GetQueueClientService(StorageType storageType)
=> _storage.GetQueueServiceClientForAccount(_storage.GetPrimaryAccount(storageType));

public async Task<bool> QueueObject<T>(string name, T obj, StorageType storageType, TimeSpan? visibilityTimeout = null, TimeSpan? timeToLive = null) {
var queueClient = await GetQueueClient(name, storageType) ?? throw new Exception($"unable to queue object, no such queue: {name}");
var queueClient = await GetQueueClient(name, storageType);
try {
var serialized = JsonSerializer.Serialize(obj, EntityConverter.GetJsonSerializerOptions());
var res = await queueClient.SendMessageAsync(serialized, visibilityTimeout: visibilityTimeout, timeToLive);
Expand All @@ -74,12 +65,14 @@ public async Task<bool> QueueObject<T>(string name, T obj, StorageType storageTy
}

public async Task<Uri> GetQueueSas(string name, StorageType storageType, QueueSasPermissions permissions, TimeSpan? duration) {
var queue = await GetQueueClient(name, storageType) ?? throw new Exception($"unable to queue object, no such queue: {name}");
var sasaBuilder = new QueueSasBuilder(permissions, DateTimeOffset.UtcNow + (duration ?? DEFAULT_DURATION));
return queue.GenerateSasUri(sasaBuilder);
var queueClient = await GetQueueClient(name, storageType);
var now = DateTimeOffset.UtcNow;
return _storage.GenerateQueueSasUri(
permissions,
queueClient,
(now, now + (duration ?? DEFAULT_DURATION)));
}


public async Async.Task CreateQueue(string name, StorageType storageType) {
var client = await GetQueueClient(name, storageType);
var resp = await client.CreateIfNotExistsAsync();
Expand Down
Loading

0 comments on commit 84b2cc9

Please sign in to comment.