in packages/type-safe-api/src/construct/type-safe-websocket-api.ts [141:479]
constructor(scope: Construct, id: string, props: TypeSafeWebsocketApiProps) {
super(scope, id);
addMetric(scope, "type-safe-websocket-api");
this._props = props;
// Create the WebSocket API
this.api = new WebSocketApi(this, id, {
...props,
routeSelectionExpression: "$request.body.route",
});
// Add the connect/disconnect routes
this.addRoute("$connect", {
integration:
props.connect?.integration ??
new WebSocketMockIntegration("ConnectIntegration"),
authorizer: props.authorizer,
});
const disconnectRoute = this.addRoute("$disconnect", {
integration:
props.connect?.integration ??
new WebSocketMockIntegration("DisconnectIntegration"),
});
NagSuppressions.addResourceSuppressions(
disconnectRoute,
["AwsPrototyping-APIGWAuthorization", "AwsSolutions-APIG4"].map(
(ruleId) => ({
id: ruleId,
reason: `Authorizers only apply to the $connect route`,
})
),
true
);
// Create a default stage
this.defaultStage = new WebSocketStage(this, "default", {
webSocketApi: this.api,
autoDeploy: true,
stageName: "default",
...props.stageProps,
});
// Enable execution logs by default
(this.defaultStage.node.defaultChild as CfnStage).defaultRouteSettings = {
loggingLevel: "INFO",
dataTraceEnabled: false,
};
// Enable access logging by default
if (!props.disableAccessLogging) {
const logGroup = new LogGroup(this, `AccessLogs`);
(this.defaultStage.node.defaultChild as CfnStage).accessLogSettings = {
destinationArn: logGroup.logGroupArn,
format: `$context.identity.sourceIp - - [$context.requestTime] "$context.httpMethod $context.routeKey $context.protocol" $context.status $context.responseLength $context.requestId`,
};
}
const lambdaHandlers: IFunction[] = [
...Object.values(props.integrations),
props.connect,
props.disconnect,
].flatMap((integration) =>
integration?.integration instanceof WebSocketLambdaIntegration &&
(integration.integration as any).handler?.grantPrincipal
? [(integration.integration as any).handler]
: []
);
const stack = Stack.of(this);
// By default, grant lambda handlers access to the management api
if (!props.disableGrantManagementAccessToLambdas) {
lambdaHandlers.forEach((fn) => {
this.defaultStage.grantManagementApiAccess(fn);
NagSuppressions.addResourceSuppressions(
fn,
["AwsPrototyping-IAMNoWildcardPermissions", "AwsSolutions-IAM5"].map(
(ruleId) => ({
id: ruleId,
reason:
"WebSocket handlers are granted permissions to manage arbitrary connections",
appliesTo: [
{
regex: `/^Resource::arn:${PDKNag.getStackPartitionRegex(
stack
)}:execute-api:${PDKNag.getStackRegionRegex(
stack
)}:${PDKNag.getStackAccountRegex(stack)}:.*\\/${
this.defaultStage.stageName
}\\/\\*\\/@connections\\/\\*$/g`,
},
],
})
),
true
);
});
}
// Where the same function is used for multiple routes, grant permission for API gateway to invoke
// the lambda for all routes
const uniqueLambdaHandlers = new Set<IFunction>();
const duplicateLambdaHandlers = new Set<IFunction>();
lambdaHandlers.forEach((fn) => {
if (uniqueLambdaHandlers.has(fn)) {
duplicateLambdaHandlers.add(fn);
}
uniqueLambdaHandlers.add(fn);
});
[...duplicateLambdaHandlers].forEach((fn, i) => {
new CfnPermission(this, `GrantRouteInvoke${i}`, {
action: "lambda:InvokeFunction",
principal: "apigateway.amazonaws.com",
functionName: fn.functionArn,
sourceArn: stack.formatArn({
service: "execute-api",
resource: this.api.apiId,
resourceName: "*",
}),
});
});
// Read and parse the spec
const spec = JSON.parse(
fs.readFileSync(props.specPath, "utf-8")
) as OpenAPIV3.Document;
// Map of route key to paths
const serverOperationPaths = Object.fromEntries(
Object.values(props.operationLookup).map((details) => [
details.path.replace(/\//g, ""),
details.path,
])
);
// Locally check that we can extract the schema for every operation
const schemas = extractWebSocketSchemas(
Object.keys(serverOperationPaths),
serverOperationPaths,
spec
);
// Check that every operation has an integration
const missingIntegrations = Object.keys(props.operationLookup).filter(
(operationId) => !props.integrations[operationId]
);
if (missingIntegrations.length > 0) {
throw new Error(
`Missing integrations for operations ${missingIntegrations.join(", ")}`
);
}
// Create an asset for the spec, which we'll read from the custom resource
const inputSpec = new Asset(this, "InputSpec", {
path: props.specPath,
});
// Function for managing schemas/models associated with routes
const schemaHandler = new LambdaFunction(this, "SchemaHandler", {
handler: "websocket-schema-handler.handler",
runtime: Runtime.NODEJS_20_X,
code: Code.fromAsset(
path.join(__dirname, "./prepare-spec-event-handler")
),
timeout: Duration.minutes(1),
});
NagSuppressions.addResourceSuppressions(
schemaHandler,
["AwsPrototyping-IAMNoManagedPolicies", "AwsSolutions-IAM4"].map(
(ruleId) => ({
id: ruleId,
reason: `AWSLambdaBasicExecutionRole grants minimal permissions required for lambda execution`,
})
),
true
);
schemaHandler.addToRolePolicy(
new PolicyStatement({
actions: ["s3:GetObject"],
resources: [inputSpec.bucket.arnForObjects(inputSpec.s3ObjectKey)],
})
);
schemaHandler.addToRolePolicy(
new PolicyStatement({
actions: [
"apigateway:DELETE",
"apigateway:PATCH",
"apigateway:POST",
"apigateway:GET",
],
resources: [
stack.formatArn({
service: "apigateway",
account: "",
resource: `/apis/${this.api.apiId}/models`,
}),
stack.formatArn({
service: "apigateway",
account: "",
resource: `/apis/${this.api.apiId}/models/*`,
}),
],
})
);
schemaHandler.addToRolePolicy(
new PolicyStatement({
actions: ["apigateway:PATCH", "apigateway:GET"],
resources: [
stack.formatArn({
service: "apigateway",
account: "",
resource: `/apis/${this.api.apiId}/routes`,
}),
stack.formatArn({
service: "apigateway",
account: "",
resource: `/apis/${this.api.apiId}/routes/*`,
}),
],
})
);
NagSuppressions.addResourceSuppressions(
schemaHandler,
["AwsPrototyping-IAMNoWildcardPermissions", "AwsSolutions-IAM5"].map(
(ruleId) => ({
id: ruleId,
reason: `Schema custom resource manages all routes and models`,
})
),
true
);
const providerRole = new Role(this, "PrepareSpecProviderRole", {
assumedBy: new ServicePrincipal("lambda.amazonaws.com"),
inlinePolicies: {
logs: new PolicyDocument({
statements: [
new PolicyStatement({
effect: Effect.ALLOW,
actions: [
"logs:CreateLogGroup",
"logs:CreateLogStream",
"logs:PutLogEvents",
],
resources: [
`arn:aws:logs:${stack.region}:${stack.account}:log-group:/aws/lambda/*`,
],
}),
],
}),
},
});
const provider = new Provider(this, "SchemaProvider", {
onEventHandler: schemaHandler,
role: providerRole,
});
NagSuppressions.addResourceSuppressions(
providerRole,
["AwsPrototyping-IAMNoWildcardPermissions", "AwsSolutions-IAM5"].map(
(ruleId) => ({
id: ruleId,
reason: `Custom resource provider may invoke arbitrary lambda versions`,
})
),
true
);
NagSuppressions.addResourceSuppressions(
provider,
["AwsPrototyping-LambdaLatestVersion", "AwsSolutions-L1"].map(
(ruleId) => ({
id: ruleId,
reason: `Provider framework lambda is managed by CDK`,
})
),
true
);
const schemaCustomResourceProperties: WebSocketSchemaResourceProperties = {
apiId: this.api.apiId,
inputSpecLocation: {
bucket: inputSpec.s3BucketName,
key: inputSpec.s3ObjectKey,
},
serverOperationPaths,
};
const schemaCustomResource = new CustomResource(
this,
"SchemaCustomResource",
{
serviceToken: provider.serviceToken,
properties: schemaCustomResourceProperties,
}
);
// Add a route for every integration
Object.entries(props.integrations).forEach(([operationId, integration]) => {
const op = props.operationLookup[operationId];
if (!op) {
throw new Error(
`Integration not found in operation lookup for operation ${operationId}`
);
}
// Add the route
const routeKey = op.path.replace(/\//g, "");
const route = this.addRoute(routeKey, {
integration: integration.integration,
});
NagSuppressions.addResourceSuppressions(
route,
["AwsPrototyping-APIGWAuthorization", "AwsSolutions-APIG4"].map(
(ruleId) => ({
id: ruleId,
reason: `Authorizers only apply to the $connect route`,
})
),
true
);
// Associate the route with its corresponding schema (which is created by the custom resource)
if (schemas[routeKey]) {
(route.node.defaultChild as CfnRoute).requestModels = {
model: routeKey,
};
(route.node.defaultChild as CfnRoute).modelSelectionExpression =
"model";
}
route.node.addDependency(schemaCustomResource);
});
}