From 8ff696c70d50ad47fce0ee5115baafd1fe289cb6 Mon Sep 17 00:00:00 2001
From: Ahmad Azizi <91204996+its-aazizi@users.noreply.github.com>
Date: Tue, 20 Aug 2024 18:15:55 +0500
Subject: [PATCH] update websocket URL with -ats when needed

---
 .../default-unauth-resources-template.yaml           | 12 ++++++++----
 extra/cloudformation/main-cf-template.yaml           | 12 ++++++++----
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/extra/cloudformation/default-unauth-resources-template.yaml b/extra/cloudformation/default-unauth-resources-template.yaml
index 2a134132..110fd770 100644
--- a/extra/cloudformation/default-unauth-resources-template.yaml
+++ b/extra/cloudformation/default-unauth-resources-template.yaml
@@ -698,13 +698,17 @@ Resources:
                                   responseData = {Error: "describeEndpoint call failed"};
                                   console.log(responseData.Error + ":\n", err);
                             } else {
+                                  const splitEndpointAddress = data.endpointAddress.split(".");
+                                  const endpointAddress = !splitEndpointAddress[0].endsWith("-ats")
+                                        ? `${splitEndpointAddress[0]}-ats.${splitEndpointAddress
+                                                .splice(1)
+                                                .join(".")}`
+                                        : data.endpointAddress;
                                   responseStatus = "SUCCESS";
                                   responseData = {
-                                        IotEndpointAddress: data.endpointAddress,
+                                        IotEndpointAddress: endpointAddress,
                                   };
-                                  console.log(
-                                        "response data: " + JSON.stringify(responseData)
-                                  );
+                                  console.log("response data: " + JSON.stringify(responseData));
                             }
 
                             sendResponse(event, context, responseStatus, responseData);
diff --git a/extra/cloudformation/main-cf-template.yaml b/extra/cloudformation/main-cf-template.yaml
index ec07d094..cd867211 100644
--- a/extra/cloudformation/main-cf-template.yaml
+++ b/extra/cloudformation/main-cf-template.yaml
@@ -557,13 +557,17 @@ Resources:
                                   responseData = {Error: "describeEndpoint call failed"};
                                   console.log(responseData.Error + ":\n", err);
                             } else {
+                                  const splitEndpointAddress = data.endpointAddress.split(".");
+                                  const endpointAddress = !splitEndpointAddress[0].endsWith("-ats")
+                                        ? `${splitEndpointAddress[0]}-ats.${splitEndpointAddress
+                                                .splice(1)
+                                                .join(".")}`
+                                        : data.endpointAddress;
                                   responseStatus = "SUCCESS";
                                   responseData = {
-                                        IotEndpointAddress: data.endpointAddress,
+                                        IotEndpointAddress: endpointAddress,
                                   };
-                                  console.log(
-                                        "response data: " + JSON.stringify(responseData)
-                                  );
+                                  console.log("response data: " + JSON.stringify(responseData));
                             }
 
                             sendResponse(event, context, responseStatus, responseData);