From e73506385a25a282ccedf391a4e579d94ed8b0ba Mon Sep 17 00:00:00 2001 From: Zexi Li Date: Fri, 14 Aug 2026 14:40:39 +0800 Subject: [PATCH] feat(aiproxy): encrypt ai_key and virtual_key secrets at rest (#25328) Store provider secrets and virtual keys encrypted with AES, add hash-based VK lookup, and migrate existing plaintext rows on init. --- pkg/aiproxy/models/ai_key_resolve.go | 6 +- pkg/aiproxy/models/ai_keys.go | 68 ++++++++++- pkg/aiproxy/models/ai_virtual_keys.go | 84 +++++++++++++- pkg/aiproxy/models/chat_upstream.go | 2 +- pkg/aiproxy/models/secret_crypt.go | 83 ++++++++++++++ pkg/aiproxy/models/secret_crypt_test.go | 145 ++++++++++++++++++++++++ pkg/apis/aiproxy/ai_key.go | 1 + 7 files changed, 379 insertions(+), 10 deletions(-) create mode 100644 pkg/aiproxy/models/secret_crypt.go create mode 100644 pkg/aiproxy/models/secret_crypt_test.go diff --git a/pkg/aiproxy/models/ai_key_resolve.go b/pkg/aiproxy/models/ai_key_resolve.go index 3a62b597ed..93c8c3f9b9 100644 --- a/pkg/aiproxy/models/ai_key_resolve.go +++ b/pkg/aiproxy/models/ai_key_resolve.go @@ -112,7 +112,7 @@ func aiKeySkipReason(k *SAiKey, modelKey string, exclude map[string]bool) string return "?: nil ai_key" } label := aiKeyLabel(k) - if strings.TrimSpace(k.Secret) == "" { + if strings.TrimSpace(k.GetSecret()) == "" { return label + ": empty secret" } if exclude != nil && exclude[k.Id] { @@ -209,7 +209,7 @@ func resolveUpstreamAPIKeyExcluding(prov *SAiProvider, modelKey string, exclude skipReasons := make([]string, 0, len(keys)) for i := range keys { k := &keys[i] - if strings.TrimSpace(k.Secret) == "" { + if strings.TrimSpace(k.GetSecret()) == "" { continue } hasSecretKey = true @@ -225,7 +225,7 @@ func resolveUpstreamAPIKeyExcluding(prov *SAiProvider, modelKey string, exclude return nil, errors.Wrap(httperrors.ErrInvalidStatus, "failed to pick ai_key") } return &resolvedUpstreamAPIKey{ - Secret: strings.TrimSpace(chosen.Secret), + Secret: strings.TrimSpace(chosen.GetSecret()), AiKeyId: chosen.Id, FromRows: true, }, nil diff --git a/pkg/aiproxy/models/ai_keys.go b/pkg/aiproxy/models/ai_keys.go index 6a8770cd10..14bc35c9b1 100644 --- a/pkg/aiproxy/models/ai_keys.go +++ b/pkg/aiproxy/models/ai_keys.go @@ -39,7 +39,7 @@ type SAiKey struct { // AiProviderId optionally associates this key with a catalog provider row. AiProviderId string `width:"128" charset:"ascii" nullable:"true" list:"user" create:"optional" update:"user"` // Secret holds raw key material; only privileged scopes should list it. - Secret string `width:"4096" charset:"ascii" nullable:"false" create:"required"` + Secret string `width:"4096" charset:"ascii" nullable:"false" create:"required" get:"user"` // Weight is used for weighted random load balancing among matching keys (default 1). Weight int `default:"1" nullable:"false" list:"user" create:"optional" update:"user"` // Routing limits which request "model" values may use this key. @@ -66,7 +66,68 @@ func init() { } func (manager *SAiKeyManager) InitializeData() error { - return backfillEmptyTenantId(manager) + if err := backfillEmptyTenantId(manager); err != nil { + return err + } + return migrateAiKeySecrets() +} + +func (k *SAiKey) BeforeInsert() { + if len(k.Id) == 0 { + k.Id = db.DefaultUUIDGenerator() + } + if err := k.sealSecret(); err != nil { + log.Errorf("ai_key sealSecret: %v", err) + } + k.SVirtualResourceBase.BeforeInsert() +} + +func (k *SAiKey) BeforeUpdate() { + if err := k.sealSecret(); err != nil { + log.Errorf("ai_key sealSecret: %v", err) + } +} + +func (k *SAiKey) GetSecret() string { + return decryptAtRest(k.Id, k.Secret) +} + +func (k *SAiKey) sealSecret() error { + if len(k.Id) == 0 { + k.Id = db.DefaultUUIDGenerator() + } + plain := strings.TrimSpace(decryptAtRest(k.Id, k.Secret)) + if plain == "" { + return nil + } + if secretLooksEncrypted(k.Secret) { + return nil + } + enc, err := encryptAtRest(k.Id, plain) + if err != nil { + return err + } + k.Secret = enc + return nil +} + +func migrateAiKeySecrets() error { + keys := make([]SAiKey, 0) + q := queryUnprefixedSecret(AiKeyManager.Query(), "secret") + if err := q.All(&keys); err != nil { + return errors.Wrap(err, "list ai_keys for secret migrate") + } + for i := range keys { + k := &keys[i] + k.SetModelManager(AiKeyManager, k) + _, err := db.Update(k, func() error { + return k.sealSecret() + }) + if err != nil { + return errors.Wrapf(err, "encrypt ai_key %s secret", k.Id) + } + } + return nil } func (manager *SAiKeyManager) ListItemFilter( @@ -134,6 +195,9 @@ func (manager *SAiKeyManager) FetchCustomizeColumns( for i := range objs { rows[i].VirtualResourceDetails = baseRows[i] k := objs[i].(*SAiKey) + plain := k.GetSecret() + k.Secret = plain + rows[i].Secret = plain providerIds[i] = k.AiProviderId } providerNames, err := db.FetchIdNameMap2(AiProviderManager, providerIds) diff --git a/pkg/aiproxy/models/ai_virtual_keys.go b/pkg/aiproxy/models/ai_virtual_keys.go index a0e30f000a..3aee9a3a9b 100644 --- a/pkg/aiproxy/models/ai_virtual_keys.go +++ b/pkg/aiproxy/models/ai_virtual_keys.go @@ -44,8 +44,10 @@ type SAiVirtualKey struct { // OwnerId is the user that owns this virtual key within the project. OwnerId string `width:"128" charset:"ascii" index:"true" list:"user" nullable:"false" create:"optional" update:"user"` - // VirtualKey is the opaque key id or prefix presented to clients (not the upstream provider secret). - VirtualKey string `width:"128" charset:"ascii" nullable:"false" list:"user" create:"optional" update:"user"` + // VirtualKey stores the client token encrypted at rest with the row id. + VirtualKey string `width:"512" charset:"ascii" nullable:"false" create:"optional" update:"user" get:"user"` + // VirtualKeyHash is SHA256 of the plaintext token for auth lookup. + VirtualKeyHash string `width:"64" charset:"ascii" nullable:"true" unique:"true"` // Limits constrains allowed providers, per-request max_tokens, and request rate. Limits *api.SAiVirtualKeyLimits `length:"medium" charset:"utf8" list:"user" create:"optional" update:"user"` } @@ -69,6 +71,73 @@ func init() { AiVirtualKeyManager.SetVirtualObject(AiVirtualKeyManager) } +func (manager *SAiVirtualKeyManager) InitializeData() error { + return migrateAiVirtualKeys() +} + +func (m *SAiVirtualKey) BeforeInsert() { + if len(m.Id) == 0 { + m.Id = db.DefaultUUIDGenerator() + } + if err := m.sealVirtualKey(); err != nil { + log.Errorf("ai_virtual_key sealVirtualKey: %v", err) + } + m.SVirtualResourceBase.BeforeInsert() +} + +func (m *SAiVirtualKey) BeforeUpdate() { + if err := m.sealVirtualKey(); err != nil { + log.Errorf("ai_virtual_key sealVirtualKey: %v", err) + } +} + +func (m *SAiVirtualKey) GetVirtualKey() string { + return decryptAtRest(m.Id, m.VirtualKey) +} + +func (m *SAiVirtualKey) sealVirtualKey() error { + if len(m.Id) == 0 { + m.Id = db.DefaultUUIDGenerator() + } + plain := strings.TrimSpace(decryptAtRest(m.Id, m.VirtualKey)) + if plain == "" { + return nil + } + m.VirtualKeyHash = virtualKeyDigest(plain) + if secretLooksEncrypted(m.VirtualKey) { + return nil + } + enc, err := encryptAtRest(m.Id, plain) + if err != nil { + return err + } + m.VirtualKey = enc + return nil +} + +func migrateAiVirtualKeys() error { + keys := make([]SAiVirtualKey, 0) + q := AiVirtualKeyManager.Query().IsNotEmpty("virtual_key") + q = q.Filter(sqlchemy.OR( + sqlchemy.NOT(sqlchemy.Startswith(q.Field("virtual_key"), atRestPrefix)), + sqlchemy.IsNullOrEmpty(q.Field("virtual_key_hash")), + )) + if err := q.All(&keys); err != nil { + return errors.Wrap(err, "list ai_virtual_keys for migrate") + } + for i := range keys { + m := &keys[i] + m.SetModelManager(AiVirtualKeyManager, m) + _, err := db.Update(m, func() error { + return m.sealVirtualKey() + }) + if err != nil { + return errors.Wrapf(err, "encrypt ai_virtual_key %s", m.Id) + } + } + return nil +} + func (m *SAiVirtualKey) GetOwnerId() mcclient.IIdentityProvider { owner := db.SOwnerId{ UserId: m.OwnerId, @@ -126,7 +195,7 @@ func (manager *SAiVirtualKeyManager) ListItemFilter( return nil, errors.Wrap(err, "SEnabledResourceBaseManager.ListItemFilter") } if v := strings.TrimSpace(query.VirtualKey); v != "" { - q = q.Equals("virtual_key", v) + q = q.Equals("virtual_key_hash", virtualKeyDigest(v)) } userId := strings.TrimSpace(query.UserId) if userId != "" { @@ -175,6 +244,9 @@ func (manager *SAiVirtualKeyManager) FetchCustomizeColumns( for i := range objs { rows[i].VirtualResourceDetails = virtRows[i] vk := objs[i].(*SAiVirtualKey) + plain := vk.GetVirtualKey() + vk.VirtualKey = plain + rows[i].VirtualKey = plain if strings.TrimSpace(vk.OwnerId) != "" { userIds[i] = vk.OwnerId } @@ -291,7 +363,11 @@ func validateAiVirtualKeyLimits(ctx context.Context, userCred mcclient.TokenCred } func aiVirtualKeyExists(virtualKey string) (bool, error) { - cnt, err := AiVirtualKeyManager.Query().Equals("virtual_key", virtualKey).CountWithError() + digest := virtualKeyDigest(virtualKey) + if digest == "" { + return false, nil + } + cnt, err := AiVirtualKeyManager.Query().Equals("virtual_key_hash", digest).CountWithError() if err != nil { return false, errors.Wrap(err, "count ai_virtual_key") } diff --git a/pkg/aiproxy/models/chat_upstream.go b/pkg/aiproxy/models/chat_upstream.go index 649dd408d3..3e09d7d422 100644 --- a/pkg/aiproxy/models/chat_upstream.go +++ b/pkg/aiproxy/models/chat_upstream.go @@ -142,7 +142,7 @@ func loadEnabledVirtualKey(virtualKey string) (*SAiVirtualKey, error) { return nil, errors.Wrap(httperrors.ErrInputParameter, "missing virtual key (Authorization: Bearer or X-Ai-Virtual-Key)") } vk := SAiVirtualKey{} - qvk := AiVirtualKeyManager.Query().Equals("virtual_key", virtualKey).Equals("enabled", true) + qvk := AiVirtualKeyManager.Query().Equals("virtual_key_hash", virtualKeyDigest(virtualKey)).Equals("enabled", true) err := qvk.First(&vk) if err != nil { if stderrors.Is(err, sql.ErrNoRows) { diff --git a/pkg/aiproxy/models/secret_crypt.go b/pkg/aiproxy/models/secret_crypt.go new file mode 100644 index 0000000000..419f2b9ace --- /dev/null +++ b/pkg/aiproxy/models/secret_crypt.go @@ -0,0 +1,83 @@ +// Copyright 2019 Yunion +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package models + +import ( + "crypto/sha256" + "encoding/hex" + "strings" + + "yunion.io/x/pkg/errors" + "yunion.io/x/pkg/utils" + "yunion.io/x/sqlchemy" +) + +const atRestPrefix = "aescfb64:" + +func encryptAtRest(id, plain string) (string, error) { + id = strings.TrimSpace(id) + plain = strings.TrimSpace(plain) + if id == "" { + return "", errors.Error("empty encrypt id") + } + if plain == "" { + return "", nil + } + enc, err := utils.EncryptAESBase64(id, plain) + if err != nil { + return "", errors.Wrap(err, "EncryptAESBase64") + } + return atRestPrefix + enc, nil +} + +func decryptAtRest(id, stored string) string { + stored = strings.TrimSpace(stored) + if stored == "" { + return "" + } + id = strings.TrimSpace(id) + if id == "" { + return stored + } + if strings.HasPrefix(stored, atRestPrefix) { + plain, err := utils.DescryptAESBase64(id, strings.TrimPrefix(stored, atRestPrefix)) + if err != nil { + return "" + } + return plain + } + plain, err := utils.DescryptAESBase64(id, stored) + if err != nil { + return stored + } + return plain +} + +func secretLooksEncrypted(stored string) bool { + return strings.HasPrefix(strings.TrimSpace(stored), atRestPrefix) +} + +func virtualKeyDigest(plain string) string { + plain = strings.TrimSpace(plain) + if plain == "" { + return "" + } + sum := sha256.Sum256([]byte(plain)) + return hex.EncodeToString(sum[:]) +} + +func queryUnprefixedSecret(q *sqlchemy.SQuery, field string) *sqlchemy.SQuery { + return q.IsNotEmpty(field).Filter(sqlchemy.NOT(sqlchemy.Startswith(q.Field(field), atRestPrefix))) +} diff --git a/pkg/aiproxy/models/secret_crypt_test.go b/pkg/aiproxy/models/secret_crypt_test.go new file mode 100644 index 0000000000..145269d42d --- /dev/null +++ b/pkg/aiproxy/models/secret_crypt_test.go @@ -0,0 +1,145 @@ +// Copyright 2019 Yunion +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package models + +import ( + "strings" + "testing" + + "yunion.io/x/pkg/utils" +) + +func TestEncryptDecryptAtRest(t *testing.T) { + id := "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" + plain := "sk-upstream-secret" + enc, err := encryptAtRest(id, plain) + if err != nil { + t.Fatal(err) + } + if !strings.HasPrefix(enc, atRestPrefix) { + t.Fatalf("expected prefix %q, got %q", atRestPrefix, enc) + } + if !secretLooksEncrypted(enc) { + t.Fatal("ciphertext should look encrypted") + } + if got := decryptAtRest(id, enc); got != plain { + t.Fatalf("decrypt got %q want %q", got, plain) + } + if got := decryptAtRest(id, plain); got != plain { + t.Fatalf("plaintext fallback got %q", got) + } + if secretLooksEncrypted(plain) { + t.Fatal("plaintext should not look encrypted") + } +} + +func TestDecryptLegacyCiphertextWithoutPrefix(t *testing.T) { + id := "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" + plain := "sk-legacy-secret" + legacy, err := utils.EncryptAESBase64(id, plain) + if err != nil { + t.Fatal(err) + } + if secretLooksEncrypted(legacy) { + t.Fatal("legacy ciphertext has no prefix") + } + if got := decryptAtRest(id, legacy); got != plain { + t.Fatalf("legacy decrypt got %q want %q", got, plain) + } + k := &SAiKey{Secret: legacy} + k.Id = id + if err := k.sealSecret(); err != nil { + t.Fatal(err) + } + if !secretLooksEncrypted(k.Secret) { + t.Fatal("reseal should add prefix") + } + if got := k.GetSecret(); got != plain { + t.Fatalf("GetSecret after reseal=%q", got) + } +} + +func TestAiKeySealAndGetSecret(t *testing.T) { + k := &SAiKey{Secret: "sk-plain-secret"} + k.Id = "key-id-1" + if err := k.sealSecret(); err != nil { + t.Fatal(err) + } + if !secretLooksEncrypted(k.Secret) { + t.Fatal("secret should be encrypted with prefix") + } + if got := k.GetSecret(); got != "sk-plain-secret" { + t.Fatalf("GetSecret=%q", got) + } + if err := k.sealSecret(); err != nil { + t.Fatal(err) + } + if got := k.GetSecret(); got != "sk-plain-secret" { + t.Fatalf("reseal changed plaintext: %q", got) + } +} + +func TestAiKeyGetSecretPlaintext(t *testing.T) { + k := &SAiKey{Secret: "still-plain"} + k.Id = "key-id-2" + if got := k.GetSecret(); got != "still-plain" { + t.Fatalf("unmigrated GetSecret=%q", got) + } +} + +func TestVirtualKeySealHashAndLookup(t *testing.T) { + plain := "sk-0123456789abcdef0123456789abcdef" + m := &SAiVirtualKey{VirtualKey: plain} + m.Id = "vk-id-1" + if err := m.sealVirtualKey(); err != nil { + t.Fatal(err) + } + if !secretLooksEncrypted(m.VirtualKey) { + t.Fatal("virtual_key should be encrypted with prefix") + } + if m.VirtualKeyHash != virtualKeyDigest(plain) { + t.Fatalf("hash=%q want %q", m.VirtualKeyHash, virtualKeyDigest(plain)) + } + if got := m.GetVirtualKey(); got != plain { + t.Fatalf("GetVirtualKey=%q", got) + } + if virtualKeyDigest("sk-other") == m.VirtualKeyHash { + t.Fatal("different keys should not share hash") + } + if err := m.sealVirtualKey(); err != nil { + t.Fatal(err) + } + if got := m.GetVirtualKey(); got != plain { + t.Fatalf("reseal changed plaintext: %q", got) + } +} + +func TestMigrateVirtualKeySetsHash(t *testing.T) { + plain := "sk-legacyplaintext0000000000000000" + m := &SAiVirtualKey{VirtualKey: plain} + m.Id = "vk-id-legacy" + if secretLooksEncrypted(m.VirtualKey) { + t.Fatal("legacy value should be plaintext") + } + if err := m.sealVirtualKey(); err != nil { + t.Fatal(err) + } + if !secretLooksEncrypted(m.VirtualKey) { + t.Fatal("migrated value should be encrypted") + } + if m.VirtualKeyHash == "" || !strings.EqualFold(m.VirtualKeyHash, virtualKeyDigest(plain)) { + t.Fatalf("migrated hash=%q", m.VirtualKeyHash) + } +} diff --git a/pkg/apis/aiproxy/ai_key.go b/pkg/apis/aiproxy/ai_key.go index e867af54cc..faf672d508 100644 --- a/pkg/apis/aiproxy/ai_key.go +++ b/pkg/apis/aiproxy/ai_key.go @@ -88,6 +88,7 @@ type AiKeyDetails struct { AiProviderId string `json:"ai_provider_id"` AiProviderName string `json:"ai_provider_name"` + Secret string `json:"secret"` Weight int `json:"weight"` Routing *SAiKeyRouting `json:"routing"` }