mirror of
https://github.com/openimsdk/open-im-server.git
synced 2025-10-25 12:42:12 +08:00
288 lines
9.8 KiB
Go
288 lines
9.8 KiB
Go
// Copyright © 2024 OpenIM. All rights reserved.
|
|
//
|
|
// 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 encryption
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/openimsdk/open-im-server/v3/internal/rpc/encryption/stores"
|
|
"github.com/openimsdk/open-im-server/v3/pkg/common/storage/model/signal"
|
|
"github.com/openimsdk/tools/log"
|
|
)
|
|
|
|
const (
|
|
// Maximum number of one-time prekeys that can be uploaded at once
|
|
MaxOneTimePreKeys = 100
|
|
)
|
|
|
|
type KeysManager struct {
|
|
identityStore stores.IdentityStoreInterface
|
|
preKeyStore stores.PreKeyStoreInterface
|
|
signedPreKeyStore stores.SignedPreKeyStoreInterface
|
|
}
|
|
|
|
func NewKeysManager(
|
|
identityStore stores.IdentityStoreInterface,
|
|
preKeyStore stores.PreKeyStoreInterface,
|
|
signedPreKeyStore stores.SignedPreKeyStoreInterface,
|
|
) *KeysManager {
|
|
return &KeysManager{
|
|
identityStore: identityStore,
|
|
preKeyStore: preKeyStore,
|
|
signedPreKeyStore: signedPreKeyStore,
|
|
}
|
|
}
|
|
|
|
// GetIdentityKey retrieves the identity key for a user/device
|
|
func (km *KeysManager) GetIdentityKey(ctx context.Context, userID string, deviceID int32) (*signal.SignalIdentityKey, error) {
|
|
return km.identityStore.Get(ctx, userID, deviceID)
|
|
}
|
|
|
|
// SetIdentityKey sets the identity key for a user/device
|
|
func (km *KeysManager) SetIdentityKey(ctx context.Context, userID string, deviceID int32, identityKey []byte, registrationID int32) error {
|
|
now := time.Now()
|
|
|
|
// Check if identity key already exists
|
|
existing, err := km.identityStore.Get(ctx, userID, deviceID)
|
|
if err == nil && existing != nil {
|
|
// Update existing identity key
|
|
existing.IdentityKey = identityKey
|
|
existing.RegistrationID = registrationID
|
|
existing.UpdatedTime = now
|
|
return km.identityStore.Update(ctx, userID, deviceID, existing)
|
|
}
|
|
|
|
// Create new identity key
|
|
identityKeyRecord := &signal.SignalIdentityKey{
|
|
UserID: userID,
|
|
DeviceID: deviceID,
|
|
IdentityKey: identityKey,
|
|
RegistrationID: registrationID,
|
|
CreatedTime: now,
|
|
UpdatedTime: now,
|
|
}
|
|
|
|
return km.identityStore.Create(ctx, identityKeyRecord)
|
|
}
|
|
|
|
// GetActiveSignedPreKey retrieves the active signed prekey for a user/device
|
|
func (km *KeysManager) GetActiveSignedPreKey(ctx context.Context, userID string, deviceID int32) (*signal.SignalSignedPreKey, error) {
|
|
return km.signedPreKeyStore.GetActive(ctx, userID, deviceID)
|
|
}
|
|
|
|
// SetSignedPreKey sets a signed prekey for a user/device
|
|
func (km *KeysManager) SetSignedPreKey(ctx context.Context, userID string, deviceID int32, signedPreKey *SignedPreKeyResponse) error {
|
|
now := time.Now()
|
|
|
|
signedPreKeyRecord := &signal.SignalSignedPreKey{
|
|
UserID: userID,
|
|
DeviceID: deviceID,
|
|
KeyID: signedPreKey.KeyId,
|
|
PublicKey: signedPreKey.PublicKey,
|
|
Signature: signedPreKey.Signature,
|
|
CreatedTime: now,
|
|
Active: true,
|
|
}
|
|
|
|
// Deactivate existing signed prekeys and set this one as active
|
|
err := km.signedPreKeyStore.SetActive(ctx, userID, deviceID, signedPreKey.KeyId)
|
|
if err != nil {
|
|
log.ZWarn(ctx, "failed to deactivate existing signed prekeys", err)
|
|
}
|
|
|
|
// Create or update the signed prekey
|
|
existing, err := km.signedPreKeyStore.GetByKeyID(ctx, userID, deviceID, signedPreKey.KeyId)
|
|
if err == nil && existing != nil {
|
|
return km.signedPreKeyStore.Update(ctx, userID, deviceID, signedPreKey.KeyId, signedPreKeyRecord)
|
|
}
|
|
|
|
return km.signedPreKeyStore.Create(ctx, signedPreKeyRecord)
|
|
}
|
|
|
|
// GetOneTimePreKey retrieves an available one-time prekey and marks it as used
|
|
func (km *KeysManager) GetOneTimePreKey(ctx context.Context, userID string, deviceID int32) (*signal.SignalPreKey, error) {
|
|
preKey, err := km.preKeyStore.GetAvailable(ctx, userID, deviceID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("no available one-time prekey: %w", err)
|
|
}
|
|
|
|
// Mark the prekey as used
|
|
err = km.preKeyStore.MarkUsed(ctx, userID, deviceID, preKey.KeyID)
|
|
if err != nil {
|
|
log.ZError(ctx, "failed to mark prekey as used", err, "userID", userID, "deviceID", deviceID, "keyID", preKey.KeyID)
|
|
// Don't fail the request, but log the error
|
|
}
|
|
|
|
return preKey, nil
|
|
}
|
|
|
|
// SetOneTimePreKeys sets multiple one-time prekeys for a user/device
|
|
func (km *KeysManager) SetOneTimePreKeys(ctx context.Context, userID string, deviceID int32, preKeys []*PreKeyResponse) (int, error) {
|
|
if len(preKeys) > MaxOneTimePreKeys {
|
|
return 0, fmt.Errorf("too many one-time prekeys: %d (max: %d)", len(preKeys), MaxOneTimePreKeys)
|
|
}
|
|
|
|
now := time.Now()
|
|
var preKeyRecords []*signal.SignalPreKey
|
|
|
|
for _, preKey := range preKeys {
|
|
preKeyRecord := &signal.SignalPreKey{
|
|
UserID: userID,
|
|
DeviceID: deviceID,
|
|
KeyID: preKey.KeyId,
|
|
PublicKey: preKey.PublicKey,
|
|
Used: false,
|
|
CreatedTime: now,
|
|
}
|
|
preKeyRecords = append(preKeyRecords, preKeyRecord)
|
|
}
|
|
|
|
err := km.preKeyStore.CreateBatch(ctx, preKeyRecords)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("failed to create one-time prekeys: %w", err)
|
|
}
|
|
|
|
return len(preKeyRecords), nil
|
|
}
|
|
|
|
// GetPreKeyCount returns the count of available one-time prekeys for a user/device
|
|
func (km *KeysManager) GetPreKeyCount(ctx context.Context, userID string, deviceID int32) (int64, error) {
|
|
return km.preKeyStore.CountAvailable(ctx, userID, deviceID)
|
|
}
|
|
|
|
// GetSignedPreKeyInfo returns information about signed prekey existence and last rotation
|
|
func (km *KeysManager) GetSignedPreKeyInfo(ctx context.Context, userID string, deviceID int32) (exists bool, lastRotation time.Time, err error) {
|
|
signedPreKey, err := km.signedPreKeyStore.GetActive(ctx, userID, deviceID)
|
|
if err != nil {
|
|
return false, time.Time{}, err
|
|
}
|
|
|
|
if signedPreKey != nil {
|
|
return true, signedPreKey.CreatedTime, nil
|
|
}
|
|
|
|
return false, time.Time{}, nil
|
|
}
|
|
|
|
// CleanupExpiredKeys removes expired keys from storage
|
|
func (km *KeysManager) CleanupExpiredKeys(ctx context.Context) error {
|
|
// Cleanup used one-time prekeys older than 7 days
|
|
usedPreKeysCleanedCount, err := km.preKeyStore.CleanupUsed(ctx, 7*24*time.Hour)
|
|
if err != nil {
|
|
log.ZError(ctx, "failed to cleanup used prekeys", err)
|
|
} else {
|
|
log.ZInfo(ctx, "cleaned up used prekeys", "count", usedPreKeysCleanedCount)
|
|
}
|
|
|
|
// Cleanup inactive signed prekeys older than 30 days
|
|
inactiveSignedPreKeysCount, err := km.signedPreKeyStore.CleanupInactive(ctx, 30*24*time.Hour)
|
|
if err != nil {
|
|
log.ZError(ctx, "failed to cleanup inactive signed prekeys", err)
|
|
} else {
|
|
log.ZInfo(ctx, "cleaned up inactive signed prekeys", "count", inactiveSignedPreKeysCount)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// ValidateSignedPreKey validates the signature of a signed prekey
|
|
func (km *KeysManager) ValidateSignedPreKey(ctx context.Context, userID string, deviceID int32, signedPreKey *SignedPreKeyResponse) error {
|
|
// Get the identity key to verify the signature
|
|
_, err := km.GetIdentityKey(ctx, userID, deviceID)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to get identity key for signature validation: %w", err)
|
|
}
|
|
|
|
// TODO: Implement actual signature validation using Signal Protocol
|
|
// This would require integrating with the Signal Protocol library
|
|
// For now, we'll skip the validation
|
|
log.ZInfo(ctx, "signed prekey signature validation skipped (not implemented)",
|
|
"userID", userID, "deviceID", deviceID, "keyID", signedPreKey.KeyId)
|
|
|
|
return nil
|
|
}
|
|
|
|
// RotateSignedPreKey creates a new signed prekey and deactivates the old one
|
|
func (km *KeysManager) RotateSignedPreKey(ctx context.Context, userID string, deviceID int32, newSignedPreKey *SignedPreKeyResponse) error {
|
|
// Validate the new signed prekey
|
|
err := km.ValidateSignedPreKey(ctx, userID, deviceID, newSignedPreKey)
|
|
if err != nil {
|
|
return fmt.Errorf("signed prekey validation failed: %w", err)
|
|
}
|
|
|
|
// Set the new signed prekey (this will automatically deactivate the old one)
|
|
return km.SetSignedPreKey(ctx, userID, deviceID, newSignedPreKey)
|
|
}
|
|
|
|
// GetPreKeyBundleForUser retrieves all necessary keys for X3DH key agreement
|
|
func (km *KeysManager) GetPreKeyBundleForUser(ctx context.Context, userID string, deviceID int32) (map[string]interface{}, error) {
|
|
// Get identity key
|
|
identityKey, err := km.GetIdentityKey(ctx, userID, deviceID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get identity key: %w", err)
|
|
}
|
|
|
|
// Get signed prekey
|
|
signedPreKey, err := km.GetActiveSignedPreKey(ctx, userID, deviceID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get signed prekey: %w", err)
|
|
}
|
|
|
|
// Get one-time prekey (optional)
|
|
oneTimePreKey, err := km.GetOneTimePreKey(ctx, userID, deviceID)
|
|
if err != nil {
|
|
log.ZWarn(ctx, "no one-time prekey available", err)
|
|
oneTimePreKey = nil
|
|
}
|
|
|
|
bundle := map[string]interface{}{
|
|
"identityKey": map[string]interface{}{
|
|
"identityKey": identityKey.IdentityKey,
|
|
"registrationId": identityKey.RegistrationID,
|
|
"createdTime": identityKey.CreatedTime.Unix(),
|
|
},
|
|
"signedPreKey": map[string]interface{}{
|
|
"keyId": signedPreKey.KeyID,
|
|
"publicKey": signedPreKey.PublicKey,
|
|
"signature": signedPreKey.Signature,
|
|
"createdTime": signedPreKey.CreatedTime.Unix(),
|
|
},
|
|
"registrationId": identityKey.RegistrationID,
|
|
}
|
|
|
|
if oneTimePreKey != nil {
|
|
bundle["oneTimePreKey"] = map[string]interface{}{
|
|
"keyId": oneTimePreKey.KeyID,
|
|
"publicKey": oneTimePreKey.PublicKey,
|
|
}
|
|
}
|
|
|
|
return bundle, nil
|
|
}
|
|
|
|
// Response structures used in key manager
|
|
type SignedPreKeyResponse struct {
|
|
KeyId uint32 `json:"keyId"`
|
|
PublicKey []byte `json:"publicKey"`
|
|
Signature []byte `json:"signature"`
|
|
}
|
|
|
|
type PreKeyResponse struct {
|
|
KeyId uint32 `json:"keyId"`
|
|
PublicKey []byte `json:"publicKey"`
|
|
}
|