fix: load uid key from manager config

2026-09-17 16:09:30 +08:00
parent 228e3a54fd
commit ba6fd7f1b6
11 changed files with 416 additions and 14 deletions
@@ -0,0 +1,83 @@
+package managerconfig
+
+import (
+	"context"
+	"crypto/hmac"
+	"crypto/sha256"
+	"encoding/hex"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"git.118im.com/im_team/im_workbench/internal/config"
+)
+
+func TestResolveUIDKeyPrefersManager(t *testing.T) {
+	const secret = "test-secret"
+	const managerKey = "manager-key"
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/internal/im-config/third-party/im" {
+			t.Fatalf("unexpected path %s", r.URL.Path)
+		}
+		if r.URL.Query().Get("keys") != "uid_encryption_key,UID_ENCRYPTION_KEY,uid_key" {
+			t.Fatalf("unexpected keys %s", r.URL.Query().Get("keys"))
+		}
+		if r.Header.Get("X-Service-Name") != "im-server" {
+			t.Fatalf("unexpected service %s", r.Header.Get("X-Service-Name"))
+		}
+		if !validSignature(t, r, secret) {
+			t.Fatal("invalid signature")
+		}
+		_, _ = w.Write([]byte(`{"code":0,"data":{"values":{"uid_encryption_key":"manager-key"}}}`))
+	}))
+	defer server.Close()
+
+	got := ResolveUIDKey(context.Background(), config.ManagerConfig{
+		BaseURL:     server.URL,
+		ServiceName: "im-server",
+		Secret:      secret,
+		UIDKeyGroup: "im",
+		UIDKeyKeys:  []string{"uid_encryption_key", "UID_ENCRYPTION_KEY", "uid_key"},
+	}, "fallback-key")
+	if got != managerKey {
+		t.Fatalf("ResolveUIDKey() = %q, want %q", got, managerKey)
+	}
+}
+
+func TestResolveUIDKeyFallsBack(t *testing.T) {
+	got := ResolveUIDKey(context.Background(), config.ManagerConfig{}, "fallback-key")
+	if got != "fallback-key" {
+		t.Fatalf("ResolveUIDKey() = %q, want fallback", got)
+	}
+}
+
+func TestParseValuesAcceptsUnwrappedResponse(t *testing.T) {
+	values, err := parseValues([]byte(`{"values":{"UID_ENCRYPTION_KEY":"key"}}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if values["UID_ENCRYPTION_KEY"] != "key" {
+		t.Fatalf("unexpected values %#v", values)
+	}
+}
+
+func validSignature(t *testing.T, r *http.Request, secret string) bool {
+	t.Helper()
+	body, err := io.ReadAll(r.Body)
+	if err != nil {
+		t.Fatal(err)
+	}
+	bodyHash := sha256.Sum256(body)
+	signingText := strings.Join([]string{
+		r.Method,
+		r.URL.RequestURI(),
+		hex.EncodeToString(bodyHash[:]),
+		r.Header.Get("X-Timestamp"),
+		r.Header.Get("X-Nonce"),
+	}, "\n")
+	mac := hmac.New(sha256.New, []byte(secret))
+	_, _ = mac.Write([]byte(signingText))
+	return hmac.Equal([]byte(r.Header.Get("X-Signature")), []byte(hex.EncodeToString(mac.Sum(nil))))
+}