From ac75caba862f3118547bc88fb2cf955fad3e0362 Mon Sep 17 00:00:00 2001 From: richjohnrambo Date: Thu, 17 Sep 2026 16:30:47 +0800 Subject: [PATCH] fix: load uid key from manager config --- server/admin_internal_config.go | 12 +++--- server/main.go | 14 +++++-- server/uid_key_config.go | 53 +++++++++++++++++++++++++ server/uid_key_config_test.go | 69 +++++++++++++++++++++++++++++++++ 4 files changed, 139 insertions(+), 9 deletions(-) create mode 100644 server/uid_key_config.go create mode 100644 server/uid_key_config_test.go diff --git a/server/admin_internal_config.go b/server/admin_internal_config.go index 8deb907..2e784c2 100644 --- a/server/admin_internal_config.go +++ b/server/admin_internal_config.go @@ -60,15 +60,15 @@ type adminInternalTenantDTO struct { UpdatedBy int64 `json:"updated_by"` } -func initAdminInternalAPI(raw json.RawMessage) error { +func initAdminInternalAPI(raw json.RawMessage) (*adminInternalClient, error) { store.SetAdminInternalConfigProvider(nil) loginIPWhitelistProvider = nil if len(raw) == 0 || string(raw) == "null" { - return nil + return nil, nil } var cfg adminInternalAPIConfig if err := json.Unmarshal(raw, &cfg); err != nil { - return err + return nil, err } cfg.BaseURL = strings.TrimRight(strings.TrimSpace(cfg.BaseURL), "/") cfg.ServiceName = strings.TrimSpace(cfg.ServiceName) @@ -79,10 +79,10 @@ func initAdminInternalAPI(raw json.RawMessage) error { cfg.TimeoutMS = 3000 } if !cfg.Enabled { - return nil + return nil, nil } if cfg.BaseURL == "" || strings.TrimSpace(cfg.Secret) == "" { - return errors.New("admin_internal_api enabled but base_url or secret is empty") + return nil, errors.New("admin_internal_api enabled but base_url or secret is empty") } client := &adminInternalClient{ config: cfg, @@ -91,7 +91,7 @@ func initAdminInternalAPI(raw json.RawMessage) error { store.SetAdminInternalConfigProvider(client) loginIPWhitelistProvider = client logs.Info.Printf("Admin internal config client enabled: base_url=%s service=%s", cfg.BaseURL, cfg.ServiceName) - return nil + return client, nil } func (c *adminInternalClient) GetTenantByCode(code string) (*types.Tenant, error) { diff --git a/server/main.go b/server/main.go index 055019a..3fd87f7 100644 --- a/server/main.go +++ b/server/main.go @@ -494,6 +494,17 @@ func main() { logs.Info.Printf("Profiling info saved to '%s.(cpu|mem)'", *pprofFile) } + adminInternalClient, err := initAdminInternalAPI(config.AdminAPI) + if err != nil { + logs.Err.Fatal("Failed to initialize Admin internal API client:", err) + } + if storeConfig, ok, err := storeConfigWithAdminUIDKey(config.Store, adminInternalClient); err != nil { + logs.Warn.Printf("Failed to load UID key from Admin internal config, falling back to store_config.uid_key: %v", err) + } else if ok { + config.Store = storeConfig + logs.Info.Println("Loaded store UID key from Admin internal config") + } + err = store.Store.Open(workerId, config.Store) logs.Info.Println("DB adapter", store.Store.GetAdapterName(), store.Store.GetAdapterVersion()) if err != nil { @@ -513,9 +524,6 @@ func main() { if err = initInternalAPI(config.Internal); err != nil { logs.Err.Fatal("Failed to initialize internal API:", err) } - if err = initAdminInternalAPI(config.AdminAPI); err != nil { - logs.Err.Fatal("Failed to initialize Admin internal API client:", err) - } // API key signing secret globals.apiKeySalt = config.APIKeySalt diff --git a/server/uid_key_config.go b/server/uid_key_config.go new file mode 100644 index 0000000..384ea9f --- /dev/null +++ b/server/uid_key_config.go @@ -0,0 +1,53 @@ +package main + +import ( + "encoding/json" + "strings" +) + +const adminUIDKeyConfigGroup = "im" + +var adminUIDKeyConfigKeys = []string{"uid_encryption_key", "UID_ENCRYPTION_KEY", "uid_key"} + +func storeConfigWithAdminUIDKey(raw json.RawMessage, client *adminInternalClient) (json.RawMessage, bool, error) { + if client == nil { + return raw, false, nil + } + values, err := client.ThirdPartyConfig(adminUIDKeyConfigGroup, adminUIDKeyConfigKeys...) + if err != nil { + return raw, false, err + } + uidKey := firstAdminUIDKeyValue(values) + if uidKey == "" { + return raw, false, nil + } + merged, err := mergeStoreUIDKey(raw, uidKey) + if err != nil { + return raw, false, err + } + return merged, true, nil +} + +func firstAdminUIDKeyValue(values map[string]string) string { + for _, key := range adminUIDKeyConfigKeys { + if value := strings.TrimSpace(values[key]); value != "" { + return value + } + } + return "" +} + +func mergeStoreUIDKey(raw json.RawMessage, uidKey string) (json.RawMessage, error) { + config := map[string]json.RawMessage{} + if len(raw) > 0 && string(raw) != "null" { + if err := json.Unmarshal(raw, &config); err != nil { + return nil, err + } + } + encoded, err := json.Marshal(strings.TrimSpace(uidKey)) + if err != nil { + return nil, err + } + config["uid_key"] = encoded + return json.Marshal(config) +} diff --git a/server/uid_key_config_test.go b/server/uid_key_config_test.go new file mode 100644 index 0000000..3b97933 --- /dev/null +++ b/server/uid_key_config_test.go @@ -0,0 +1,69 @@ +package main + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" +) + +func TestStoreConfigWithAdminUIDKey(t *testing.T) { + const managerUIDKey = "MTIzNDU2Nzg5MGFiY2RlZg==" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/internal/im-config/third-party/im" { + t.Errorf("unexpected path: %s", r.URL.Path) + return + } + if r.Header.Get(internalHeaderSignature) == "" { + t.Error("internal request must be signed") + return + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "code": 0, + "data": map[string]any{ + "values": map[string]string{ + "uid_encryption_key": managerUIDKey, + }, + }, + }) + })) + defer server.Close() + + client := &adminInternalClient{ + config: adminInternalAPIConfig{BaseURL: server.URL, ServiceName: "im-server", Secret: "test-secret"}, + http: server.Client(), + } + merged, ok, err := storeConfigWithAdminUIDKey(json.RawMessage(`{"uid_key":"bG9jYWwtZmFsbGJhY2s=","max_results":1024}`), client) + if err != nil { + t.Fatal(err) + } + if !ok { + t.Fatal("expected uid key to be loaded from manager") + } + var config struct { + UIDKey []byte `json:"uid_key"` + MaxResults int `json:"max_results"` + } + if err := json.Unmarshal(merged, &config); err != nil { + t.Fatal(err) + } + if string(config.UIDKey) != "1234567890abcdef" { + t.Fatalf("unexpected uid key: %q", string(config.UIDKey)) + } + if config.MaxResults != 1024 { + t.Fatalf("max_results should be preserved, got %d", config.MaxResults) + } +} + +func TestStoreConfigWithAdminUIDKeyFallsBackWhenMissing(t *testing.T) { + merged, ok, err := storeConfigWithAdminUIDKey(json.RawMessage(`{"uid_key":"bG9jYWw="}`), nil) + if err != nil { + t.Fatal(err) + } + if ok { + t.Fatal("nil admin client should not override uid key") + } + if string(merged) != `{"uid_key":"bG9jYWw="}` { + t.Fatalf("unexpected config: %s", string(merged)) + } +}