From 354a579e15b2afc757243d3f7e813884d0f22029 Mon Sep 17 00:00:00 2001
From: mashiike <ikeda-masashi@kayac.com>
Date: Thu, 27 Jan 2022 17:54:44 +0900
Subject: [PATCH] if run on Lambda, the arguments passed at startup are
 prepended at each execution.

---
 cmd/flexentry/main.go | 21 +++++++++++-
 flexentry.go          | 74 +++++++++++++++++++++++--------------------
 2 files changed, 60 insertions(+), 35 deletions(-)

diff --git a/cmd/flexentry/main.go b/cmd/flexentry/main.go
index 838397d..147d951 100644
--- a/cmd/flexentry/main.go
+++ b/cmd/flexentry/main.go
@@ -2,6 +2,8 @@ package main
 
 import (
 	"context"
+	"flag"
+	"fmt"
 	"log"
 	"os"
 	"strings"
@@ -17,6 +19,15 @@ var (
 )
 
 func main() {
+	showVersion := flag.Bool("v", false, "show version")
+	showHelp := flag.Bool("h", false, "show help")
+	flag.Usage = func() {
+		fmt.Fprintln(flag.CommandLine.Output(), "Usage of flexenty:")
+		fmt.Fprintln(flag.CommandLine.Output(), "flexentry [opts] [commands...]")
+		flag.PrintDefaults()
+	}
+	flag.Parse()
+
 	logLevel := "info"
 	if l := os.Getenv("FLEXENTRY_LOG_LEVEL"); l != "" {
 		logLevel = l
@@ -34,13 +45,21 @@ func main() {
 	}
 	log.SetOutput(filter)
 	log.Println("[debug] flexentry version:", Version)
+	if *showVersion {
+		fmt.Println("flexentry:", Version)
+		return
+	}
+	if *showHelp {
+		flag.Usage()
+		return
+	}
 	entrypoint := flexentry.Entrypoint{
 		Executer: flexentry.NewSSMWrapExecuter(
 			flexentry.NewShellExecuter(),
 			time.Minute,
 		),
 	}
-	if err := entrypoint.Run(context.Background()); err != nil {
+	if err := entrypoint.Run(context.Background(), flag.Args()...); err != nil {
 		log.Fatalln("[error] ", err)
 	}
 }
diff --git a/flexentry.go b/flexentry.go
index 95e040a..26dc9f0 100644
--- a/flexentry.go
+++ b/flexentry.go
@@ -24,17 +24,17 @@ type Entrypoint struct {
 	query *gojq.Query
 }
 
-func (e *Entrypoint) Run(ctx context.Context) error {
+func (e *Entrypoint) Run(ctx context.Context, args ...string) error {
 	if e.isLambda() {
 		log.Println("[debug] start lambda handler")
-		lambda.Start(e.handleRequest)
+		lambda.Start(e.getHandler(args...))
 		return nil
 	}
 	return e.Execute(ctx, Pipe{
 		Stdin:  os.Stdin,
 		Stdout: os.Stdout,
 		Stderr: os.Stderr,
-	}, os.Args[1:]...)
+	}, args...)
 }
 
 func (e *Entrypoint) isLambda() bool {
@@ -51,39 +51,45 @@ func (e *Entrypoint) Execute(ctx context.Context, pipe Pipe, commands ...string)
 	return e.Executer.Execute(ctx, pipe, commands...)
 }
 
-func (e *Entrypoint) handleRequest(ctx context.Context, event Event) (interface{}, error) {
-	commands, err := e.DetectCommand(ctx, event)
-	if err != nil {
-		log.Println("[error] ", err)
-		return nil, err
-	}
-	pipe := Pipe{
-		Stderr: os.Stderr,
-		Stdout: os.Stdout,
-	}
-	var bufInput, bufOutput bytes.Buffer
-	if err := json.NewEncoder(&bufInput).Encode(event); err != nil {
-		log.Println("[warn] failed event encode", err)
-	} else {
-		pipe.Stdin = &bufInput
-	}
-	if os.Getenv("FLEXENTRY_FUNCTION_OUTPUT") == "enable" {
-		pipe.Stdout = io.MultiWriter(os.Stdout, &bufOutput)
-	}
-	err = e.Execute(ctx, pipe, commands...)
-	if err != nil {
-		log.Println("[error] ", err)
-		return nil, err
-	}
-	if bufOutput.Len() > 0 {
-		var functionOutput interface{}
-		if err := json.NewDecoder(&bufOutput).Decode(&functionOutput); err != nil {
-			log.Println("[warn] failed output encode", err)
-			return nil, nil
+func (e *Entrypoint) getHandler(args ...string) func(ctx context.Context, event Event) (interface{}, error) {
+	return func(ctx context.Context, event Event) (interface{}, error) {
+		commands, err := e.DetectCommand(ctx, event)
+		if err != nil {
+			log.Println("[error] ", err)
+			return nil, err
+		}
+		pipe := Pipe{
+			Stderr: os.Stderr,
+			Stdout: os.Stdout,
+		}
+		var bufInput, bufOutput bytes.Buffer
+		if err := json.NewEncoder(&bufInput).Encode(event); err != nil {
+			log.Println("[warn] failed event encode", err)
+		} else {
+			pipe.Stdin = &bufInput
+		}
+		if os.Getenv("FLEXENTRY_FUNCTION_OUTPUT") == "enable" {
+			pipe.Stdout = io.MultiWriter(os.Stdout, &bufOutput)
+		}
+		executeCommand := make([]string, 0, len(args)+len(commands))
+		executeCommand = append(executeCommand, args...)
+		executeCommand = append(executeCommand, commands...)
+
+		err = e.Execute(ctx, pipe, executeCommand...)
+		if err != nil {
+			log.Println("[error] ", err)
+			return nil, err
+		}
+		if bufOutput.Len() > 0 {
+			var functionOutput interface{}
+			if err := json.NewDecoder(&bufOutput).Decode(&functionOutput); err != nil {
+				log.Println("[warn] failed output encode", err)
+				return nil, nil
+			}
+			return functionOutput, nil
 		}
-		return functionOutput, nil
+		return nil, nil
 	}
-	return nil, nil
 }
 
 func (e *Entrypoint) DetectCommand(ctx context.Context, event Event) ([]string, error) {