优化结构

2026-08-27 16:20:37 +08:00
parent 001e39da50
commit 653414a1da
54 changed files with 1257 additions and 3450 deletions
@@ -1,41 +1,12 @@
 package config

-// Clone returns a deep copy of a configuration snapshot. The copy helpers are
-// intentionally kept in the config package so data/integration code does not
-// need protobuf cloning or storage-specific serializers.
-func Clone(value *Config) *Config { return cloneConfig(value) }
-
-func CloneData(value *Data) *Data {
-	if value == nil {
-		return nil
-	}
-	return cloneConfig(&Config{Data: value}).Data
-}
-
-func CloneAdmin(value *Admin) *Admin {
-	if value == nil {
-		return nil
-	}
-	return cloneConfig(&Config{Admin: value}).Admin
-}
-
-func CloneDatabase(value *Database) *Database {
-	if value == nil {
-		return nil
-	}
-	return cloneConfig(&Config{Data: &Data{Database: value}}).Data.Database
-}
-
-func CloneStorage(value *Storage) *Storage {
-	if value == nil {
-		return nil
-	}
-	return cloneConfig(&Config{Admin: &Admin{Storage: value}}).Admin.Storage
-}
-
-func CloneEmail(value *Email) *Email {
-	if value == nil {
-		return nil
-	}
-	return cloneConfig(&Config{Admin: &Admin{Email: value}}).Admin.Email
-}
+// The exported clone surface. Deep copying lives in the config package so
+// data/integration code needs neither protobuf cloning nor a storage-specific
+// serializer to hand out a snapshot a caller may safely mutate. Each function
+// returns nil for a nil input.
+func Clone(value *Config) *Config             { return cloneConfig(value) }
+func CloneData(value *Data) *Data             { return cloneData(value) }
+func CloneAdmin(value *Admin) *Admin          { return cloneAdmin(value) }
+func CloneDatabase(value *Database) *Database { return clonePtr(value) }
+func CloneStorage(value *Storage) *Storage    { return cloneStorage(value) }
+func CloneEmail(value *Email) *Email          { return clonePtr(value) }
@@ -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
+	}
+}
@@ -0,0 +1,39 @@
+package config
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+)
+
+func TestDocumentPreservesUnknownKeysAndUpdatesHTTPPort(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "config.yaml")
+	input := "# keep this comment\nserver:\n  http:\n    addr: 127.0.0.1:8000\ncustom:\n  enabled: true\n"
+	if err := os.WriteFile(path, []byte(input), 0o600); err != nil {
+		t.Fatal(err)
+	}
+	document, err := OpenDocument(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err = document.SetServerHTTPPort(9000); err != nil {
+		t.Fatal(err)
+	}
+	if err = document.Set("admin", map[string]any{"router_prefix": "/api"}); err != nil {
+		t.Fatal(err)
+	}
+	if err = document.Save(); err != nil {
+		t.Fatal(err)
+	}
+	loaded, err := os.ReadFile(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	text := string(loaded)
+	for _, want := range []string{"127.0.0.1:9000", "custom:", "enabled: true", "admin:", "router_prefix: /api"} {
+		if !strings.Contains(text, want) {
+			t.Fatalf("saved document missing %q:\n%s", want, text)
+		}
+	}
+}
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"slices"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -128,7 +129,7 @@ func (r *Store) UpdateDatabase(database *Database) {
 	if config.Data == nil {
 		config.Data = &Data{}
 	}
-	config.Data.Database = cloneDatabase(database)
+	config.Data.Database = clonePtr(database)
 	r.replace(config, false)
 }

@@ -358,162 +359,115 @@ func decode(v *viper.Viper) (*Config, error) {
 }

 func setConfigPath(config *Config, path string) {
-	if config != nil && config.Admin != nil {
-		config.Admin.ConfigPath = path
+	if config == nil {
+		return
 	}
+	if config.Admin == nil {
+		config.Admin = &Admin{}
+	}
+	config.Admin.ConfigPath = path
+}
+
+// Every config type is plain data: scalars, pointers to structs, and slices of
+// those. clonePtr and cloneSlice cover both shapes, so a clone function below is
+// one line per pointer or slice field and nothing else. A field holding only
+// scalars needs no line at all — the struct assignment already copied it.
+func clonePtr[T any](value *T) *T {
+	if value == nil {
+		return nil
+	}
+	copied := *value
+	return &copied
+}
+
+func cloneSlice[T any](values []*T, clone func(*T) *T) []*T {
+	if values == nil {
+		return nil
+	}
+	copied := make([]*T, len(values))
+	for i, item := range values {
+		copied[i] = clone(item)
+	}
+	return copied
 }

 func cloneConfig(value *Config) *Config {
 	if value == nil {
 		return nil
 	}
-	// Viper's mapstructure output contains only value fields and pointers;
-	// round-tripping through a temporary map is unnecessarily lossy for
-	// durations. Explicit copy via YAML gives us a robust deep copy while
-	// preserving time.Duration values through the text hook.
-	copy := *value
-	if value.Server != nil {
-		server := *value.Server
-		copy.Server = &server
-		if value.Server.HTTP != nil {
-			http := *value.Server.HTTP
-			copy.Server.HTTP = &http
-		}
+	copied := *value
+	copied.Server = clonePtr(value.Server)
+	if copied.Server != nil {
+		copied.Server.HTTP = clonePtr(value.Server.HTTP)
 	}
-	if value.Data != nil {
-		data := *value.Data
-		copy.Data = &data
-		data.Database = cloneDatabase(value.Data.Database)
-		data.Redis = cloneRedis(value.Data.Redis)
-		data.Mongo = cloneMongo(value.Data.Mongo)
-		data.DatabaseList = make([]*Database, len(value.Data.DatabaseList))
-		for i, item := range value.Data.DatabaseList {
-			data.DatabaseList[i] = cloneDatabase(item)
-		}
-		data.RedisList = make([]*Redis, len(value.Data.RedisList))
-		for i, item := range value.Data.RedisList {
-			data.RedisList[i] = cloneRedis(item)
-		}
-	}
-	if value.Admin != nil {
-		copy.Admin = cloneAdmin(value.Admin)
-	}
-	return &copy
+	copied.Data = cloneData(value.Data)
+	copied.Admin = cloneAdmin(value.Admin)
+	return &copied
 }

-func cloneDatabase(value *Database) *Database {
+func cloneData(value *Data) *Data {
 	if value == nil {
 		return nil
 	}
-	copy := *value
-	return &copy
+	copied := *value
+	copied.Database = clonePtr(value.Database)
+	copied.Redis = cloneRedis(value.Redis)
+	copied.Mongo = cloneMongo(value.Mongo)
+	copied.DatabaseList = cloneSlice(value.DatabaseList, clonePtr)
+	copied.RedisList = cloneSlice(value.RedisList, cloneRedis)
+	return &copied
 }
+
 func cloneRedis(value *Redis) *Redis {
-	if value == nil {
-		return nil
+	copied := clonePtr(value)
+	if copied != nil {
+		copied.ClusterAddrs = slices.Clone(value.ClusterAddrs)
 	}
-	copy := *value
-	copy.ClusterAddrs = append([]string(nil), value.ClusterAddrs...)
-	return &copy
+	return copied
 }
+
 func cloneMongo(value *Mongo) *Mongo {
-	if value == nil {
-		return nil
+	copied := clonePtr(value)
+	if copied != nil {
+		copied.Hosts = cloneSlice(value.Hosts, clonePtr)
 	}
-	copy := *value
-	copy.Hosts = make([]*MongoHost, len(value.Hosts))
-	for i, item := range value.Hosts {
-		if item != nil {
-			host := *item
-			copy.Hosts[i] = &host
-		}
-	}
-	return &copy
+	return copied
 }

 func cloneAdmin(value *Admin) *Admin {
 	if value == nil {
 		return nil
 	}
-	copy := *value
-	if value.JWT != nil {
-		item := *value.JWT
-		copy.JWT = &item
+	copied := *value
+	copied.JWT = clonePtr(value.JWT)
+	copied.Captcha = clonePtr(value.Captcha)
+	copied.Local = clonePtr(value.Local)
+	copied.Email = clonePtr(value.Email)
+	copied.Media = clonePtr(value.Media)
+	copied.System = clonePtr(value.System)
+	copied.App = clonePtr(value.App)
+	copied.DiskList = cloneSlice(value.DiskList, clonePtr)
+	copied.Storage = cloneStorage(value.Storage)
+	if copied.Zap = clonePtr(value.Zap); copied.Zap != nil {
+		copied.Zap.FileOnlyModules = slices.Clone(value.Zap.FileOnlyModules)
 	}
-	if value.Captcha != nil {
-		item := *value.Captcha
-		copy.Captcha = &item
+	if copied.CORS = clonePtr(value.CORS); copied.CORS != nil {
+		copied.CORS.Whitelist = cloneSlice(value.CORS.Whitelist, clonePtr)
 	}
-	if value.Local != nil {
-		item := *value.Local
-		copy.Local = &item
-	}
-	if value.Email != nil {
-		item := *value.Email
-		copy.Email = &item
-	}
-	if value.Media != nil {
-		item := *value.Media
-		copy.Media = &item
-	}
-	if value.System != nil {
-		item := *value.System
-		copy.System = &item
-	}
-	if value.Zap != nil {
-		item := *value.Zap
-		item.FileOnlyModules = append([]string(nil), value.Zap.FileOnlyModules...)
-		copy.Zap = &item
-	}
-	if value.App != nil {
-		item := *value.App
-		copy.App = &item
-	}
-	if value.CORS != nil {
-		item := *value.CORS
-		item.Whitelist = make([]*CORSRule, len(value.CORS.Whitelist))
-		for i, rule := range value.CORS.Whitelist {
-			if rule != nil {
-				next := *rule
-				item.Whitelist[i] = &next
-			}
-		}
-		copy.CORS = &item
-	}
-	if value.Storage != nil {
-		copy.Storage = cloneStorage(value.Storage)
-	}
-	copy.DiskList = make([]*Disk, len(value.DiskList))
-	for i, item := range value.DiskList {
-		if item != nil {
-			next := *item
-			copy.DiskList[i] = &next
-		}
-	}
-	return &copy
+	return &copied
 }

 func cloneStorage(value *Storage) *Storage {
 	if value == nil {
 		return nil
 	}
-	copy := *value
-	if value.Qiniu != nil {
-		item := *value.Qiniu
-		copy.Qiniu = &item
-	}
-	copy.AliyunOSS = cloneObjectStore(value.AliyunOSS)
-	copy.HuaweiOBS = cloneObjectStore(value.HuaweiOBS)
-	copy.TencentCOS = cloneObjectStore(value.TencentCOS)
-	copy.AWSS3 = cloneObjectStore(value.AWSS3)
-	copy.CloudflareR2 = cloneObjectStore(value.CloudflareR2)
-	copy.Minio = cloneObjectStore(value.Minio)
-	return &copy
-}
-func cloneObjectStore(value *ObjectStore) *ObjectStore {
-	if value == nil {
-		return nil
-	}
-	copy := *value
-	return &copy
+	copied := *value
+	copied.Qiniu = clonePtr(value.Qiniu)
+	copied.AliyunOSS = clonePtr(value.AliyunOSS)
+	copied.HuaweiOBS = clonePtr(value.HuaweiOBS)
+	copied.TencentCOS = clonePtr(value.TencentCOS)
+	copied.AWSS3 = clonePtr(value.AWSS3)
+	copied.CloudflareR2 = clonePtr(value.CloudflareR2)
+	copied.Minio = clonePtr(value.Minio)
+	return &copied
 }