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

@@ -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

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)

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")
 	}

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 <vk> 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) {

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)))
+}

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)
+	}
+}

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"`
 }