Repository: aws/chalice

Issue:
botocore's LambdaContext returns wrong stage map on custom domains
Hello,

This is not related to chalice but botocore but I'm posting it here so we can have it documented and also figure out if chalice should implement something around it.

When the apigateway is published on a custom domain name and the stage mapped to something different, for example "dev" to "dev-aws" (dev-aws being the mapping on the custom domain) the event context still displays "dev" as the stage and so chalice websocket configure doesn't work. Mapping for the stage is optional also so the url formatting also fails.

Thanks!

Discussion:

Patch:
diff --git a/.changes/next-release/30774531330-bugfix-Websocket-824.json b/.changes/next-release/30774531330-bugfix-Websocket-824.json
new file mode 100644
index 0000000..b7ad7dd
--- /dev/null
+++ b/.changes/next-release/30774531330-bugfix-Websocket-824.json
@@ -0,0 +1,5 @@
+{
+  "type": "bugfix",
+  "category": "Websocket",
+  "description": "Fix websocket client configuration when using a custom domain (#1503)"
+}
diff --git a/chalice/app.py b/chalice/app.py
index 8dad18f..959bb98 100644
--- a/chalice/app.py
+++ b/chalice/app.py
@@ -579,11 +579,15 @@ class APIGateway(object):
 
 class WebsocketAPI(object):
     _WEBSOCKET_ENDPOINT_TEMPLATE = 'https://{domain_name}/{stage}'
+    _REGION_ENV_VARS = ['AWS_REGION', 'AWS_DEFAULT_REGION']
 
-    def __init__(self):
+    def __init__(self, env=None):
         self.session = None
         self._endpoint = None
         self._client = None
+        if env is None:
+            env = os.environ
+        self._env = env
 
     def configure(self, domain_name, stage):
         if self._endpoint is not None:
@@ -593,6 +597,37 @@ class WebsocketAPI(object):
             stage=stage,
         )
 
+    def configure_from_api_id(self, api_id, stage):
+        if self._endpoint is not None:
+            return
+        region_name = self._get_region()
+        domain_name = '{api_id}.execute-api.{region}.amazonaws.com'.format(
+            api_id=api_id, region=region_name)
+        self.configure(domain_name, stage)
+
+    def _get_region(self):
+        # Attempt to get the region so we can configure the
+        # apigatewaymanagementapi client.  We'll first try
+        # retrieving this value from env vars because these should
+        # always be set in the Lambda runtime environment.
+        for varname in self._REGION_ENV_VARS:
+            if varname in self._env:
+                return self._env[varname]
+        # As a last attempt we'll try to retrieve the region
+        # from the currently configured region.  If the session
+        # isn't configured or we can't get the region, we have
+        # no choice but to error out.
+        if self.session is not None:
+            region_name = self.session.region_name
+            if region_name is not None:
+                return region_name
+        raise ValueError(
+            "Unable to retrieve the region name when configuring the "
+            "websocket client.  Either set the 'AWS_REGION' environment "
+            "variable or assign 'app.websocket_api.session' to a boto3 "
+            "session."
+        )
+
     def _get_client(self):
         if self.session is None:
             raise ValueError(
@@ -1515,8 +1550,8 @@ class WebsocketEventSourceHandler(EventSourceHandler):
         self.handler = None
 
     def __call__(self, event, context):
-        self.websocket_api.configure(
-            event['requestContext']['domainName'],
+        self.websocket_api.configure_from_api_id(
+            event['requestContext']['apiId'],
             event['requestContext']['stage'],
         )
         super(WebsocketEventSourceHandler, self).__call__(event, context)


Test patch:
diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py
index dbe77cd..e1e663f 100644
--- a/tests/unit/conftest.py
+++ b/tests/unit/conftest.py
@@ -187,9 +187,10 @@ def create_websocket_event():
         return {
             'requestContext': {
                 'routeKey': route_key,
-                'domainName': 'abcd1234.us-west-2.amazonaws.com',
+                'domainName': 'abcd1234.execute-api.us-west-2.amazonaws.com',
                 'stage': 'api',
                 'connectionId': 'ABCD1234=',
+                'apiId': 'abcd1234',
             },
             'body': body,
         }
diff --git a/tests/unit/test_app.py b/tests/unit/test_app.py
index 8a2a261..1c59ea7 100644
--- a/tests/unit/test_app.py
+++ b/tests/unit/test_app.py
@@ -28,6 +28,7 @@ from chalice.app import (
     WebsocketDisconnectedError,
     WebsocketEventSourceHandler,
     ConvertToMiddleware,
+    WebsocketAPI
 )
 from chalice import __version__ as chalice_version
 from chalice.deploy.validate import ExperimentalFeatureError
@@ -139,9 +140,10 @@ class FakeClient(object):
 
 
 class FakeSession(object):
-    def __init__(self, client=None):
+    def __init__(self, client=None, region_name='us-west-2'):
         self.calls = []
         self._client = client
+        self.region_name = region_name
 
     def client(self, name, endpoint_url=None):
         self.calls.append((name, endpoint_url))
@@ -2633,7 +2635,7 @@ def test_can_route_websocket_connect_message(sample_websocket_app,
     assert calls[0][0] == 'connect'
     event = calls[0][1]
     assert isinstance(event, WebsocketEvent)
-    assert event.domain_name == 'abcd1234.us-west-2.amazonaws.com'
+    assert event.domain_name == 'abcd1234.execute-api.us-west-2.amazonaws.com'
     assert event.stage == 'api'
     assert event.connection_id == 'ABCD1234='
 
@@ -2652,7 +2654,7 @@ def test_can_route_websocket_disconnect_message(sample_websocket_app,
     assert calls[0][0] == 'disconnect'
     event = calls[0][1]
     assert isinstance(event, WebsocketEvent)
-    assert event.domain_name == 'abcd1234.us-west-2.amazonaws.com'
+    assert event.domain_name == 'abcd1234.execute-api.us-west-2.amazonaws.com'
     assert event.stage == 'api'
     assert event.connection_id == 'ABCD1234='
 
@@ -2671,7 +2673,7 @@ def test_can_route_websocket_default_message(sample_websocket_app,
     assert calls[0][0] == 'default'
     event = calls[0][1]
     assert isinstance(event, WebsocketEvent)
-    assert event.domain_name == 'abcd1234.us-west-2.amazonaws.com'
+    assert event.domain_name == 'abcd1234.execute-api.us-west-2.amazonaws.com'
     assert event.stage == 'api'
     assert event.connection_id == 'ABCD1234='
     assert event.body == 'foo bar'
@@ -2688,7 +2690,44 @@ def test_can_configure_client_on_connect(sample_websocket_app,
 
     assert demo.websocket_api.session.calls == [
         ('apigatewaymanagementapi',
-         'https://abcd1234.us-west-2.amazonaws.com/api'),
+         'https://abcd1234.execute-api.us-west-2.amazonaws.com/api'),
+    ]
+
+
+def test_uses_api_id_not_domain_name(sample_websocket_app,
+                                     create_websocket_event):
+    demo, calls = sample_websocket_app
+    client = FakeClient()
+    demo.websocket_api.session = FakeSession(client)
+    event = create_websocket_event('$connect')
+    # If you configure a custom domain name, we should still use the
+    # original domainName generated from API gateway when configuring
+    # the apigatewaymanagementapi client.
+    event['requestContext']['domainName'] = 'api.custom-domain-name.com'
+    handler = websocket_handler_for_route('$connect', demo)
+    handler(event, context=None)
+    assert demo.websocket_api.session.calls == [
+        ('apigatewaymanagementapi',
+         'https://abcd1234.execute-api.us-west-2.amazonaws.com/api'),
+    ]
+
+
+def test_fallsback_to_session_if_needed(sample_websocket_app,
+                                        create_websocket_event):
+    demo, calls = sample_websocket_app
+    client = FakeClient()
+    demo.websocket_api = WebsocketAPI(env={})
+    demo.websocket_api.session = FakeSession(client, region_name='us-east-2')
+    event = create_websocket_event('$connect')
+    # If you configure a custom domain name, we should still use the
+    # original domainName generated from API gateway when configuring
+    # the apigatewaymanagementapi client.
+    event['requestContext']['domainName'] = 'api.custom-domain-name.com'
+    handler = websocket_handler_for_route('$connect', demo)
+    handler(event, context=None)
+    assert demo.websocket_api.session.calls == [
+        ('apigatewaymanagementapi',
+         'https://abcd1234.execute-api.us-east-2.amazonaws.com/api'),
     ]
 
 
@@ -2703,7 +2742,7 @@ def test_can_configure_client_on_disconnect(sample_websocket_app,
 
     assert demo.websocket_api.session.calls == [
         ('apigatewaymanagementapi',
-         'https://abcd1234.us-west-2.amazonaws.com/api'),
+         'https://abcd1234.execute-api.us-west-2.amazonaws.com/api'),
     ]
 
 
@@ -2719,7 +2758,7 @@ def test_can_configure_client_on_message(sample_websocket_app,
 
     assert demo.websocket_api.session.calls == [
         ('apigatewaymanagementapi',
-         'https://abcd1234.us-west-2.amazonaws.com/api'),
+         'https://abcd1234.execute-api.us-west-2.amazonaws.com/api'),
     ]
 
 
@@ -2736,7 +2775,7 @@ def test_does_only_configure_client_once(sample_websocket_app,
 
     assert demo.websocket_api.session.calls == [
         ('apigatewaymanagementapi',
-         'https://abcd1234.us-west-2.amazonaws.com/api'),
+         'https://abcd1234.execute-api.us-west-2.amazonaws.com/api'),
     ]
 
 
@@ -3239,6 +3278,7 @@ class TestMiddleware:
                     'domainName': 'example.com',
                     'stage': 'dev',
                     'connectionId': 'abcd',
+                    'apiId': 'abcd1234',
                 },
                 'body': "body"
             }
@@ -3325,6 +3365,7 @@ class TestMiddleware:
                     'domainName': 'example.com',
                     'stage': 'dev',
                     'connectionId': 'abcd',
+                    'apiId': 'abcd1234',
                 },
                 'body': "body"
             })