From ccaf7eb26775e087257515f0d8c2a101b45bef10 Mon Sep 17 00:00:00 2001
From: Mikhail Shilkov <github@mikhail.io>
Date: Tue, 17 Sep 2024 16:15:11 +0200
Subject: [PATCH] Apply secrets to function result properties marked as such

---
 examples/credentials/consumer/Pulumi.yaml | 11 ++++
 examples/credentials/main.go              | 28 +++++++++
 infer/function.go                         |  2 +-
 tests/invoke_test.go                      | 69 +++++++++++++++++++++++
 4 files changed, 109 insertions(+), 1 deletion(-)
 create mode 100644 tests/invoke_test.go

diff --git a/examples/credentials/consumer/Pulumi.yaml b/examples/credentials/consumer/Pulumi.yaml
index f1ba9ac1..d631690d 100644
--- a/examples/credentials/consumer/Pulumi.yaml
+++ b/examples/credentials/consumer/Pulumi.yaml
@@ -6,6 +6,16 @@ plugins:
     - name: credentials
       path: ..
 
+variables:
+  helloworld:
+    fn::invoke:
+      function: credentials:sign
+      arguments:
+        message: "Hello, World!"
+      return: out
+      options:
+        provider: ${provider}
+
 resources:
   provider:
     type: pulumi:providers:credentials
@@ -22,3 +32,4 @@ outputs:
   user: ${user.name}
   password: ${user.password}
   rawPassword: ${provider.password}
+  helloworld: ${helloworld}
diff --git a/examples/credentials/main.go b/examples/credentials/main.go
index 0eea0d1a..d8b6b768 100644
--- a/examples/credentials/main.go
+++ b/examples/credentials/main.go
@@ -27,6 +27,9 @@ func provider() p.Provider {
 			"credentials": "index",
 		},
 		Config: infer.Config[*Config](),
+		Functions: []infer.InferredFunction{
+			infer.Function[*Sign](),
+		},
 	})
 }
 
@@ -110,3 +113,28 @@ func (*User) Diff(ctx context.Context, id string, olds UserState, news UserArgs)
 	}
 	return p.DiffResponse{}, nil
 }
+
+type Sign struct{}
+
+func (Sign) Call(ctx context.Context, args SignArgs) (SignRes, error) {
+	config := infer.GetConfig[Config](ctx)
+	return SignRes{
+		Out: fmt.Sprintf("%s by %s", args.Message, config.User),
+	}, nil
+}
+
+func (r *Sign) Annotate(a infer.Annotator) {
+	a.Describe(r, "Signs the message with the user name and returns the result as a secret.")
+}
+
+type SignArgs struct {
+	Message string `pulumi:"message"`
+}
+
+func (ra *SignArgs) Annotate(a infer.Annotator) {
+	a.Describe(&ra.Message, "Message to sign.")
+}
+
+type SignRes struct {
+	Out string `pulumi:"out" provider:"secret"`
+}
diff --git a/infer/function.go b/infer/function.go
index b5067d97..2147260d 100644
--- a/infer/function.go
+++ b/infer/function.go
@@ -154,6 +154,6 @@ func (r *derivedInvokeController[F, I, O]) Invoke(ctx context.Context, req p.Inv
 		return p.InvokeResponse{}, err
 	}
 	return p.InvokeResponse{
-		Return: m,
+		Return: applySecrets[O](m),
 	}, nil
 }
diff --git a/tests/invoke_test.go b/tests/invoke_test.go
new file mode 100644
index 00000000..83ff9e79
--- /dev/null
+++ b/tests/invoke_test.go
@@ -0,0 +1,69 @@
+// Copyright 2024, Pulumi Corporation.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tests
+
+import (
+	"context"
+	"testing"
+
+	"github.com/blang/semver"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	p "github.com/pulumi/pulumi-go-provider"
+	"github.com/pulumi/pulumi-go-provider/infer"
+	"github.com/pulumi/pulumi-go-provider/integration"
+	"github.com/pulumi/pulumi/sdk/v3/go/common/resource"
+)
+
+type inv struct{}
+
+type invInput struct {
+	Field string `pulumi:"field"`
+}
+
+type invOutput struct {
+	Out string `pulumi:"out" provider:"secret"`
+}
+
+func (inv) Call(ctx context.Context, args invInput) (invOutput, error) {
+	return invOutput{
+		Out: args.Field + "-secret",
+	}, nil
+}
+
+var _ infer.Annotated = inv{}
+
+func (c inv) Annotate(a infer.Annotator) { a.SetToken("index", "inv") }
+
+func TestInferInvokeSecrets(t *testing.T) {
+	t.Parallel()
+
+	resp, err := integration.NewServer("test", semver.MustParse("0.0.0"), infer.Provider(infer.Options{
+		Functions: []infer.InferredFunction{
+			infer.Function[inv, invInput, invOutput](),
+		},
+	})).Invoke(p.InvokeRequest{
+		Token: "test:index:inv",
+		Args: map[resource.PropertyKey]resource.PropertyValue{
+			"field": resource.NewProperty("value"),
+		},
+	})
+	require.NoError(t, err)
+	require.Empty(t, resp.Failures)
+	assert.Equal(t, resource.PropertyMap{
+		"out": resource.MakeSecret(resource.NewProperty("value-secret")),
+	}, resp.Return)
+}