feat(apigatewayv2): websocket api: grant manage connections (aws#16872)

tmokmss · TikiTDO · commit 5dc3931ed42d · 2022-02-21T18:26:39.000-05:00
closes aws#14828 By this PR, we can allow access to management API by the following code. ```ts const api = new WebSocketApi(stack, 'Api'); const defaultStage = new WebSocketStage(stack, 'Stage', { webSocketApi: api, stageName: 'dev', }); const principal = new User(stack, 'User'); api.grantManagementApiAccess(principal); // allow access to the management API for all the stage defaultStage.grantManagementApiAccess(principal); // allow access to the management API for a specific stage ``` We use WebSocket API Management API to send messages to a WebSocket API. [(doc)](https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-how-to-call-websocket-api-connections.html) To use the API, we must set IAM statement as below [(doc)](https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-websocket-control-access-iam.html): ```json { "Effect": "Allow", "Action": [ "execute-api:ManageConnections" ], "Resource": [ "arn:aws:execute-api:us-east-1:account-id:api-id/stage-name/POST/@connections/*" ] } ``` We need `/*` at the end of resource ARN because there will be arbitrary strings (`connectionId`). i.e. `{apiArn}/{stageName}/POST/@connections/{connectionId}` ---- *By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
diff --git a/packages/@aws-cdk/aws-apigatewayv2/README.md b/packages/@aws-cdk/aws-apigatewayv2/README.md
@@ -40,6 +40,7 @@ Higher level constructs for Websocket APIs | ![Experimental](https://img.shields
   - [VPC Link](#vpc-link)
   - [Private Integration](#private-integration)
 - [WebSocket API](#websocket-api)
+  - [Manage Connections Permission](#manage-connections-permission)
 
 ## Introduction
 
@@ -403,3 +404,22 @@ webSocketApi.addRoute('sendmessage', {
   }),
 });
 ```
+
+### Manage Connections Permission
+
+Grant permission to use API Gateway Management API of a WebSocket API by calling the `grantManageConnections` API.
+You can use Management API to send a callback message to a connected client, get connection information, or disconnect the client. Learn more at [Use @connections commands in your backend service](https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-how-to-call-websocket-api-connections.html).
+
+```ts
+const lambda = new lambda.Function(this, 'lambda', { /* ... */ });
+
+const webSocketApi = new WebSocketApi(stack, 'mywsapi');
+const stage = new WebSocketStage(stack, 'mystage', {
+  webSocketApi,
+  stageName: 'dev',
+});
+// per stage permission
+stage.grantManageConnections(lambda);
+// for all the stages permission
+webSocketApi.grantManageConnections(lambda);
+```
diff --git a/packages/@aws-cdk/aws-apigatewayv2/lib/websocket/api.ts b/packages/@aws-cdk/aws-apigatewayv2/lib/websocket/api.ts
@@ -1,3 +1,5 @@
+import { Grant, IGrantable } from '@aws-cdk/aws-iam';
+import { Stack } from '@aws-cdk/core';
 import { Construct } from 'constructs';
 import { CfnApi } from '../apigatewayv2.generated';
 import { IApi } from '../common/api';
@@ -127,4 +129,23 @@ export class WebSocketApi extends ApiBase implements IWebSocketApi {
       ...options,
     });
   }
+
+  /**
+   * Grant access to the API Gateway management API for this WebSocket API to an IAM
+   * principal (Role/Group/User).
+   *
+   * @param identity The principal
+   */
+  public grantManageConnections(identity: IGrantable): Grant {
+    const arn = Stack.of(this).formatArn({
+      service: 'execute-api',
+      resource: this.apiId,
+    });
+
+    return Grant.addToPrincipal({
+      grantee: identity,
+      actions: ['execute-api:ManageConnections'],
+      resourceArns: [`${arn}/*/POST/@connections/*`],
+    });
+  }
 }
diff --git a/packages/@aws-cdk/aws-apigatewayv2/lib/websocket/stage.ts b/packages/@aws-cdk/aws-apigatewayv2/lib/websocket/stage.ts
@@ -1,3 +1,4 @@
+import { Grant, IGrantable } from '@aws-cdk/aws-iam';
 import { Stack } from '@aws-cdk/core';
 import { Construct } from 'constructs';
 import { CfnStage } from '../apigatewayv2.generated';
@@ -114,4 +115,23 @@ export class WebSocketStage extends StageBase implements IWebSocketStage {
     const urlPath = this.stageName;
     return `https://${this.api.apiId}.execute-api.${s.region}.${s.urlSuffix}/${urlPath}`;
   }
+
+  /**
+   * Grant access to the API Gateway management API for this WebSocket API Stage to an IAM
+   * principal (Role/Group/User).
+   *
+   * @param identity The principal
+   */
+  public grantManagementApiAccess(identity: IGrantable): Grant {
+    const arn = Stack.of(this.api).formatArn({
+      service: 'execute-api',
+      resource: this.api.apiId,
+    });
+
+    return Grant.addToPrincipal({
+      grantee: identity,
+      actions: ['execute-api:ManageConnections'],
+      resourceArns: [`${arn}/${this.stageName}/POST/@connections/*`],
+    });
+  }
 }
diff --git a/packages/@aws-cdk/aws-apigatewayv2/test/websocket/api.test.ts b/packages/@aws-cdk/aws-apigatewayv2/test/websocket/api.test.ts
@@ -1,4 +1,5 @@
-import { Template } from '@aws-cdk/assertions';
+import { Match, Template } from '@aws-cdk/assertions';
+import { User } from '@aws-cdk/aws-iam';
 import { Stack } from '@aws-cdk/core';
 import {
   IWebSocketRouteIntegration, WebSocketApi, WebSocketIntegrationType,
@@ -80,6 +81,49 @@ describe('WebSocketApi', () => {
       RouteKey: '$default',
     });
   });
+
+  describe('grantManageConnections', () => {
+    test('adds an IAM policy to the principal', () => {
+      // GIVEN
+      const stack = new Stack();
+      const api = new WebSocketApi(stack, 'api');
+      const principal = new User(stack, 'user');
+
+      // WHEN
+      api.grantManageConnections(principal);
+
+      // THEN
+      Template.fromStack(stack).hasResourceProperties('AWS::IAM::Policy', {
+        PolicyDocument: {
+          Statement: Match.arrayWith([{
+            Action: 'execute-api:ManageConnections',
+            Effect: 'Allow',
+            Resource: {
+              'Fn::Join': ['', [
+                'arn:',
+                {
+                  Ref: 'AWS::Partition',
+                },
+                ':execute-api:',
+                {
+                  Ref: 'AWS::Region',
+                },
+                ':',
+                {
+                  Ref: 'AWS::AccountId',
+                },
+                ':',
+                {
+                  Ref: 'apiC8550315',
+                },
+                '/*/POST/@connections/*',
+              ]],
+            },
+          }]),
+        },
+      });
+    });
+  });
 });
 
 class DummyIntegration implements IWebSocketRouteIntegration {
diff --git a/packages/@aws-cdk/aws-apigatewayv2/test/websocket/stage.test.ts b/packages/@aws-cdk/aws-apigatewayv2/test/websocket/stage.test.ts
@@ -1,4 +1,5 @@
-import { Template } from '@aws-cdk/assertions';
+import { Match, Template } from '@aws-cdk/assertions';
+import { User } from '@aws-cdk/aws-iam';
 import { Stack } from '@aws-cdk/core';
 import { WebSocketApi, WebSocketStage } from '../../lib';
 
@@ -59,4 +60,51 @@ describe('WebSocketStage', () => {
     expect(defaultStage.callbackUrl.endsWith('/dev')).toBe(true);
     expect(defaultStage.callbackUrl.startsWith('https://')).toBe(true);
   });
+
+  describe('grantManageConnections', () => {
+    test('adds an IAM policy to the principal', () => {
+      // GIVEN
+      const stack = new Stack();
+      const api = new WebSocketApi(stack, 'Api');
+      const defaultStage = new WebSocketStage(stack, 'Stage', {
+        webSocketApi: api,
+        stageName: 'dev',
+      });
+      const principal = new User(stack, 'User');
+
+      // WHEN
+      defaultStage.grantManagementApiAccess(principal);
+
+      // THEN
+      Template.fromStack(stack).hasResourceProperties('AWS::IAM::Policy', {
+        PolicyDocument: {
+          Statement: Match.arrayWith([{
+            Action: 'execute-api:ManageConnections',
+            Effect: 'Allow',
+            Resource: {
+              'Fn::Join': ['', [
+                'arn:',
+                {
+                  Ref: 'AWS::Partition',
+                },
+                ':execute-api:',
+                {
+                  Ref: 'AWS::Region',
+                },
+                ':',
+                {
+                  Ref: 'AWS::AccountId',
+                },
+                ':',
+                {
+                  Ref: 'ApiF70053CD',
+                },
+                `/${defaultStage.stageName}/POST/@connections/*`,
+              ]],
+            },
+          }]),
+        },
+      });
+    });
+  });
 });
