werf/werf

package helm

import (
    "bytes"
    "context"
    "fmt"
    "os"
    "reflect"
    "sort"
    "strings"

    yaml_v3 "gopkg.in/yaml.v3"
    "helm.sh/helm/v3/pkg/releaseutil"
    "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
    "sigs.k8s.io/yaml"

    "github.com/werf/logboek"
    "github.com/werf/werf/pkg/werf"
    "github.com/werf/werf/pkg/werf/global_warnings"
)

var WerfRuntimeAnnotations = map[string]string{
    "werf.io/version": werf.Version,
}

var WerfRuntimeLabels = map[string]string{}

func NewExtraAnnotationsAndLabelsPostRenderer(extraAnnotations, extraLabels map[string]string, ignoreInvalidAnnotationsAndLabels bool) *ExtraAnnotationsAndLabelsPostRenderer {
    return &ExtraAnnotationsAndLabelsPostRenderer{
        ExtraAnnotations:                  extraAnnotations,
        ExtraLabels:                       extraLabels,
        IgnoreInvalidAnnotationsAndLabels: ignoreInvalidAnnotationsAndLabels,
        globalWarnings:                    &defaultGlobalWarnings{},
    }
}

type ExtraAnnotationsAndLabelsPostRenderer struct {
    ExtraAnnotations                  map[string]string
    ExtraLabels                       map[string]string
    IgnoreInvalidAnnotationsAndLabels bool

    globalWarnings globalWarnings
}

type defaultGlobalWarnings struct{}

func (gw *defaultGlobalWarnings) GlobalWarningLn(ctx context.Context, msg string) {
    global_warnings.GlobalWarningLn(ctx, msg)
}

type globalWarnings interface {
    GlobalWarningLn(ctx context.Context, msg string)
}

func replaceNodeByKey(node *yaml_v3.Node, key string, value *yaml_v3.Node) {
    content := node.Content
    end := len(content)

    for pos := 0; pos < end; pos += 2 {
        keyNode := content[pos]

        if keyNode.Tag != "!!str" {
            continue
        }

        var k string
        if err := keyNode.Decode(&k); err != nil {
            continue
        }

        if k == key {
            content[pos+1] = value
            return
        }
    }
}

func findNodeByKey(node *yaml_v3.Node, key string) *yaml_v3.Node {
    content := node.Content
    end := len(content)

    for pos := 0; pos < end; pos += 2 {
        keyNode := content[pos]
        valueNode := content[pos+1]

        if keyNode.Tag != "!!str" {
            continue
        }

        var k string
        if err := keyNode.Decode(&k); err != nil {
            continue
        }

        if k == key {
            return valueNode
        }
    }

    return nil
}

func dereferenceAliasNode(node *yaml_v3.Node) (*yaml_v3.Node, error) {
    dereferencedNode, err := createNode(map[string]interface{}{})
    if err != nil {
        return nil, err
    }

    dereferencedNode.Content = append(dereferencedNode.Content, node.Alias.Content...)

    return dereferencedNode, nil
}

func getMapNode(docNode *yaml_v3.Node) *yaml_v3.Node {
    if docNode.Kind == yaml_v3.DocumentNode {
        if len(docNode.Content) > 0 {
            n := docNode.Content[0]
            if n.Tag == "!!map" {
                return n
            }
        }
    }
    return nil
}

func createNode(v interface{}) (*yaml_v3.Node, error) {
    newNode := &yaml_v3.Node{}
    if err := newNode.Encode(v); err != nil {
        return nil, fmt.Errorf("unable to encode value %#v into yaml node: %w", v, err)
    }
    return newNode, nil
}

func appendToNode(node *yaml_v3.Node, data interface{}) (*yaml_v3.Node, error) {
    newNode, err := createNode(data)
    if err != nil {
        return nil, err
    }

    node.Content = append(node.Content, newNode.Content...)
    return node, nil
}

func validateStringNode(node *yaml_v3.Node) error {
    var v interface{}
    if err := node.Decode(&v); err != nil {
        return fmt.Errorf("unable to decode value %q: %w", node.Value, err)
    }
    if _, ok := v.(string); !ok {
        typeOf := reflect.TypeOf(v)
        if typeOf != nil {
            return fmt.Errorf("invalid node %q: expected string, got %s", node.Value, reflect.TypeOf(v).String())
        } else {
            return fmt.Errorf("invalid node %q: expected string, got null value", node.Value)
        }
    }
    return nil
}

func validateKeyAndValue(keyNode, valueNode *yaml_v3.Node) (errors []error) {
    keyErr := validateStringNode(keyNode)
    if keyErr != nil {
        errors = append(errors, keyErr)
    }

    var valueErr error
    if valueNode.Kind == yaml_v3.AliasNode {
        valueErr = validateStringNode(valueNode.Alias)
    } else {
        valueErr = validateStringNode(valueNode)
    }
    if valueErr != nil {
        errors = append(errors, valueErr)
    }

    return
}

func validateMapStringStringNode(node *yaml_v3.Node) ([]*yaml_v3.Node, []error) {
    content := node.Content
    end := len(content)

    var validValues []*yaml_v3.Node
    var errors []error

    for pos := 0; pos < end; pos += 2 {
        keyNode := content[pos]
        valueNode := content[pos+1]

        if keyNode.Tag == "!!merge" && valueNode.Kind == yaml_v3.AliasNode {
            for pos := 0; pos < len(valueNode.Alias.Content); pos += 2 {
                keyNode := valueNode.Alias.Content[pos]
                valueNode := valueNode.Alias.Content[pos+1]

                newErrs := validateKeyAndValue(keyNode, valueNode)
                if len(newErrs) > 0 {
                    errors = append(errors, newErrs...)
                } else {
                    validValues = append(validValues, keyNode, valueNode)
                }
            }

            continue
        }

        newErrs := validateKeyAndValue(keyNode, valueNode)
        if len(newErrs) > 0 {
            errors = append(errors, newErrs...)
        } else {
            validValues = append(validValues, keyNode, valueNode)
        }
    }

    return validValues, errors
}

func appendExtraData(node *yaml_v3.Node, key string, data interface{}) error {
    if targetNode := findNodeByKey(node, key); targetNode != nil {
        if targetNode.Kind == yaml_v3.AliasNode {
            dereferencedTargetNode, err := dereferenceAliasNode(targetNode)
            if err != nil {
                return err
            }

            appendToNode(dereferencedTargetNode, data)
            replaceNodeByKey(node, key, dereferencedTargetNode)
        } else {
            appendToNode(targetNode, data)
        }

        if targetNode.Kind != yaml_v3.AliasNode && targetNode.Kind != yaml_v3.MappingNode {
            targetNode.Kind = yaml_v3.MappingNode
        }
        targetNode.Tag = "!!map"
    } else {
        appendToNode(node, map[string]interface{}{key: data})
    }

    return nil
}

func (pr *ExtraAnnotationsAndLabelsPostRenderer) Run(renderedManifests *bytes.Buffer) (*bytes.Buffer, error) {
    extraAnnotations := map[string]string{}
    for k, v := range WerfRuntimeAnnotations {
        extraAnnotations[k] = v
    }
    for k, v := range pr.ExtraAnnotations {
        extraAnnotations[k] = v
    }

    extraLabels := map[string]string{}
    for k, v := range WerfRuntimeLabels {
        extraLabels[k] = v
    }
    for k, v := range pr.ExtraLabels {
        extraLabels[k] = v
    }

    splitManifestsByKeys := releaseutil.SplitManifests(renderedManifests.String())

    manifestsKeys := make([]string, 0, len(splitManifestsByKeys))
    for k := range splitManifestsByKeys {
        manifestsKeys = append(manifestsKeys, k)
    }
    sort.Sort(releaseutil.BySplitManifestsOrder(manifestsKeys))

    splitModifiedManifests := make([]string, 0)

    for _, manifestKey := range manifestsKeys {
        manifestContent := splitManifestsByKeys[manifestKey]

        if os.Getenv("WERF_HELM_V3_EXTRA_ANNOTATIONS_AND_LABELS_DEBUG") == "1" {
            fmt.Printf("ExtraAnnotationsAndLabelsPostRenderer -- original manifest BEGIN\n")
            fmt.Printf("%s\n", manifestContent)
            fmt.Printf("ExtraAnnotationsAndLabelsPostRenderer -- original manifest END\n")
        }

        var obj unstructured.Unstructured
        if err := yaml.Unmarshal([]byte(manifestContent), &obj); err != nil {
            logboek.Warn().LogF("Unable to decode yaml manifest as unstructured object: %s: will not add extra annotations and labels to this object:\n%s\n---\n", err, manifestContent)
            splitModifiedManifests = append(splitModifiedManifests, manifestContent)
            continue
        }
        if obj.GetKind() == "" {
            logboek.Debug().LogF("Skipping empty object\n")
            continue
        }

        var objNode yaml_v3.Node
        if err := yaml_v3.Unmarshal([]byte(manifestContent), &objNode); err != nil {
            logboek.Warn().LogF("Unable to decode yaml manifest as map slice: %s: will not add extra annotations and labels to this object:\n%s\n---\n", err, manifestContent)
            splitModifiedManifests = append(splitModifiedManifests, manifestContent)
            continue
        }

        if os.Getenv("WERF_HELM_V3_EXTRA_ANNOTATIONS_AND_LABELS_DEBUG") == "1" {
            fmt.Printf("Unpacket obj annotations: %#v\n", obj.GetAnnotations())
        }

        if objMapNode := getMapNode(&objNode); objMapNode != nil {
            if metadataNode := findNodeByKey(objMapNode, "metadata"); metadataNode != nil {

                if obj.IsList() && len(extraAnnotations) > 0 {
                    pr.globalWarnings.GlobalWarningLn(context.Background(), fmt.Sprintf("werf annotations won't be applied to *List resource Kinds, including %s. We advise to replace *List resources with multiple separate resources of the same Kind", obj.GetKind()))
                } else if len(extraAnnotations) > 0 {
                    if err := appendExtraData(metadataNode, "annotations", extraAnnotations); err != nil {
                        pr.globalWarnings.GlobalWarningLn(context.Background(), fmt.Sprintf("werf annotations won't be applied to the %s/%s: an error have occurred during annotations injection: %s\n", strings.ToLower(obj.GetKind()), obj.GetName(), err))
                    }
                }

                if obj.IsList() && len(extraLabels) > 0 {
                    pr.globalWarnings.GlobalWarningLn(context.Background(), fmt.Sprintf("werf labels won't be applied to *List resource Kinds, including %s. We advise to replace *List resources with multiple separate resources of the same Kind", obj.GetKind()))
                } else if len(extraLabels) > 0 {
                    if err := appendExtraData(metadataNode, "labels", extraLabels); err != nil {
                        pr.globalWarnings.GlobalWarningLn(context.Background(), fmt.Sprintf("werf labels won't be applied to the %s/%s: an error have occurred during labels injection: %s\n", strings.ToLower(obj.GetKind()), obj.GetName(), err))
                    }
                }

                if annotationsNode := findNodeByKey(metadataNode, "annotations"); annotationsNode != nil {
                    validNodes, errors := validateMapStringStringNode(annotationsNode)
                    for _, err := range errors {
                        pr.globalWarnings.GlobalWarningLn(context.Background(), fmt.Sprintf("%s/%s annotations validation: %s", strings.ToLower(obj.GetKind()), obj.GetName(), err.Error()))
                    }
                    if pr.IgnoreInvalidAnnotationsAndLabels {
                        annotationsNode.Content = validNodes
                    }
                }

                if labelsNode := findNodeByKey(metadataNode, "labels"); labelsNode != nil {
                    validNodes, errors := validateMapStringStringNode(labelsNode)
                    for _, err := range errors {
                        pr.globalWarnings.GlobalWarningLn(context.Background(), fmt.Sprintf("%s/%s labels validation: %s\n", strings.ToLower(obj.GetKind()), obj.GetName(), err.Error()))
                    }
                    if pr.IgnoreInvalidAnnotationsAndLabels {
                        labelsNode.Content = validNodes
                    }
                }

            }
        }

        var modifiedManifestContent bytes.Buffer
        yamlEncoder := yaml_v3.NewEncoder(&modifiedManifestContent)
        yamlEncoder.SetIndent(2)

        if err := yamlEncoder.Encode(&objNode); err != nil {
            return nil, fmt.Errorf("unable to modify manifest: %w\n%s\n---\n", err, manifestContent)
        } else {
            splitModifiedManifests = append(splitModifiedManifests, modifiedManifestContent.String())

            if os.Getenv("WERF_HELM_V3_EXTRA_ANNOTATIONS_AND_LABELS_DEBUG") == "1" {
                fmt.Printf("ExtraAnnotationsAndLabelsPostRenderer -- modified manifest BEGIN\n")
                fmt.Printf("%s", modifiedManifestContent.String())
                fmt.Printf("ExtraAnnotationsAndLabelsPostRenderer -- modified manifest END\n")
            }
        }
    }

    modifiedManifests := bytes.NewBufferString(strings.Join(splitModifiedManifests, "---\n"))
    if os.Getenv("WERF_HELM_V3_EXTRA_ANNOTATIONS_AND_LABELS_DEBUG") == "1" {
        fmt.Printf("ExtraAnnotationsAndLabelsPostRenderer -- modified manifests RESULT BEGIN\n")
        fmt.Printf("%s\n", modifiedManifests.String())
        fmt.Printf("ExtraAnnotationsAndLabelsPostRenderer -- modified manifests RESULT END\n")
    }

    return modifiedManifests, nil
}

func (pr *ExtraAnnotationsAndLabelsPostRenderer) Add(extraAnnotations, extraLabels map[string]string) {
    if len(extraAnnotations) > 0 {
        if pr.ExtraAnnotations == nil {
            pr.ExtraAnnotations = make(map[string]string)
        }
        for k, v := range extraAnnotations {
            pr.ExtraAnnotations[k] = v
        }
    }

    if len(extraLabels) > 0 {
        if pr.ExtraLabels == nil {
            pr.ExtraLabels = make(map[string]string)
        }
        for k, v := range extraLabels {
            pr.ExtraLabels[k] = v
        }
    }
}