fix: load uid key from manager config
This commit is contained in:
richjohnrambo
@@ -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
|
||||
}
|
||||
@@ -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))))
|
||||
}
|
||||
Reference in New Issue
Block a user