Merge pull request #476 from ericzbeard/forecast-lambda

Forecast issues with lambda functions

internal/aws/s3/s3.go

@@ -1,6 +1,7 @@
 package s3
 
 import (
+	"archive/zip"
 	"bytes"
 	"context"
 	"crypto/sha256"
@@ -246,6 +247,81 @@ func GetObject(bucketName string, key string) ([]byte, error) {
 	return body, nil
 }
 
+// GetUnzippedObjectSize gets the uncompressed length in bytes of an object.
+// Calling this on a large object will be slow!
+func GetUnzippedObjectSize(bucketName string, key string) (int64, error) {
+	result, err := getClient().GetObject(context.Background(),
+		&s3.GetObjectInput{
+			Bucket: &bucketName,
+			Key:    &key,
+		})
+	if err != nil {
+		return 0, err
+	}
+	var size int64 = 0
+
+	body, err := io.ReadAll(result.Body)
+	if err != nil {
+		return 0, err
+	}
+
+	// Unzip the archive and count the total bytes of all files
+	zipReader, err := zip.NewReader(bytes.NewReader(body), int64(len(body)))
+	if err != nil {
+		// TODO: What if it's not a zip file? Maybe return something like -1?
+		return 0, err
+	}
+
+	// Read all the files from zip archive and count total size
+	for _, zipFile := range zipReader.File {
+		config.Debugf("Reading file from zip archive: %s", zipFile.Name)
+
+		f, err := zipFile.Open()
+		if err != nil {
+			config.Debugf("Error opening zip file %s: %v", zipFile.Name, err)
+			return 0, err
+		}
+		defer f.Close()
+
+		bytesRead := 0
+		buf := make([]byte, 256)
+		for {
+			bytesRead, err = f.Read(buf)
+			if err != nil {
+				config.Debugf("Error reading from zip file %s: %v", zipFile.Name, err)
+			}
+			if bytesRead == 0 {
+				break
+			}
+			size += int64(bytesRead)
+		}
+	}
+
+	config.Debugf("Total size for %s/%s is %d", bucketName, key, size)
+
+	return size, nil
+}
+
+type S3ObjectInfo struct {
+	SizeBytes int64
+}
+
+// HeadObject gets information about an object without downloading it
+func HeadObject(bucketName string, key string) (*S3ObjectInfo, error) {
+	result, err := getClient().HeadObject(context.Background(),
+		&s3.HeadObjectInput{
+			Bucket: &bucketName,
+			Key:    &key,
+		})
+	if err != nil {
+		return nil, err
+	}
+	retval := &S3ObjectInfo{
+		SizeBytes: *result.ContentLength,
+	}
+	return retval, nil
+}
+
 // PutObject puts an object into a bucket
 func PutObject(bucketName string, key string, body []byte) error {
 	_, err := getClient().PutObject(context.Background(),

internal/cmd/forecast/README.md

@@ -54,6 +54,9 @@ These can be ignored with the `--ignore` argument.
 | F0016 | Lambda function role exists                                                    |
 | F0017 | Lambda function role can be assumed                                            |
 | F0018 | SageMaker Notebook quota limit has not been reached                            |
+| F0019 | Lambda S3Bucket exists                                                         |
+| F0020 | Lambda S3Key exists                                                            |
+| F0021 | Lambda zip file has a valid size                                               |
 
 ## Estimates
 

internal/cmd/forecast/forecast.go

@@ -95,6 +95,12 @@ func (input *PredictionInput) GetPropertyNode(name string) *yaml.Node {
 	return nil
 }
 
+// GetNode is a simplified version of s11n.GetMapValue that returns the value only
+func GetNode(prop *yaml.Node, name string) *yaml.Node {
+	_, n, _ := s11n.GetMapValue(prop, name)
+	return n
+}
+
 // LineNumber is the current line number in the template
 var LineNumber int
 

internal/cmd/forecast/forecast_integ.sh

@@ -4,6 +4,10 @@
 #
 # Run this from the root directory
 #
+# TODO: We have a bunch of ad-hoc integ tests in 
+# /test/templates/forecast/ that we run manually.
+# Add them here. 
+#
 # ./internal/cmd/forecast/forecast_integ.sh
 
 set -x

internal/cmd/forecast/lambda.go

@@ -1,35 +1,33 @@
 package forecast
 
 import (
+	"fmt"
+
 	"github.com/aws-cloudformation/rain/internal/aws/iam"
+	"github.com/aws-cloudformation/rain/internal/aws/s3"
 	"github.com/aws-cloudformation/rain/internal/config"
-	"github.com/aws-cloudformation/rain/internal/s11n"
+	"github.com/aws-cloudformation/rain/internal/console/spinner"
 	"gopkg.in/yaml.v3"
 )
 
-// checkLambdaFunction checks for potential stack failures related to functions
-func checkLambdaFunction(input PredictionInput) Forecast {
-
-	forecast := makeForecast(input.typeName, input.logicalId)
+func checkLambdaRole(input *PredictionInput, forecast *Forecast) {
 
-	_, props, _ := s11n.GetMapValue(input.resource, "Properties")
-	if props == nil {
-		config.Debugf("No Properties found for %s", input.logicalId)
-		return forecast
-	}
-	_, roleProp, _ := s11n.GetMapValue(props, "Role")
+	roleProp := input.GetPropertyNode("Role")
 
 	// If the role is specified, and it's a scalar, check if it exists
 	if roleProp != nil && roleProp.Kind == yaml.ScalarNode {
+		spin(input.typeName, input.logicalId, "Checking if lambda role exists")
 		roleArn := roleProp.Value
 		LineNumber = roleProp.Line
 		if !iam.RoleExists(roleArn) {
 			forecast.Add(F0016, false, "Role does not exist")
 		} else {
 			forecast.Add(F0016, true, "Role exists")
 		}
+		spinner.Pop()
 
 		// Check to make sure the iam role can be assumed by the lambda function
+		spin(input.typeName, input.logicalId, "Checking if lambda role can be assumed")
 		canAssume, err := iam.CanAssumeRole(roleArn, "lambda.amazonaws.com")
 		if err != nil {
 			config.Debugf("Error checking role: %s", err)
@@ -40,7 +38,91 @@ func checkLambdaFunction(input PredictionInput) Forecast {
 				forecast.Add(F0017, true, "Role can be assumed")
 			}
 		}
+		spinner.Pop()
 	}
+}
+
+func checkLambdaS3Bucket(input *PredictionInput, forecast *Forecast) {
+	// If the lambda function has an s3 bucket and key, make sure they exist
+	codeProp := input.GetPropertyNode("Code")
+	if codeProp != nil {
+		s3Bucket := GetNode(codeProp, "S3Bucket")
+		s3Key := GetNode(codeProp, "S3Key")
+		if s3Bucket != nil && s3Key != nil {
+			spin(input.typeName, input.logicalId,
+				fmt.Sprintf("Checking to see if S3 object %s/%s exists",
+					s3Bucket.Value, s3Key.Value))
+
+			// See if the bucket exists
+			exists, err := s3.BucketExists(s3Bucket.Value)
+			if err != nil {
+				config.Debugf("Unable to check if S3 bucket exists: %v", err)
+			}
+			if !exists {
+				forecast.Add(F0019, false, "S3 bucket does not exist")
+			} else {
+				forecast.Add(F0019, true, "S3 bucket exists")
+
+				// If the bucket exists, check to see if the object exists
+				obj, err := s3.HeadObject(s3Bucket.Value, s3Key.Value)
+
+				if err != nil || obj == nil {
+					forecast.Add(F0020, false, "S3 object does not exist")
+				} else {
+					forecast.Add(F0020, true, "S3 object exists")
+
+					config.Debugf("S3 Object %s/%s SizeBytes: %v",
+						s3Bucket.Value, s3Key.Value, obj.SizeBytes)
+
+					// Make sure it's less than 50Mb and greater than 0
+					// We are not downloading it and unzipping to check total size,
+					// since that would take too long for large files.
+					var max int64 = 50 * 1024 * 1024
+					if obj.SizeBytes > 0 && obj.SizeBytes <= max {
+
+						if obj.SizeBytes < 256 {
+							// This is suspiciously small. Download it and decompress
+							// to see if it's a zero byte file. A simple "Hello" python
+							// handler will zip down to 207b but an empty file has a
+							// similar zip size, so we can't know from the zip size itself.
+							unzippedSize, err := s3.GetUnzippedObjectSize(s3Bucket.Value, s3Key.Value)
+							if err != nil {
+								config.Debugf("Unable to unzip object: %v", err)
+							} else if unzippedSize == 0 {
+								forecast.Add(F0021, false, "S3 object has a zero byte unzipped size")
+							} else {
+								forecast.Add(F0021, true, "S3 object has a non-zero unzipped size")
+							}
+						} else {
+							forecast.Add(F0021, true, "S3 object has a non-zero length less than 50Mb")
+						}
+					} else {
+						if obj.SizeBytes == 0 {
+							forecast.Add(F0021, false, "S3 object has zero bytes")
+						} else {
+							forecast.Add(F0021, false, "S3 object is greater than 50Mb")
+						}
+					}
+				}
+			}
+
+			spinner.Pop()
+		} else {
+			config.Debugf("%s does not have S3Bucket and S3Key", input.logicalId)
+		}
+	} else {
+		config.Debugf("Unexpected missing Code property from %s", input.logicalId)
+	}
+}
+
+// checkLambdaFunction checks for potential stack failures related to functions
+func checkLambdaFunction(input PredictionInput) Forecast {
+
+	forecast := makeForecast(input.typeName, input.logicalId)
+
+	checkLambdaRole(&input, &forecast)
+
+	checkLambdaS3Bucket(&input, &forecast)
 
 	return forecast
 }

test/templates/forecast/lambda-fail.yaml

@@ -39,3 +39,52 @@ Resources:
       Runtime: nodejs20.x
       Timeout: 30
       MemorySize: 128
+
+  LambdaFunction3:
+    Type: AWS::Lambda::Function
+    Properties:
+      FunctionName: InvalidS3Bucket
+      Handler: index.handler
+      Code:
+        S3Bucket: does-not-exist-0123456789012345awsedrf
+        S3Key: abc
+      Runtime: nodejs20.x
+      Timeout: 30
+      MemorySize: 128
+
+  LambdaFunction4:
+    Type: AWS::Lambda::Function
+    Properties:
+      FunctionName: InvalidS3Key
+      Handler: index.handler
+      Code:
+        S3Bucket: rain-artifacts-755952356119-us-east-1 
+        S3Key: does-not-exist-123
+      Runtime: nodejs20.x
+      Timeout: 30
+      MemorySize: 128
+
+  LambdaFunction5:
+    Type: AWS::Lambda::Function
+    Properties:
+      FunctionName: ZeroByteObject
+      Handler: index.handler
+      Code:
+        S3Bucket: ezbeard-rain-notempty
+        S3Key: zero.zip
+      Runtime: nodejs20.x
+      Timeout: 30
+      MemorySize: 128
+
+  LambdaFunction6:
+    Type: AWS::Lambda::Function
+    Properties:
+      FunctionName: ObjectTooLarge
+      Handler: index.handler
+      Code:
+        S3Bucket: ezbeard-rain-notempty
+        S3Key: cdk.zip 
+      Runtime: nodejs20.x
+      Timeout: 30
+      MemorySize: 128
+

test/templates/forecast/lambda-succeed.yaml

@@ -20,3 +20,28 @@ Resources:
       Runtime: nodejs20.x
       Timeout: 30
       MemorySize: 128
+
+  LambdaFunction2:
+    Type: AWS::Lambda::Function
+    Properties:
+      FunctionName: Rain-Forecast-S3Exists
+      Handler: index.handler
+      Role: "arn:aws:iam::755952356119:role/lambda-basic"
+      Code:
+        S3Bucket: ezbeard-rain-notempty 
+        S3Key: README.md 
+      Runtime: nodejs20.x
+      Timeout: 30
+      MemorySize: 128
+
+  LambdaFunction3:
+    Type: AWS::Lambda::Function
+    Properties:
+      FunctionName: SmallObjectNonZero
+      Handler: index.handler
+      Code:
+        S3Bucket: ezbeard-rain-notempty
+        S3Key: small.zip
+      Runtime: nodejs20.x
+      Timeout: 30
+      MemorySize: 128