fix: correctly parse ECR URIs and tags to OSMLContainer

.eslintrc.yaml

@@ -32,7 +32,6 @@ rules:
   require-await: off
   "@typescript-eslint/no-unused-expressions":
     - error
-    - allowShortCircuit: true
     - allowTernary: true
   "@typescript-eslint/no-unsafe-assignment": warn
   "@typescript-eslint/interface-name-prefix": off

lib/osml/model_endpoint/me_http_endpoint.ts

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023-2024 Amazon.com, Inc. or its affiliates.
+ * Copyright 2023-2025 Amazon.com, Inc. or its affiliates.
  */
 
 import { Duration, RemovalPolicy } from "aws-cdk-lib";
@@ -8,6 +8,7 @@ import {
   Cluster,
   Compatibility,
   ContainerImage,
+  ContainerInsights,
   LogDriver,
   Protocol,
   TaskDefinition
@@ -168,7 +169,9 @@ export class MEHTTPEndpoint extends Construct {
     const httpEndpointCluster = new Cluster(this, props.clusterName, {
       clusterName: props.clusterName,
       vpc: props.osmlVpc.vpc,
-      containerInsights: props.account.prodLike
+      containerInsightsV2: props.account.prodLike
+        ? ContainerInsights.ENABLED
+        : ContainerInsights.DISABLED
     });
 
     // Determine the removal policy based on the environment

lib/osml/model_runner/mr_dataplane.ts

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023-2024 Amazon.com, Inc. or its affiliates.
+ * Copyright 2023-2025 Amazon.com, Inc. or its affiliates.
  */
 
 import { EcsIsoServiceAutoscaler } from "@cdklabs/cdk-enterprise-iac";
@@ -18,6 +18,7 @@ import {
   Cluster,
   Compatibility,
   ContainerDefinition,
+  ContainerInsights,
   FargateService,
   Protocol,
   TaskDefinition
@@ -662,7 +663,9 @@ export class MRDataplane extends Construct {
     this.cluster = new Cluster(this, "MRCluster", {
       clusterName: this.config.ECS_CLUSTER_NAME,
       vpc: props.osmlVpc.vpc,
-      containerInsights: props.account.prodLike
+      containerInsightsV2: props.account.prodLike
+        ? ContainerInsights.ENABLED
+        : ContainerInsights.DISABLED
     });
 
     // Define our ECS task

lib/osml/osml_container.ts

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023-2024 Amazon.com, Inc. or its affiliates.
+ * Copyright 2023-2025 Amazon.com, Inc. or its affiliates.
  */
 
 import { RemovalPolicy, SymlinkFollowMode } from "aws-cdk-lib";
@@ -21,6 +21,7 @@ import { RegionalConfig, RegionConfig } from "./utils/regional_config";
 export class OSMLContainerConfig extends BaseConfig {
   /**
    * The default container image.
+   * Can only be specified if REPOSITORY_ARN is not set.
    */
   public CONTAINER_URI?: string;
 
@@ -39,12 +40,25 @@ export class OSMLContainerConfig extends BaseConfig {
    */
   public CONTAINER_DOCKERFILE?: string;
 
+  /**
+   * (Optional) ARN of an existing container to be used.
+   * Can only be specified if CONTAINER_URI is not set.
+   */
+  public ECR_REPOSITORY_ARN?: string;
+
+  /**
+   * (Optional) ARN of an existing container to be used.
+   * Only used if REPOSITORY_ARN is set and will default to "latest".
+   */
+  public ECR_REPOSITORY_TAG?: string;
+
   /**
    * Creates an instance of OSMLContainerConfig.
    * @param config - The configuration object for OSMLContainer.
    */
   constructor(config: ConfigType = {}) {
     super({
+      REPOSITORY_TAG: "latest",
       ...config
     });
   }
@@ -212,95 +226,161 @@ export class OSMLContainer extends Construct {
   }
 
   /**
-   * Builds the container from an existing repository using the specified repository URI in the config.
-   * Validates the configuration to ensure that the necessary parameters are provided.
+   * Builds the container from an existing repository using the specified repository URI or ARN in the config.
+   * Ensures that only one of `CONTAINER_URI` or `CONTAINER_ARN` is provided.
    *
-   * @throws {Error} If CONTAINER_URI is not set in the configuration.
+   * @throws {Error} If both `CONTAINER_URI` and `CONTAINER_ARN` are provided or neither is set.
    */
   private buildFromRepository(): void {
-    // Validate that the CONTAINER_URI is provided in the configuration.
-    if (!this.config.CONTAINER_URI) {
+    const { CONTAINER_URI, ECR_REPOSITORY_ARN } = this.config;
+
+    this.validateContainerConfig(CONTAINER_URI, ECR_REPOSITORY_ARN);
+
+    if (ECR_REPOSITORY_ARN) {
+      this.importRepositoryByArn(ECR_REPOSITORY_ARN);
+    } else {
+      this.importContainerByUri(CONTAINER_URI!);
+    }
+  }
+
+  /**
+   * Validates the configuration to ensure only one of CONTAINER_URI or CONTAINER_ARN is provided.
+   *
+   * @param {string | undefined} uri - The container URI.
+   * @param {string | undefined} arn - The container ARN.
+   * @throws {Error} If both or neither are provided.
+   */
+  private validateContainerConfig(uri?: string, arn?: string): void {
+    if (uri && arn) {
       throw new Error(
-        "CONTAINER_URI must be set in the configuration to use a pre-built container image."
+        "Only one of CONTAINER_URI or ECR_REPOSITORY_ARN can be set."
       );
     }
-
-    // Determine the tag to use for the container image, defaulting to "latest" if not provided.
-    const tag = "latest";
-
-    // Check if the CONTAINER_URI indicates an Amazon ECR repository by checking the ARN format.
-    const ecrArnRegex =
-      /^arn:([^:]+):ecr:([^:]+):[^:]+:repository\/[a-zA-Z0-9-._]+(:[a-zA-Z0-9-._]+)?$/;
-
-    if (ecrArnRegex.test(this.config.CONTAINER_URI)) {
-      let repositoryArn = this.config.CONTAINER_URI;
-      let extractedTag = tag;
-
-      // Extract tag if present
-      const tagSeparatorIndex = this.config.CONTAINER_URI.lastIndexOf(":");
-      if (
-        tagSeparatorIndex > this.config.CONTAINER_URI.indexOf("repository/")
-      ) {
-        repositoryArn = this.config.CONTAINER_URI.substring(
-          0,
-          tagSeparatorIndex
-        );
-        extractedTag = this.config.CONTAINER_URI.substring(
-          tagSeparatorIndex + 1
-        );
-      }
-
-      // Import the existing ECR repository using the ARN provided in the CONTAINER_URI.
-      this.repository = Repository.fromRepositoryArn(
-        this,
-        "ImportedECRRepo",
-        repositoryArn
+    if (!uri && !arn) {
+      throw new Error(
+        "Either CONTAINER_URI or ECR_REPOSITORY_ARN must be set in the configuration."
       );
+    }
+  }
 
-      // Create a ContainerImage object from the imported ECR repository and specified tag.
-      this.containerImage = ContainerImage.fromEcrRepository(
-        this.repository,
-        extractedTag
-      );
+  /**
+   * Imports an ECR repository by ARN.
+   *
+   * @param {string} arn - The ECR repository ARN.
+   */
+  private importRepositoryByArn(arn: string): void {
+    this.repository = Repository.fromRepositoryArn(
+      this,
+      "ImportedECRRepo",
+      arn
+    );
+    this.containerImage = ContainerImage.fromEcrRepository(
+      this.repository,
+      this.config.ECR_REPOSITORY_TAG
+    );
+    this.containerUri = this.repository.repositoryUri;
 
-      // Set the containerUri to the full URI of the container image in ECR.
-      this.containerUri = this.repository.repositoryUriForTag(extractedTag);
+    if (this.buildDockerImageCode) {
+      this.dockerImageCode = DockerImageCode.fromEcr(this.repository);
+    }
+  }
 
-      if (this.buildDockerImageCode) {
-        // Create a DockerImageCode object for Lambda using the imported ECR repository.
-        this.dockerImageCode = DockerImageCode.fromEcr(this.repository, {
-          tagOrDigest: extractedTag
-        });
-      }
+  /**
+   * Imports a container by URI, handling ECR repositories and external registries.
+   *
+   * @param {string} uri - The container URI.
+   */
+  private importContainerByUri(uri: string): void {
+    const ecrUriRegex =
+      /^[0-9]{12}\.dkr\.ecr\.[a-z0-9-]+\.amazonaws\.com\/[a-zA-Z0-9-_]+/;
+
+    if (ecrUriRegex.test(uri)) {
+      this.importEcrContainer(uri);
     } else {
-      // If the CONTAINER_URI does not indicate an ECR repository, assume it is a public or private Docker registry.
-      this.repositoryAccessMode = "Vpc";
+      this.importExternalContainer(uri);
+    }
+  }
 
-      // Create a ContainerImage object from the provided Docker registry URI and tag.
-      this.containerImage = ContainerImage.fromRegistry(
-        this.config.CONTAINER_URI
-      );
+  /**
+   * Imports a container from an ECR repository.
+   *
+   * @param {string} uri - The container URI.
+   */
+  private importEcrContainer(uri: string): void {
+    // Extract the repository name and tag from the URI
+    const tag = this.extractTagFromUri(uri);
+    const repoName = uri.split("/").pop()?.split(":")[0] ?? "";
+
+    this.repository = Repository.fromRepositoryName(
+      this,
+      "ImportedECRRepo",
+      repoName
+    );
+    this.containerImage = ContainerImage.fromEcrRepository(
+      this.repository,
+      tag
+    );
+    this.containerUri = `${this.repository.repositoryUri}:${tag}`;
 
-      // Set the containerUri to the full URI of the container image in the Docker registry.
-      this.containerUri = this.config.CONTAINER_URI;
+    if (this.buildDockerImageCode) {
+      this.dockerImageCode = DockerImageCode.fromEcr(this.repository, {
+        tagOrDigest: tag
+      });
+    }
+  }
 
-      if (this.buildDockerImageCode) {
-        // Define the Dockerfile content dynamically based on containerUri
-        const dockerfileContent = `FROM ${this.config.CONTAINER_URI}`;
+  /**
+   * Imports a container from an external registry such as Docker Hub.
+   *
+   * @param {string} uri - The container URI.
+   */
+  private importExternalContainer(uri: string): void {
+    this.repositoryAccessMode = "Vpc";
+    this.containerImage = ContainerImage.fromRegistry(uri);
+    this.containerUri = uri;
 
-        // Create a temporary Dockerfile to build the Docker image with
-        const tmpDockerfile = "Dockerfile.tmp";
+    if (this.buildDockerImageCode) {
+      this.createDockerImageFromRegistry(uri);
+    }
+  }
+
+  /**
+   * Creates a temporary Dockerfile to use an external registry image.
+   *
+   * @param {string} uri - The container URI.
+   */
+  private createDockerImageFromRegistry(uri: string): void {
+    const tmpDockerfile = path.join(__dirname, "Dockerfile.tmp");
+    fs.writeFileSync(tmpDockerfile, `FROM ${uri}`);
 
-        // Write the temp Dockerfile to the build directory
-        const dockerfilePath = path.join(__dirname, tmpDockerfile);
-        fs.writeFileSync(dockerfilePath, dockerfileContent);
+    this.dockerImageCode = DockerImageCode.fromImageAsset(__dirname, {
+      file: "Dockerfile.tmp",
+      followSymlinks: SymlinkFollowMode.ALWAYS
+    });
+  }
 
-        // Create a DockerImageCode object for Lambda using the DockerImageAsset
-        this.dockerImageCode = DockerImageCode.fromImageAsset(__dirname, {
-          file: tmpDockerfile,
-          followSymlinks: SymlinkFollowMode.ALWAYS
-        });
-      }
+  /**
+   * Extracts the tag from a container URI.
+   * If no tag is found, returns `"latest"` as the default.
+   *
+   * @param {string} uri - The container URI.
+   * @returns {string} The extracted tag or `"latest"` if none is found.
+   */
+  private extractTagFromUri(uri: string): string {
+    if (!uri) {
+      throw new Error("Container URI is undefined or empty.");
     }
+
+    // Find the last `:` in the URI to check for a tag
+    const tagSeparatorIndex = uri.lastIndexOf(":");
+    const lastSlashIndex = uri.lastIndexOf("/");
+
+    // Ensure the colon comes after the last slash (indicating a tag, not a protocol separator)
+    if (tagSeparatorIndex > lastSlashIndex) {
+      return uri.substring(tagSeparatorIndex + 1);
+    }
+
+    // Default if no tag is found
+    return "latest";
   }
 }

lib/osml/tile_server/ts_dataplane.ts

@@ -39,7 +39,6 @@ import {
   NetworkLoadBalancer,
   Protocol as elbv2_protocol
 } from "aws-cdk-lib/aws-elasticloadbalancingv2";
-import { IpTarget } from "aws-cdk-lib/aws-elasticloadbalancingv2-targets";
 import {
   AnyPrincipal,
   IRole,
@@ -617,7 +616,9 @@ export class TSDataplane extends Construct {
       "TSServiceApplicationLoadBalancer",
       {
         vpc: props.osmlVpc.vpc,
-        vpcSubnets: props.osmlVpc.selectedSubnets
+        vpcSubnets: props.osmlVpc.selectedSubnets,
+        securityGroup: this.securityGroup,
+        internetFacing: false
       }
     );
 
@@ -632,12 +633,10 @@ export class TSDataplane extends Construct {
         securityGroups: this.securityGroup ? [this.securityGroup] : [],
         taskSubnets: props.osmlVpc.selectedSubnets,
         assignPublicIp: false,
-        publicLoadBalancer: false,
         loadBalancer: this.alb
       }
     );
     this.fargateService.node.addDependency(this.tsContainer);
-    this.fargateService.node.addDependency(this.alb);
 
     // Allow access to EFS from Fargate ECS
     this.fileSystem.grantRootAccess(
@@ -707,10 +706,8 @@ export class TSDataplane extends Construct {
       });
 
       nlbListener.addTargets("TSNlbTargetGroup", {
-        targets: [
-          new IpTarget(this.fargateService.loadBalancer.loadBalancerDnsName)
-        ],
-        port: this.config.ECS_NETWORK_LOAD_BALANCER_PORT // Target port
+        targets: [this.fargateService.service],
+        port: this.config.ECS_NETWORK_LOAD_BALANCER_PORT
       });
 
       const vpcLink = new VpcLink(this, "TSVpcLink", {