kra-new/internal/config/document.go

package config

import (
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strconv"

	"gopkg.in/yaml.v3"
)

// Document is config.yaml as an editable YAML tree. Persisting through the tree
// rather than re-marshalling a Config keeps two things the typed model cannot
// represent: keys this program does not know about, and the user's comments.
// Loading lives next to it in this package so one place owns the file format.
type Document struct {
	path string
	node yaml.Node
}

func OpenDocument(path string) (*Document, error) {
	raw, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	document := &Document{path: path}
	if err = yaml.Unmarshal(raw, &document.node); err != nil {
		return nil, err
	}
	return document, nil
}

// Set replaces a top-level section, merging into whatever is already there.
// The value is normalized through a generic map so a typed struct and a
// hand-built map produce the same tree.
func (d *Document) Set(key string, value any) error {
	normalized, err := normalizeYAML(value)
	if err != nil {
		return err
	}
	return setYAMLMapping(&d.node, key, normalized)
}

// Delete removes the nested key path if present.
func (d *Document) Delete(keys ...string) { deleteYAMLMapping(&d.node, keys...) }

// Has reports whether the nested key path exists.
func (d *Document) Has(keys ...string) bool { return yamlMappingValue(&d.node, keys...) != nil }

// SetServerHTTPPort rewrites only the port of server.http.addr, keeping the
// host the user configured. A non-positive port is ignored.
func (d *Document) SetServerHTTPPort(port int) error {
	if port <= 0 {
		return nil
	}
	host := "0.0.0.0"
	if addr := yamlMappingValue(&d.node, "server", "http", "addr"); addr != nil {
		if currentHost, _, err := net.SplitHostPort(addr.Value); err == nil && currentHost != "" {
			host = currentHost
		}
	}
	return d.Set("server", map[string]any{"http": map[string]any{"addr": net.JoinHostPort(host, strconv.Itoa(port))}})
}

// Save writes the document through a temporary file so a crash mid-write can
// never leave a truncated config.yaml behind.
func (d *Document) Save() error {
	output, err := yaml.Marshal(&d.node)
	if err != nil {
		return err
	}
	if err = os.MkdirAll(filepath.Dir(d.path), 0o755); err != nil {
		return err
	}
	temporary, err := os.CreateTemp(filepath.Dir(d.path), ".kra-config-*.yaml")
	if err != nil {
		return err
	}
	tempName := temporary.Name()
	defer os.Remove(tempName)
	if _, err = temporary.Write(output); err != nil {
		_ = temporary.Close()
		return err
	}
	if err = temporary.Chmod(0o600); err != nil {
		_ = temporary.Close()
		return err
	}
	if err = temporary.Close(); err != nil {
		return err
	}
	return os.Rename(tempName, d.path)
}

func normalizeYAML(input any) (map[string]any, error) {
	raw, err := yaml.Marshal(input)
	if err != nil {
		return nil, err
	}
	var output map[string]any
	if err = yaml.Unmarshal(raw, &output); err != nil {
		return nil, err
	}
	return output, nil
}

func setYAMLMapping(node *yaml.Node, key string, value any) error {
	if node.Kind == yaml.DocumentNode {
		node = node.Content[0]
	}
	if node.Kind != yaml.MappingNode {
		return fmt.Errorf("configuration root is not a mapping")
	}
	raw, err := yaml.Marshal(value)
	if err != nil {
		return err
	}
	var replacement yaml.Node
	if err = yaml.Unmarshal(raw, &replacement); err != nil {
		return err
	}
	for i := 0; i < len(node.Content); i += 2 {
		if node.Content[i].Value == key {
			mergeYAMLNode(node.Content[i+1], replacement.Content[0])
			return nil
		}
	}
	node.Content = append(node.Content, &yaml.Node{Kind: yaml.ScalarNode, Value: key}, replacement.Content[0])
	return nil
}

// mergeYAMLNode updates values produced from the runtime configuration while
// retaining keys that are not represented by the typed configuration and
// comments that were already present in the user's config file.
func mergeYAMLNode(dst, src *yaml.Node) {
	if dst.Kind == yaml.MappingNode && src.Kind == yaml.MappingNode {
		for i := 0; i+1 < len(src.Content); i += 2 {
			key := src.Content[i].Value
			found := false
			for j := 0; j+1 < len(dst.Content); j += 2 {
				if dst.Content[j].Value == key {
					mergeYAMLNode(dst.Content[j+1], src.Content[i+1])
					found = true
					break
				}
			}
			if !found {
				dst.Content = append(dst.Content, src.Content[i], src.Content[i+1])
			}
		}
		return
	}

	// Keep comments attached to a scalar/sequence node when its value changes.
	head, line, foot := dst.HeadComment, dst.LineComment, dst.FootComment
	*dst = *src
	dst.HeadComment, dst.LineComment, dst.FootComment = head, line, foot
}

func yamlMappingValue(node *yaml.Node, keys ...string) *yaml.Node {
	if node == nil {
		return nil
	}
	if node.Kind == yaml.DocumentNode && len(node.Content) > 0 {
		node = node.Content[0]
	}
	for _, key := range keys {
		if node.Kind != yaml.MappingNode {
			return nil
		}
		var next *yaml.Node
		for i := 0; i+1 < len(node.Content); i += 2 {
			if node.Content[i].Value == key {
				next = node.Content[i+1]
				break
			}
		}
		if next == nil {
			return nil
		}
		node = next
	}
	return node
}

func deleteYAMLMapping(node *yaml.Node, keys ...string) {
	if len(keys) == 0 || node == nil {
		return
	}
	if node.Kind == yaml.DocumentNode && len(node.Content) > 0 {
		node = node.Content[0]
	}
	if node.Kind != yaml.MappingNode {
		return
	}
	key := keys[0]
	for i := 0; i+1 < len(node.Content); i += 2 {
		if node.Content[i].Value != key {
			continue
		}
		if len(keys) == 1 {
			node.Content = append(node.Content[:i], node.Content[i+2:]...)
			return
		}
		deleteYAMLMapping(node.Content[i+1], keys[1:]...)
		return
	}
}