feat: add identity.assumeRoleChainRef for a dynamic assume-role chain

README.md

@@ -104,14 +104,23 @@ skipQueryWhenTargetHasData: true        # optional
 queryIntervalMinutes: 10                # optional throttle
 identity:                               # optional; defaults to Secret
   source: Secret                        # Secret | IRSA | WebIdentity | PodIdentity
-  assumeRoleChain:
+  assumeRoleChain:                      # or assumeRoleChainRef: <path> (status./context./spec.)
   - roleARN: arn:aws:iam::222222222222:role/crossplane-readonly
 ```
 
-**Dynamic queries:** `regionRef`, `filtersRef`, and `parametersRef` resolve their
-value from a `status.`/`context.`/`spec.` path at runtime (overriding the static
-field), so a prior pipeline step or the XR can drive the query — the AWS analog
-of azresourcegraph's `queryRef`/`subscriptionsRef`.
+**Dynamic queries:** `regionRef`, `filtersRef`, `parametersRef`, and
+`identity.assumeRoleChainRef` resolve their value from a
+`status.`/`context.`/`spec.` path at runtime (overriding the static field), so a
+prior pipeline step or the XR can drive the query — the AWS analog of
+azresourcegraph's `queryRef`/`subscriptionsRef`.
+
+The function pod's base identity is shared by every composition on the control
+plane, so prefer a `context.` path fed from platform-controlled config (e.g. an
+EnvironmentConfig) for `assumeRoleChainRef`: a `spec.` path lets any XR author
+pick the role. The ref never falls back to the base identity: an empty or
+malformed chain is fatal, and a path that doesn't resolve uses the static
+`assumeRoleChain` if one is set and is fatal otherwise. The target roles' trust
+policies remain the real control. See `example/composition-assume-role-chain-ref.yaml`.
 
 ## Authentication
 

aws_test.go

@@ -983,6 +983,117 @@ func TestRunFunctionResolvesRefs(t *testing.T) {
 	}
 }
 
+// TestRunFunctionResolvesAssumeRoleChainRef proves identity.assumeRoleChainRef
+// replaces the static chain before the query runs (from context, spec or
+// status), keeps it when the path is absent, and is fatal for a malformed value
+// or for an absent path with no static chain to fall back to.
+func TestRunFunctionResolvesAssumeRoleChainRef(t *testing.T) {
+	const envRef = "context.[apiextensions.crossplane.io/environment].roleChain"
+	static := []v1beta1.AssumeRole{{RoleARN: "arn:aws:iam::111111111111:role/static"}}
+	runWith := func(t *testing.T, identity, ctx, xr string) (*fnv1.RunFunctionResponse, *v1beta1.Input) {
+		t.Helper()
+		var seen *v1beta1.Input
+		f := &Function{log: logging.NewNopLogger(), awsQuery: &MockAWSQuery{fn: func(_ context.Context, _ map[string][]byte, in *v1beta1.Input) (any, error) {
+			seen = in
+			return map[string]any{}, nil
+		}}}
+		req := &fnv1.RunFunctionRequest{
+			Meta: &fnv1.RequestMeta{Tag: "test"},
+			Input: resource.MustStructJSON(fmt.Sprintf(`{
+				"apiVersion":"aws.fn.crossplane.io/v1beta1","kind":"Input",
+				"queryType":"GetCallerIdentity","target":"status.id",
+				"identity":%s
+			}`, identity)),
+			Observed: &fnv1.State{Composite: &fnv1.Resource{Resource: resource.MustStructJSON(xr)}},
+			Context:  resource.MustStructJSON(ctx),
+		}
+		rsp, err := f.RunFunction(context.Background(), req)
+		if err != nil {
+			t.Fatalf("unexpected err: %v", err)
+		}
+		return rsp, seen
+	}
+	run := func(t *testing.T, ref, ctx, xr string) (*fnv1.RunFunctionResponse, *v1beta1.Input) {
+		t.Helper()
+		return runWith(t, fmt.Sprintf(`{"assumeRoleChain":[{"roleARN":"arn:aws:iam::111111111111:role/static"}],"assumeRoleChainRef":%q}`, ref), ctx, xr)
+	}
+
+	chain := `[
+		{"roleARN":"arn:aws:iam::222222222222:role/hop1","externalID":"ext-1","roleSessionName":"fn-aws-query",
+		 "tags":{"team":"platform"},"transitiveTagKeys":["team"]},
+		{"roleARN":"arn:aws:iam::333333333333:role/hop2"}]`
+	want := []v1beta1.AssumeRole{
+		{
+			RoleARN:           "arn:aws:iam::222222222222:role/hop1",
+			ExternalID:        aws.String("ext-1"),
+			RoleSessionName:   aws.String("fn-aws-query"),
+			Tags:              map[string]string{"team": "platform"},
+			TransitiveTagKeys: []string{"team"},
+		},
+		{RoleARN: "arn:aws:iam::333333333333:role/hop2"},
+	}
+	sources := map[string]struct{ ref, ctx, xr string }{
+		"ContextOverridesStatic": {ref: envRef, ctx: `{"apiextensions.crossplane.io/environment":{"roleChain":` + chain + `}}`, xr: observedXR},
+		"SpecOverridesStatic":    {ref: "spec.roleChain", ctx: `{}`, xr: `{"apiVersion":"example.org/v1","kind":"XR","metadata":{"name":"test"},"spec":{"roleChain":` + chain + `}}`},
+		"StatusOverridesStatic":  {ref: "status.roleChain", ctx: `{}`, xr: `{"apiVersion":"example.org/v1","kind":"XR","metadata":{"name":"test"},"status":{"roleChain":` + chain + `}}`},
+	}
+	for name, tc := range sources {
+		t.Run(name, func(t *testing.T) {
+			_, seen := run(t, tc.ref, tc.ctx, tc.xr)
+			if seen == nil {
+				t.Fatal("awsQuery was not called")
+			}
+			if diff := cmp.Diff(want, seen.Identity.AssumeRoleChain); diff != "" {
+				t.Errorf("-want +got:\n%s", diff)
+			}
+		})
+	}
+
+	t.Run("AbsentKeepsStatic", func(t *testing.T) {
+		_, seen := run(t, envRef, `{}`, observedXR)
+		if seen == nil {
+			t.Fatal("awsQuery was not called")
+		}
+		if diff := cmp.Diff(static, seen.Identity.AssumeRoleChain); diff != "" {
+			t.Errorf("-want +got:\n%s", diff)
+		}
+	})
+
+	// Without a static chain, an unresolved ref must not silently run the query
+	// as the function's base identity.
+	t.Run("AbsentWithoutStaticIsFatal", func(t *testing.T) {
+		rsp, seen := runWith(t, fmt.Sprintf(`{"assumeRoleChainRef":%q}`, envRef), `{}`, observedXR)
+		if seen != nil {
+			t.Error("awsQuery should not be called when assumeRoleChainRef does not resolve and there is no static chain")
+		}
+		if !hasFatal(rsp) || !strings.Contains(rsp.GetResults()[0].GetMessage(), envRef) {
+			t.Errorf("want a fatal result naming %q, got: %v", envRef, rsp.GetResults())
+		}
+	})
+
+	cases := map[string]struct{ ref, chain, wantMsg string }{
+		"UnrecognizedPrefix": {ref: "bogus.roleChain", chain: `[]`, wantMsg: "unrecognized reference"},
+		"NotAList":           {ref: envRef, chain: `{"roleARN":"arn:aws:iam::222222222222:role/x"}`, wantMsg: "expected a list"},
+		"NotAnObject":        {ref: envRef, chain: `["arn:aws:iam::222222222222:role/x"]`, wantMsg: "role 0 is not an object"},
+		"MissingRoleARN":     {ref: envRef, chain: `[{"roleARN":"arn:aws:iam::222222222222:role/x"},{"externalID":"e"}]`, wantMsg: "role 1 is missing"},
+		"WrongFieldType":     {ref: envRef, chain: `[{"roleARN":"arn:aws:iam::222222222222:role/x","tags":"team"}]`, wantMsg: "role 0"},
+		"UnknownField":       {ref: envRef, chain: `[{"roleARN":"arn:aws:iam::222222222222:role/x","externalId":"e"}]`, wantMsg: "role 0"},
+		"EmptyList":          {ref: envRef, chain: `[]`, wantMsg: "empty list"},
+	}
+	for name, tc := range cases {
+		t.Run(name, func(t *testing.T) {
+			rsp, seen := run(t, tc.ref, fmt.Sprintf(`{"apiextensions.crossplane.io/environment":{"roleChain":%s}}`, tc.chain), observedXR)
+			if seen != nil {
+				t.Error("awsQuery should not be called for a malformed assumeRoleChainRef")
+			}
+			msg := rsp.GetResults()[0].GetMessage()
+			if !hasFatal(rsp) || !strings.Contains(msg, tc.ref) || !strings.Contains(msg, tc.wantMsg) {
+				t.Errorf("want a fatal result naming %q and %q, got: %v", tc.ref, tc.wantMsg, rsp.GetResults())
+			}
+		})
+	}
+}
+
 func TestRunFunctionIntervalSkip(t *testing.T) {
 	called := false
 	f := &Function{log: logging.NewNopLogger(), awsQuery: &MockAWSQuery{fn: func(_ context.Context, _ map[string][]byte, _ *v1beta1.Input) (any, error) {

example/Makefile

@@ -5,6 +5,8 @@
 #   make commands      # print the copy-pasteable render commands (for manual inspection)
 #   make serve         # run the function in the foreground (for manual `crossplane render`)
 #   make render-irsa   # render the IRSA example (needs real IRSA/role; excluded from `render`)
+#   make render-assume-role ROLE_ARN=arn:aws:iam::<acct>:role/<name>
+#                      # render the assumeRoleChainRef example (needs a role your creds can assume)
 #   make clean
 #
 # Requires: go, the crossplane CLI, and working AWS credentials in
@@ -46,7 +48,7 @@ dynamic-context|xr.yaml|composition-dynamic-context.yaml|--context-values='$(CTX
 endef
 export EXAMPLES
 
-.PHONY: render commands serve render-irsa build clean help
+.PHONY: render commands serve render-irsa render-assume-role build clean help
 
 build: $(BIN)
 
@@ -73,6 +75,9 @@ commands:
 	done <<< "$$EXAMPLES"
 	@echo "# irsa (needs real IRSA/Pod Identity + a valid assumeRoleChain role)"
 	@echo "crossplane render xr.yaml composition-irsa.yaml $(FUNCS) -rc"
+	@echo
+	@echo "# assume-role-chain-ref (needs a role your credentials can assume)"
+	@echo "crossplane render xr.yaml composition-assume-role-chain-ref.yaml $(FUNCS) --function-credentials=$(CREDS) --context-values='$$ROLE_CTX' -rc"
 
 render: build
 	@set -uo pipefail
@@ -110,6 +115,21 @@ render-irsa: build
 	for i in $$(seq 1 40); do (exec 3<>/dev/tcp/127.0.0.1/$(PORT)) 2>/dev/null && break; sleep 0.25; done
 	crossplane render xr.yaml composition-irsa.yaml $(FUNCS) -rc
 
+# The assumeRoleChainRef example takes its chain from the environment context,
+# which a real control plane fills from an EnvironmentConfig. The role must be
+# assumable by the credentials in $(CREDS), so it is excluded from `render`.
+ROLE_ARN ?= arn:aws:iam::222222222222:role/crossplane-readonly
+ROLE_CTX := apiextensions.crossplane.io/environment={"awsQuery":{"roleChain":[{"roleARN":"$(ROLE_ARN)"}]}}
+export ROLE_CTX
+
+render-assume-role: build
+	@echo ">> assumeRoleChainRef example: assuming $(ROLE_ARN)"
+	$(BIN) --insecure --address=$(ADDR) >/tmp/fn-render.log 2>&1 &
+	FN_PID=$$!
+	trap 'kill $$FN_PID 2>/dev/null || true' EXIT
+	for i in $$(seq 1 40); do (exec 3<>/dev/tcp/127.0.0.1/$(PORT)) 2>/dev/null && break; sleep 0.25; done
+	crossplane render xr.yaml composition-assume-role-chain-ref.yaml $(FUNCS) --function-credentials=$(CREDS) --context-values='$(ROLE_CTX)' -rc
+
 clean:
 	rm -f $(BIN) /tmp/render-*.out /tmp/fn-render.log
 
@@ -119,5 +139,6 @@ help:
 	@echo "  commands     print the copy-pasteable render commands for manual inspection"
 	@echo "  serve        run the function in the foreground for manual 'crossplane render'"
 	@echo "  render-irsa  render the IRSA example (needs real IRSA/role)"
+	@echo "  render-assume-role ROLE_ARN=...  render the assumeRoleChainRef example"
 	@echo "  build        build the function binary ($(BIN))"
 	@echo "  clean        remove the built binary and render logs"

example/README.md

@@ -115,6 +115,28 @@ existing AWS provider secret. The `Identity.source` selects the mechanism:
 An optional `identity.assumeRoleChain` assumes one or more roles on top of the
 base credentials (e.g. cross-account reads). See `composition-irsa.yaml`.
 
+### Choosing the role at runtime (`assumeRoleChainRef`)
+
+`identity.assumeRoleChainRef` reads the chain from a `context.`/`status.`/`spec.`
+path instead, so one Composition can assume a different role per control plane.
+`composition-assume-role-chain-ref.yaml` takes it from the environment context,
+which a `function-environment-configs` step fills from a platform-owned
+EnvironmentConfig; `render` seeds it with `--context-values`. Use a role your
+credentials can assume:
+
+```shell
+make render-assume-role ROLE_ARN=arn:aws:iam::222222222222:role/crossplane-readonly
+# or, with `make serve` running:
+crossplane render xr.yaml composition-assume-role-chain-ref.yaml functions.yaml \
+  --function-credentials=secrets/aws-creds.yaml \
+  --context-values='apiextensions.crossplane.io/environment={"awsQuery":{"roleChain":[{"roleARN":"arn:aws:iam::222222222222:role/crossplane-readonly"}]}}' \
+  -rc
+```
+
+`status.callerIdentity.arn` shows the assumed role. Without the context value the
+step is fatal: the example has no static `assumeRoleChain` to fall back to, and
+the function never falls back to its own base identity.
+
 ### Creating the Secret from `~/.aws/credentials`
 
 The function reads the `[default]` profile of an AWS shared-credentials INI, so

example/composition-assume-role-chain-ref.yaml

@@ -0,0 +1,36 @@
+apiVersion: apiextensions.crossplane.io/v1
+kind: Composition
+metadata:
+  name: function-aws-query-assume-role-chain-ref
+spec:
+  compositeTypeRef:
+    apiVersion: example.io/v1alpha1
+    kind: XAccount
+  mode: Pipeline
+  pipeline:
+  # The role to assume is chosen at runtime from the pipeline context, so one
+  # Composition can target a different AWS account per control plane. In a real
+  # control plane a function-environment-configs step loads it from a
+  # platform-owned EnvironmentConfig into the environment context key; `crossplane
+  # render` seeds that key with --context-values (see `make render-assume-role`).
+  # Prefer a context path over spec.*: a spec path lets any XR author pick the role.
+  #
+  # There is no static assumeRoleChain, so if the path doesn't resolve the step is
+  # fatal instead of silently querying as the function's base identity.
+  # GetCallerIdentity makes the assumed role visible in status.callerIdentity.arn.
+  - step: caller-identity
+    functionRef:
+      name: function-aws-query
+    input:
+      apiVersion: aws.fn.crossplane.io/v1beta1
+      kind: Input
+      queryType: GetCallerIdentity
+      target: status.callerIdentity
+      identity:
+        assumeRoleChainRef: context.[apiextensions.crossplane.io/environment].awsQuery.roleChain
+    credentials:
+    - name: aws-creds
+      source: Secret
+      secretRef:
+        namespace: crossplane-system
+        name: aws-creds

fn.go

@@ -8,6 +8,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/go-json-experiment/json"
 	"github.com/upbound/function-aws-query/input/v1beta1"
 	"google.golang.org/protobuf/types/known/structpb"
 
@@ -47,8 +48,8 @@ func (f *Function) RunFunction(ctx context.Context, req *fnv1.RunFunctionRequest
 		return rsp, nil //nolint:nilerr // err is surfaced via rsp; the RPC itself must not fail
 	}
 
-	// Resolve region, filters, and parameters from references if specified
-	// (dynamic queries driven by spec/status/context).
+	// Resolve region, filters, parameters, and the assume-role chain from
+	// references if specified (dynamic queries driven by spec/status/context).
 	if err := f.resolveRegionRef(req, in, rsp); err != nil {
 		return rsp, nil //nolint:nilerr // err is surfaced via rsp; the RPC itself must not fail
 	}
@@ -58,6 +59,9 @@ func (f *Function) RunFunction(ctx context.Context, req *fnv1.RunFunctionRequest
 	if err := f.resolveParametersRef(req, in, rsp); err != nil {
 		return rsp, nil //nolint:nilerr // err is surfaced via rsp; the RPC itself must not fail
 	}
+	if err := f.resolveAssumeRoleChainRef(req, in, rsp); err != nil {
+		return rsp, nil //nolint:nilerr // err is surfaced via rsp; the RPC itself must not fail
+	}
 
 	if in.QueryType == "" {
 		response.Fatal(rsp, errors.New("queryType is required"))
@@ -206,6 +210,38 @@ func (f *Function) resolveParametersRef(req *fnv1.RunFunctionRequest, in *v1beta
 	return nil
 }
 
+// resolveAssumeRoleChainRef resolves Identity.AssumeRoleChainRef into
+// Identity.AssumeRoleChain (overriding any static chain), so the role to assume
+// can be chosen at runtime. An unresolved path keeps the static chain; with no
+// static chain it is fatal, so the query never silently runs as the base identity.
+func (f *Function) resolveAssumeRoleChainRef(req *fnv1.RunFunctionRequest, in *v1beta1.Input, rsp *fnv1.RunFunctionResponse) error {
+	if in.Identity == nil || in.Identity.AssumeRoleChainRef == nil || *in.Identity.AssumeRoleChainRef == "" {
+		return nil
+	}
+	root, path, err := f.refRoot(req, *in.Identity.AssumeRoleChainRef)
+	if err != nil {
+		response.Fatal(rsp, err)
+		return err
+	}
+	v, ok := GetNestedValue(root, path)
+	if !ok {
+		if len(in.Identity.AssumeRoleChain) > 0 {
+			return nil
+		}
+		err := errors.Errorf("assumeRoleChainRef %q did not resolve and no static assumeRoleChain is set; refusing to run as the function's base identity", *in.Identity.AssumeRoleChainRef)
+		response.Fatal(rsp, err)
+		return err
+	}
+	chain, err := toAssumeRoleChain(v)
+	if err != nil {
+		err = errors.Wrapf(err, "cannot resolve assumeRoleChainRef %q", *in.Identity.AssumeRoleChainRef)
+		response.Fatal(rsp, err)
+		return err
+	}
+	in.Identity.AssumeRoleChain = chain
+	return nil
+}
+
 // toFilters converts a resolved reference value into a []v1beta1.Filter. The
 // value must be a list of objects, each with a "name" string and a "values" list.
 func toFilters(v any) ([]v1beta1.Filter, error) {
@@ -234,6 +270,42 @@ func toFilters(v any) ([]v1beta1.Filter, error) {
 	return out, nil
 }
 
+// toAssumeRoleChain converts a resolved reference value into a
+// []v1beta1.AssumeRole. The value must be a non-empty list of objects using the
+// AssumeRole field names, each with a "roleARN" string. Entries are decoded as
+// strictly as the static Input (exact field names, unknown fields rejected).
+func toAssumeRoleChain(v any) ([]v1beta1.AssumeRole, error) {
+	list, ok := v.([]any)
+	if !ok {
+		return nil, errors.New("expected a list of {roleARN, ...} objects")
+	}
+	if len(list) == 0 {
+		// An empty chain would silently drop the static chain and run the
+		// query as the function's base identity.
+		return nil, errors.New("expected at least one role, got an empty list")
+	}
+	out := make([]v1beta1.AssumeRole, len(list))
+	for i, e := range list {
+		m, ok := e.(map[string]any)
+		if !ok {
+			return nil, errors.Errorf("role %d is not an object", i)
+		}
+		if arn, _ := m["roleARN"].(string); arn == "" {
+			return nil, errors.Errorf("role %d is missing a string 'roleARN'", i)
+		}
+		// Round-trip through JSON with the decoder options the SDK uses for the
+		// Input, so the AssumeRole field names and types apply.
+		b, err := json.Marshal(m)
+		if err != nil {
+			return nil, errors.Wrapf(err, "role %d", i)
+		}
+		if err := json.Unmarshal(b, &out[i], json.RejectUnknownMembers(true)); err != nil {
+			return nil, errors.Wrapf(err, "role %d", i)
+		}
+	}
+	return out, nil
+}
+
 // toParameters converts a resolved reference value into a map[string]string.
 func toParameters(v any) (map[string]string, error) {
 	m, ok := v.(map[string]any)

go.mod

@@ -13,6 +13,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/servicequotas v1.35.1
 	github.com/aws/aws-sdk-go-v2/service/sts v1.42.2
 	github.com/crossplane/function-sdk-go v0.7.1
+	github.com/go-json-experiment/json v0.0.0-20240815175050-ebd3a8989ca1
 	github.com/google/go-cmp v0.7.0
 	google.golang.org/protobuf v1.36.11
 	gopkg.in/ini.v1 v1.67.2
@@ -41,7 +42,6 @@ require (
 	github.com/evanphx/json-patch/v5 v5.9.11 // indirect
 	github.com/fatih/color v1.18.0 // indirect
 	github.com/fxamacker/cbor/v2 v2.9.0 // indirect
-	github.com/go-json-experiment/json v0.0.0-20240815175050-ebd3a8989ca1 // indirect
 	github.com/go-logr/logr v1.4.3 // indirect
 	github.com/go-logr/zapr v1.3.0 // indirect
 	github.com/go-openapi/jsonpointer v0.22.5 // indirect