BT-711 Refresh SAS token for filesystem on expiry #6831
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,139 @@ | ||
| package cromwell.filesystems.blob | ||
|
|
||
| import com.azure.core.credential.AzureSasCredential | ||
| import com.azure.core.management.AzureEnvironment | ||
| import com.azure.core.management.profile.AzureProfile | ||
| import com.azure.identity.DefaultAzureCredentialBuilder | ||
| import com.azure.resourcemanager.AzureResourceManager | ||
| import com.azure.resourcemanager.storage.models.{StorageAccount, StorageAccountKey} | ||
| import com.azure.storage.blob.nio.AzureFileSystem | ||
| import com.azure.storage.blob.sas.{BlobContainerSasPermission, BlobServiceSasSignatureValues} | ||
| import com.azure.storage.blob.{BlobContainerClient, BlobContainerClientBuilder} | ||
| import com.azure.storage.common.StorageSharedKeyCredential | ||
|
|
||
| import java.net.URI | ||
| import java.nio.file.{FileSystem, FileSystemNotFoundException, FileSystems} | ||
| import java.time.temporal.ChronoUnit | ||
| import java.time.{Duration, Instant, OffsetDateTime} | ||
| import scala.jdk.CollectionConverters._ | ||
| import scala.util.{Failure, Success, Try} | ||
|
|
||
| case class FileSystemAPI() { | ||
| def getFileSystem(uri: URI): Try[FileSystem] = Try(FileSystems.getFileSystem(uri)) | ||
| def newFileSystem(uri: URI, config: Map[String, Object]): FileSystem = FileSystems.newFileSystem(uri, config.asJava) | ||
| def closeFileSystem(uri: URI): Option[Unit] = getFileSystem(uri).toOption.map(_.close) | ||
| } | ||
|
|
||
| object BlobFileSystemManager { | ||
| def parseTokenExpiry(token: AzureSasCredential): Option[Instant] = for { | ||
| expiryString <- token.getSignature.split("&").find(_.startsWith("se")).map(_.replaceFirst("se=","")).map(_.replace("%3A", ":")) | ||
| instant = Instant.parse(expiryString) | ||
| } yield instant | ||
|
|
||
| def buildConfigMap(credential: AzureSasCredential, container: BlobContainerName): Map[String, Object] = { | ||
| Map((AzureFileSystem.AZURE_STORAGE_SAS_TOKEN_CREDENTIAL, credential), | ||
| (AzureFileSystem.AZURE_STORAGE_FILE_STORES, container.value), | ||
| (AzureFileSystem.AZURE_STORAGE_SKIP_INITIAL_CONTAINER_CHECK, java.lang.Boolean.TRUE)) | ||
| } | ||
| def hasTokenExpired(tokenExpiry: Instant, buffer: Duration): Boolean = Instant.now.plus(buffer).isAfter(tokenExpiry) | ||
| def uri(endpoint: EndpointURL) = new URI("azb://?endpoint=" + endpoint) | ||
| } | ||
| case class BlobFileSystemManager( | ||
| container: BlobContainerName, | ||
| endpoint: EndpointURL, | ||
| expiryBufferMinutes: Long, | ||
| blobTokenGenerator: BlobTokenGenerator, | ||
| fileSystemAPI: FileSystemAPI = FileSystemAPI(), | ||
| private val initialExpiration: Option[Instant] = None) { | ||
| private var expiry: Option[Instant] = initialExpiration | ||
| val buffer: Duration = Duration.of(expiryBufferMinutes, ChronoUnit.MINUTES) | ||
|
|
||
| def getExpiry: Option[Instant] = expiry | ||
| def uri: URI = BlobFileSystemManager.uri(endpoint) | ||
| def isTokenExpired: Boolean = expiry.exists(BlobFileSystemManager.hasTokenExpired(_, buffer)) | ||
| def shouldReopenFilesystem: Boolean = isTokenExpired || expiry.isEmpty | ||
| def retrieveFilesystem(): Try[FileSystem] = { | ||
| synchronized { | ||
| shouldReopenFilesystem match { | ||
| case false => fileSystemAPI.getFileSystem(uri).recoverWith { | ||
| // If no filesystem already exists, this will create a new connection, with the provided configs | ||
| case _: FileSystemNotFoundException => blobTokenGenerator.generateAccessToken.flatMap(generateFilesystem(uri, container, _)) | ||
| } | ||
| // If the token has expired, OR there is no token record, try to close the FS and regenerate | ||
| case true => | ||
| fileSystemAPI.closeFileSystem(uri) | ||
| blobTokenGenerator.generateAccessToken.flatMap(generateFilesystem(uri, container, _)) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| private def generateFilesystem(uri: URI, container: BlobContainerName, token: AzureSasCredential): Try[FileSystem] = { | ||
| expiry = BlobFileSystemManager.parseTokenExpiry(token) | ||
| if (expiry.isEmpty) return Failure(new Exception("Could not reopen filesystem, no expiration found")) | ||
| Try(fileSystemAPI.newFileSystem(uri, BlobFileSystemManager.buildConfigMap(token, container))) | ||
| } | ||
|
|
||
| } | ||
|
|
||
| sealed trait BlobTokenGenerator {def generateAccessToken: Try[AzureSasCredential]} | ||
| object BlobTokenGenerator { | ||
| def createBlobTokenGenerator(container: BlobContainerName, endpoint: EndpointURL, subscription: Option[String]): BlobTokenGenerator = { | ||
| createBlobTokenGenerator(container, endpoint, None, None, subscription) | ||
| } | ||
| def createBlobTokenGenerator(container: BlobContainerName, endpoint: EndpointURL, workspaceId: Option[WorkspaceId], workspaceManagerURL: Option[WorkspaceManagerURL], subscription: Option[String]): BlobTokenGenerator = { | ||
| (container: BlobContainerName, endpoint: EndpointURL, workspaceId, workspaceManagerURL) match { | ||
| case (container, endpoint, None, None) => | ||
| NativeBlobTokenGenerator(container, endpoint, subscription) | ||
| case (container, endpoint, Some(workspaceId), Some(workspaceManagerURL)) => | ||
| WSMBlobTokenGenerator(container, endpoint, workspaceId, workspaceManagerURL) | ||
| case _ => | ||
| throw new Exception("Arguments provided do not match any available BlobTokenGenerator implementation.") | ||
| } | ||
| } | ||
| def createBlobTokenGenerator(container: BlobContainerName, endpoint: EndpointURL): BlobTokenGenerator = createBlobTokenGenerator(container, endpoint, None) | ||
| def createBlobTokenGenerator(container: BlobContainerName, endpoint: EndpointURL, workspaceId: Option[WorkspaceId], workspaceManagerURL: Option[WorkspaceManagerURL]): BlobTokenGenerator = | ||
| createBlobTokenGenerator(container, endpoint, workspaceId, workspaceManagerURL, None) | ||
|
|
||
| } | ||
|
|
||
| case class WSMBlobTokenGenerator(container: BlobContainerName, endpoint: EndpointURL, workspaceId: WorkspaceId, workspaceManagerURL: WorkspaceManagerURL) extends BlobTokenGenerator { | ||
| def generateAccessToken: Try[AzureSasCredential] = Failure(new NotImplementedError) | ||
| } | ||
|
|
||
| case class NativeBlobTokenGenerator(container: BlobContainerName, endpoint: EndpointURL, subscription: Option[String] = None) extends BlobTokenGenerator { | ||
|
|
||
| private val azureProfile = new AzureProfile(AzureEnvironment.AZURE) | ||
| private def azureCredentialBuilder = new DefaultAzureCredentialBuilder() | ||
| .authorityHost(azureProfile.getEnvironment.getActiveDirectoryEndpoint) | ||
| .build | ||
| private def authenticateWithSubscription(sub: String) = AzureResourceManager.authenticate(azureCredentialBuilder, azureProfile).withSubscription(sub) | ||
| private def authenticateWithDefaultSubscription = AzureResourceManager.authenticate(azureCredentialBuilder, azureProfile).withDefaultSubscription() | ||
| private def azure = subscription.map(authenticateWithSubscription(_)).getOrElse(authenticateWithDefaultSubscription) | ||
|
|
||
| private def findAzureStorageAccount(name: StorageAccountName) = azure.storageAccounts.list.asScala.find(_.name.equals(name.value)) | ||
| .fold[Try[StorageAccount]](Failure(new Exception("Azure Storage Account not found")))(Success(_)) | ||
| private def buildBlobContainerClient(credential: StorageSharedKeyCredential, endpoint: EndpointURL, container: BlobContainerName): BlobContainerClient = { | ||
| new BlobContainerClientBuilder() | ||
| .credential(credential) | ||
| .endpoint(endpoint.value) | ||
| .containerName(container.value) | ||
| .buildClient() | ||
| } | ||
| private val bcsp = new BlobContainerSasPermission() | ||
| .setReadPermission(true) | ||
| .setCreatePermission(true) | ||
| .setListPermission(true) | ||
|
|
||
|
|
||
| def generateAccessToken: Try[AzureSasCredential] = for { | ||
| configuredAccount <- BlobPathBuilder.parseStorageAccount(BlobPathBuilder.parseURI(endpoint.value)) | ||
| azureAccount <- findAzureStorageAccount(configuredAccount) | ||
| keys = azureAccount.getKeys.asScala | ||
| key <- keys.headOption.fold[Try[StorageAccountKey]](Failure(new Exception("Storage account has no keys")))(Success(_)) | ||
| first = key.value | ||
| sskc = new StorageSharedKeyCredential(configuredAccount.value, first) | ||
| bcc = buildBlobContainerClient(sskc, endpoint, container) | ||
| bsssv = new BlobServiceSasSignatureValues(OffsetDateTime.now.plusDays(1), bcsp) | ||
| asc = new AzureSasCredential(bcc.generateSas(bsssv)) | ||
| } yield asc | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,114 +1,35 @@ | ||
| package cromwell.filesystems.blob | ||
|
There was a problem hiding this comment. Now is a good time to take another look at this file and see if there's anything we want to move out into other files in this package. There are a lot of top-level classes in here now, the whole group would probably be easier to understand with a little division. |
||
|
|
||
| import akka.actor.ActorSystem | ||
| import com.typesafe.config.Config | ||
| import cromwell.core.WorkflowOptions | ||
| import cromwell.core.path.PathBuilderFactory | ||
| import net.ceedubs.ficus.Ficus._ | ||
|
|
||
| import scala.concurrent.{ExecutionContext, Future} | ||
|
|
||
| final case class BlobFileSystemConfig(config: Config) | ||
|
|
||
| final case class BlobContainerName(value: String) {override def toString: String = value} | ||
| final case class StorageAccountName(value: String) {override def toString: String = value} | ||
| final case class EndpointURL(value: String) {override def toString: String = value} | ||
| final case class WorkspaceId(value: String) {override def toString: String = value} | ||
| final case class WorkspaceManagerURL(value: String) {override def toString: String = value} | ||
| final case class BlobPathBuilderFactory(globalConfig: Config, instanceConfig: Config, singletonConfig: BlobFileSystemConfig) extends PathBuilderFactory { | ||
|
There was a problem hiding this comment. We have a LOT of important string identifiers floating around. Not necessarily in scope for this PR, but we should consider replacing them with simple case classes so that the compiler can check that we're passing them around correctly. Otherwise it's too easy to swap their ordering in a method call and introduce a very hard-to-find bug. ...and so on. There was a problem hiding this comment. I think I will give this a try, It will likely help me get the tests right without accidentally swapping strings |
||
| val subscription: Option[String] = instanceConfig.as[Option[String]]("subscription") | ||
| val container: BlobContainerName = BlobContainerName(instanceConfig.as[String]("container")) | ||
| val endpoint: EndpointURL = EndpointURL(instanceConfig.as[String]("endpoint")) | ||
| val workspaceId: Option[WorkspaceId] = instanceConfig.as[Option[String]]("workspace-id").map(WorkspaceId(_)) | ||
| val expiryBufferMinutes: Long = instanceConfig.as[Option[Long]]("expiry-buffer-minutes").getOrElse(10) | ||
|
There was a problem hiding this comment. Thinking out loud, once we have this base config checked into |
||
| val workspaceManagerURL: Option[WorkspaceManagerURL] = singletonConfig.config.as[Option[String]]("workspace-manager-url").map(WorkspaceManagerURL(_)) | ||
|
|
||
| val blobTokenGenerator: BlobTokenGenerator = BlobTokenGenerator.createBlobTokenGenerator( | ||
| container, endpoint, workspaceId, workspaceManagerURL, subscription) | ||
| val fsm: BlobFileSystemManager = BlobFileSystemManager(container, endpoint, expiryBufferMinutes, blobTokenGenerator) | ||
|
|
||
| override def withOptions(options: WorkflowOptions)(implicit as: ActorSystem, ec: ExecutionContext): Future[BlobPathBuilder] = { | ||
| Future { | ||
| new BlobPathBuilder(container, endpoint)(fsm) | ||
| } | ||
| } | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
IntelliJ is telling me that
.map(StorageAccountName(_))can be just.map(StorageAccountName). But I'm also seeing it say that doesn't work in other cases and I don't understand why. 🤷♂️There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah, I was wondering about this shorthand. I may leave this to be consistent