优化结构

2026-08-27 16:20:37 +08:00
parent 001e39da50
commit 653414a1da
54 changed files with 1257 additions and 3450 deletions
@@ -0,0 +1,209 @@
+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
+	}
+}