384 lines
13 KiB
Swift
384 lines
13 KiB
Swift
//
|
|
// Copyright 2023 Signal Messenger, LLC
|
|
// SPDX-License-Identifier: AGPL-3.0-only
|
|
//
|
|
|
|
public import LibSignalClient
|
|
|
|
public struct KyberPreKeyRecord: Codable {
|
|
|
|
enum CodingKeys: String, CodingKey {
|
|
case keyData
|
|
case isLastResort
|
|
case replacedAt
|
|
}
|
|
|
|
public let signature: Data
|
|
public let generatedAt: Date
|
|
public var replacedAt: Date?
|
|
public let id: UInt32
|
|
public let keyPair: KEMKeyPair
|
|
public let isLastResort: Bool
|
|
|
|
public init(
|
|
_ id: UInt32,
|
|
keyPair: KEMKeyPair,
|
|
signature: Data,
|
|
generatedAt: Date,
|
|
replacedAt: Date?,
|
|
isLastResort: Bool
|
|
) {
|
|
self.id = id
|
|
self.keyPair = keyPair
|
|
self.signature = signature
|
|
self.generatedAt = generatedAt
|
|
self.replacedAt = replacedAt
|
|
self.isLastResort = isLastResort
|
|
}
|
|
|
|
public func encode(to encoder: Encoder) throws {
|
|
var container = encoder.container(keyedBy: CodingKeys.self)
|
|
|
|
try container.encode(isLastResort, forKey: .isLastResort)
|
|
try container.encode(asLSCRecord().serialize(), forKey: .keyData)
|
|
try container.encodeIfPresent(replacedAt, forKey: .replacedAt)
|
|
}
|
|
|
|
public init(from decoder: Decoder) throws {
|
|
let container = try decoder.container(keyedBy: CodingKeys.self)
|
|
|
|
let isLastResort = try container.decode(Bool.self, forKey: .isLastResort)
|
|
let keyData = try container.decode(Data.self, forKey: .keyData)
|
|
let replacedAt = try container.decodeIfPresent(Date.self, forKey: .replacedAt)
|
|
|
|
let record = try LibSignalClient.KyberPreKeyRecord(bytes: keyData)
|
|
self.id = record.id
|
|
self.keyPair = try record.keyPair()
|
|
self.signature = record.signature
|
|
self.generatedAt = Date(millisecondsSince1970: record.timestamp)
|
|
self.replacedAt = replacedAt
|
|
self.isLastResort = isLastResort
|
|
}
|
|
}
|
|
|
|
public class KyberPreKeyStoreImpl: LibSignalClient.KyberPreKeyStore {
|
|
|
|
internal enum Constants {
|
|
internal enum ACI {
|
|
static let keyStoreCollection = "SSKKyberPreKeyStoreACIKeyStore"
|
|
static let metadataStoreCollection = "SSKKyberPreKeyStoreACIMetadataStore"
|
|
}
|
|
|
|
internal enum PNI {
|
|
static let keyStoreCollection = "SSKKyberPreKeyStorePNIKeyStore"
|
|
static let metadataStoreCollection = "SSKKyberPreKeyStorePNIMetadataStore"
|
|
}
|
|
|
|
static let lastKeyId = "lastKeyId"
|
|
|
|
static let lastKeyRotationDate = "lastKeyRotationDate"
|
|
}
|
|
|
|
let identity: OWSIdentity
|
|
|
|
// Store both one-time and last resort keys
|
|
private let keyStore: KeyValueStore
|
|
|
|
// Store current last resort ID
|
|
private let metadataStore: KeyValueStore
|
|
|
|
private let dateProvider: DateProvider
|
|
|
|
public init(
|
|
for identity: OWSIdentity,
|
|
dateProvider: @escaping DateProvider,
|
|
) {
|
|
self.identity = identity
|
|
self.dateProvider = dateProvider
|
|
|
|
switch identity {
|
|
case .aci:
|
|
self.keyStore = KeyValueStore(collection: Constants.ACI.keyStoreCollection)
|
|
self.metadataStore = KeyValueStore(collection: Constants.ACI.metadataStoreCollection)
|
|
case .pni:
|
|
self.keyStore = KeyValueStore(collection: Constants.PNI.keyStoreCollection)
|
|
self.metadataStore = KeyValueStore(collection: Constants.PNI.metadataStoreCollection)
|
|
}
|
|
}
|
|
|
|
private func nextKyberPreKeyId(minimumCapacity: UInt32 = 1, tx: DBReadTransaction) -> UInt32 {
|
|
return PreKeyId.nextPreKeyId(
|
|
lastPreKeyId: UInt32(exactly: metadataStore.getInt32(Constants.lastKeyId, defaultValue: 0, transaction: tx)) ?? 0,
|
|
minimumCapacity: minimumCapacity
|
|
)
|
|
}
|
|
|
|
private func generateKyberPreKeyRecord(
|
|
id: UInt32,
|
|
signedBy identityKeyPair: ECKeyPair,
|
|
isLastResort: Bool
|
|
) -> KyberPreKeyRecord {
|
|
let keyPair = KEMKeyPair.generate()
|
|
let signature = identityKeyPair.keyPair.privateKey.generateSignature(message: keyPair.publicKey.serialize())
|
|
|
|
let record = KyberPreKeyRecord(
|
|
id,
|
|
keyPair: keyPair,
|
|
signature: signature,
|
|
generatedAt: dateProvider(),
|
|
replacedAt: nil,
|
|
isLastResort: isLastResort
|
|
)
|
|
|
|
return record
|
|
}
|
|
|
|
// One-Time prekeys
|
|
|
|
public func generateKyberPreKeyRecords(
|
|
count: Int,
|
|
signedBy keyPair: ECKeyPair,
|
|
tx: DBWriteTransaction
|
|
) -> [KyberPreKeyRecord] {
|
|
var nextKeyId = nextKyberPreKeyId(minimumCapacity: UInt32(count), tx: tx)
|
|
let records = (0..<count).map { _ in
|
|
let record = generateKyberPreKeyRecord(
|
|
id: nextKeyId,
|
|
signedBy: keyPair,
|
|
isLastResort: false
|
|
)
|
|
nextKeyId += 1
|
|
return record
|
|
}
|
|
metadataStore.setInt32(Int32(nextKeyId - 1), key: Constants.lastKeyId, transaction: tx)
|
|
return records
|
|
}
|
|
|
|
// Last Resort Keys
|
|
|
|
public func generateLastResortKyberPreKey(
|
|
signedBy keyPair: ECKeyPair,
|
|
tx: DBWriteTransaction
|
|
) -> SignalServiceKit.KyberPreKeyRecord {
|
|
let keyId = nextKyberPreKeyId(tx: tx)
|
|
let record = generateKyberPreKeyRecord(
|
|
id: keyId,
|
|
signedBy: keyPair,
|
|
isLastResort: true
|
|
)
|
|
metadataStore.setInt32(Int32(keyId), key: Constants.lastKeyId, transaction: tx)
|
|
return record
|
|
}
|
|
|
|
/// Keys returned by this method should not be stored in the local
|
|
/// KyberPreKeyStore since there is no guarantee the key ID is unique.
|
|
public func generateLastResortKyberPreKeyForLinkedDevice(
|
|
signedBy keyPair: ECKeyPair
|
|
) -> SignalServiceKit.KyberPreKeyRecord {
|
|
return generateKyberPreKeyRecord(
|
|
id: PreKeyId.random(),
|
|
signedBy: keyPair,
|
|
isLastResort: true
|
|
)
|
|
}
|
|
|
|
public func storeLastResortPreKey(record: KyberPreKeyRecord, tx: DBWriteTransaction) throws {
|
|
try storeKyberPreKey(record: record, tx: tx)
|
|
}
|
|
|
|
public func storeLastResortPreKeyFromLinkedDevice(record: KyberPreKeyRecord, tx: DBWriteTransaction) throws {
|
|
try storeLastResortPreKey(record: record, tx: tx)
|
|
metadataStore.setInt32(Int32(record.id), key: Constants.lastKeyId, transaction: tx)
|
|
}
|
|
|
|
// MARK: - LibSignalClient.KyberPreKeyStore conformance
|
|
|
|
public func loadKyberPreKey(id: UInt32, tx: DBReadTransaction) -> SignalServiceKit.KyberPreKeyRecord? {
|
|
try? self.keyStore.getCodableValue(forKey: key(for: id), transaction: tx)
|
|
}
|
|
|
|
public func storeKyberPreKeyRecords(records: [KyberPreKeyRecord], tx: DBWriteTransaction) throws {
|
|
for record in records {
|
|
try storeKyberPreKey(record: record, tx: tx)
|
|
}
|
|
}
|
|
|
|
private func storeKyberPreKey(record: SignalServiceKit.KyberPreKeyRecord, tx: DBWriteTransaction) throws {
|
|
try self.keyStore.setCodable(record, key: key(for: record.id), transaction: tx)
|
|
}
|
|
|
|
public func markKyberPreKeyUsed(id: UInt32, tx: DBWriteTransaction) throws {
|
|
// fetch the key, see if it's was a last resort.
|
|
// if not, remove the key from the list of uses (or mark it as used?)
|
|
guard let record = loadKyberPreKey(id: id, tx: tx) else { throw Error.noKyberPreKeyWithId(id) }
|
|
if !record.isLastResort {
|
|
self.keyStore.removeValue(forKey: key(for: id), transaction: tx)
|
|
}
|
|
}
|
|
|
|
public func removeAll(tx: DBWriteTransaction) {
|
|
self.keyStore.removeAll(transaction: tx)
|
|
self.metadataStore.removeAll(transaction: tx)
|
|
}
|
|
|
|
internal func key(for id: UInt32) -> String { "\(id)" }
|
|
|
|
enum Error: Swift.Error {
|
|
case noKyberPreKeyWithId(UInt32)
|
|
case noKyberLastResortKey
|
|
}
|
|
|
|
public func loadKyberPreKey(
|
|
id: UInt32,
|
|
context: LibSignalClient.StoreContext
|
|
) throws -> LibSignalClient.KyberPreKeyRecord {
|
|
guard let preKey = self.loadKyberPreKey(
|
|
id: id,
|
|
tx: context.asTransaction
|
|
) else {
|
|
throw Error.noKyberPreKeyWithId(id)
|
|
}
|
|
|
|
return try LibSignalClient.KyberPreKeyRecord(
|
|
id: preKey.id,
|
|
timestamp: preKey.generatedAt.ows_millisecondsSince1970,
|
|
keyPair: preKey.keyPair,
|
|
signature: preKey.signature
|
|
)
|
|
}
|
|
|
|
// This method isn't used in practice, so it doesn't matter if this is a
|
|
// one time or last resort Kyber key. It also doesn't set replacedAt.
|
|
public func storeKyberPreKey(
|
|
_ record: LibSignalClient.KyberPreKeyRecord,
|
|
id: UInt32,
|
|
context: LibSignalClient.StoreContext
|
|
) throws {
|
|
owsFailDebug("This method doesn't behave correctly.")
|
|
let record = SignalServiceKit.KyberPreKeyRecord(
|
|
id,
|
|
keyPair: try record.keyPair(),
|
|
signature: record.signature,
|
|
generatedAt: Date(millisecondsSince1970: record.timestamp),
|
|
replacedAt: nil,
|
|
isLastResort: false
|
|
)
|
|
try self.storeKyberPreKey(record: record, tx: context.asTransaction)
|
|
}
|
|
|
|
public func markKyberPreKeyUsed(id: UInt32, context: StoreContext) throws {
|
|
try self.markKyberPreKeyUsed(id: id, tx: context.asTransaction)
|
|
}
|
|
|
|
public func setLastSuccessfulRotationDate(_ date: Date, tx: DBWriteTransaction) {
|
|
self.metadataStore.setDate(date, key: Constants.lastKeyRotationDate, transaction: tx)
|
|
}
|
|
|
|
public func getLastSuccessfulRotationDate(tx: DBReadTransaction) -> Date? {
|
|
self.metadataStore.getDate(Constants.lastKeyRotationDate, transaction: tx)
|
|
}
|
|
|
|
public func setOneTimePreKeysReplacedAtToNowIfNil(
|
|
exceptFor justUploadedOneTimePreKey: [KyberPreKeyRecord],
|
|
tx: DBWriteTransaction
|
|
) throws {
|
|
try setReplacedAtToNowIfNil(exceptFor: justUploadedOneTimePreKey, isLastResort: false, tx: tx)
|
|
}
|
|
|
|
public func cullOneTimePreKeyRecords(gracePeriod: TimeInterval, tx: DBWriteTransaction) {
|
|
cullPreKeys(isLastResort: false, gracePeriod: gracePeriod, tx: tx)
|
|
}
|
|
|
|
public func setLastResortPreKeysReplacedAtToNowIfNil(
|
|
exceptFor justUploadedLastResortPreKey: KyberPreKeyRecord,
|
|
tx: DBWriteTransaction
|
|
) throws {
|
|
try setReplacedAtToNowIfNil(exceptFor: [justUploadedLastResortPreKey], isLastResort: true, tx: tx)
|
|
}
|
|
|
|
public func cullLastResortPreKeyRecords(gracePeriod: TimeInterval, tx: DBWriteTransaction) {
|
|
cullPreKeys(isLastResort: true, gracePeriod: gracePeriod, tx: tx)
|
|
}
|
|
|
|
private func setReplacedAtToNowIfNil(exceptFor preKeys: [KyberPreKeyRecord], isLastResort: Bool, tx: DBWriteTransaction) throws {
|
|
let exceptForIds = Set(preKeys.map(\.id))
|
|
for key in self.keyStore.allKeys(transaction: tx) {
|
|
let record: KyberPreKeyRecord? = try self.keyStore.getCodableValue(forKey: key, transaction: tx)
|
|
guard
|
|
var record,
|
|
record.isLastResort == isLastResort,
|
|
!exceptForIds.contains(record.id),
|
|
record.replacedAt == nil
|
|
else {
|
|
continue
|
|
}
|
|
record.replacedAt = dateProvider()
|
|
try self.keyStore.setCodable(record, key: key, transaction: tx)
|
|
}
|
|
}
|
|
|
|
private func cullPreKeys(isLastResort: Bool, gracePeriod: TimeInterval, tx: DBWriteTransaction) {
|
|
for key in self.keyStore.allKeys(transaction: tx) {
|
|
let record: KyberPreKeyRecord?
|
|
do {
|
|
record = try self.keyStore.getCodableValue(forKey: key, transaction: tx)
|
|
} catch {
|
|
owsFailDebug("Couldn't decode KyberPreKeyRecord: \(error)")
|
|
record = nil
|
|
}
|
|
guard let record else {
|
|
self.keyStore.removeValue(forKey: key, transaction: tx)
|
|
continue
|
|
}
|
|
guard
|
|
record.isLastResort == isLastResort,
|
|
let replacedAt = record.replacedAt,
|
|
dateProvider().timeIntervalSince(replacedAt) > (PreKeyManagerImpl.Constants.maxUnacknowledgedSessionAge + gracePeriod)
|
|
else {
|
|
continue
|
|
}
|
|
self.keyStore.removeValue(forKey: key, transaction: tx)
|
|
}
|
|
}
|
|
|
|
public func removeLastResortPreKey(
|
|
record: SignalServiceKit.KyberPreKeyRecord,
|
|
tx: DBWriteTransaction
|
|
) {
|
|
self.keyStore.removeValue(forKey: key(for: record.id), transaction: tx)
|
|
}
|
|
|
|
#if TESTABLE_BUILD
|
|
|
|
func count(isLastResort: Bool, tx: DBReadTransaction) -> Int {
|
|
let records: [KyberPreKeyRecord] = try! self.keyStore.allCodableValues(transaction: tx)
|
|
return records.count(where: { $0.isLastResort == isLastResort })
|
|
}
|
|
|
|
#endif
|
|
}
|
|
|
|
extension LibSignalClient.KyberPreKeyRecord {
|
|
func asSSKLastResortRecord() throws -> SignalServiceKit.KyberPreKeyRecord {
|
|
return SignalServiceKit.KyberPreKeyRecord(
|
|
self.id,
|
|
keyPair: try self.keyPair(),
|
|
signature: self.signature,
|
|
generatedAt: Date(millisecondsSince1970: self.timestamp),
|
|
replacedAt: nil,
|
|
isLastResort: true
|
|
)
|
|
}
|
|
}
|
|
|
|
extension SignalServiceKit.KyberPreKeyRecord {
|
|
func asLSCRecord() throws -> LibSignalClient.KyberPreKeyRecord {
|
|
try LibSignalClient.KyberPreKeyRecord(
|
|
id: self.id,
|
|
timestamp: self.generatedAt.ows_millisecondsSince1970,
|
|
keyPair: self.keyPair,
|
|
signature: self.signature
|
|
)
|
|
}
|
|
}
|