CloudFront + Lambda Web Adapter でセキュアなストリーミングレスポンスAPIを実現

# 注意事項（CloudFormationのtemplateコードについて）
# 1. templateコードはAWS CloudFormationのIaCジェネレータで自動生成されたものを、独自にカスタマイズして作成しています。
# 2. Lambda@Edge部分のPythonコードについては、以下のブログ記事で公開されていたNode.jsコードを参考にさせていただき、機能要件を満たすようにPythonで新たに実装したものです。
#    https://dev.classmethod.jp/articles/cloudfront-lambda-url-with-post-put-request/


AWSTemplateFormatVersion: '2010-09-09'
Description: 'Bedrock Proxy API with CloudFront Distribution and WAF'
Parameters:
  BastionIpAddress:
    Type: String
    Description: 'Bastion host IP address with /32 CIDR (e.g. 10.0.0.1/32)'
  ServiceName:
    Type: String
    Default: bedrock-proxy-api
Resources:
  # ===== IAM Roles and Policies =====
  LambdaBasicExecutionRolePolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Description: Lambda basic execution role policy
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Action:
              - logs:CreateLogGroup
              - logs:CreateLogStream
              - logs:PutLogEvents
            Resource: !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:*
  IAMRoleBedrockProxyLambda:
    Type: AWS::IAM::Role
    Properties:
      Path: /service-role/
      ManagedPolicyArns:
        - !Ref LambdaBasicExecutionRolePolicy
        - arn:aws:iam::aws:policy/AmazonSSMFullAccess
        - arn:aws:iam::aws:policy/AmazonBedrockFullAccess
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Action: sts:AssumeRole
            Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
  IAMRoleEdgeHeaderControlLambda:
    Type: AWS::IAM::Role
    Properties:
      Path: /service-role/
      ManagedPolicyArns:
        - !Ref LambdaBasicExecutionRolePolicy
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Action: sts:AssumeRole
            Effect: Allow
            Principal:
              Service:
                - edgelambda.amazonaws.com
                - lambda.amazonaws.com
  # ===== Lambda Functions =====
  LambdaBedrockProxyFunction:
    Type: AWS::Lambda::Function
    DependsOn: IAMRoleBedrockProxyLambda
    Properties:
      MemorySize: 512
      Timeout: 60
      Code:
        ImageUri: !Sub ${AWS::AccountId}.dkr.ecr.${AWS::Region}.${AWS::URLSuffix}/${ServiceName}:latest
      Role: !GetAtt IAMRoleBedrockProxyLambda.Arn
      FunctionName: !Ref ServiceName
      PackageType: Image
      LoggingConfig:
        LogFormat: Text
        LogGroup: !Sub /aws/lambda/${ServiceName}
      Environment:
        Variables:
          API_KEY_PARAM_NAME: BedrockProxyAPIKey
          DEBUG: false
          DEFAULT_MODEL: us.anthropic.claude-3-5-sonnet-20241022-v2:0
          DEFAULT_EMBEDDING_MODEL: cohere.embed-multilingual-v3
          ENABLE_CROSS_REGION_INFERENCE: true
          AWS_LWA_ASYNC_INIT: true
          AWS_LWA_INVOKE_MODE: response_stream
          AWS_LWA_AUTHORIZATION_SOURCE: authorization-custom
      EphemeralStorage:
        Size: 512
      Architectures:
        - x86_64
  LambdaEdgeHeaderControlFunction:
    Type: AWS::Lambda::Function
    DependsOn: IAMRoleEdgeHeaderControlLambda
    Properties:
      MemorySize: 128
      Timeout: 5
      Handler: index.lambda_handler
      Role: !GetAtt IAMRoleEdgeHeaderControlLambda.Arn
      FunctionName: edge-HeaderControl-for-lambda
      Runtime: python3.13
      Code:
        ZipFile: |
          import hashlib
          import base64
          def lambda_handler(event, context):
              request = event['Records'][0]['cf']['request']
              headers = request['headers']
              # x-amz-content-sha256ヘッダーへのペイロードハッシュ値登録
              body = ''
              if 'body' in request:
                  encoding = request['body'].get('encoding', '')
                  data = request['body'].get('data', '')
                  if encoding == 'base64':
                      body = base64.b64decode(data).decode('utf-8')
                  else:
                      body = data
              sha256_hash = hashlib.sha256(body.encode('utf-8')).hexdigest()
              headers['x-amz-content-sha256'] = [{
                  'key': 'x-amz-content-sha256',
                  'value': sha256_hash
              }]
              # authorization header replace for Lambda-web-adapter
              if 'authorization' in headers:
                  headers['authorization-custom'] = [{
                      'key': 'authorization-custom',
                      'value': headers['authorization'][0]['value']
                  }]
                  del headers['authorization']
              return request
      LoggingConfig:
        LogFormat: Text
        LogGroup: /aws/lambda/edge-HeaderControl-for-lambda
      EphemeralStorage:
        Size: 512
      Architectures:
        - x86_64
  # ===== Lambda Versions and Permissions =====
  LambdaVersionfunctionedgeHeaderControlforlambda:
    Type: AWS::Lambda::Version
    DependsOn: LambdaEdgeHeaderControlFunction
    Properties:
      FunctionName: !Ref LambdaEdgeHeaderControlFunction
      Description: Version published for Lambda@Edge
      RuntimePolicy:
        UpdateRuntimeOn: Auto
  LambdaUrlBedrockProxyFunction:
    Type: AWS::Lambda::Url
    DependsOn: LambdaBedrockProxyFunction
    Properties:
      AuthType: AWS_IAM
      InvokeMode: RESPONSE_STREAM
      TargetFunctionArn: !GetAtt LambdaBedrockProxyFunction.Arn
  LambdaPermissionBedrockProxyFunction:
    Type: AWS::Lambda::Permission
    DependsOn: LambdaBedrockProxyFunction
    Properties:
      FunctionName: !GetAtt LambdaBedrockProxyFunction.Arn
      Action: lambda:InvokeFunctionUrl
      Principal: cloudfront.amazonaws.com
      SourceArn: !Sub arn:aws:cloudfront::${AWS::AccountId}:distribution/${CloudFrontDistributionBedrockProxy}
  LambdaPermissionfunctionedgeHeaderControlLambda:
    Type: AWS::Lambda::Permission
    DependsOn: LambdaVersionfunctionedgeHeaderControlforlambda
    Properties:
      FunctionName: !Ref LambdaVersionfunctionedgeHeaderControlforlambda
      Action: lambda:GetFunction
      Principal: replicator.lambda.amazonaws.com
  # ===== WAF Resources =====
  WAFv2IPSet00CLOUDFRONTBedrockProxy:
    Type: AWS::WAFv2::IPSet
    Properties:
      Addresses:
        - !Ref BastionIpAddress
      Description: IP restriction for API access
      IPAddressVersion: IPV4
      Scope: CLOUDFRONT
      Name: !Sub ${AWS::StackName}-allowed-ips
  WAFv2WebACLCLOUDFRONTBedrockProxy:
    Type: AWS::WAFv2::WebACL
    DependsOn: WAFv2IPSet00CLOUDFRONTBedrockProxy
    Properties:
      Description: WAF rules for Bedrock Proxy API
      DefaultAction:
        Allow: {}
      Scope: CLOUDFRONT
      Rules:
        - Action:
            Allow: {}
          Name: ipRestrict
          Priority: 0
          Statement:
            IPSetReferenceStatement:
              Arn: !GetAtt WAFv2IPSet00CLOUDFRONTBedrockProxy.Arn
          VisibilityConfig:
            SampledRequestsEnabled: true
            CloudWatchMetricsEnabled: true
            MetricName: ipRestrict
        - Name: AWS-AWSManagedRulesAmazonIpReputationList
          Priority: 1
          OverrideAction:
            None: {}
          Statement:
            ManagedRuleGroupStatement:
              VendorName: AWS
              Name: AWSManagedRulesAmazonIpReputationList
          VisibilityConfig:
            SampledRequestsEnabled: true
            CloudWatchMetricsEnabled: true
            MetricName: AWS-AWSManagedRulesAmazonIpReputationList
        - Name: AWS-AWSManagedRulesCommonRuleSet
          Priority: 2
          OverrideAction:
            None: {}
          Statement:
            ManagedRuleGroupStatement:
              VendorName: AWS
              Name: AWSManagedRulesCommonRuleSet
          VisibilityConfig:
            SampledRequestsEnabled: true
            CloudWatchMetricsEnabled: true
            MetricName: AWS-AWSManagedRulesCommonRuleSet
        - Name: AWS-AWSManagedRulesKnownBadInputsRuleSet
          Priority: 3
          OverrideAction:
            None: {}
          Statement:
            ManagedRuleGroupStatement:
              VendorName: AWS
              Name: AWSManagedRulesKnownBadInputsRuleSet
          VisibilityConfig:
            SampledRequestsEnabled: true
            CloudWatchMetricsEnabled: true
            MetricName: AWS-AWSManagedRulesKnownBadInputsRuleSet
      VisibilityConfig:
        SampledRequestsEnabled: true
        CloudWatchMetricsEnabled: true
        MetricName: !Sub ${ServiceName}-webacl-metrics
  # ===== CloudFront Resources =====
  CloudFrontOACBedrockProxy:
    Type: AWS::CloudFront::OriginAccessControl
    Properties:
      OriginAccessControlConfig:
        Name: oac-for-lambda-url
        SigningBehavior: always
        SigningProtocol: sigv4
        OriginAccessControlOriginType: lambda
  CloudFrontDistributionBedrockProxy:
    Type: AWS::CloudFront::Distribution
    DependsOn:
      - LambdaUrlBedrockProxyFunction
      - WAFv2WebACLCLOUDFRONTBedrockProxy
      - CloudFrontOACBedrockProxy
    Properties:
      DistributionConfig:
        Enabled: true
        Comment: Distribution for Bedrock Proxy API
        DefaultRootObject: ""
        IPV6Enabled: true
        HttpVersion: http2
        Origins:
          - DomainName: !Select [2, !Split ["/", !GetAtt LambdaUrlBedrockProxyFunction.FunctionUrl]] 
            Id: CloudfrontToLambda
            ConnectionTimeout: 10
            ConnectionAttempts: 3
            OriginAccessControlId: !Ref CloudFrontOACBedrockProxy
            CustomOriginConfig:
              OriginProtocolPolicy: https-only
              OriginSSLProtocols: [TLSv1.2]
              OriginKeepaliveTimeout: 5
              OriginReadTimeout: 60
        DefaultCacheBehavior:
          TargetOriginId: CloudfrontToLambda
          ViewerProtocolPolicy: allow-all
          Compress: true
          LambdaFunctionAssociations:
            - EventType: origin-request
              IncludeBody: true
              LambdaFunctionARN: !Ref LambdaVersionfunctionedgeHeaderControlforlambda
          AllowedMethods:
            - HEAD
            - DELETE
            - POST
            - GET
            - OPTIONS
            - PUT
            - PATCH
          CachedMethods:
            - HEAD
            - GET
          CachePolicyId: !Ref CloudFrontCachePolicyBedrockProxy
          OriginRequestPolicyId: !Ref CloudFrontOriginRequestPolicyBedrockProxy
        WebACLId: !GetAtt WAFv2WebACLCLOUDFRONTBedrockProxy.Arn
        ViewerCertificate:
          CloudFrontDefaultCertificate: true
          MinimumProtocolVersion: TLSv1
        PriceClass: PriceClass_All
  CloudFrontCachePolicyBedrockProxy:
    Type: AWS::CloudFront::CachePolicy
    Properties:
      CachePolicyConfig:
        Name: Custom-CachingDisabled
        Comment: Policy with caching disabled
        DefaultTTL: 0
        MaxTTL: 0
        MinTTL: 0
        ParametersInCacheKeyAndForwardedToOrigin:
          QueryStringsConfig:
            QueryStringBehavior: none
          HeadersConfig:
            HeaderBehavior: none
          CookiesConfig:
            CookieBehavior: none
          EnableAcceptEncodingGzip: false
          EnableAcceptEncodingBrotli: false
  CloudFrontOriginRequestPolicyBedrockProxy:
    Type: AWS::CloudFront::OriginRequestPolicy
    Properties:
      OriginRequestPolicyConfig:
        Name: Custom-AllViewerExceptHostHeader
        Comment: Policy to forward all parameters in viewer requests except for the Host header
        QueryStringsConfig:
          QueryStringBehavior: all
        HeadersConfig:
          HeaderBehavior: allExcept
          Headers:
            - host
        CookiesConfig:
          CookieBehavior: all
Outputs:
  ApiEndpoint:
    Description: API Endpoint URL
    Value: !Sub https://${CloudFrontDistributionBedrockProxy.DomainName}/api/v1