pro/olons_server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: load uid key from manager config

Browse Source
This commit is contained in:
richjohnrambo
2026-09-17 16:09:30 +08:00
parent 228e3a54fd
commit ba6fd7f1b6
11 changed files with 416 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/app/app.go
+5
View File
@@ -9,6 +9,7 @@ import (
"git.118im.com/im_team/im_workbench/internal/config"
"git.118im.com/im_team/im_workbench/internal/events"
"git.118im.com/im_team/im_workbench/internal/httpapi"
"git.118im.com/im_team/im_workbench/internal/managerconfig"
"git.118im.com/im_team/im_workbench/internal/store"
"github.com/prometheus/client_golang/prometheus/promhttp"
"github.com/redis/go-redis/v9"
@@ -22,6 +23,10 @@ type App struct {
}
func New(cfg *config.Config) (*App, error) {
uidKey := managerconfig.ResolveUIDKey(context.Background(), cfg.Manager, cfg.Auth.UIDKey)
if err := httpapi.SetUIDKey(uidKey); err != nil {
return nil, err
}
st, err := store.Open(cfg.MySQL)
if err != nil {
return nil, err
internal/config/config.go
+37
View File
@@ -16,6 +16,7 @@ type Config struct {
Redis RedisConfig `json:"redis"`
Kafka KafkaConfig `json:"kafka"`
Auth AuthConfig `json:"auth"`
Manager ManagerConfig `json:"manager"`
ReadTimeout time.Duration `json:"-"`
WriteTimeout time.Duration `json:"-"`
}
@@ -46,6 +47,7 @@ type KafkaConfig struct {
type AuthConfig struct {
APIKeys []string `json:"api_keys"`
SessionKeyPattern string `json:"session_key_pattern"`
UIDKey string `json:"uid_key"`
AllowDevHeaderAuth bool `json:"allow_dev_header_auth"`
RequiredTenantCode bool `json:"required_tenant_code"`
RequiredAPIKey bool `json:"required_api_key"`
@@ -55,6 +57,15 @@ type AuthConfig struct {
InternalSecret string `json:"internal_secret"`
}
type ManagerConfig struct {
BaseURL string `json:"base_url"`
ServiceName string `json:"service_name"`
Secret string `json:"secret"`
UIDKeyGroup string `json:"uid_key_group"`
UIDKeyKeys []string `json:"uid_key_keys"`
TimeoutMS int `json:"timeout_ms"`
}
func Load() (*Config, error) {
cfg := defaultConfig()
if path := strings.TrimSpace(os.Getenv("IM_WORKBENCH_CONFIG")); path != "" {
@@ -89,9 +100,16 @@ func defaultConfig() *Config {
},
Auth: AuthConfig{
SessionKeyPattern: "im:session:token:%s",
UIDKey: "la6YsO+bNX/+XIkOqc5Svw==",
RequiredTenantCode: true,
RequiredAPIKey: true,
},
Manager: ManagerConfig{
ServiceName: "im-server",
UIDKeyGroup: "im",
UIDKeyKeys: []string{"uid_encryption_key", "UID_ENCRYPTION_KEY", "uid_key"},
TimeoutMS: 3000,
},
ReadTimeout: 8 * time.Second,
WriteTimeout: 8 * time.Second,
}
@@ -110,26 +128,45 @@ func applyEnv(cfg *Config) {
setString(&cfg.Kafka.SASL, "IM_WORKBENCH_KAFKA_SASL")
setString(&cfg.Kafka.ClientID, "IM_WORKBENCH_KAFKA_CLIENT_ID")
setString(&cfg.Auth.SessionKeyPattern, "IM_WORKBENCH_SESSION_KEY_PATTERN")
setString(&cfg.Auth.UIDKey, "IM_WORKBENCH_UID_KEY")
setString(&cfg.Auth.TrustedInternalAddr, "IM_WORKBENCH_TRUSTED_INTERNAL_ADDR")
setString(&cfg.Auth.InternalServiceName, "IM_WORKBENCH_INTERNAL_SERVICE_NAME")
setString(&cfg.Auth.InternalSecret, "IM_WORKBENCH_INTERNAL_SECRET")
setString(&cfg.Manager.BaseURL, "IM_WORKBENCH_MANAGER_BASE_URL")
setString(&cfg.Manager.ServiceName, "IM_WORKBENCH_MANAGER_SERVICE_NAME")
setString(&cfg.Manager.Secret, "IM_WORKBENCH_MANAGER_SECRET")
setString(&cfg.Manager.UIDKeyGroup, "IM_WORKBENCH_MANAGER_UID_KEY_GROUP")
if strings.TrimSpace(cfg.Auth.InternalSecret) == "" {
setString(&cfg.Auth.InternalSecret, "IM_INTERNAL_API_SECRET")
}
if strings.TrimSpace(cfg.Auth.InternalSecret) == "" {
setString(&cfg.Auth.InternalSecret, "IM_ADMIN_API_SECRET")
}
if strings.TrimSpace(cfg.Manager.Secret) == "" {
setString(&cfg.Manager.Secret, "IM_CONFIG_API_SECRET")
}
if strings.TrimSpace(cfg.Manager.Secret) == "" {
setString(&cfg.Manager.Secret, "IM_INTERNAL_API_SECRET")
}
if brokers := strings.TrimSpace(os.Getenv("IM_WORKBENCH_KAFKA_BROKERS")); brokers != "" {
cfg.Kafka.Brokers = splitCSV(brokers)
}
if keys := strings.TrimSpace(os.Getenv("IM_WORKBENCH_API_KEYS")); keys != "" {
cfg.Auth.APIKeys = splitCSV(keys)
}
if keys := strings.TrimSpace(os.Getenv("IM_WORKBENCH_MANAGER_UID_KEY_KEYS")); keys != "" {
cfg.Manager.UIDKeyKeys = splitCSV(keys)
}
if value := strings.TrimSpace(os.Getenv("IM_WORKBENCH_KAFKA_BATCH_SIZE")); value != "" {
if parsed, err := strconv.Atoi(value); err == nil {
cfg.Kafka.BatchSize = parsed
}
}
if value := strings.TrimSpace(os.Getenv("IM_WORKBENCH_MANAGER_TIMEOUT_MS")); value != "" {
if parsed, err := strconv.Atoi(value); err == nil {
cfg.Manager.TimeoutMS = parsed
}
}
if value := strings.TrimSpace(os.Getenv("IM_WORKBENCH_REDIS_DB")); value != "" {
if parsed, err := strconv.Atoi(value); err == nil {
cfg.Redis.DB = parsed
internal/httpapi/uid.go
+54 -6
View File
@@ -7,10 +7,41 @@ import (
"fmt"
"strconv"
"strings"
"sync"
"golang.org/x/crypto/xtea"
)
type userID int64
const defaultUIDKey = "la6YsO+bNX/+XIkOqc5Svw=="
var (
uidCodecMu sync.RWMutex
uidCipher *xtea.Cipher
)
func init() {
if err := SetUIDKey(defaultUIDKey); err != nil {
panic(err)
}
}
func SetUIDKey(encoded string) error {
key, err := base64.StdEncoding.DecodeString(strings.TrimSpace(encoded))
if err != nil {
return fmt.Errorf("invalid uid key: %w", err)
}
cipher, err := xtea.NewCipher(key)
if err != nil {
return fmt.Errorf("invalid uid key: %w", err)
}
uidCodecMu.Lock()
uidCipher = cipher
uidCodecMu.Unlock()
return nil
}
func (id *userID) UnmarshalJSON(raw []byte) error {
value, err := parseUserIDJSON(raw)
if err != nil {
@@ -63,7 +94,7 @@ func parseUserID(value string) (int64, error) {
if err != nil || len(decoded) < 8 {
return 0, fmt.Errorf("invalid user id")
}
return int64(binary.LittleEndian.Uint64(decoded[:8])), nil
return decryptUID(decoded[:8]), nil
}
return strconv.ParseInt(value, 10, 64)
}
@@ -72,11 +103,8 @@ func formatUserID(id int64) string {
if id == 0 {
return ""
}
src := make([]byte, 8)
dst := make([]byte, base64.URLEncoding.WithPadding(base64.NoPadding).EncodedLen(8))
binary.LittleEndian.PutUint64(src, uint64(id))
base64.URLEncoding.WithPadding(base64.NoPadding).Encode(dst, src)
return "usr" + string(dst)
encoded := encryptUID(id)
return "usr" + base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString(encoded)
}
func parseUserIDList(values []userID) []int64 {
@@ -92,3 +120,23 @@ func parseUserIDList(values []userID) []int64 {
}
return items
}
func decryptUID(encoded []byte) int64 {
dst := make([]byte, 8)
uidCodecMu.RLock()
cipher := uidCipher
uidCodecMu.RUnlock()
cipher.Decrypt(dst, encoded[:8])
return int64(binary.LittleEndian.Uint64(dst))
}
func encryptUID(id int64) []byte {
src := make([]byte, 8)
dst := make([]byte, 8)
binary.LittleEndian.PutUint64(src, uint64(id))
uidCodecMu.RLock()
cipher := uidCipher
uidCodecMu.RUnlock()
cipher.Encrypt(dst, src)
return dst
}
internal/httpapi/uid_test.go
+17
View File
@@ -22,6 +22,23 @@ func TestUserIDJSONAcceptsIMUserID(t *testing.T) {
}
}
func TestUserIDCodecMatchesIMServer(t *testing.T) {
const (
publicID = "usrOj2NWDNdlMU"
dbID = int64(1574592702383656960)
)
got, err := parseUserID(publicID)
if err != nil {
t.Fatalf("parseUserID: %v", err)
}
if got != dbID {
t.Fatalf("parseUserID(%q)=%d, want %d", publicID, got, dbID)
}
if formatted := formatUserID(dbID); formatted != publicID {
t.Fatalf("formatUserID(%d)=%q, want %q", dbID, formatted, publicID)
}
}
func TestUserIDJSONRejectsNumericID(t *testing.T) {
var id userID
if err := json.Unmarshal([]byte(`1573862062872858624`), &id); err == nil {
internal/managerconfig/client.go
+170
View File
@@ -0,0 +1,170 @@
package managerconfig
import (
"bytes"
"context"
"crypto/hmac"
cryptorand "crypto/rand"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"strconv"
"strings"
"time"
"git.118im.com/im_team/im_workbench/internal/config"
)
const maxResponseBody = 1 << 20
type Client struct {
cfg config.ManagerConfig
http *http.Client
}
func ResolveUIDKey(ctx context.Context, cfg config.ManagerConfig, fallback string) string {
key, err := New(cfg).UIDKey(ctx)
if err != nil || strings.TrimSpace(key) == "" {
return fallback
}
return key
}
func New(cfg config.ManagerConfig) *Client {
timeout := time.Duration(cfg.TimeoutMS) * time.Millisecond
if timeout <= 0 {
timeout = 3 * time.Second
}
return &Client{
cfg: cfg,
http: &http.Client{Timeout: timeout},
}
}
func (c *Client) UIDKey(ctx context.Context) (string, error) {
baseURL := strings.TrimRight(strings.TrimSpace(c.cfg.BaseURL), "/")
if baseURL == "" {
return "", fmt.Errorf("manager base url is empty")
}
if strings.TrimSpace(c.cfg.Secret) == "" {
return "", fmt.Errorf("manager secret is empty")
}
keys := cleanList(c.cfg.UIDKeyKeys)
if len(keys) == 0 {
keys = []string{"uid_encryption_key", "UID_ENCRYPTION_KEY", "uid_key"}
}
group := strings.TrimSpace(c.cfg.UIDKeyGroup)
if group == "" {
group = "im"
}
values := url.Values{}
values.Set("keys", strings.Join(keys, ","))
path := "/internal/im-config/third-party/" + url.PathEscape(group) + "?" + values.Encode()
req, err := c.newRequest(ctx, http.MethodGet, baseURL+path, nil)
if err != nil {
return "", err
}
resp, err := c.http.Do(req)
if err != nil {
return "", err
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return "", fmt.Errorf("manager status %d", resp.StatusCode)
}
body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBody))
if err != nil {
return "", err
}
valuesMap, err := parseValues(body)
if err != nil {
return "", err
}
for _, key := range keys {
if value := strings.TrimSpace(valuesMap[key]); value != "" {
return value, nil
}
}
return "", fmt.Errorf("uid key is empty")
}
func (c *Client) newRequest(ctx context.Context, method string, rawURL string, body []byte) (*http.Request, error) {
req, err := http.NewRequestWithContext(ctx, method, rawURL, bytes.NewReader(body))
if err != nil {
return nil, err
}
timestamp := strconv.FormatInt(time.Now().UTC().Unix(), 10)
nonce, err := randomNonce()
if err != nil {
return nil, err
}
bodyHash := sha256.Sum256(body)
signingText := strings.Join([]string{
method,
req.URL.RequestURI(),
hex.EncodeToString(bodyHash[:]),
timestamp,
nonce,
}, "\n")
mac := hmac.New(sha256.New, []byte(c.cfg.Secret))
_, _ = mac.Write([]byte(signingText))
serviceName := strings.TrimSpace(c.cfg.ServiceName)
if serviceName == "" {
serviceName = "im-server"
}
req.Header.Set("Accept", "application/json")
req.Header.Set("X-Service-Name", serviceName)
req.Header.Set("X-Timestamp", timestamp)
req.Header.Set("X-Nonce", nonce)
req.Header.Set("X-Signature", hex.EncodeToString(mac.Sum(nil)))
return req, nil
}
func parseValues(body []byte) (map[string]string, error) {
var envelope struct {
Code int `json:"code"`
Data json.RawMessage `json:"data"`
}
if err := json.Unmarshal(body, &envelope); err == nil && len(envelope.Data) > 0 {
if envelope.Code != 0 {
return nil, fmt.Errorf("manager code %d", envelope.Code)
}
body = envelope.Data
}
var payload struct {
Values map[string]string `json:"values"`
}
if err := json.Unmarshal(body, &payload); err != nil {
return nil, err
}
if payload.Values == nil {
return map[string]string{}, nil
}
return payload.Values, nil
}
func cleanList(values []string) []string {
items := make([]string, 0, len(values))
seen := make(map[string]bool, len(values))
for _, value := range values {
value = strings.TrimSpace(value)
if value == "" || seen[value] {
continue
}
seen[value] = true
items = append(items, value)
}
return items
}
func randomNonce() (string, error) {
var b [16]byte
if _, err := cryptorand.Read(b[:]); err != nil {
return "", err
}
return hex.EncodeToString(b[:]), nil
}
internal/managerconfig/client_test.go
+83
View File
@@ -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))))
}