aes-go: AWS Lambda Support + Workflow Test

Signed-off-by: Alan Nair <[email protected]>
manyiw99 · May 23, 2023 · 7556ba9 · 7556ba9
1 parent 99d166a
commit 7556ba9
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 62 deletions.
diff --git a/.github/workflows/e2e-aes.yml b/.github/workflows/e2e-aes.yml
@@ -46,7 +46,8 @@ jobs:
             aes-python,
             aes-nodejs,
             aes-python-lambda,
-            aes-nodejs-lambda
+            aes-nodejs-lambda,
+            aes-go-lambda
           ]
 
     steps:
@@ -237,7 +238,8 @@ jobs:
         service:
           [
             aes-python-lambda,
-            aes-nodejs-lambda
+            aes-nodejs-lambda,
+            aes-go-lambda
           ]
     steps:
       - name: Check out code

diff --git a/benchmarks/aes/Makefile b/benchmarks/aes/Makefile
@@ -80,6 +80,14 @@ aes-nodejs-lambda-image: docker/Dockerfile.Lambda nodejs/server.js
 	-f docker/Dockerfile.Lambda \
 	$(ROOT) --load
 
+
+aes-go-lambda-image: docker/Dockerfile.Lambda go/server.go
+	DOCKER_BUILDKIT=1 docker buildx build \
+	--tag $(AWS_ACCOUNT_ID).dkr.ecr.$(AWS_REGION).amazonaws.com/aes-go-lambda:latest \
+	--target aesGoLambda \
+	-f docker/Dockerfile.Lambda \
+	$(ROOT) --load
+
 ## Push images
 
 push-%: %-image

diff --git a/benchmarks/aes/docker/Dockerfile.Lambda b/benchmarks/aes/docker/Dockerfile.Lambda
@@ -47,3 +47,25 @@ RUN npm install
 
 # Set the CMD to your handler
 CMD [ "server.lambda_handler" ]
+
+#---------- GoLang -----------#
+# First stage (Builder):
+FROM vhiveease/golang-builder:latest AS aesGoLambdaBuilder
+WORKDIR /app/app/
+RUN apt-get install git ca-certificates
+
+COPY ./utils/tracing/go ../../utils/tracing/go
+COPY ./benchmarks/aes/go/go.mod ./
+COPY ./benchmarks/aes/go/go.sum ./
+COPY ./benchmarks/aes/go/server.go ./
+
+RUN go mod tidy
+RUN CGO_ENABLED=0 GOOS=linux go build -v -o ./server server.go
+
+# Second stage (Runner):
+FROM scratch as aesGoLambda
+WORKDIR /app/
+COPY --from=aesGoLambdaBuilder /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/
+COPY --from=aesGoLambdaBuilder /app/app/server .
+
+ENTRYPOINT [ "/app/server" ]
diff --git a/benchmarks/aes/go/server.go b/benchmarks/aes/go/server.go
@@ -30,17 +30,19 @@ import (
 	"encoding/hex"
 	"flag"
 	"fmt"
-
 	"net"
+	"os"
+	"strings"
 
 	log "github.com/sirupsen/logrus"
-
 	pb "github.com/vhive-serverless/vSwarm-proto/proto/aes"
-
 	tracing "github.com/vhive-serverless/vSwarm/utils/tracing/go"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/reflection"
+
+	"github.com/aws/aws-lambda-go/events"
+	"github.com/aws/aws-lambda-go/lambda"
 )
 
 var (
@@ -50,38 +52,6 @@ var (
 	default_plaintext_message = flag.String("default-plaintext", "defaultplaintext", "Default plaintext when the function is called with the plaintext_message world")
 )
 
-// func AESModeCBC(plaintext []byte) []byte {
-// 	// Reference: cipher documentation
-// 	// https://golang.org/pkg/crypto/cipher/#BlockMode
-
-// 	key, _ := hex.DecodeString(*key_string)
-
-// 	// CBC mode works on blocks so plaintexts may need to be padded to the
-// 	// next whole block. For an example of such padding, see
-// 	// https://tools.ietf.org/html/rfc5246#section-6.2.3.2.
-// 	var padding [aes.BlockSize]byte
-// 	if len(plaintext)%aes.BlockSize != 0 {
-// 		plaintext = append(plaintext, padding[(len(plaintext)%aes.BlockSize):]...)
-// 	}
-
-// 	block, err := aes.NewCipher(key)
-// 	if err != nil {
-// 		panic(err)
-// 	}
-
-// 	// The IV needs to be unique, but not secure. Therefore it's common to
-// 	// include it at the beginning of the ciphertext.
-// 	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
-// 	iv := ciphertext[:aes.BlockSize]
-// 	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
-// 		panic(err)
-// 	}
-
-// 	mode := cipher.NewCBCEncrypter(block, iv)
-// 	mode.CryptBlocks(ciphertext[aes.BlockSize:], plaintext)
-
-// 	return ciphertext
-// }
 
 func AESModeCTR(plaintext []byte) []byte {
 	// Reference: cipher documentation
@@ -111,8 +81,6 @@ type server struct {
 
 // ShowEncryption implements aes.AesServer
 func (s *server) ShowEncryption(ctx context.Context, in *pb.PlainTextMessage) (*pb.ReturnEncryptionInfo, error) {
-	// log.Printf("Received: %v", in.GetPlaintextMessage())
-
 	var plaintext, ciphertext []byte
 	if in.GetPlaintextMessage() == "" || in.GetPlaintextMessage() == "world" {
 		plaintext = []byte(*default_plaintext_message)
@@ -125,33 +93,54 @@ func (s *server) ShowEncryption(ctx context.Context, in *pb.PlainTextMessage) (*
 	return &pb.ReturnEncryptionInfo{EncryptionInfo: resp}, nil
 }
 
-func main() {
-	flag.Parse()
-	if tracing.IsTracingEnabled() {
-		log.Printf("Start tracing on : %s\n", *zipkin)
-		shutdown, err := tracing.InitBasicTracer(*zipkin, "aes function")
-		if err != nil {
-			log.Warn(err)
-		}
-		defer shutdown()
+func HandleRequest(ctx context.Context, request events.APIGatewayProxyRequest) (string, error) {
+	var plaintext, ciphertext []byte
+	plaintext_string := request.QueryStringParameters["plaintext"]
+	if plaintext_string == "" || plaintext_string == "world" {
+		plaintext = []byte(*default_plaintext_message)
+	} else {
+		plaintext = []byte(plaintext_string)
 	}
 
-	lis, err := net.Listen("tcp", *address)
-	if err != nil {
-		log.Fatalf("failed to listen: %v", err)
-	}
-	log.Printf("Start AES-go server. Addr: %s\n", *address)
+	ciphertext = AESModeCTR(plaintext)
+	responsemsg := fmt.Sprintf("fn: AES | plaintext: %s | ciphertext: %x | runtime: golang | platform: AWS Lambda", plaintext, ciphertext)
+	return responsemsg, nil
+}
+
+func main() {
+	val, ok := os.LookupEnv("IS_LAMBDA");
+	LAMBDA := (ok && (strings.ToLower(val) == "true"	|| strings.ToLower(val) == "yes" || strings.ToLower(val) == "1"))
 
-	var grpcServer *grpc.Server
-	if tracing.IsTracingEnabled() {
-		grpcServer = tracing.GetGRPCServerWithUnaryInterceptor()
+	if LAMBDA {
+		lambda.Start(HandleRequest)
 	} else {
-		grpcServer = grpc.NewServer()
-	}
-	pb.RegisterAesServer(grpcServer, &server{})
-	reflection.Register(grpcServer)
+		flag.Parse()
+		if tracing.IsTracingEnabled() {
+			log.Printf("Start tracing on : %s\n", *zipkin)
+			shutdown, err := tracing.InitBasicTracer(*zipkin, "aes function")
+			if err != nil {
+				log.Warn(err)
+			}
+			defer shutdown()
+		}
+
+		lis, err := net.Listen("tcp", *address)
+		if err != nil {
+			log.Fatalf("failed to listen: %v", err)
+		}
+		log.Printf("Start AES-go server. Addr: %s\n", *address)
 
-	if err := grpcServer.Serve(lis); err != nil {
-		log.Fatalf("failed to serve: %v", err)
+		var grpcServer *grpc.Server
+		if tracing.IsTracingEnabled() {
+			grpcServer = tracing.GetGRPCServerWithUnaryInterceptor()
+		} else {
+			grpcServer = grpc.NewServer()
+		}
+		pb.RegisterAesServer(grpcServer, &server{})
+		reflection.Register(grpcServer)
+
+		if err := grpcServer.Serve(lis); err != nil {
+			log.Fatalf("failed to serve: %v", err)
+		}
 	}
 }