fn.go

package main

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/crossplane/crossplane-runtime/v2/pkg/errors"
	"github.com/crossplane/crossplane-runtime/v2/pkg/logging"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/durationpb"
	"google.golang.org/protobuf/types/known/structpb"
	"k8s.io/apimachinery/pkg/runtime"
	"kcl-lang.io/krm-kcl/pkg/api"
	"kcl-lang.io/krm-kcl/pkg/api/v1alpha1"
	"kcl-lang.io/krm-kcl/pkg/kio"
	remoteauth "oras.land/oras-go/v2/registry/remote/auth"

	fnv1 "github.com/crossplane/function-sdk-go/proto/v1"
	"github.com/crossplane/function-sdk-go/request"
	"github.com/crossplane/function-sdk-go/response"

	fkcl "github.com/crossplane-contrib/function-kcl/input/v1alpha1"
	pkgresource "github.com/crossplane-contrib/function-kcl/pkg/resource"
	"sigs.k8s.io/yaml"
)

var defaultSource = os.Getenv("FUNCTION_KCL_DEFAULT_SOURCE")

const ociCacheMaxAge = 30 * time.Minute

var (
	ociCacheMu      sync.Mutex
	ociCacheCreated time.Time
)

// Function returns whatever response you ask it to.
type Function struct {
	fnv1.UnimplementedFunctionRunnerServiceServer

	log          logging.Logger
	dependencies string
	recycler     *recycler
	cache        *renderCache
}

// RunFunction runs the Function.
func (f *Function) RunFunction(ctx context.Context, req *fnv1.RunFunctionRequest) (*fnv1.RunFunctionResponse, error) {
	// Reject new work while the process is draining for a memory recycle so
	// Crossplane retries this reconcile elsewhere instead of having it cut off
	// mid-render. begin() also bounds the recycle drain to in-flight calls.
	if !f.recycler.begin() {
		return nil, status.Error(codes.Unavailable, "function-kcl is recycling to release native memory; retry")
	}
	defer f.recycler.end()

	log := f.log.WithValues("tag", req.GetMeta().GetTag())
	log.Debug("Running Function")

	rsp := response.To(req, response.DefaultTTL)
	in := &fkcl.KCLInput{}
	if err := request.GetInput(req, in); err != nil {
		response.Fatal(rsp, errors.Wrapf(err, "cannot get Function input from %T", req))
		return rsp, nil
	}
	// Allow callers to override the function response TTL via spec.ttl. When
	// unset we keep the Crossplane default so existing compositions behave the
	// same. Negative durations are clamped to 0 — Crossplane treats 0 as
	// "requeue immediately after the current reconcile finishes", which is
	// useful for waiting on conditions this function has not yet observed.
	if in.Spec.TTL != nil {
		ttl := in.Spec.TTL.Duration
		if ttl < 0 {
			ttl = 0
		}
		log.Debug("using custom function response TTL", "ttl", ttl.String())
		rsp.Meta.Ttl = durationpb.New(ttl)
	}
	// Set default source
	if in.Spec.Source == "" {
		in.Spec.Source = defaultSource
	}
	if strings.HasPrefix(in.Spec.Source, "oci://") {
		resetOCITokenCacheIfNeeded(log)
	}
	// Set default target
	if in.Spec.Target == "" {
		in.Spec.Target = pkgresource.Default
	}
	// Set default params
	if in.Spec.Params == nil {
		in.Spec.Params = make(map[string]runtime.RawExtension)
	}
	// Add base dependencies
	if f.dependencies != "" {
		in.Spec.Dependencies = f.dependencies + "\n" + in.Spec.Dependencies
	}
	// Add credentials
	if creds, ok := req.Credentials["kcl-registry"]; ok {
		data := creds.GetCredentialData()
		if data != nil {
			if password, ok := data.Data["password"]; ok {
				in.Spec.Credentials.Password = string(password)
				if username, ok := data.Data["username"]; ok {
					in.Spec.Credentials.Username = string(username)
				}
				if url, ok := data.Data["url"]; ok {
					in.Spec.Credentials.Url = string(url)
				}
			} else if _, hasProvider := data.Data["provider"]; hasProvider {
				// No password is required when a provider is
				// configured: the provider gets the credential
				// (e.g. GCP Workload Identity).
				if url, ok := data.Data["url"]; ok {
					in.Spec.Credentials.Url = string(url)
				}
			} else {
				log.Info("Warning: required password not found in the credentials")
			}
			if provider, ok := data.Data["provider"]; ok {
				in.Spec.Credentials.Provider = string(provider)
			}
		}
	}
	// The "aws" provider gets a short-lived ECR credential from the pod's AWS identity
	// (IRSA / Pod Identity), the same way the "gcp" provider uses Workload Identity. We write it into a
	// docker config the pull reads, then CLEAR the credentials so krm-kcl does not call kpm login -
	// kpm refuses to STORE plaintext credentials for an HTTPS registry, but READING them for the pull
	// is fine. No static credential is stored.
	if strings.EqualFold(in.Spec.Credentials.Provider, "aws") {
		username, password, err := getECRCredential(ctx, in.Spec.Credentials.Url)
		if err != nil {
			response.Fatal(rsp, errors.Wrap(err, "cannot get ECR credential"))
			return rsp, nil
		}
		if err := writeECRDockerConfig(in.Spec.Credentials.Url, username, password); err != nil {
			response.Fatal(rsp, errors.Wrap(err, "cannot write ECR docker config"))
			return rsp, nil
		}
		in.Spec.Credentials = fkcl.CredSpec{}
	}
	if err := in.Validate(); err != nil {
		response.Fatal(rsp, errors.Wrap(err, "invalid function input"))
		return rsp, nil
	}
	// The composite resource that actually exists.
	oxr, err := request.GetObservedCompositeResource(req)
	if err != nil {
		response.Fatal(rsp, errors.Wrap(err, "cannot get observed composite resource"))
		return rsp, nil
	}
	// Set option("params").oxr
	in.Spec.Params["oxr"], err = pkgresource.UnstructuredToRawExtension(&oxr.Resource.Unstructured)
	if err != nil {
		response.Fatal(rsp, err)
		return rsp, nil
	}
	log = log.WithValues(
		"xr-version", oxr.Resource.GetAPIVersion(),
		"xr-kind", oxr.Resource.GetKind(),
		"xr-name", oxr.Resource.GetName(),
		"target", in.Spec.Target,
	)

	// The composite resource desired by previous functions in the pipeline.
	dxr, err := request.GetDesiredCompositeResource(req)
	if err != nil {
		response.Fatal(rsp, errors.Wrap(err, "cannot get desired composite resource"))
		return rsp, nil
	}
	// Set option("params").dxr
	dxr.Resource.SetAPIVersion(oxr.Resource.GetAPIVersion())
	dxr.Resource.SetKind(oxr.Resource.GetKind())
	in.Spec.Params["dxr"], err = pkgresource.UnstructuredToRawExtension(&dxr.Resource.Unstructured)
	if err != nil {
		response.Fatal(rsp, err)
		return rsp, nil
	}
	// The composed resources desired by any previous Functions in the pipeline.
	desired, err := request.GetDesiredComposedResources(req)
	if err != nil {
		response.Fatal(rsp, errors.Wrapf(err, "cannot get desired composed resources from %T", req))
		return rsp, nil
	}
	log.Debug(fmt.Sprintf("DesiredComposed resources: %d", len(desired)))
	in.Spec.Params["dcds"], err = pkgresource.ObjToRawExtension(desired)
	if err != nil {
		response.Fatal(rsp, err)
		return rsp, nil
	}

	// The composed resources desired by any previous Functions in the pipeline.
	observed, err := request.GetObservedComposedResources(req)
	if err != nil {
		response.Fatal(rsp, errors.Wrapf(err, "cannot get observed composed resources from %T", req))
		return rsp, nil
	}
	log.Debug(fmt.Sprintf("ObservedComposed resources: %d", len(observed)))
	in.Spec.Params["ocds"], err = pkgresource.ObjToRawExtension(observed)
	if err != nil {
		response.Fatal(rsp, err)
		return rsp, nil
	}
	// Set function context
	ctxByte, err := req.Context.MarshalJSON()
	if err != nil {
		response.Fatal(rsp, err)
		return rsp, nil
	}
	ctxObj, err := pkgresource.JsonByteToRawExtension(ctxByte)
	if err != nil {
		response.Fatal(rsp, err)
		return rsp, nil
	}
	in.Spec.Params["ctx"] = ctxObj
	// The extra resources by myself or any previous Functions in the pipeline.
	extras, err := request.GetExtraResources(req)
	if err != nil {
		response.Fatal(rsp, errors.Wrapf(err, "cannot get extra resources from %T", req))
		return rsp, nil
	}
	log.Debug(fmt.Sprintf("Extra resources: %d", len(extras)))
	in.Spec.Params["extraResources"], err = pkgresource.ObjToRawExtension(extras)
	if err != nil {
		response.Fatal(rsp, err)
		return rsp, nil
	}
	// The required resources by myself or any previous Functions in the pipeline.
	required, err := request.GetRequiredResources(req)
	if err != nil {
		response.Fatal(rsp, errors.Wrapf(err, "cannot get required resources from %T", req))
		return rsp, nil
	}
	log.Debug(fmt.Sprintf("Required resources: %d", len(required)))
	in.Spec.Params["requiredResources"], err = pkgresource.ObjToRawExtension(required)
	if err != nil {
		response.Fatal(rsp, err)
		return rsp, nil
	}
	// Convert the function-kcl KCLInput to the KRM-KCL spec and run function pipelines.
	// Input Example: https://github.com/kcl-lang/krm-kcl/blob/main/examples/mutation/set-annotations/suite/good.yaml
	in.APIVersion = v1alpha1.KCLRunAPIVersion
	in.Kind = api.KCLRunKind
	// The KCL render is deterministic over the input (source + dependencies +
	// all params + config), so renderKey identifies it. Reuse the previous output
	// for byte-identical inputs to skip recompiling the module — this avoids both
	// the CPU cost and a native memory-leak increment on no-op re-syncs. See
	// rendercache.go.
	var key []byte
	if f.cache.enabled() {
		if key, err = renderKey(in); err != nil {
			response.Fatal(rsp, errors.Wrap(err, "cannot derive render cache key"))
			return rsp, nil
		}
	}

	var outputData []byte
	if cached, ok := f.cache.lookup(key); ok {
		outputData = cached
		hits, misses := f.cache.stats()
		log.Debug("render cache hit", "hits", hits, "misses", misses)
	} else {
		// Fast path: feed the KCL runtime the JSON we already hold, skipping the
		// JSON -> YAML -> RNode -> JSON round trip. Falls back to the krm-kcl
		// pipeline for non-inline sources (oci://, git, http, local path).
		out, ok, err := renderInline(in)
		if err != nil {
			response.Fatal(rsp, errors.Wrap(err, "failed to run kcl function pipelines"))
			return rsp, nil
		}
		if !ok {
			// Note use "sigs.k8s.io/yaml" here.
			kclRunBytes, err := yaml.Marshal(in)
			if err != nil {
				response.Fatal(rsp, errors.Wrap(err, "cannot marshal input to yaml"))
				return rsp, nil
			}
			inputBytes, outputBytes := bytes.NewBuffer(kclRunBytes), bytes.NewBuffer([]byte{})
			// Run pipeline to get the result mutated or validated by the KCL source.
			pipeline := kio.NewPipeline(inputBytes, outputBytes, false)
			if err := pipeline.Execute(); err != nil {
				response.Fatal(rsp, errors.Wrap(err, "failed to run kcl function pipelines"))
				return rsp, nil
			}
			out = outputBytes.Bytes()
		}
		outputData = out
		f.cache.store(key, outputData)
	}
	log.Debug(fmt.Sprintf("Pipeline output: %v", string(outputData)))
	data, err := pkgresource.DataResourcesFromYaml(outputData)
	if err != nil {
		response.Fatal(rsp, errors.Wrapf(err, "cannot parse data resources from the pipeline output in %T", rsp))
		return rsp, nil
	}
	log.Debug(fmt.Sprintf("Pipeline data: %v", data))

	var resources pkgresource.ResourceList
	for _, r := range in.Spec.Resources {
		base, err := pkgresource.JsonByteToUnstructured(r.Base.Raw)
		if err != nil {
			response.Fatal(rsp, errors.Wrapf(err, "cannot parse data resources from the pipeline output in %T", rsp))
			return rsp, nil
		}
		resources = append(resources, pkgresource.Resource{
			Name: r.Name,
			Base: *base,
		})
	}
	log.Debug(fmt.Sprintf("Input resources: %v", resources))
	extraResources := map[string]*fnv1.ResourceSelector{}
	requiredResources := map[string]*fnv1.ResourceSelector{}
	var conditions pkgresource.ConditionResources
	var events pkgresource.EventResources
	contextData := make(map[string]interface{})
	result, err := pkgresource.ProcessResources(dxr, oxr, desired, observed, extraResources, requiredResources, &conditions, &events, &contextData, in.Spec.Target, resources, &pkgresource.AddResourcesOptions{
		Basename:  in.Name,
		Data:      data,
		Overwrite: true,
	})
	if err != nil {
		response.Fatal(rsp, errors.Wrapf(err, "cannot process xr and state with the pipeline output in %T", rsp))
		return rsp, nil
	}
	if len(extraResources) > 0 || len(requiredResources) > 0 {
		for n, d := range extraResources {
			log.Debug(fmt.Sprintf("Requesting ExtraResources from %s named %s", d.String(), n))
		}
		for n, d := range requiredResources {
			log.Debug(fmt.Sprintf("Requesting RequiredResources from %s named %s", d.String(), n))
		}
		rsp.Requirements = &fnv1.Requirements{ExtraResources: extraResources, Resources: requiredResources}
	}

	if len(conditions) > 0 {
		err := pkgresource.SetConditions(rsp, conditions, log)
		if err != nil {
			return rsp, nil
		}
	}

	if len(events) > 0 {
		err := pkgresource.SetEvents(rsp, events)
		if err != nil {
			return rsp, nil
		}
	}

	if len(contextData) > 0 {
		mergedCtx, err := pkgresource.MergeContext(req, contextData)
		if err != nil {
			response.Fatal(rsp, errors.Wrapf(err, "cannot merge Context"))
			return rsp, nil
		}
		for key, v := range mergedCtx {
			vv, err := structpb.NewValue(v)
			if err != nil {
				response.Fatal(rsp, errors.Wrap(err, "cannot convert value to structpb.Value"))
				return rsp, nil
			}
			f.log.Debug("Updating Composition environment", "key", key, "data", v)
			response.SetContextKey(rsp, key, vv)
		}

	}

	log.Debug(fmt.Sprintf("Set %d resource(s) to the desired state", result.MsgCount))
	// Set dxr and desired state
	log.Debug(fmt.Sprintf("Setting desired XR state to %+v", dxr.Resource))
	if err := response.SetDesiredCompositeResource(rsp, dxr); err != nil {
		response.Fatal(rsp, errors.Wrapf(err, "cannot set desired composite resource in %T", rsp))
		return rsp, nil
	}
	for n, d := range desired {
		log.Debug(fmt.Sprintf("Setting DesiredComposed state to %+v named %s", d.Resource, n))
	}
	if err := response.SetDesiredComposedResources(rsp, desired); err != nil {
		response.Fatal(rsp, errors.Wrapf(err, "cannot set desired composed resources in %T", rsp))
		return rsp, nil
	}
	log.Debug("Successfully processed crossplane KCL function resources", "input", in.Name)
	return rsp, nil
}

// resetOCITokenCacheIfNeeded replaces the global ORAS auth token cache when
// it is older than ociCacheMaxAge.  This prevents long-lived gRPC servers from
// using stale Bearer tokens that expired server-side.
//
// ociCacheCreated starts as zero-value, so the very first call always resets.
func resetOCITokenCacheIfNeeded(log logging.Logger) {
	ociCacheMu.Lock()
	defer ociCacheMu.Unlock()

	age := time.Since(ociCacheCreated)
	if ociCacheCreated.IsZero() || age >= ociCacheMaxAge {
		remoteauth.DefaultCache = remoteauth.NewCache()
		ociCacheCreated = time.Now()
		log.Debug("Reset ORAS OCI token cache", "age", age.Round(time.Second), "maxAge", ociCacheMaxAge)
	}
}