Skip to content

Instantly share code, notes, and snippets.

@prabhakhar
Created November 9, 2025 20:45
Show Gist options
  • Select an option

  • Save prabhakhar/de062ac88c16ea0e69c7a52999197d9d to your computer and use it in GitHub Desktop.

Select an option

Save prabhakhar/de062ac88c16ea0e69c7a52999197d9d to your computer and use it in GitHub Desktop.
package main
import (
"context"
"encoding/json"
"fmt"
"strings"
ackv1alpha1 "github.com/aws-controllers-k8s/iam-controller/apis/v1alpha1"
"k8s.io/apimachinery/pkg/api/errors"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/runtime"
"k8s.io/apimachinery/pkg/types"
ctrl "sigs.k8s.io/controller-runtime"
"sigs.k8s.io/controller-runtime/pkg/client"
"sigs.k8s.io/controller-runtime/pkg/log"
)
// ==================================================================
// 1. NEW CUSTOM RESOURCE DEFINITIONS (The Inputs)
// ==================================================================
// In a real project, these would be in api/v1alpha1/xxx_types.go
// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
type EKSWorkloadIAM struct {
metav1.TypeMeta `json:",inline"`
metav1.ObjectMeta `json:"metadata,omitempty"`
// This spec matches ActorConfig from the previous script
Spec ActorConfig `json:"spec,omitempty"`
}
// +kubebuilder:object:root=true
type EKSWorkloadIAMList struct {
metav1.TypeMeta `json:",inline"`
metav1.ListMeta `json:"metadata,omitempty"`
Items []EKSWorkloadIAM `json:"items"`
}
// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
type AuroraAccessIAM struct {
metav1.TypeMeta `json:",inline"`
metav1.ObjectMeta `json:"metadata,omitempty"`
// This spec matches TargetConfig from the previous script
Spec TargetConfig `json:"spec,omitempty"`
}
// +kubebuilder:object:root=true
type AuroraAccessIAMList struct {
metav1.TypeMeta `json:",inline"`
metav1.ListMeta `json:"metadata,omitempty"`
Items []AuroraAccessIAM `json:"items"`
}
// Re-using the config structs as Kubernetes Specs
type ActorConfig struct {
// AccountID is the AWS Account ID where this role should live
AccountID string `json:"accountId"`
// RoleName is the desired IAM Role name
RoleName string `json:"roleName"`
// --- Trust Method Selection ---
// Option A: EKS Pod Identity (PIA)
// e.g., "pods.eks.amazonaws.com"
TrustPrincipal string `json:"trustPrincipal,omitempty"`
// Option B: IRSA (OIDC)
// OIDCProvider is the OIDC issuer URL without https://
// e.g., "oidc.eks.us-west-2.amazonaws.com/id/EXAMPLED539D4633E53DE1B71EXAMPLE"
OIDCProvider string `json:"oidcProvider,omitempty"`
// ServiceAccountName for IRSA condition
ServiceAccountName string `json:"serviceAccountName,omitempty"`
// ServiceAccountNamespace for IRSA condition
ServiceAccountNamespace string `json:"serviceAccountNamespace,omitempty"`
// ------------------------------
// TargetRoleARNs is a list of DB roles this pod can assume
TargetRoleARNs []string `json:"targetRoleArns,omitempty"`
}
type TargetConfig struct {
AccountID string `json:"accountId"`
RoleName string `json:"roleName"`
Region string `json:"region"`
DBClusterID string `json:"dbClusterId"`
DBUser string `json:"dbUser"`
TrustedRoleARNs []string `json:"trustedRoleArns,omitempty"`
}
// ==================================================================
// 2. RECONCILERS (The Controller Logic)
// ==================================================================
// EKSWorkloadIAMReconciler watches EKSWorkloadIAM inputs and creates ACK Roles
type EKSWorkloadIAMReconciler struct {
client.Client
Scheme *runtime.Scheme
}
func (r *EKSWorkloadIAMReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
l := log.FromContext(ctx)
// 1. Fetch Input (EKSWorkloadIAM)
var input EKSWorkloadIAM
if err := r.Get(ctx, req.NamespacedName, &input); err != nil {
return ctrl.Result{}, client.IgnoreNotFound(err)
}
// 2. Generate the desired ACK Role
desiredACKRole, err := r.generateACKRole(&input)
if err != nil {
l.Error(err, "failed to generate ACK role definition")
// Don't requeue if it's a configuration error in the CR
if strings.Contains(err.Error(), "invalid configuration") {
return ctrl.Result{}, nil
}
return ctrl.Result{}, err // Retry for other errors
}
// 3. Create or Update the ACK Role on the cluster
var currentACKRole ackv1alpha1.Role
err = r.Get(ctx, types.NamespacedName{Name: desiredACKRole.Name, Namespace: desiredACKRole.Namespace}, &currentACKRole)
if errors.IsNotFound(err) {
l.Info("Creating new ACK Role", "role", desiredACKRole.Name)
if err := r.Create(ctx, desiredACKRole); err != nil {
return ctrl.Result{}, err
}
return ctrl.Result{}, nil
} else if err != nil {
return ctrl.Result{}, err
}
// Simple update logic: overwrite spec if it exists
// (In prod, you might compare specs to avoid unnecessary API calls)
currentACKRole.Spec = desiredACKRole.Spec
l.Info("Updating existing ACK Role", "role", currentACKRole.Name)
if err := r.Update(ctx, &currentACKRole); err != nil {
return ctrl.Result{}, err
}
return ctrl.Result{}, nil
}
// generateACKRole adapts the old 'GenerateActorRole' function for K8s
func (r *EKSWorkloadIAMReconciler) generateACKRole(input *EKSWorkloadIAM) (*ackv1alpha1.Role, error) {
var trustPolicy IAMPolicy
// --- Trust Policy Generation Logic ---
if input.Spec.OIDCProvider != "" {
// === OPTION B: IRSA ===
if input.Spec.ServiceAccountName == "" || input.Spec.ServiceAccountNamespace == "" {
return nil, fmt.Errorf("invalid configuration: OIDCProvider requires ServiceAccountName and ServiceAccountNamespace")
}
oidcARN := fmt.Sprintf("arn:aws:iam::%s:oidc-provider/%s", input.Spec.AccountID, input.Spec.OIDCProvider)
subCondition := fmt.Sprintf("system:serviceaccount:%s:%s", input.Spec.ServiceAccountNamespace, input.Spec.ServiceAccountName)
trustPolicy = IAMPolicy{
Version: "2012-10-17",
Statement: []IAMStatement{{
Sid: "TrustIRSA",
Effect: "Allow",
Principal: map[string]string{
"Federated": oidcARN,
},
Action: "sts:AssumeRoleWithWebIdentity",
Condition: map[string]interface{}{
"StringEquals": map[string]string{
// "oidc.eks.REGION.amazonaws.com/id/XXX:sub": "system:serviceaccount:ns:sa"
fmt.Sprintf("%s:sub", input.Spec.OIDCProvider): subCondition,
// "oidc.eks.REGION.amazonaws.com/id/XXX:aud": "sts.amazonaws.com"
fmt.Sprintf("%s:aud", input.Spec.OIDCProvider): "sts.amazonaws.com",
},
},
}},
}
} else if input.Spec.TrustPrincipal != "" {
// === OPTION A: EKS Pod Identity ===
trustPolicy = IAMPolicy{
Version: "2012-10-17",
Statement: []IAMStatement{{
Sid: "TrustEKSPodIdentity",
Effect: "Allow",
Principal: map[string]string{
"Service": input.Spec.TrustPrincipal,
},
Action: []string{"sts:AssumeRole", "sts:TagSession"},
}},
}
} else {
return nil, fmt.Errorf("invalid configuration: must specify either trustPrincipal (PIA) or oidcProvider (IRSA)")
}
trustJSON, _ := json.Marshal(trustPolicy)
// --- Permissions Policy Generation Logic ---
permPolicy := IAMPolicy{
Version: "2012-10-17",
Statement: []IAMStatement{{
Sid: "AllowCrossAccountAssume",
Effect: "Allow",
Action: "sts:AssumeRole",
Resource: input.Spec.TargetRoleARNs,
}},
}
permJSON, _ := json.Marshal(permPolicy)
// ------------------------------------------------
ackRole := &ackv1alpha1.Role{
ObjectMeta: metav1.ObjectMeta{
Name: "ack-" + strings.ToLower(input.Spec.RoleName),
Namespace: input.Namespace,
},
Spec: ackv1alpha1.RoleSpec{
Name: &input.Spec.RoleName,
AssumeRolePolicyDocument: string(trustJSON),
InlinePolicies: map[string]*string{
"CrossAccountHops": awsString(string(permJSON)),
},
},
}
// CRITICAL: Set OwnerReference so deleting the Input deletes the ACK Role
if err := ctrl.SetControllerReference(input, ackRole, r.Scheme); err != nil {
return nil, err
}
return ackRole, nil
}
// AuroraAccessIAMReconciler watches AuroraAccessIAM inputs and creates ACK Roles
type AuroraAccessIAMReconciler struct {
client.Client
Scheme *runtime.Scheme
}
func (r *AuroraAccessIAMReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
l := log.FromContext(ctx)
var input AuroraAccessIAM
if err := r.Get(ctx, req.NamespacedName, &input); err != nil {
return ctrl.Result{}, client.IgnoreNotFound(err)
}
desiredACKRole, err := r.generateACKRole(&input)
if err != nil {
l.Error(err, "failed to generate ACK role definition")
return ctrl.Result{}, err
}
var currentACKRole ackv1alpha1.Role
err = r.Get(ctx, types.NamespacedName{Name: desiredACKRole.Name, Namespace: desiredACKRole.Namespace}, &currentACKRole)
if errors.IsNotFound(err) {
l.Info("Creating new Aurora Target ACK Role", "role", desiredACKRole.Name)
if err := r.Create(ctx, desiredACKRole); err != nil {
return ctrl.Result{}, err
}
return ctrl.Result{}, nil
} else if err != nil {
return ctrl.Result{}, err
}
currentACKRole.Spec = desiredACKRole.Spec
l.Info("Updating existing Aurora Target ACK Role", "role", currentACKRole.Name)
if err := r.Update(ctx, &currentACKRole); err != nil {
return ctrl.Result{}, err
}
return ctrl.Result{}, nil
}
func (r *AuroraAccessIAMReconciler) generateACKRole(input *AuroraAccessIAM) (*ackv1alpha1.Role, error) {
// --- Re-using IAM generation logic ---
trustPolicy := IAMPolicy{
Version: "2012-10-17",
Statement: []IAMStatement{{
Sid: "TrustExternalPods",
Effect: "Allow",
Principal: map[string][]string{"AWS": input.Spec.TrustedRoleARNs},
Action: "sts:AssumeRole",
}},
}
trustJSON, _ := json.Marshal(trustPolicy)
dbResourceARN := fmt.Sprintf("arn:aws:rds-db:%s:%s:dbuser:%s/%s",
input.Spec.Region, input.Spec.AccountID, input.Spec.DBClusterID, input.Spec.DBUser)
permPolicy := IAMPolicy{
Version: "2012-10-17",
Statement: []IAMStatement{{
Sid: "AllowDBConnect",
Effect: "Allow",
Action: "rds-db:connect",
Resource: dbResourceARN,
}},
}
permJSON, _ := json.Marshal(permPolicy)
// ------------------------------------
ackRole := &ackv1alpha1.Role{
ObjectMeta: metav1.ObjectMeta{
Name: "ack-" + strings.ToLower(input.Spec.RoleName),
Namespace: input.Namespace,
},
Spec: ackv1alpha1.RoleSpec{
Name: &input.Spec.RoleName,
AssumeRolePolicyDocument: string(trustJSON),
InlinePolicies: map[string]*string{
"DBConnectPermissions": awsString(string(permJSON)),
},
},
}
if err := ctrl.SetControllerReference(input, ackRole, r.Scheme); err != nil {
return nil, err
}
return ackRole, nil
}
// ==================================================================
// 3. HELPER STRUCTS (from original script)
// ==================================================================
type IAMPolicy struct {
Version string `json:"Version"`
Statement []IAMStatement `json:"Statement"`
}
type IAMStatement struct {
Sid string `json:"Sid,omitempty"`
Effect string `json:"Effect"`
Principal interface{} `json:"Principal,omitempty"`
Action interface{} `json:"Action"`
Resource interface{} `json:"Resource,omitempty"`
Condition map[string]interface{} `json:"Condition,omitempty"`
}
func awsString(v string) *string { return &v }
// ==================================================================
// 4. MAIN (Setup)
// ==================================================================
// In a real operator, this would be in main.go
func Setup(mgr ctrl.Manager) error {
// Register our 2 new controllers with the manager
if err := (&EKSWorkloadIAMReconciler{
Client: mgr.GetClient(),
Scheme: mgr.GetScheme(),
}).SetupWithManager(mgr); err != nil {
return err
}
if err := (&AuroraAccessIAMReconciler{
Client: mgr.GetClient(),
Scheme: mgr.GetScheme(),
}).SetupWithManager(mgr); err != nil {
return err
}
return nil
}
// Boilerplate to attach reconciler to manager
func (r *EKSWorkloadIAMReconciler) SetupWithManager(mgr ctrl.Manager) error {
return ctrl.NewControllerManagedBy(mgr).
For(&EKSWorkloadIAM{}).
Owns(&ackv1alpha1.Role{}). // Watch ACK roles we own
Complete(r)
}
func (r *AuroraAccessIAMReconciler) SetupWithManager(mgr ctrl.Manager) error {
return ctrl.NewControllerManagedBy(mgr).
For(&AuroraAccessIAM{}).
Owns(&ackv1alpha1.Role{}).
Complete(r)
}
@prabhakhar
Copy link
Author

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	ackv1alpha1 "github.com/aws-controllers-k8s/iam-controller/apis/v1alpha1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

// ==================================================================
// 1. NEW CUSTOM RESOURCE DEFINITIONS (The Inputs)
// ==================================================================

// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
type WorkloadIdentity struct {
	metav1.TypeMeta   `json:",inline"`
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Spec for the actor (pod) role
	Spec WorkloadIdentitySpec `json:"spec,omitempty"`
}

// +kubebuilder:object:root=true
type WorkloadIdentityList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata,omitempty"`
	Items           []WorkloadIdentity `json:"items"`
}

// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
type ResourceAccessRole struct {
	metav1.TypeMeta   `json:",inline"`
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Spec for the target resource role
	Spec ResourceAccessRoleSpec `json:"spec,omitempty"`
}

// +kubebuilder:object:root=true
type ResourceAccessRoleList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata,omitempty"`
	Items           []ResourceAccessRole `json:"items"`
}

// --- Capability Specs ---

type WorkloadIdentitySpec struct {
	AccountID string `json:"accountId"`
	RoleName  string `json:"roleName"`

	// Trust Method (PIA vs IRSA)
	TrustPrincipal          string `json:"trustPrincipal,omitempty"`
	OIDCProvider            string `json:"oidcProvider,omitempty"`
	ServiceAccountName      string `json:"serviceAccountName,omitempty"`
	ServiceAccountNamespace string `json:"serviceAccountNamespace,omitempty"`

	TargetRoleARNs []string `json:"targetRoleArns,omitempty"`
}

type ResourceType string

const (
	ResourceTypeDatabase ResourceType = "Database"
	ResourceTypeS3       ResourceType = "S3"
	ResourceTypeDynamoDB ResourceType = "DynamoDB"
	// Add more as needed
)

type ResourceAccessRoleSpec struct {
	AccountID       string       `json:"accountId"`
	RoleName        string       `json:"roleName"`
	Type            ResourceType `json:"type"`
	TrustedRoleARNs []string     `json:"trustedRoleArns,omitempty"`

	// --- Resource-Specific Fields ---

	// For Type=Database
	Region      string `json:"region,omitempty"`
	DBClusterID string `json:"dbClusterId,omitempty"`
	DBUser      string `json:"dbUser,omitempty"`

	// For Type=S3, DynamoDB, etc. (generic fallback)
	GenericAction      string `json:"genericAction,omitempty"`      // e.g., "s3:GetObject"
	GenericResourceARN string `json:"genericResourceArn,omitempty"` // e.g., "arn:aws:s3:::my-bucket/*"
}

// ==================================================================
// 2. RECONCILERS
// ==================================================================

// WorkloadIdentityReconciler
type WorkloadIdentityReconciler struct {
	client.Client
	Scheme *runtime.Scheme
}

func (r *WorkloadIdentityReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var input WorkloadIdentity
	if err := r.Get(ctx, req.NamespacedName, &input); err != nil {
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}
	desiredACKRole, err := r.generateACKRole(&input)
	if err != nil {
		l.Error(err, "failed to generate ACK role")
		if strings.Contains(err.Error(), "invalid configuration") {
			return ctrl.Result{}, nil
		}
		return ctrl.Result{}, err
	}
	return r.reconcileACKRole(ctx, desiredACKRole)
}

func (r *WorkloadIdentityReconciler) reconcileACKRole(ctx context.Context, desired *ackv1alpha1.Role) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var current ackv1alpha1.Role
	err := r.Get(ctx, types.NamespacedName{Name: desired.Name, Namespace: desired.Namespace}, &current)
	if errors.IsNotFound(err) {
		l.Info("Creating new Workload ACK Role", "role", desired.Name)
		if err := r.Create(ctx, desired); err != nil {
			return ctrl.Result{}, err
		}
		return ctrl.Result{}, nil
	} else if err != nil {
		return ctrl.Result{}, err
	}
	current.Spec = desired.Spec
	l.Info("Updating Workload ACK Role", "role", current.Name)
	if err := r.Update(ctx, &current); err != nil {
		return ctrl.Result{}, err
	}
	return ctrl.Result{}, nil
}

func (r *WorkloadIdentityReconciler) generateACKRole(input *WorkloadIdentity) (*ackv1alpha1.Role, error) {
	var trustPolicy IAMPolicy
	if input.Spec.OIDCProvider != "" {
		if input.Spec.ServiceAccountName == "" || input.Spec.ServiceAccountNamespace == "" {
			return nil, fmt.Errorf("invalid configuration: IRSA requires ServiceAccountName and Namespace")
		}
		oidcARN := fmt.Sprintf("arn:aws:iam::%s:oidc-provider/%s", input.Spec.AccountID, input.Spec.OIDCProvider)
		trustPolicy = IAMPolicy{
			Version: "2012-10-17",
			Statement: []IAMStatement{{
				Sid:       "TrustIRSA",
				Effect:    "Allow",
				Principal: map[string]string{"Federated": oidcARN},
				Action:    "sts:AssumeRoleWithWebIdentity",
				Condition: map[string]interface{}{
					"StringEquals": map[string]string{
						fmt.Sprintf("%s:sub", input.Spec.OIDCProvider): fmt.Sprintf("system:serviceaccount:%s:%s", input.Spec.ServiceAccountNamespace, input.Spec.ServiceAccountName),
						fmt.Sprintf("%s:aud", input.Spec.OIDCProvider): "sts.amazonaws.com",
					},
				},
			}},
		}
	} else if input.Spec.TrustPrincipal != "" {
		trustPolicy = IAMPolicy{
			Version: "2012-10-17",
			Statement: []IAMStatement{{
				Sid:       "TrustPIA",
				Effect:    "Allow",
				Principal: map[string]string{"Service": input.Spec.TrustPrincipal},
				Action:    []string{"sts:AssumeRole", "sts:TagSession"},
			}},
		}
	} else {
		return nil, fmt.Errorf("invalid configuration: must specify trustPrincipal or oidcProvider")
	}

	trustJSON, _ := json.Marshal(trustPolicy)
	permPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:      "AllowCrossAccountAssume",
			Effect:   "Allow",
			Action:   "sts:AssumeRole",
			Resource: input.Spec.TargetRoleARNs,
		}},
	}
	permJSON, _ := json.Marshal(permPolicy)

	ackRole := &ackv1alpha1.Role{
		ObjectMeta: metav1.ObjectMeta{Name: "ack-" + strings.ToLower(input.Spec.RoleName), Namespace: input.Namespace},
		Spec: ackv1alpha1.RoleSpec{
			Name:                     &input.Spec.RoleName,
			AssumeRolePolicyDocument: string(trustJSON),
			InlinePolicies:           map[string]*string{"CrossAccountHops": awsString(string(permJSON))},
		},
	}
	if err := ctrl.SetControllerReference(input, ackRole, r.Scheme); err != nil {
		return nil, err
	}
	return ackRole, nil
}

// ResourceAccessRoleReconciler
type ResourceAccessRoleReconciler struct {
	client.Client
	Scheme *runtime.Scheme
}

func (r *ResourceAccessRoleReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var input ResourceAccessRole
	if err := r.Get(ctx, req.NamespacedName, &input); err != nil {
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}
	desiredACKRole, err := r.generateACKRole(&input)
	if err != nil {
		l.Error(err, "failed to generate ACK role")
		return ctrl.Result{}, err
	}
	return r.reconcileACKRole(ctx, desiredACKRole)
}

func (r *ResourceAccessRoleReconciler) reconcileACKRole(ctx context.Context, desired *ackv1alpha1.Role) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var current ackv1alpha1.Role
	err := r.Get(ctx, types.NamespacedName{Name: desired.Name, Namespace: desired.Namespace}, &current)
	if errors.IsNotFound(err) {
		l.Info("Creating new Resource ACK Role", "role", desired.Name)
		if err := r.Create(ctx, desired); err != nil {
			return ctrl.Result{}, err
		}
		return ctrl.Result{}, nil
	} else if err != nil {
		return ctrl.Result{}, err
	}
	current.Spec = desired.Spec
	l.Info("Updating Resource ACK Role", "role", current.Name)
	if err := r.Update(ctx, &current); err != nil {
		return ctrl.Result{}, err
	}
	return ctrl.Result{}, nil
}

func (r *ResourceAccessRoleReconciler) generateACKRole(input *ResourceAccessRole) (*ackv1alpha1.Role, error) {
	trustPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:       "TrustExternalPods",
			Effect:    "Allow",
			Principal: map[string][]string{"AWS": input.Spec.TrustedRoleARNs},
			Action:    "sts:AssumeRole",
		}},
	}
	trustJSON, _ := json.Marshal(trustPolicy)

	// --- Generic Resource Logic ---
	var action string
	var resourceARN string

	switch input.Spec.Type {
	case ResourceTypeDatabase:
		action = "rds-db:connect"
		resourceARN = fmt.Sprintf("arn:aws:rds-db:%s:%s:dbuser:%s/%s",
			input.Spec.Region, input.Spec.AccountID, input.Spec.DBClusterID, input.Spec.DBUser)
	default:
		// Fallback for S3, DynamoDB, etc.
		if input.Spec.GenericAction == "" || input.Spec.GenericResourceARN == "" {
			return nil, fmt.Errorf("generic resource type requires genericAction and genericResourceArn")
		}
		action = input.Spec.GenericAction
		resourceARN = input.Spec.GenericResourceARN
	}
	// ------------------------------

	permPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:      "AllowResourceAccess",
			Effect:   "Allow",
			Action:   action,
			Resource: resourceARN,
		}},
	}
	permJSON, _ := json.Marshal(permPolicy)

	ackRole := &ackv1alpha1.Role{
		ObjectMeta: metav1.ObjectMeta{Name: "ack-" + strings.ToLower(input.Spec.RoleName), Namespace: input.Namespace},
		Spec: ackv1alpha1.RoleSpec{
			Name:                     &input.Spec.RoleName,
			AssumeRolePolicyDocument: string(trustJSON),
			InlinePolicies:           map[string]*string{"ResourcePermissions": awsString(string(permJSON))},
		},
	}
	if err := ctrl.SetControllerReference(input, ackRole, r.Scheme); err != nil {
		return nil, err
	}
	return ackRole, nil
}

// --- Helpers ---
type IAMPolicy struct {
	Version   string         `json:"Version"`
	Statement []IAMStatement `json:"Statement"`
}
type IAMStatement struct {
	Sid       string                 `json:"Sid,omitempty"`
	Effect    string                 `json:"Effect"`
	Principal interface{}            `json:"Principal,omitempty"`
	Action    interface{}            `json:"Action"`
	Resource  interface{}            `json:"Resource,omitempty"`
	Condition map[string]interface{} `json:"Condition,omitempty"`
}

func awsString(v string) *string { return &v }

func Setup(mgr ctrl.Manager) error {
	if err := (&WorkloadIdentityReconciler{Client: mgr.GetClient(), Scheme: mgr.GetScheme()}).SetupWithManager(mgr); err != nil {
		return err
	}
	if err := (&ResourceAccessRoleReconciler{Client: mgr.GetClient(), Scheme: mgr.GetScheme()}).SetupWithManager(mgr); err != nil {
		return err
	}
	return nil
}
func (r *WorkloadIdentityReconciler) SetupWithManager(mgr ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(mgr).For(&WorkloadIdentity{}).Owns(&ackv1alpha1.Role{}).Complete(r)
}
func (r *ResourceAccessRoleReconciler) SetupWithManager(mgr ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(mgr).For(&ResourceAccessRole{}).Owns(&ackv1alpha1.Role{}).Complete(r)
}

@prabhakhar
Copy link
Author

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	ackv1alpha1 "github.com/aws-controllers-k8s/iam-controller/apis/v1alpha1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

// ==================================================================
// 1. NEW CUSTOM RESOURCE DEFINITIONS (The Inputs)
// ==================================================================

// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
type WorkloadIdentity struct {
	metav1.TypeMeta   `json:",inline"`
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Spec for the actor (pod) role
	Spec WorkloadIdentitySpec `json:"spec,omitempty"`
}

// +kubebuilder:object:root=true
type WorkloadIdentityList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata,omitempty"`
	Items           []WorkloadIdentity `json:"items"`
}

// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
type ResourceAccessRole struct {
	metav1.TypeMeta   `json:",inline"`
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Spec for the target resource role
	Spec ResourceAccessRoleSpec `json:"spec,omitempty"`
}

// +kubebuilder:object:root=true
type ResourceAccessRoleList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata,omitempty"`
	Items           []ResourceAccessRole `json:"items"`
}

// --- Capability Specs ---

type WorkloadIdentitySpec struct {
	AccountID string `json:"accountId"`
	RoleName  string `json:"roleName"`

	// --- Trust Method Selection ---

	// Option A: AWS Managed (Pod Identity Agent)
	TrustPrincipal string `json:"trustPrincipal,omitempty"` // e.g. "pods.eks.amazonaws.com"

	// Option B: OIDC (Standard EKS IRSA OR SPIFFE/SPIRE)
	OIDCProvider string `json:"oidcProvider,omitempty"` // e.g. "oidc.eks.us-west-2.amazonaws.com/id/XXX" or "spire.example.com"
	OIDCAudience string `json:"oidcAudience,omitempty"` // Defaults to "sts.amazonaws.com" if unset

	// OIDC sub option 1: Standard Kubernetes SA (IRSA helper)
	ServiceAccountName      string `json:"serviceAccountName,omitempty"`
	ServiceAccountNamespace string `json:"serviceAccountNamespace,omitempty"`

	// OIDC sub option 2: Raw Subject (required for SPIFFE)
	// e.g. "spiffe://example.org/ns/default/sa/myservice"
	OIDCSubject string `json:"oidcSubject,omitempty"`

	// ------------------------------

	TargetRoleARNs []string `json:"targetRoleArns,omitempty"`
}

type ResourceType string

const (
	ResourceTypeDatabase ResourceType = "Database"
	ResourceTypeS3       ResourceType = "S3"
	ResourceTypeDynamoDB ResourceType = "DynamoDB"
	// Add more as needed
)

type ResourceAccessRoleSpec struct {
	AccountID       string       `json:"accountId"`
	RoleName        string       `json:"roleName"`
	Type            ResourceType `json:"type"`
	TrustedRoleARNs []string     `json:"trustedRoleArns,omitempty"`

	// --- Resource-Specific Fields ---

	// For Type=Database
	Region      string `json:"region,omitempty"`
	DBClusterID string `json:"dbClusterId,omitempty"`
	DBUser      string `json:"dbUser,omitempty"`

	// For Type=S3, DynamoDB, etc. (generic fallback)
	GenericAction      string `json:"genericAction,omitempty"`      // e.g., "s3:GetObject"
	GenericResourceARN string `json:"genericResourceArn,omitempty"` // e.g., "arn:aws:s3:::my-bucket/*"
}

// ==================================================================
// 2. RECONCILERS
// ==================================================================

// WorkloadIdentityReconciler
type WorkloadIdentityReconciler struct {
	client.Client
	Scheme *runtime.Scheme
}

func (r *WorkloadIdentityReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var input WorkloadIdentity
	if err := r.Get(ctx, req.NamespacedName, &input); err != nil {
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}
	desiredACKRole, err := r.generateACKRole(&input)
	if err != nil {
		l.Error(err, "failed to generate ACK role")
		if strings.Contains(err.Error(), "invalid configuration") {
			return ctrl.Result{}, nil
		}
		return ctrl.Result{}, err
	}
	return r.reconcileACKRole(ctx, desiredACKRole)
}

func (r *WorkloadIdentityReconciler) reconcileACKRole(ctx context.Context, desired *ackv1alpha1.Role) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var current ackv1alpha1.Role
	err := r.Get(ctx, types.NamespacedName{Name: desired.Name, Namespace: desired.Namespace}, &current)
	if errors.IsNotFound(err) {
		l.Info("Creating new Workload ACK Role", "role", desired.Name)
		if err := r.Create(ctx, desired); err != nil {
			return ctrl.Result{}, err
		}
		return ctrl.Result{}, nil
	} else if err != nil {
		return ctrl.Result{}, err
	}
	current.Spec = desired.Spec
	l.Info("Updating Workload ACK Role", "role", current.Name)
	if err := r.Update(ctx, &current); err != nil {
		return ctrl.Result{}, err
	}
	return ctrl.Result{}, nil
}

func (r *WorkloadIdentityReconciler) generateACKRole(input *WorkloadIdentity) (*ackv1alpha1.Role, error) {
	var trustPolicy IAMPolicy

	if input.Spec.OIDCProvider != "" {
		// === OPTION B: OIDC (IRSA or SPIFFE) ===
		oidcARN := fmt.Sprintf("arn:aws:iam::%s:oidc-provider/%s", input.Spec.AccountID, input.Spec.OIDCProvider)

		// 1. Determine Subject (sub)
		var sub string
		if input.Spec.OIDCSubject != "" {
			// User provided raw subject (use this for SPIFFE)
			sub = input.Spec.OIDCSubject
		} else if input.Spec.ServiceAccountName != "" && input.Spec.ServiceAccountNamespace != "" {
			// User provided K8s helpers (standard IRSA)
			sub = fmt.Sprintf("system:serviceaccount:%s:%s", input.Spec.ServiceAccountNamespace, input.Spec.ServiceAccountName)
		} else {
			return nil, fmt.Errorf("invalid configuration: OIDC requires either oidcSubject OR serviceAccountName+Namespace")
		}

		// 2. Determine Audience (aud)
		aud := input.Spec.OIDCAudience
		if aud == "" {
			aud = "sts.amazonaws.com" // Default for standard EKS IRSA
		}

		trustPolicy = IAMPolicy{
			Version: "2012-10-17",
			Statement: []IAMStatement{{
				Sid:       "TrustOIDC",
				Effect:    "Allow",
				Principal: map[string]string{"Federated": oidcARN},
				Action:    "sts:AssumeRoleWithWebIdentity",
				Condition: map[string]interface{}{
					"StringEquals": map[string]string{
						fmt.Sprintf("%s:sub", input.Spec.OIDCProvider): sub,
						fmt.Sprintf("%s:aud", input.Spec.OIDCProvider): aud,
					},
				},
			}},
		}

	} else if input.Spec.TrustPrincipal != "" {
		// === OPTION A: AWS Managed (Pod Identity Agent) ===
		trustPolicy = IAMPolicy{
			Version: "2012-10-17",
			Statement: []IAMStatement{{
				Sid:       "TrustPIA",
				Effect:    "Allow",
				Principal: map[string]string{"Service": input.Spec.TrustPrincipal},
				Action:    []string{"sts:AssumeRole", "sts:TagSession"},
			}},
		}
	} else {
		return nil, fmt.Errorf("invalid configuration: must specify trustPrincipal or oidcProvider")
	}

	trustJSON, _ := json.Marshal(trustPolicy)
	permPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:      "AllowCrossAccountAssume",
			Effect:   "Allow",
			Action:   "sts:AssumeRole",
			Resource: input.Spec.TargetRoleARNs,
		}},
	}
	permJSON, _ := json.Marshal(permPolicy)

	ackRole := &ackv1alpha1.Role{
		ObjectMeta: metav1.ObjectMeta{Name: "ack-" + strings.ToLower(input.Spec.RoleName), Namespace: input.Namespace},
		Spec: ackv1alpha1.RoleSpec{
			Name:                     &input.Spec.RoleName,
			AssumeRolePolicyDocument: string(trustJSON),
			InlinePolicies:           map[string]*string{"CrossAccountHops": awsString(string(permJSON))},
		},
	}
	if err := ctrl.SetControllerReference(input, ackRole, r.Scheme); err != nil {
		return nil, err
	}
	return ackRole, nil
}

// ResourceAccessRoleReconciler
type ResourceAccessRoleReconciler struct {
	client.Client
	Scheme *runtime.Scheme
}

func (r *ResourceAccessRoleReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var input ResourceAccessRole
	if err := r.Get(ctx, req.NamespacedName, &input); err != nil {
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}
	desiredACKRole, err := r.generateACKRole(&input)
	if err != nil {
		l.Error(err, "failed to generate ACK role")
		return ctrl.Result{}, err
	}
	return r.reconcileACKRole(ctx, desiredACKRole)
}

func (r *ResourceAccessRoleReconciler) reconcileACKRole(ctx context.Context, desired *ackv1alpha1.Role) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var current ackv1alpha1.Role
	err := r.Get(ctx, types.NamespacedName{Name: desired.Name, Namespace: desired.Namespace}, &current)
	if errors.IsNotFound(err) {
		l.Info("Creating new Resource ACK Role", "role", desired.Name)
		if err := r.Create(ctx, desired); err != nil {
			return ctrl.Result{}, err
		}
		return ctrl.Result{}, nil
	} else if err != nil {
		return ctrl.Result{}, err
	}
	current.Spec = desired.Spec
	l.Info("Updating Resource ACK Role", "role", current.Name)
	if err := r.Update(ctx, &current); err != nil {
		return ctrl.Result{}, err
	}
	return ctrl.Result{}, nil
}

func (r *ResourceAccessRoleReconciler) generateACKRole(input *ResourceAccessRole) (*ackv1alpha1.Role, error) {
	trustPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:       "TrustExternalPods",
			Effect:    "Allow",
			Principal: map[string][]string{"AWS": input.Spec.TrustedRoleARNs},
			Action:    "sts:AssumeRole",
		}},
	}
	trustJSON, _ := json.Marshal(trustPolicy)

	// --- Generic Resource Logic ---
	var action string
	var resourceARN string

	switch input.Spec.Type {
	case ResourceTypeDatabase:
		action = "rds-db:connect"
		resourceARN = fmt.Sprintf("arn:aws:rds-db:%s:%s:dbuser:%s/%s",
			input.Spec.Region, input.Spec.AccountID, input.Spec.DBClusterID, input.Spec.DBUser)
	default:
		// Fallback for S3, DynamoDB, etc.
		if input.Spec.GenericAction == "" || input.Spec.GenericResourceARN == "" {
			return nil, fmt.Errorf("generic resource type requires genericAction and genericResourceArn")
		}
		action = input.Spec.GenericAction
		resourceARN = input.Spec.GenericResourceARN
	}
	// ------------------------------

	permPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:      "AllowResourceAccess",
			Effect:   "Allow",
			Action:   action,
			Resource: resourceARN,
		}},
	}
	permJSON, _ := json.Marshal(permPolicy)

	ackRole := &ackv1alpha1.Role{
		ObjectMeta: metav1.ObjectMeta{Name: "ack-" + strings.ToLower(input.Spec.RoleName), Namespace: input.Namespace},
		Spec: ackv1alpha1.RoleSpec{
			Name:                     &input.Spec.RoleName,
			AssumeRolePolicyDocument: string(trustJSON),
			InlinePolicies:           map[string]*string{"ResourcePermissions": awsString(string(permJSON))},
		},
	}
	if err := ctrl.SetControllerReference(input, ackRole, r.Scheme); err != nil {
		return nil, err
	}
	return ackRole, nil
}

// --- Helpers ---
type IAMPolicy struct {
	Version   string         `json:"Version"`
	Statement []IAMStatement `json:"Statement"`
}
type IAMStatement struct {
	Sid       string                 `json:"Sid,omitempty"`
	Effect    string                 `json:"Effect"`
	Principal interface{}            `json:"Principal,omitempty"`
	Action    interface{}            `json:"Action"`
	Resource  interface{}            `json:"Resource,omitempty"`
	Condition map[string]interface{} `json:"Condition,omitempty"`
}

func awsString(v string) *string { return &v }

func Setup(mgr ctrl.Manager) error {
	if err := (&WorkloadIdentityReconciler{Client: mgr.GetClient(), Scheme: mgr.GetScheme()}).SetupWithManager(mgr); err != nil {
		return err
	}
	if err := (&ResourceAccessRoleReconciler{Client: mgr.GetClient(), Scheme: mgr.GetScheme()}).SetupWithManager(mgr); err != nil {
		return err
	}
	return nil
}
func (r *WorkloadIdentityReconciler) SetupWithManager(mgr ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(mgr).For(&WorkloadIdentity{}).Owns(&ackv1alpha1.Role{}).Complete(r)
}
func (r *ResourceAccessRoleReconciler) SetupWithManager(mgr ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(mgr).For(&ResourceAccessRole{}).Owns(&ackv1alpha1.Role{}).Complete(r)
}

@prabhakhar
Copy link
Author

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	// You will need to import the ACK IAM Role type definition.
	// The exact path may vary based on your Go modules.
	// e.g., "github.com/aws-controllers-k8s/iam-controller/apis/v1alpha1"
	ackv1alpha1 "github.com/aws-controllers-k8s/iam-controller/apis/v1alpha1"

	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

// ==================================================================
// 1. CUSTOM RESOURCE DEFINITIONS (The Inputs)
// ==================================================================

// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
// WorkloadIdentity defines the desired IAM role for a pod (the "actor").
type WorkloadIdentity struct {
	metav1.TypeMeta   `json:",inline"`
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Spec for the actor (pod) role
	Spec WorkloadIdentitySpec `json:"spec,omitempty"`
}

// +kubebuilder:object:root=true
type WorkloadIdentityList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata,omitempty"`
	Items           []WorkloadIdentity `json:"items"`
}

// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
// ResourceAccessRole defines the IAM role for a target resource (e.g., DB).
type ResourceAccessRole struct {
	metav1.TypeMeta   `json:",inline"`
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Spec for the target resource role
	Spec ResourceAccessRoleSpec `json:"spec,omitempty"`
}

// +kubebuilder:object:root=true
type ResourceAccessRoleList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata,omitempty"`
	Items           []ResourceAccessRole `json:"items"`
}

// --- Capability Specs ---

type WorkloadIdentitySpec struct {
	AccountID string `json:"accountId"`
	RoleName  string `json:"roleName"`

	// --- Trust Method Selection ---

	// Option A: AWS Managed (Pod Identity Agent)
	TrustPrincipal string `json:"trustPrincipal,omitempty"` // e.g. "pods.eks.amazonaws.com"

	// Option B: OIDC (Standard EKS IRSA OR SPIFFE/SPIRE)
	OIDCProvider string `json:"oidcProvider,omitempty"` // e.g. "oidc.eks.us-west-2.amazonaws.com/id/XXX" or "spire.example.com"
	OIDCAudience string `json:"oidcAudience,omitempty"` // Defaults to "sts.amazonaws.com" if unset

	// OIDC sub option 1: Standard Kubernetes SA (IRSA helper)
	ServiceAccountName      string `json:"serviceAccountName,omitempty"`
	ServiceAccountNamespace string `json:"serviceAccountNamespace,omitempty"`

	// OIDC sub option 2: Raw Subject (required for SPIFFE)
	// e.g. "spiffe://example.org/ns/default/sa/myservice"
	OIDCSubject string `json:"oidcSubject,omitempty"`

	// ------------------------------

	TargetRoleARNs []string `json:"targetRoleArns,omitempty"`
}

type ResourceType string

const (
	ResourceTypeDatabase ResourceType = "Database"
	ResourceTypeS3       ResourceType = "S3"
	ResourceTypeDynamoDB ResourceType = "DynamoDB"
	// Add more as needed
)

type ResourceAccessRoleSpec struct {
	AccountID       string       `json:"accountId"`
	RoleName        string       `json:"roleName"`
	Type            ResourceType `json:"type"`
	TrustedRoleARNs []string     `json:"trustedRoleArns,omitempty"`

	// --- Resource-Specific Fields ---

	// For Type=Database
	Region      string `json:"region,omitempty"`
	DBClusterID string `json:"dbClusterId,omitempty"`
	DBUser      string `json:"dbUser,omitempty"`

	// For Type=S3, DynamoDB, etc. (generic fallback)
	GenericAction      string `json:"genericAction,omitempty"`      // e.g., "s3:GetObject"
	GenericResourceARN string `json:"genericResourceArn,omitempty"` // e.g., "arn:aws:s3:::my-bucket/*"
}

// ==================================================================
// 2. RECONCILERS
// ==================================================================

// WorkloadIdentityReconciler
type WorkloadIdentityReconciler struct {
	client.Client
	Scheme *runtime.Scheme
}

// Reconcile is the main loop for the WorkloadIdentity controller
func (r *WorkloadIdentityReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var input WorkloadIdentity
	if err := r.Get(ctx, req.NamespacedName, &input); err != nil {
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}

	// Generate the ACK Role definition based on the input CR
	desiredACKRole, err := r.generateACKRole(&input)
	if err != nil {
		l.Error(err, "failed to generate ACK role")
		if strings.Contains(err.Error(), "invalid configuration") {
			// Don't requeue if the spec is invalid
			return ctrl.Result{}, nil
		}
		return ctrl.Result{}, err
	}

	// Create or Update the underlying ACK Role
	return r.reconcileACKRole(ctx, desiredACKRole)
}

// reconcileACKRole handles the Create/Update logic for the ACK Role
func (r *WorkloadIdentityReconciler) reconcileACKRole(ctx context.Context, desired *ackv1alpha1.Role) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var current ackv1alpha1.Role
	err := r.Get(ctx, types.NamespacedName{Name: desired.Name, Namespace: desired.Namespace}, &current)

	if errors.IsNotFound(err) {
		l.Info("Creating new Workload ACK Role", "role", desired.Name)
		if err := r.Create(ctx, desired); err != nil {
			return ctrl.Result{}, err
		}
		return ctrl.Result{}, nil
	} else if err != nil {
		return ctrl.Result{}, err
	}

	// Role exists, update its spec
	current.Spec = desired.Spec
	l.Info("Updating Workload ACK Role", "role", current.Name)
	if err := r.Update(ctx, &current); err != nil {
		return ctrl.Result{}, err
	}
	return ctrl.Result{}, nil
}

// generateACKRole translates the WorkloadIdentity CR into an ACK Role
func (r *WorkloadIdentityReconciler) generateACKRole(input *WorkloadIdentity) (*ackv1alpha1.Role, error) {
	var trustPolicy IAMPolicy

	if input.Spec.OIDCProvider != "" {
		// === OPTION B: OIDC (IRSA or SPIFFE) ===
		oidcARN := fmt.Sprintf("arn:aws:iam::%s:oidc-provider/%s", input.Spec.AccountID, input.Spec.OIDCProvider)

		// 1. Determine Subject (sub)
		var sub string
		if input.Spec.OIDCSubject != "" {
			// User provided raw subject (use this for SPIFFE)
			sub = input.Spec.OIDCSubject
		} else if input.Spec.ServiceAccountName != "" && input.Spec.ServiceAccountNamespace != "" {
			// User provided K8s helpers (standard IRSA)
			sub = fmt.Sprintf("system:serviceaccount:%s:%s", input.Spec.ServiceAccountNamespace, input.Spec.ServiceAccountName)
		} else {
			return nil, fmt.Errorf("invalid configuration: OIDC requires either oidcSubject OR serviceAccountName+Namespace")
		}

		// 2. Determine Audience (aud)
		aud := input.Spec.OIDCAudience
		if aud == "" {
			aud = "sts.amazonaws.com" // Default for standard EKS IRSA
		}

		trustPolicy = IAMPolicy{
			Version: "2012-10-17",
			Statement: []IAMStatement{{
				Sid:       "TrustOIDC",
				Effect:    "Allow",
				Principal: map[string]string{"Federated": oidcARN},
				Action:    "sts:AssumeRoleWithWebIdentity",
				Condition: map[string]interface{}{
					"StringEquals": map[string]string{
						fmt.Sprintf("%s:sub", input.Spec.OIDCProvider): sub,
						fmt.Sprintf("%s:aud", input.Spec.OIDCProvider): aud,
					},
				},
			}},
		}

	} else if input.Spec.TrustPrincipal != "" {
		// === OPTION A: AWS Managed (Pod Identity Agent) ===
		trustPolicy = IAMPolicy{
			Version: "2012-10-17",
			Statement: []IAMStatement{{
				Sid:       "TrustPIA",
				Effect:    "Allow",
				Principal: map[string]string{"Service": input.Spec.TrustPrincipal},
				Action:    []string{"sts:AssumeRole", "sts:TagSession"},
			}},
		}
	} else {
		return nil, fmt.Errorf("invalid configuration: must specify trustPrincipal or oidcProvider")
	}

	trustJSON, _ := json.Marshal(trustPolicy)
	permPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:      "AllowCrossAccountAssume",
			Effect:   "Allow",
			Action:   "sts:AssumeRole",
			Resource: input.Spec.TargetRoleARNs,
		}},
	}
	permJSON, _ := json.Marshal(permPolicy)

	ackRole := &ackv1alpha1.Role{
		ObjectMeta: metav1.ObjectMeta{Name: "ack-" + strings.ToLower(input.Spec.RoleName), Namespace: input.Namespace},
		Spec: ackv1alpha1.RoleSpec{
			Name:                     &input.Spec.RoleName,
			AssumeRolePolicyDocument: string(trustJSON),
			InlinePolicies:           map[string]*string{"CrossAccountHops": awsString(string(permJSON))},
		},
	}
	if err := ctrl.SetControllerReference(input, ackRole, r.Scheme); err != nil {
		return nil, err
	}
	return ackRole, nil
}

// ResourceAccessRoleReconciler
type ResourceAccessRoleReconciler struct {
	client.Client
	Scheme *runtime.Scheme
}

// Reconcile is the main loop for the ResourceAccessRole controller
func (r *ResourceAccessRoleReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var input ResourceAccessRole
	if err := r.Get(ctx, req.NamespacedName, &input); err != nil {
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}
	desiredACKRole, err := r.generateACKRole(&input)
	if err != nil {
		l.Error(err, "failed to generate ACK role")
		return ctrl.Result{}, err
	}
	return r.reconcileACKRole(ctx, desiredACKRole)
}

// reconcileACKRole handles the Create/Update logic for the ACK Role
func (r *ResourceAccessRoleReconciler) reconcileACKRole(ctx context.Context, desired *ackv1alpha1.Role) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var current ackv1alpha1.Role
	err := r.Get(ctx, types.NamespacedName{Name: desired.Name, Namespace: desired.Namespace}, &current)
	if errors.IsNotFound(err) {
		l.Info("Creating new Resource ACK Role", "role", desired.Name)
		if err := r.Create(ctx, desired); err != nil {
			return ctrl.Result{}, err
		}
		return ctrl.Result{}, nil
	} else if err != nil {
		return ctrl.Result{}, err
	}
	current.Spec = desired.Spec
	l.Info("Updating Resource ACK Role", "role", current.Name)
	if err := r.Update(ctx, &current); err != nil {
		return ctrl.Result{}, err
	}
	return ctrl.Result{}, nil
}

// generateACKRole translates the ResourceAccessRole CR into an ACK Role
func (r *ResourceAccessRoleReconciler) generateACKRole(input *ResourceAccessRole) (*ackv1alpha1.Role, error) {
	trustPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:       "TrustExternalPods",
			Effect:    "Allow",
			Principal: map[string][]string{"AWS": input.Spec.TrustedRoleARNs},
			Action:    "sts:AssumeRole",
		}},
	}
	trustJSON, _ := json.Marshal(trustPolicy)

	// --- Generic Resource Logic ---
	var action string
	var resourceARN string

	switch input.Spec.Type {
	case ResourceTypeDatabase:
		action = "rds-db:connect"
		resourceARN = fmt.Sprintf("arn:aws:rds-db:%s:%s:dbuser:%s/%s",
			input.Spec.Region, input.Spec.AccountID, input.Spec.DBClusterID, input.Spec.DBUser)
	default:
		// Fallback for S3, DynamoDB, etc.
		if input.Spec.GenericAction == "" || input.Spec.GenericResourceARN == "" {
			return nil, fmt.Errorf("generic resource type requires genericAction and genericResourceArn")
		}
		action = input.Spec.GenericAction
		resourceARN = input.Spec.GenericResourceARN
	}
	// ------------------------------

	permPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:      "AllowResourceAccess",
			Effect:   "Allow",
			Action:   action,
			Resource: resourceARN,
		}},
	}
	permJSON, _ := json.Marshal(permPolicy)

	ackRole := &ackv1alpha1.Role{
		ObjectMeta: metav1.ObjectMeta{Name: "ack-" + strings.ToLower(input.Spec.RoleName), Namespace: input.Namespace},
		Spec: ackv1alpha1.RoleSpec{
			Name:                     &input.Spec.RoleName,
			AssumeRolePolicyDocument: string(trustJSON),
			InlinePolicies:           map[string]*string{"ResourcePermissions": awsString(string(permJSON))},
		},
	}
	if err := ctrl.SetControllerReference(input, ackRole, r.Scheme); err != nil {
		return nil, err
	}
	return ackRole, nil
}

// ==================================================================
// 3. HELPERS
// ==================================================================

// IAMPolicy is a helper struct for JSON marshaling
type IAMPolicy struct {
	Version   string         `json:"Version"`
	Statement []IAMStatement `json:"Statement"`
}

// IAMStatement is a helper struct for JSON marshaling
type IAMStatement struct {
	Sid       string                 `json:"Sid,omitempty"`
	Effect    string                 `json:"Effect"`
	Principal interface{}            `json:"Principal,omitempty"`
	Action    interface{}            `json:"Action"`
	Resource  interface{}            `json:"Resource,omitempty"`
	Condition map[string]interface{} `json:"Condition,omitempty"`
}

// awsString is a helper to convert a string to a *string
func awsString(v string) *string {
	return &v
}

// ==================================================================
// 4. SETUP (in main.go)
// ==================================================================

// Setup registers the controllers with the manager
func Setup(mgr ctrl.Manager) error {
	if err := (&WorkloadIdentityReconciler{
		Client: mgr.GetClient(),
		Scheme: mgr.GetScheme(),
	}).SetupWithManager(mgr); err != nil {
		return err
	}

	if err := (&ResourceAccessRoleReconciler{
		Client: mgr.GetClient(),
		Scheme: mgr.GetScheme(),
	}).SetupWithManager(mgr); err != nil {
		return err
	}
	return nil
}

// SetupWithManager sets up the controller with the Manager.
func (r *WorkloadIdentityReconciler) SetupWithManager(mgr ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(mgr).
		For(&WorkloadIdentity{}).
		Owns(&ackv1alpha1.Role{}). // Automatically watch ACK Roles it creates
		Complete(r)
}

// SetupWithManager sets up the controller with the Manager.
func (r *ResourceAccessRoleReconciler) SetupWithManager(mgr ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(mgr).
		For(&ResourceAccessRole{}).
		Owns(&ackv1alpha1.Role{}). // Automatically watch ACK Roles it creates
		Complete(r)
}

@prabhakhar
Copy link
Author

prabhakhar commented Nov 9, 2025

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	ackv1alpha1 "github.com/aws-controllers-k8s/iam-controller/apis/v1alpha1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

// ==================================================================
// 1. IMPROVED CRD MODELING (Discriminated Union)
// ==================================================================

type WorkloadIdentityType string

const (
	WorkloadTypePodIdentity WorkloadIdentityType = "PodIdentity"
	WorkloadTypeIRSA        WorkloadIdentityType = "IRSA"
	WorkloadTypeSPIFFE      WorkloadIdentityType = "SPIFFE"
)

// WorkloadIdentitySpec uses a discriminated union pattern.
type WorkloadIdentitySpec struct {
	AccountID string `json:"accountId"`
	RoleName  string `json:"roleName"`

	// Type explicitly tells the controller which sub-struct to use.
	// +kubebuilder:validation:Enum=PodIdentity;IRSA;SPIFFE
	Type WorkloadIdentityType `json:"type"`

	// --- mutually exclusive configuration blocks ---
	PodIdentity *PodIdentitySpec `json:"podIdentity,omitempty"`
	IRSA        *IRSASpec        `json:"irsa,omitempty"`
	SPIFFE      *SPIFFESpec      `json:"spiffe,omitempty"`
	// -----------------------------------------------

	TargetRoleARNs []string `json:"targetRoleArns,omitempty"`
}

type PodIdentitySpec struct {
	TrustPrincipal string `json:"trustPrincipal"` // e.g. "pods.eks.amazonaws.com"
}

type IRSASpec struct {
	OIDCProvider      string `json:"oidcProvider"`
	OIDCAudience      string `json:"oidcAudience,omitempty"` // Optional, defaults to sts.amazonaws.com
	ServiceAccount    string `json:"serviceAccount"`
	ServiceNamespace  string `json:"serviceNamespace"`
}

type SPIFFESpec struct {
	OIDCProvider string `json:"oidcProvider"`
	OIDCAudience string `json:"oidcAudience,omitempty"` // Optional
	SPIFFEID     string `json:"spiffeId"`               // The exact subject
}

// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
type WorkloadIdentity struct {
	metav1.TypeMeta   `json:",inline"`
	metav1.ObjectMeta `json:"metadata,omitempty"`
	Spec              WorkloadIdentitySpec `json:"spec,omitempty"`
}

// +kubebuilder:object:root=true
type WorkloadIdentityList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata,omitempty"`
	Items           []WorkloadIdentity `json:"items"`
}

// ==================================================================
// 2. RECONCILER UPDATES
// ==================================================================

type WorkloadIdentityReconciler struct {
	client.Client
	Scheme *runtime.Scheme
}

// Reconcile is the main loop for the WorkloadIdentity controller
func (r *WorkloadIdentityReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var input WorkloadIdentity
	if err := r.Get(ctx, req.NamespacedName, &input); err != nil {
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}

	// Generate the ACK Role definition based on the input CR
	desiredACKRole, err := r.generateACKRole(&input)
	if err != nil {
		l.Error(err, "failed to generate ACK role")
		if strings.Contains(err.Error(), "invalid configuration") || strings.Contains(err.Error(), "config missing") || strings.Contains(err.Error(), "incomplete") {
			// Don't requeue if the spec is invalid to avoid hot loops
			return ctrl.Result{}, nil
		}
		return ctrl.Result{}, err
	}

	// Create or Update the underlying ACK Role
	return r.reconcileACKRole(ctx, desiredACKRole)
}

// reconcileACKRole handles the Create/Update logic for the ACK Role
func (r *WorkloadIdentityReconciler) reconcileACKRole(ctx context.Context, desired *ackv1alpha1.Role) (ctrl.Result, error) {
	l := log.FromContext(ctx)
	var current ackv1alpha1.Role
	err := r.Get(ctx, types.NamespacedName{Name: desired.Name, Namespace: desired.Namespace}, &current)

	if errors.IsNotFound(err) {
		l.Info("Creating new Workload ACK Role", "role", desired.Name)
		if err := r.Create(ctx, desired); err != nil {
			return ctrl.Result{}, err
		}
		return ctrl.Result{}, nil
	} else if err != nil {
		return ctrl.Result{}, err
	}

	// Role exists, update its spec
	current.Spec = desired.Spec
	l.Info("Updating Workload ACK Role", "role", current.Name)
	if err := r.Update(ctx, &current); err != nil {
		return ctrl.Result{}, err
	}
	return ctrl.Result{}, nil
}

// generateACKRole now uses a clean switch statement based on Type
func (r *WorkloadIdentityReconciler) generateACKRole(input *WorkloadIdentity) (*ackv1alpha1.Role, error) {
	var trustPolicy IAMPolicy
	var err error

	switch input.Spec.Type {
	case WorkloadTypePodIdentity:
		trustPolicy, err = r.generatePodIdentityTrust(input.Spec.PodIdentity)
	case WorkloadTypeIRSA:
		trustPolicy, err = r.generateIRSATrust(input.Spec.AccountID, input.Spec.IRSA)
	case WorkloadTypeSPIFFE:
		trustPolicy, err = r.generateSPIFFETrust(input.Spec.AccountID, input.Spec.SPIFFE)
	default:
		return nil, fmt.Errorf("unknown workload identity type: %s", input.Spec.Type)
	}

	if err != nil {
		return nil, err
	}

	trustJSON, _ := json.Marshal(trustPolicy)
	permPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:      "AllowCrossAccountAssume",
			Effect:   "Allow",
			Action:   "sts:AssumeRole",
			Resource: input.Spec.TargetRoleARNs,
		}},
	}
	permJSON, _ := json.Marshal(permPolicy)

	ackRole := &ackv1alpha1.Role{
		ObjectMeta: metav1.ObjectMeta{Name: "ack-" + strings.ToLower(input.Spec.RoleName), Namespace: input.Namespace},
		Spec: ackv1alpha1.RoleSpec{
			Name:                     &input.Spec.RoleName,
			AssumeRolePolicyDocument: string(trustJSON),
			InlinePolicies:           map[string]*string{"CrossAccountHops": awsString(string(permJSON))},
		},
	}
	if err := ctrl.SetControllerReference(input, ackRole, r.Scheme); err != nil {
		return nil, err
	}
	return ackRole, nil
}

// --- Dedicated helper functions for each type ---

func (r *WorkloadIdentityReconciler) generatePodIdentityTrust(spec *PodIdentitySpec) (IAMPolicy, error) {
	if spec == nil || spec.TrustPrincipal == "" {
		return IAMPolicy{}, fmt.Errorf("PodIdentity config missing")
	}
	return IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:       "TrustPIA",
			Effect:    "Allow",
			Principal: map[string]string{"Service": spec.TrustPrincipal},
			Action:    []string{"sts:AssumeRole", "sts:TagSession"},
		}},
	}, nil
}

func (r *WorkloadIdentityReconciler) generateIRSATrust(accountID string, spec *IRSASpec) (IAMPolicy, error) {
	if spec == nil || spec.OIDCProvider == "" || spec.ServiceAccount == "" || spec.ServiceNamespace == "" {
		return IAMPolicy{}, fmt.Errorf("IRSA config incomplete")
	}
	aud := spec.OIDCAudience
	if aud == "" {
		aud = "sts.amazonaws.com"
	}
	sub := fmt.Sprintf("system:serviceaccount:%s:%s", spec.ServiceNamespace, spec.ServiceAccount)
	return r.buildOIDCTrust(accountID, spec.OIDCProvider, sub, aud), nil
}

func (r *WorkloadIdentityReconciler) generateSPIFFETrust(accountID string, spec *SPIFFESpec) (IAMPolicy, error) {
	if spec == nil || spec.OIDCProvider == "" || spec.SPIFFEID == "" {
		return IAMPolicy{}, fmt.Errorf("SPIFFE config incomplete")
	}
	aud := spec.OIDCAudience
	if aud == "" {
		aud = "sts.amazonaws.com"
	}
	return r.buildOIDCTrust(accountID, spec.OIDCProvider, spec.SPIFFEID, aud), nil
}

// Shared helper for OIDC since IRSA and SPIFFE both use it
func (r *WorkloadIdentityReconciler) buildOIDCTrust(accountID, provider, sub, aud string) IAMPolicy {
	oidcARN := fmt.Sprintf("arn:aws:iam::%s:oidc-provider/%s", accountID, provider)
	return IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:       "TrustOIDC",
			Effect:    "Allow",
			Principal: map[string]string{"Federated": oidcARN},
			Action:    "sts:AssumeRoleWithWebIdentity",
			Condition: map[string]interface{}{
				"StringEquals": map[string]string{
					fmt.Sprintf("%s:sub", provider): sub,
					fmt.Sprintf("%s:aud", provider): aud,
				},
			},
		}},
	}
}

// --- Helpers (same as before) ---
type IAMPolicy struct {
	Version   string         `json:"Version"`
	Statement []IAMStatement `json:"Statement"`
}
type IAMStatement struct {
	Sid       string                 `json:"Sid,omitempty"`
	Effect    string                 `json:"Effect"`
	Principal interface{}            `json:"Principal,omitempty"`
	Action    interface{}            `json:"Action"`
	Resource  interface{}            `json:"Resource,omitempty"`
	Condition map[string]interface{} `json:"Condition,omitempty"`
}

func awsString(v string) *string { return &v }

// SetupWithManager sets up the controller with the Manager.
func (r *WorkloadIdentityReconciler) SetupWithManager(mgr ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(mgr).
		For(&WorkloadIdentity{}).
		Owns(&ackv1alpha1.Role{}).
		Complete(r)
}

Example Usage with New Model

Now the YAML is much clearer because it's explicit.

Option 1: IRSA

spec:
  type: IRSA
  accountId: "111122223333"
  roleName: "MyRole"
  irsa:
    oidcProvider: "oidc.eks.us-west-2..."
    serviceAccount: "my-sa"
    serviceNamespace: "default"

@prabhakhar
Copy link
Author

apiVersion: myorg.com/v1alpha1
kind: WorkloadIdentity
metadata:
  name: payment-service-irsa
  namespace: ack-system
spec:
  type: IRSA
  accountId: "111122223333"
  roleName: "PaymentServiceRole"
  irsa:
    # The OIDC provider ID for your EKS cluster
    oidcProvider: "oidc.eks.us-west-2.amazonaws.com/id/EXAMPLED539D4633E53DE1B716D3041E"
    # The Kubernetes ServiceAccount details
    serviceAccount: "payment-sa"
    serviceNamespace: "default"
  targetRoleArns:
    - "arn:aws:iam::333344445555:role/AuroraAccessRole"

---
apiVersion: myorg.com/v1alpha1
kind: WorkloadIdentity
metadata:
  name: analytics-service-pia
  namespace: ack-system
spec:
  type: PodIdentity
  accountId: "111122223333"
  roleName: "AnalyticsServiceRole"
  podIdentity:
    # Standard principal for the EKS Pod Identity service
    trustPrincipal: "pods.eks.amazonaws.com"
  targetRoleArns:
    - "arn:aws:iam::333344445555:role/RedshiftAccessRole"

---
apiVersion: myorg.com/v1alpha1
kind: WorkloadIdentity
metadata:
  name: legacy-app-spiffe
  namespace: ack-system
spec:
  type: SPIFFE
  accountId: "111122223333"
  roleName: "LegacyAppRole"
  spiffe:
    # Your SPIRE server's public OIDC discovery endpoint domain
    oidcProvider: "spire.example.com"
    # The exact SPIFFE ID issued to the workload
    spiffeId: "spiffe://example.org/ns/legacy/sa/monolith"
    # Optional: If your SPIRE agent requests a specific audience
    # oidcAudience: "myspire-cluster"
  targetRoleArns:
    - "arn:aws:iam::333344445555:role/S3AccessRole"

@prabhakhar
Copy link
Author

package iamgen

// ==================================================================
// ACTOR CONFIGURATION (The Pod/Workload)
// ==================================================================

type ActorType string

const (
	ActorTypePodIdentity ActorType = "PodIdentity"
	ActorTypeIRSA        ActorType = "IRSA"
	ActorTypeSPIFFE      ActorType = "SPIFFE"
)

type ActorConfig struct {
	// K8s Metadata for the resulting ACK Role
	Name      string
	Namespace string

	// AWS Specifics
	AccountID string
	RoleName  string

	// Trust Configuration (Discriminated Union)
	Type        ActorType
	PodIdentity *PodIdentityConfig
	IRSA        *IRSAConfig
	SPIFFE      *SPIFFEConfig

	// List of target role ARNs this actor needs to assume
	TargetRoleARNs []string
}

type PodIdentityConfig struct {
	TrustPrincipal string // e.g., "pods.eks.amazonaws.com"
}

type IRSAConfig struct {
	OIDCProvider     string
	OIDCAudience     string // defaults to sts.amazonaws.com
	ServiceAccount   string
	ServiceNamespace string
}

type SPIFFEConfig struct {
	OIDCProvider string
	OIDCAudience string // defaults to sts.amazonaws.com
	SPIFFEID     string
}

// ==================================================================
// TARGET CONFIGURATION (The Database/Resource)
// ==================================================================

type TargetType string

const (
	TargetTypeDatabase TargetType = "Database"
	TargetTypeGeneric  TargetType = "Generic"
)

type TargetConfig struct {
	// K8s Metadata
	Name      string
	Namespace string

	// AWS Specifics
	AccountID string
	RoleName  string

	// The Actor roles allowed to assume this role
	TrustedRoleARNs []string

	// Resource Configuration
	Type     TargetType
	Database *DatabaseConfig
	Generic  *GenericConfig
}

type DatabaseConfig struct {
	Region      string
	DBClusterID string
	DBUser      string
}

type GenericConfig struct {
	Action      string // e.g., "s3:GetObject"
	ResourceARN string // e.g., "arn:aws:s3:::my-bucket/*"
}

@prabhakhar
Copy link
Author

package iamgen

import (
	"encoding/json"
	"fmt"

	ackv1alpha1 "github.com/aws-controllers-k8s/iam-controller/apis/v1alpha1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// GenerateActorRole creates an ACK Role for a workload
func GenerateActorRole(cfg ActorConfig) (*ackv1alpha1.Role, error) {
	var trustPolicy IAMPolicy
	var err error

	switch cfg.Type {
	case ActorTypePodIdentity:
		trustPolicy, err = generatePodIdentityTrust(cfg.PodIdentity)
	case ActorTypeIRSA:
		trustPolicy, err = generateIRSATrust(cfg.AccountID, cfg.IRSA)
	case ActorTypeSPIFFE:
		trustPolicy, err = generateSPIFFETrust(cfg.AccountID, cfg.SPIFFE)
	default:
		return nil, fmt.Errorf("unknown actor type: %s", cfg.Type)
	}

	if err != nil {
		return nil, err
	}

	return buildACKRole(
		cfg.Name, cfg.Namespace, cfg.RoleName,
		trustPolicy,
		"sts:AssumeRole", cfg.TargetRoleARNs,
		"CrossAccountHops",
	), nil
}

// GenerateTargetRole creates an ACK Role for a resource
func GenerateTargetRole(cfg TargetConfig) (*ackv1alpha1.Role, error) {
	// 1. Build Trust Policy (trusting the actors)
	trustPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:       "TrustActors",
			Effect:    "Allow",
			Principal: map[string][]string{"AWS": cfg.TrustedRoleARNs},
			Action:    "sts:AssumeRole",
		}},
	}

	// 2. Determine Permissions
	var action, resource string
	switch cfg.Type {
	case TargetTypeDatabase:
		if cfg.Database == nil {
			return nil, fmt.Errorf("database config missing")
		}
		action = "rds-db:connect"
		resource = fmt.Sprintf("arn:aws:rds-db:%s:%s:dbuser:%s/%s",
			cfg.Database.Region, cfg.AccountID, cfg.Database.DBClusterID, cfg.Database.DBUser)
	case TargetTypeGeneric:
		if cfg.Generic == nil {
			return nil, fmt.Errorf("generic config missing")
		}
		action = cfg.Generic.Action
		resource = cfg.Generic.ResourceARN
	}

	return buildACKRole(
		cfg.Name, cfg.Namespace, cfg.RoleName,
		trustPolicy,
		action, []string{resource},
		"ResourcePermissions",
	), nil
}

// --- INTERNAL HELPERS ---

func buildACKRole(k8sName, ns, awsRoleName string, trust IAMPolicy, action string, resources []string, policyName string) *ackv1alpha1.Role {
	trustJSON, _ := json.Marshal(trust)

	permPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:      "AllowAction",
			Effect:   "Allow",
			Action:   action,
			Resource: resources,
		}},
	}
	permJSON, _ := json.Marshal(permPolicy)

	return &ackv1alpha1.Role{
		ObjectMeta: metav1.ObjectMeta{
			Name:      k8sName,
			Namespace: ns,
		},
		Spec: ackv1alpha1.RoleSpec{
			Name:                     &awsRoleName,
			AssumeRolePolicyDocument: string(trustJSON),
			InlinePolicies:           map[string]*string{policyName: awsString(string(permJSON))},
		},
	}
}

func generatePodIdentityTrust(cfg *PodIdentityConfig) (IAMPolicy, error) {
	if cfg == nil {
		return IAMPolicy{}, fmt.Errorf("PodIdentity config missing")
	}
	return IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:       "TrustPIA",
			Effect:    "Allow",
			Principal: map[string]string{"Service": cfg.TrustPrincipal},
			Action:    []string{"sts:AssumeRole", "sts:TagSession"},
		}},
	}, nil
}

func generateIRSATrust(accountID string, cfg *IRSAConfig) (IAMPolicy, error) {
	if cfg == nil {
		return IAMPolicy{}, fmt.Errorf("IRSA config missing")
	}
	aud := cfg.OIDCAudience
	if aud == "" {
		aud = "sts.amazonaws.com"
	}
	sub := fmt.Sprintf("system:serviceaccount:%s:%s", cfg.ServiceNamespace, cfg.ServiceAccount)
	return buildOIDCTrust(accountID, cfg.OIDCProvider, sub, aud), nil
}

func generateSPIFFETrust(accountID string, cfg *SPIFFEConfig) (IAMPolicy, error) {
	if cfg == nil {
		return IAMPolicy{}, fmt.Errorf("SPIFFE config missing")
	}
	aud := cfg.OIDCAudience
	if aud == "" {
		aud = "sts.amazonaws.com"
	}
	return buildOIDCTrust(accountID, cfg.OIDCProvider, cfg.SPIFFEID, aud), nil
}

func buildOIDCTrust(accountID, provider, sub, aud string) IAMPolicy {
	oidcARN := fmt.Sprintf("arn:aws:iam::%s:oidc-provider/%s", accountID, provider)
	return IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:       "TrustOIDC",
			Effect:    "Allow",
			Principal: map[string]string{"Federated": oidcARN},
			Action:    "sts:AssumeRoleWithWebIdentity",
			Condition: map[string]interface{}{
				"StringEquals": map[string]string{
					fmt.Sprintf("%s:sub", provider): sub,
					fmt.Sprintf("%s:aud", provider): aud,
				},
			},
		}},
	}
}

type IAMPolicy struct {
	Version   string         `json:"Version"`
	Statement []IAMStatement `json:"Statement"`
}
type IAMStatement struct {
	Sid       string                 `json:"Sid,omitempty"`
	Effect    string                 `json:"Effect"`
	Principal interface{}            `json:"Principal,omitempty"`
	Action    interface{}            `json:"Action"`
	Resource  interface{}            `json:"Resource,omitempty"`
	Condition map[string]interface{} `json:"Condition,omitempty"`
}

func awsString(v string) *string { return &v }

@prabhakhar
Copy link
Author

// In your AccessSubscription controller's Reconcile loop...

func (r *AccessSubscriptionReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
    // 1. Fetch AccessSubscription
    var sub AccessSubscription
    if err := r.Get(ctx, req.NamespacedName, &sub); err != nil { return ctrl.Result{}, client.IgnoreNotFound(err) }

    // 2. CALCULATE CONFIG (Mocked here - this comes from your ApplicationInstance CRD)
    // Assume you found the app runs in 2 accounts.
    actorConfigs := []iamgen.ActorConfig{
        {
            Name: sub.Name + "-actor-acct1", Namespace: sub.Namespace,
            AccountID: "111111111111", RoleName: "my-app-role",
            Type: iamgen.ActorTypeIRSA,
            IRSA: &iamgen.IRSAConfig{OIDCProvider: "oidc...", ServiceAccount: "sa", ServiceNamespace: "ns"},
             // We know the target ARN ahead of time deterministically
            TargetRoleARNs: []string{fmt.Sprintf("arn:aws:iam::%s:role/%s", "333333333333", "target-db-role")},
        },
        {
             Name: sub.Name + "-actor-acct2", Namespace: sub.Namespace,
             AccountID: "222222222222", RoleName: "my-app-role",
             // ... similar config for second account ...
             TargetRoleARNs: []string{fmt.Sprintf("arn:aws:iam::%s:role/%s", "333333333333", "target-db-role")},
        },
    }

    targetConfig := iamgen.TargetConfig{
        Name: sub.Name + "-target", Namespace: sub.Namespace,
        AccountID: "333333333333", RoleName: "target-db-role",
        Type: iamgen.TargetTypeDatabase,
        Database: &iamgen.DatabaseConfig{Region: "us-west-2", DBClusterID: "mydb", DBUser: "appuser"},
        // Trust the known ARNs of the actors
        TrustedRoleARNs: []string{
             "arn:aws:iam::111111111111:role/my-app-role",
             "arn:aws:iam::222222222222:role/my-app-role",
        },
    }

    // 3. RECONCILE ACTORS
    allActorsSynced := true
    for _, cfg := range actorConfigs {
        desiredRole, _ := iamgen.GenerateActorRole(cfg)
        // Set owner ref so deleting Subscription deletes Roles
        ctrl.SetControllerReference(&sub, desiredRole, r.Scheme)
        
        synced, err := r.applyAndCheckStatus(ctx, desiredRole)
        if err != nil { return ctrl.Result{}, err }
        if !synced { allActorsSynced = false }
    }

    // 4. RECONCILE TARGET
    desiredTarget, _ := iamgen.GenerateTargetRole(targetConfig)
    ctrl.SetControllerReference(&sub, desiredTarget, r.Scheme)
    targetSynced, err := r.applyAndCheckStatus(ctx, desiredTarget)
    if err != nil { return ctrl.Result{}, err }

    // 5. AGGREGATE STATUS
    newStatus := AccessSubscriptionStatus{
        IAMReady: allActorsSynced && targetSynced,
        // ... other status fields ...
    }
    
    // Update status if changed
    if sub.Status.IAMReady != newStatus.IAMReady {
        sub.Status = newStatus
        r.Status().Update(ctx, &sub)
    }

    return ctrl.Result{}, nil
}

// Helper to Apply ACK Role and check its 'ResourceSynced' condition
func (r *AccessSubscriptionReconciler) applyAndCheckStatus(ctx context.Context, desired *ackv1alpha1.Role) (bool, error) {
    var current ackv1alpha1.Role
    err := r.Get(ctx, types.NamespacedName{Name: desired.Name, Namespace: desired.Namespace}, &current)
    
    if errors.IsNotFound(err) {
        if err := r.Create(ctx, desired); err != nil { return false, err }
        return false, nil // Created, not yet synced
    } else if err != nil { return false, err }

    // Update if spec changed
    // Note: In prod, do a deep comparison of Spec before updating to avoid noise
    current.Spec = desired.Spec
    if err := r.Update(ctx, &current); err != nil { return false, err }

    // CHECK ACK STATUS
    for _, cond := range current.Status.Conditions {
        // ACK standard condition for "AWS resource is created and matches spec"
        if cond.Type == ackv1alpha1.ConditionTypeResourceSynced && cond.Status == corev1.ConditionTrue {
            return true, nil
        }
    }
    return false, nil
}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment