Go secrets (#240)

pkg/core/persist.go

@@ -7,13 +7,9 @@ import (
 
 type (
 	Persist struct {
-		Name string
-		Kind PersistKind
-	}
-
-	Fs struct {
-		Persist
-		GenerateNewFs bool
+		Name        string
+		Kind        PersistKind
+		GenerateNew bool
 	}
 
 	Secrets struct {

pkg/lang/golang/arguments.go

@@ -3,6 +3,7 @@ package golang
 import (
 	"fmt"
 
+	"github.com/klothoplatform/klotho/pkg/query"
 	sitter "github.com/smacker/go-tree-sitter"
 )
 
@@ -12,27 +13,36 @@ type Argument struct {
 }
 
 // GetArguements is passed a tree-sitter node, which is of type argument_list, and returns a list of in order Arguments
-func GetArguements(args *sitter.Node) []Argument {
-
-	arguments := []Argument{}
-	nextMatch := doQuery(args, findArgs)
+func getArguements(args *sitter.Node) (arguments []Argument, found bool) {
+	fnName := ""
+	nextMatch := doQuery(args, findFunctionCall)
 	for {
 		match, found := nextMatch()
 		if !found {
 			break
 		}
-
+		fn := match["function"]
 		arg := match["arg"]
+
+		if fnName != "" && !query.NodeContentEquals(fn, fnName) {
+			break
+		}
+
+		fnName = fn.Content()
+
 		if arg == nil {
 			continue
 		}
 
 		arguments = append(arguments, Argument{Content: arg.Content(), Type: arg.Type()})
 	}
-	return arguments
+	if fnName != "" {
+		found = true
+	}
+	return
 }
 
-func ArgumentListToString(args []Argument) string {
+func argumentListToString(args []Argument) string {
 	result := "("
 	for index, arg := range args {
 		if index < len(args)-1 {

pkg/lang/golang/arguments_test.go

@@ -0,0 +1,58 @@
+package golang
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/klothoplatform/klotho/pkg/core"
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_GetArguements(t *testing.T) {
+	tests := []struct {
+		name      string
+		source    string
+		want      []Argument
+		wantFound bool
+	}{
+		{
+			name: "finds next function Name and args",
+			source: `
+			x = s.my_func("val")
+			y = s.other_func("something_else)
+			`,
+			want: []Argument{
+				{Content: `"val"`, Type: "interpreted_string_literal"},
+			},
+			wantFound: true,
+		},
+		{
+			name:      "args not required",
+			source:    `v, err := s.someFunc()`,
+			wantFound: false,
+		},
+		{
+			name:   "a call containing other function calls as args",
+			source: `v, err := runtimevar.OpenVariable(context.TODO(), fmt.Sprintf("file://%s?decoder=string", path))`,
+			want: []Argument{
+				{Content: "context.TODO()", Type: "call_expression"},
+				{Content: `fmt.Sprintf("file://%s?decoder=string", path)`, Type: "call_expression"},
+			},
+			wantFound: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert := assert.New(t)
+
+			f, err := core.NewSourceFile("", strings.NewReader(tt.source), Language)
+			if !assert.NoError(err) {
+				return
+			}
+			args, found := getArguements(f.Tree().RootNode())
+
+			assert.ElementsMatch(tt.want, args)
+			assert.Equal(tt.wantFound, found)
+		})
+	}
+}

pkg/lang/golang/aws_runtime/Lambda_Dockerfile

@@ -1,13 +1,13 @@
-FROM public.ecr.aws/lambda/provided:al2 
-
-WORKDIR ${LAMBDA_TASK_ROOT}
+FROM golang:1.20 as builder
 
-RUN yum install -y golang
-RUN go env -w GOPROXY=https://proxy.golang.org,direct
+WORKDIR /usr/src/app
+ENV GOOS=linux GOARCH=amd64 CGO_ENABLED=0
+COPY go.mod ./
+RUN go mod tidy && go mod download && go mod verify
 
 COPY . .
-RUN env GOOS=linux GOARCH=amd64 CGO_ENABLED=0 
-RUN go mod tidy
-RUN go build -o=/main
+RUN go build -o /usr/local/bin/app
 
-ENTRYPOINT ["/main"]
+FROM public.ecr.aws/lambda/provided:al2 
+COPY --from=builder /usr/local/bin/app main
+ENTRYPOINT ["/main"]

pkg/lang/golang/aws_runtime/aws.go

@@ -65,3 +65,12 @@ func (r *AwsRuntime) GetFsImports() []golang.Import {
 		{Package: "gocloud.dev/blob/s3blob", Alias: "_"},
 	}
 }
+
+func (r *AwsRuntime) GetSecretsImports() []golang.Import {
+	return []golang.Import{
+		{Package: "os"},
+		{Package: "strings"},
+		{Package: "gocloud.dev/runtimevar"},
+		{Package: "gocloud.dev/runtimevar/awssecretsmanager", Alias: "_"},
+	}
+}

pkg/lang/golang/plugin_fs.go

@@ -1,8 +1,8 @@
 package golang
 
 import (
-	"errors"
 	"fmt"
+	"strings"
 
 	"github.com/klothoplatform/klotho/pkg/annotation"
 	"github.com/klothoplatform/klotho/pkg/core"
@@ -81,47 +81,45 @@ func (p *PersistFsPlugin) transformFS(f *core.SourceFile, cap *core.Annotation,
 
 	unit.EnvironmentVariables = append(unit.EnvironmentVariables, fsEnvVar)
 
-	args := GetArguements(result.args)
+	args, _ := getArguements(result.expression)
+	// Generate the new node content before replacing the node. We just set it so we can compile correctly
+	newNodeContent := `var _ = ` + args[1].Content + "\n"
 
-	// We need to check to make sure the path supplied to the original node content is a static string. This is because it will get erased and we dont want to leave os level orphaned code
-	if !args[0].IsString() {
-		return nil, errors.New("must supply static string for secret path")
-	}
-
-	args[0].Content = "nil"
 	args[1].Content = fmt.Sprintf(`"s3://" + os.Getenv("%s") + "?region=" + os.Getenv("AWS_REGION")`, fsEnvVar.Name)
 
-	err := f.ReplaceNodeContent(result.args, ArgumentListToString(args))
-	if err != nil {
-		return nil, err
-	}
-	err = f.ReplaceNodeContent(result.operator, "blob")
+	newArgContent := argumentListToString(args)
+
+	newExpressionContent := strings.ReplaceAll(result.expression.Content(), result.args.Content(), newArgContent)
+	newNodeContent += newExpressionContent
+
+	err := f.ReplaceNodeContent(result.expression, newNodeContent)
 	if err != nil {
 		return nil, err
 	}
 
-	err = UpdateImportsInFile(f, p.runtime.GetFsImports(), []Import{{Package: "gocloud.dev/blob/fileblob"}})
+	err = UpdateImportsInFile(f, p.runtime.GetFsImports(), []Import{})
 	if err != nil {
 		return nil, err
 	}
 
 	persist := &core.Persist{
-		Kind: core.PersistFileKind,
-		Name: cap.Capability.ID,
+		Kind:        core.PersistFileKind,
+		Name:        cap.Capability.ID,
+		GenerateNew: true,
 	}
 	return persist, nil
 }
 
 type persistResult struct {
-	varName  string
-	operator *sitter.Node
-	args     *sitter.Node
+	varName    string
+	expression *sitter.Node
+	args       *sitter.Node
 }
 
 func queryFS(file *core.SourceFile, annotation *core.Annotation) *persistResult {
 	log := zap.L().With(logging.FileField(file), logging.AnnotationField(annotation))
 
-	fileBlobImport := GetNamedImportInFile(file, "gocloud.dev/blob/fileblob")
+	fileBlobImport := GetNamedImportInFile(file, "gocloud.dev/blob")
 
 	nextMatch := doQuery(annotation.Node, fileBucket)
 
@@ -138,7 +136,7 @@ func queryFS(file *core.SourceFile, annotation *core.Annotation) *persistResult
 				return nil
 			}
 		} else {
-			if !query.NodeContentEquals(id, "fileblob") {
+			if !query.NodeContentEquals(id, "blob") {
 				return nil
 			}
 		}
@@ -150,8 +148,8 @@ func queryFS(file *core.SourceFile, annotation *core.Annotation) *persistResult
 	}
 
 	return &persistResult{
-		varName:  varName.Content(),
-		operator: id,
-		args:     args,
+		varName:    varName.Content(),
+		expression: match["expression"],
+		args:       args,
 	}
 }