Skip to content

Commit

Permalink
feat: Add remote hashing for AWS ECR
Browse files Browse the repository at this point in the history
  • Loading branch information
sberss committed Jun 26, 2024
1 parent c089e4e commit 9ebc305
Show file tree
Hide file tree
Showing 5 changed files with 93 additions and 1 deletion.
7 changes: 7 additions & 0 deletions core/src/main/resources/reference.conf
Original file line number Diff line number Diff line change
Expand Up @@ -416,6 +416,13 @@ docker {
max-retries = 3

// Supported registries (Docker Hub, Google, Quay) can have additional configuration set separately
aws {
throttle {
number-of-requests = 1000
per = 100 seconds
}
num-threads = 10
}
azure {
// Worst case `ReadOps per minute` value from official docs
// https://github.com/MicrosoftDocs/azure-docs/blob/main/includes/container-registry-limits.md
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package cromwell.docker

import cromwell.docker.registryv2.flows.aws.AwsElasticContainerRegistry.isEcr
import cromwell.docker.registryv2.flows.azure.AzureContainerRegistry

import scala.util.{Failure, Success, Try}
Expand All @@ -19,7 +20,9 @@ sealed trait DockerImageIdentifier {
lazy val nameWithDefaultRepository =
// In ACR, the repository is part of the registry domain instead of the path
// e.g. `terrabatchdev.azurecr.io`
if (host.exists(_.contains(AzureContainerRegistry.domain)))
// In ECR, an image with no repository is supported
// e.g. 123456790.dkr.ecr.eu-west-2.amazonaws.com/example-tool
if (host.exists(_.contains(AzureContainerRegistry.domain)) || host.exists(isEcr))
image
else
repository.getOrElse("library") + s"/$image"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import cromwell.core.actor.StreamIntegration.{BackPressure, StreamContext}
import cromwell.core.{Dispatcher, DockerConfiguration}
import cromwell.docker.DockerInfoActor._
import cromwell.docker.registryv2.DockerRegistryV2Abstract
import cromwell.docker.registryv2.flows.aws.AwsElasticContainerRegistry
import cromwell.docker.registryv2.flows.azure.AzureContainerRegistry
import cromwell.docker.registryv2.flows.dockerhub.DockerHubRegistry
import cromwell.docker.registryv2.flows.google.GoogleRegistry
Expand Down Expand Up @@ -239,6 +240,7 @@ object DockerInfoActor {

// To add a new registry, simply add it to that list
List(
("aws", { c: DockerRegistryConfig => new AwsElasticContainerRegistry(c) }),
("azure", { c: DockerRegistryConfig => new AzureContainerRegistry(c) }),
("dockerhub", { c: DockerRegistryConfig => new DockerHubRegistry(c) }),
("google", { c: DockerRegistryConfig => new GoogleRegistry(c) }),
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
package cromwell.docker.registryv2.flows.aws

import cats.effect.IO
import cromwell.docker.{DockerImageIdentifier, DockerInfoActor, DockerRegistryConfig}
import cromwell.docker.registryv2.DockerRegistryV2Abstract
import cromwell.docker.registryv2.flows.aws.AwsElasticContainerRegistry.{isEcr, isPublicEcr}
import org.http4s.{AuthScheme, Header}
import org.http4s.client.Client
import software.amazon.awssdk.regions.Region
import software.amazon.awssdk.services.ecr.EcrClient
import software.amazon.awssdk.services.ecrpublic.EcrPublicClient
import software.amazon.awssdk.services.ecrpublic.model.GetAuthorizationTokenRequest

import scala.compat.java8.OptionConverters.RichOptionalGeneric

class AwsElasticContainerRegistry(config: DockerRegistryConfig) extends DockerRegistryV2Abstract(config) {

private lazy val ecrClient = EcrClient.create()
private lazy val ecrPublicClient = EcrPublicClient.builder().region(Region.US_EAST_1).build()

override def getAuthorizationScheme(dockerImageIdentifier: DockerImageIdentifier): AuthScheme =
if (isPublicEcr(dockerImageIdentifier.hostAsString)) AuthScheme.Bearer else AuthScheme.Basic

override def accepts(dockerImageIdentifier: DockerImageIdentifier): Boolean =
isEcr(dockerImageIdentifier.hostAsString)

/**
* (e.g registry-1.docker.io)
*/
override def registryHostName(dockerImageIdentifier: DockerImageIdentifier): String =
dockerImageIdentifier.host.getOrElse("")

/**
* (e.g auth.docker.io)
*/
override def authorizationServerHostName(dockerImageIdentifier: DockerImageIdentifier): String =
dockerImageIdentifier.host.getOrElse("")

override def getToken(
dockerInfoContext: DockerInfoActor.DockerInfoContext
)(implicit client: Client[IO]): IO[Option[String]] =
if (isPublicEcr(dockerInfoContext.dockerImageID.hostAsString)) getPublicEcrToken
else getPrivateEcrToken

/**
* Builds the list of headers for the token request
*/
override def buildTokenRequestHeaders(dockerInfoContext: DockerInfoActor.DockerInfoContext): List[Header] =
List.empty

private def getPublicEcrToken: IO[Option[String]] =
IO(
Option(
ecrPublicClient
.getAuthorizationToken(GetAuthorizationTokenRequest.builder().build())
.authorizationData()
.authorizationToken()
)
)

private def getPrivateEcrToken: IO[Option[String]] =
IO(
ecrClient
.getAuthorizationToken()
.authorizationData()
.stream()
.findFirst()
.asScala
.map(_.authorizationToken())
)

}

object AwsElasticContainerRegistry {
def isEcr(host: String): Boolean = isPublicEcr(host) || isPrivateEcr(host)
def isPublicEcr(host: String): Boolean = host.contains("public.ecr.aws")
def isPrivateEcr(host: String): Boolean = host.contains("amazonaws.com")
}
2 changes: 2 additions & 0 deletions project/Dependencies.scala
Original file line number Diff line number Diff line change
Expand Up @@ -402,6 +402,8 @@ object Dependencies {
"batch",
"core",
"cloudwatchlogs",
"ecr",
"ecrpublic",
"s3",
"sts",
).map(artifactName => "software.amazon.awssdk" % artifactName % awsSdkV)
Expand Down

0 comments on commit 9ebc305

Please sign in to comment.