From 61e2b6af5b92db384a598ddfc03e935ee77e4dd9 Mon Sep 17 00:00:00 2001
From: Sayali Gaikawad <gaiksaya@amazon.com>
Date: Tue, 12 Mar 2024 15:22:24 -0700
Subject: [PATCH] Allow flexible port mapping on load balancers

Signed-off-by: Sayali Gaikawad <gaiksaya@amazon.com>
---
 lib/infra/infra-stack.ts            | 43 ++++++++++++++++++-------
 test/opensearch-cluster-cdk.test.ts | 49 ++++++++++++++++++++++++++++-
 2 files changed, 80 insertions(+), 12 deletions(-)

diff --git a/lib/infra/infra-stack.ts b/lib/infra/infra-stack.ts
index cd2fbf9f151..6b1bf4908cd 100644
--- a/lib/infra/infra-stack.ts
+++ b/lib/infra/infra-stack.ts
@@ -29,7 +29,8 @@ import {
   SubnetType,
 } from 'aws-cdk-lib/aws-ec2';
 import {
-  ListenerCertificate, NetworkListener, NetworkLoadBalancer, Protocol,
+  ListenerCertificate,
+  NetworkListener, NetworkLoadBalancer, Protocol,
 } from 'aws-cdk-lib/aws-elasticloadbalancingv2';
 import { InstanceTarget } from 'aws-cdk-lib/aws-elasticloadbalancingv2-targets';
 import {
@@ -128,6 +129,10 @@ export interface InfraProps extends StackProps {
   readonly enableMonitoring?: boolean,
    /** Certificate ARN to attach to the listener */
    readonly certificateArn ?: string
+   /** Map opensearch port on load balancer to */
+   readonly mapOpensearchPortTo ?: number
+   /** Map opensearch-dashboards port on load balancer to */
+   readonly mapOpensearchDashboardsPortTo ?: number
 }
 
 export class InfraStack extends Stack {
@@ -187,9 +192,12 @@ export class InfraStack extends Stack {
 
   private enableMonitoring: boolean;
 
+  private opensearchPortMapping: number;
+
+  private opensearchDashboardsPortMapping: number;
+
   constructor(scope: Stack, id: string, props: InfraProps) {
     super(scope, id, props);
-    let opensearchListener: NetworkListener;
     let dashboardsListener: NetworkListener;
     let managerAsgCapacity: number;
     let dataAsgCapacity: number;
@@ -393,24 +401,37 @@ export class InfraStack extends Stack {
       crossZoneEnabled: true,
     });
 
+    const opensearchPortMap = `${props?.mapOpensearchPortTo ?? scope.node.tryGetContext('mapOpensearchPortTo')}`;
+    if (opensearchPortMap === 'undefined') {
+      if (!this.securityDisabled && !this.minDistribution) {
+        this.opensearchPortMapping = 443;
+      } else {
+        this.opensearchPortMapping = 80;
+      }
+    } else {
+      this.opensearchPortMapping = parseInt(opensearchPortMap, 10);
+    }
+
+    const opensearchListener = nlb.addListener('opensearch', {
+      port: this.opensearchPortMapping,
+      protocol: Protocol.TCP,
+    });
     if (!this.securityDisabled && !this.minDistribution) {
-      opensearchListener = nlb.addListener('opensearch', {
-        port: 443,
-        protocol: Protocol.TCP,
-      });
       if (certificateArn !== 'undefined') {
         opensearchListener.addCertificates('cert', [ListenerCertificate.fromArn(certificateArn)]);
       }
+    }
+
+    const opensearchDashboardsPortMap = `${props?.mapOpensearchDashboardsPortTo ?? scope.node.tryGetContext('mapOpensearchDashboardsPortTo')}`;
+    if (opensearchDashboardsPortMap === 'undefined') {
+      this.opensearchDashboardsPortMapping = 8443;
     } else {
-      opensearchListener = nlb.addListener('opensearch', {
-        port: 80,
-        protocol: Protocol.TCP,
-      });
+      this.opensearchDashboardsPortMapping = parseInt(opensearchDashboardsPortMap, 10);
     }
 
     if (this.dashboardsUrl !== 'undefined') {
       dashboardsListener = nlb.addListener('dashboards', {
-        port: 8443,
+        port: this.opensearchDashboardsPortMapping,
         protocol: Protocol.TCP,
       });
     }
diff --git a/test/opensearch-cluster-cdk.test.ts b/test/opensearch-cluster-cdk.test.ts
index 511efa39254..cc68c077203 100644
--- a/test/opensearch-cluster-cdk.test.ts
+++ b/test/opensearch-cluster-cdk.test.ts
@@ -828,7 +828,7 @@ test('Test additionalConfig overriding values', () => {
   });
 });
 
-test('Test certificate addition', () => {
+test('Test certificate addition and port mapping', () => {
   const app = new App({
     context: {
       securityDisabled: false,
@@ -841,6 +841,8 @@ test('Test certificate addition', () => {
       serverAccessType: 'ipv4',
       restrictServerAccessTo: 'all',
       certificateArn: 'arn:1234',
+      mapOpensearchPortTo: '8440',
+      mapOpensearchDashboardsPortTo: '443',
     },
   });
 
@@ -859,10 +861,55 @@ test('Test certificate addition', () => {
   // THEN
   const infraTemplate = Template.fromStack(infraStack);
   infraTemplate.hasResourceProperties('AWS::ElasticLoadBalancingV2::Listener', {
+    Port: 8440,
+    Protocol: 'TCP',
     Certificates: [
       {
         CertificateArn: 'arn:1234',
       },
     ],
   });
+  infraTemplate.hasResourceProperties('AWS::ElasticLoadBalancingV2::Listener', {
+    Port: 443,
+    Protocol: 'TCP',
+  });
+});
+
+test('Test default port mapping', () => {
+  const app = new App({
+    context: {
+      securityDisabled: false,
+      minDistribution: false,
+      distributionUrl: 'www.example.com',
+      cpuArch: 'x64',
+      singleNodeCluster: false,
+      dashboardsUrl: 'www.example.com',
+      distVersion: '1.0.0',
+      serverAccessType: 'ipv4',
+      restrictServerAccessTo: 'all',
+    },
+  });
+
+  // WHEN
+  const networkStack = new NetworkStack(app, 'opensearch-network-stack', {
+    env: { account: 'test-account', region: 'us-east-1' },
+  });
+
+  // @ts-ignore
+  const infraStack = new InfraStack(app, 'opensearch-infra-stack', {
+    vpc: networkStack.vpc,
+    securityGroup: networkStack.osSecurityGroup,
+    env: { account: 'test-account', region: 'us-east-1' },
+  });
+
+  // THEN
+  const infraTemplate = Template.fromStack(infraStack);
+  infraTemplate.hasResourceProperties('AWS::ElasticLoadBalancingV2::Listener', {
+    Port: 443,
+    Protocol: 'TCP',
+  });
+  infraTemplate.hasResourceProperties('AWS::ElasticLoadBalancingV2::Listener', {
+    Port: 8443,
+    Protocol: 'TCP',
+  });
 });