Inital draft of refresh token

broadinstitute · Aug 16, 2022 · 928e57e · 928e57e
1 parent 19d4fdb
commit 928e57e
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 34 deletions.
diff --git a/filesystems/blob/src/main/scala/cromwell/filesystems/blob/BlobPathBuilder.scala b/filesystems/blob/src/main/scala/cromwell/filesystems/blob/BlobPathBuilder.scala
@@ -7,7 +7,7 @@ import cromwell.core.path.{NioPath, Path, PathBuilder}
 import cromwell.filesystems.blob.BlobPathBuilder._
 
 import java.net.{MalformedURLException, URI}
-import java.nio.file.{FileSystem, FileSystemNotFoundException, FileSystems}
+import java.nio.file._
 import scala.jdk.CollectionConverters._
 import scala.language.postfixOps
 import scala.util.{Failure, Try}
@@ -57,35 +57,43 @@ object BlobPathBuilder {
 
 class BlobPathBuilder(blobTokenGenerator: BlobTokenGenerator, container: String, endpoint: String) extends PathBuilder {
 
-  val credential: AzureSasCredential = new AzureSasCredential(blobTokenGenerator.getAccessToken)
-  val fileSystemConfig: Map[String, Object] = Map((AzureFileSystem.AZURE_STORAGE_SAS_TOKEN_CREDENTIAL, credential),
-                                                  (AzureFileSystem.AZURE_STORAGE_FILE_STORES, container),
-                                                  (AzureFileSystem.AZURE_STORAGE_SKIP_INITIAL_CONTAINER_CHECK, java.lang.Boolean.TRUE))
-
-  def retrieveFilesystem(uri: URI): Try[FileSystem] = {
-    Try(FileSystems.getFileSystem(uri)) recover {
-      // If no filesystem already exists, this will create a new connection, with the provided configs
-      case _: FileSystemNotFoundException => FileSystems.newFileSystem(uri, fileSystemConfig.asJava)
-    }
-  }
-
   def build(string: String): Try[BlobPath] = {
     validateBlobPath(string, container, endpoint) match {
-      case ValidBlobPath(path) => for {
-            fileSystem <- retrieveFilesystem(new URI("azb://?endpoint=" + endpoint))
-            nioPath <- Try(fileSystem.getPath(path))
-            blobPath = BlobPath(nioPath, endpoint, container)
-          } yield blobPath
+      case ValidBlobPath(path) => Try(BlobPath(path, endpoint, container, blobTokenGenerator))
       case UnparsableBlobPath(errorMessage: Throwable) => Failure(errorMessage)
     }
   }
-
   override def name: String = "Azure Blob Storage"
 }
 
-// Add args for container, storage account name
-case class BlobPath private[blob](nioPath: NioPath, endpoint: String, container: String) extends Path {
-  override protected def newPath(nioPath: NioPath): Path = BlobPath(nioPath, endpoint, container)
+object BlobPath {
+  def buildConfigMap(credential: AzureSasCredential, container: String): Map[String, Object] = {
+    Map((AzureFileSystem.AZURE_STORAGE_SAS_TOKEN_CREDENTIAL, credential),
+      (AzureFileSystem.AZURE_STORAGE_FILE_STORES, container),
+      (AzureFileSystem.AZURE_STORAGE_SKIP_INITIAL_CONTAINER_CHECK, java.lang.Boolean.TRUE))
+  }
+
+  def findNioPath(path: String, endpoint: String, container: String, blobTokenGenerator: BlobTokenGenerator): NioPath = (for {
+      fileSystem <- retrieveFilesystem(new URI("azb://?endpoint=" + endpoint), container, blobTokenGenerator)
+      nioPath <- Try(fileSystem.getPath(path))
+    } yield nioPath).get // Ideally we would unwrap this to a NioPath on success and on a access failure try to recover
+
+  def retrieveFilesystem(uri: URI, container: String, blobTokenGenerator: BlobTokenGenerator): Try[FileSystem] = {
+    Try(FileSystems.getFileSystem(uri)) recover {
+      // If no filesystem already exists, this will create a new connection, with the provided configs
+      case _: FileSystemNotFoundException => {
+        val fileSystemConfig = buildConfigMap(blobTokenGenerator.getAccessToken, container)
+        FileSystems.newFileSystem(uri, fileSystemConfig.asJava)
+      }
+    }
+  }
+}
+case class BlobPath private[blob](pathString: String, endpoint: String, container: String, blobTokenGenerator: BlobTokenGenerator) extends Path {
+  //var token = blobTokenGenerator.getAccessToken
+  //var expiry = token.getSignature.split("&").filter(_.startsWith("se")).headOption.map(_.replaceFirst("se=",""))
+  override def nioPath: NioPath = BlobPath.findNioPath(path = pathString, endpoint, container, blobTokenGenerator)
+
+  override protected def newPath(nioPath: NioPath): Path = BlobPath(pathString, endpoint, container, blobTokenGenerator)
 
   override def pathAsString: String = List(endpoint, container, nioPath.toString()).mkString("/")
 

diff --git a/filesystems/blob/src/main/scala/cromwell/filesystems/blob/BlobPathBuilderFactory.scala b/filesystems/blob/src/main/scala/cromwell/filesystems/blob/BlobPathBuilderFactory.scala
@@ -1,6 +1,7 @@
 package cromwell.filesystems.blob
 
 import akka.actor.ActorSystem
+import com.azure.core.credential.AzureSasCredential
 import com.azure.core.management.AzureEnvironment
 import com.azure.core.management.profile.AzureProfile
 import com.azure.identity.DefaultAzureCredentialBuilder
@@ -14,8 +15,7 @@ import cromwell.core.path.PathBuilderFactory
 import net.ceedubs.ficus.Ficus._
 
 import java.time.OffsetDateTime
-import scala.concurrent.ExecutionContext
-import scala.concurrent.Future
+import scala.concurrent.{ExecutionContext, Future}
 import scala.jdk.CollectionConverters._
 
 final case class BlobFileSystemConfig(config: Config)
@@ -37,7 +37,7 @@ final case class BlobPathBuilderFactory(globalConfig: Config, instanceConfig: Co
 }
 
 sealed trait BlobTokenGenerator {
-  def getAccessToken: String
+  def getAccessToken: AzureSasCredential
 }
 
 object BlobTokenGenerator {
@@ -57,13 +57,13 @@ object BlobTokenGenerator {
 }
 
 case class WSMBlobTokenGenerator(container: String, endpoint: String, workspaceId: String, workspaceManagerURL: String) extends BlobTokenGenerator {
-  def getAccessToken: String = {
+  def getAccessToken: AzureSasCredential = {
     throw new NotImplementedError
   }
 }
 
 case class NativeBlobTokenGenerator(container: String, endpoint: String) extends BlobTokenGenerator {
-  def getAccessToken: String = {
+  def getAccessToken: AzureSasCredential = {
     val storageAccountName = BlobPathBuilder.parseStorageAccount(BlobPathBuilder.parseURI(endpoint)) match {
       case Some(storageAccountName) => storageAccountName
       case _ => throw new Exception("Storage account could not be parsed from endpoint")
@@ -73,7 +73,7 @@ case class NativeBlobTokenGenerator(container: String, endpoint: String) extends
     val azureCredential = new DefaultAzureCredentialBuilder()
       .authorityHost(profile.getEnvironment.getActiveDirectoryEndpoint)
       .build
-    val azure = AzureResourceManager.authenticate(azureCredential, profile).withDefaultSubscription
+    val azure = AzureResourceManager.authenticate(azureCredential, profile).withSubscription("62b22893-6bc1-46d9-8a90-806bb3cce3c9")
 
     val storageAccounts = azure.storageAccounts()
     val storageAccount = storageAccounts
@@ -110,6 +110,6 @@ case class NativeBlobTokenGenerator(container: String, endpoint: String) extends
       blobContainerSasPermission
     )
 
-    blobContainerClient.generateSas(blobServiceSasSignatureValues)
+    new AzureSasCredential(blobContainerClient.generateSas(blobServiceSasSignatureValues))
   }
 }
diff --git a/filesystems/blob/src/test/scala/cromwell/filesystems/blob/BlobPathBuilderSpec.scala b/filesystems/blob/src/test/scala/cromwell/filesystems/blob/BlobPathBuilderSpec.scala
@@ -1,7 +1,7 @@
 package cromwell.filesystems.blob
 import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatest.matchers.should.Matchers
-import java.nio.file.Files
+
 
 object BlobPathBuilderSpec {
   def buildEndpoint(storageAccount: String) = s"https://$storageAccount.blob.core.windows.net"
@@ -42,7 +42,7 @@ class BlobPathBuilderSpec extends AnyFlatSpec with Matchers{
     }
   }
 
-  ignore should "build a blob path from a test string and read a file" in {
+  it should "build a blob path from a test string and read a file" in {
     val endpoint = BlobPathBuilderSpec.buildEndpoint("coaexternalstorage")
     val endpointHost = BlobPathBuilder.parseURI(endpoint).getHost
     val store = "inputs"
@@ -55,20 +55,23 @@ class BlobPathBuilderSpec extends AnyFlatSpec with Matchers{
     blobPath.endpoint should equal(endpoint)
     blobPath.pathAsString should equal(testString)
     blobPath.pathWithoutScheme should equal(endpointHost + "/" + store + evalPath)
-
-    val is = Files.newInputStream(blobPath.nioPath)
+    val is = blobPath.newInputStream()
     val fileText = (is.readAllBytes.map(_.toChar)).mkString
     fileText should include ("This is my test file!!!! Did it work?")
   }
 
-  ignore should "build duplicate blob paths in the same filesystem" in {
+  it should "build duplicate blob paths in the same filesystem" in {
     val endpoint = BlobPathBuilderSpec.buildEndpoint("coaexternalstorage")
     val store = "inputs"
     val evalPath = "/test/inputFile.txt"
     val blobTokenGenerator: BlobTokenGenerator = BlobTokenGenerator.createBlobTokenGenerator(store, endpoint)
     val testString = endpoint + "/" + store + evalPath
     val blobPath1: BlobPath = new BlobPathBuilder(blobTokenGenerator, store, endpoint) build testString getOrElse fail()
+    blobPath1.nioPath.getFileSystem().close()
     val blobPath2: BlobPath = new BlobPathBuilder(blobTokenGenerator, store, endpoint) build testString getOrElse fail()
     blobPath1 should equal(blobPath2)
+    val is = blobPath1.newInputStream()
+    val fileText = (is.readAllBytes.map(_.toChar)).mkString
+    fileText should include ("This is my test file!!!! Did it work?")
   }
 }