287 lines
10 KiB
Swift
287 lines
10 KiB
Swift
//
|
|
// Copyright 2020 Signal Messenger, LLC
|
|
// SPDX-License-Identifier: AGPL-3.0-only
|
|
//
|
|
|
|
import LibSignalClient
|
|
|
|
public final class SSKSessionStore: SignalSessionStore {
|
|
fileprivate typealias SessionsByDeviceDictionary = [Int32: AnyObject]
|
|
|
|
private let keyValueStore: KeyValueStore
|
|
private let recipientIdFinder: RecipientIdFinder
|
|
|
|
public init(
|
|
for identity: OWSIdentity,
|
|
keyValueStoreFactory: KeyValueStoreFactory,
|
|
recipientIdFinder: RecipientIdFinder
|
|
) {
|
|
LegacySessionRecord.setUpKeyedArchiverSubstitutions()
|
|
|
|
self.keyValueStore = keyValueStoreFactory.keyValueStore(collection: {
|
|
switch identity {
|
|
case .aci:
|
|
return "TSStorageManagerSessionStoreCollection"
|
|
case .pni:
|
|
return "TSStorageManagerPNISessionStoreCollection"
|
|
}
|
|
}())
|
|
self.recipientIdFinder = recipientIdFinder
|
|
}
|
|
|
|
fileprivate func loadSerializedSession(
|
|
for serviceId: ServiceId,
|
|
deviceId: UInt32,
|
|
tx: DBReadTransaction
|
|
) throws -> Data? {
|
|
switch recipientIdFinder.recipientId(for: serviceId, tx: tx) {
|
|
case .none:
|
|
return nil
|
|
case .some(.success(let recipientId)):
|
|
return loadSerializedSession(for: recipientId, deviceId: deviceId, tx: tx)
|
|
case .some(.failure(let error)):
|
|
switch error {
|
|
case .mustNotUsePniBecauseAciExists:
|
|
throw error
|
|
}
|
|
}
|
|
}
|
|
|
|
private func serializedSession(fromDatabaseRepresentation entry: Any) -> Data? {
|
|
switch entry {
|
|
case let data as Data:
|
|
return data
|
|
case let record as LegacySessionRecord:
|
|
do {
|
|
return try record.serializeProto()
|
|
} catch {
|
|
owsFailDebug("failed to serialize AxolotlKit session: \(error)")
|
|
return nil
|
|
}
|
|
default:
|
|
owsFailDebug("unexpected entry in session store: \(entry)")
|
|
return nil
|
|
}
|
|
}
|
|
|
|
private func loadSerializedSession(
|
|
for recipientId: String,
|
|
deviceId: UInt32,
|
|
tx: DBReadTransaction
|
|
) -> Data? {
|
|
owsAssertDebug(!recipientId.isEmpty)
|
|
owsAssertDebug(deviceId > 0)
|
|
|
|
let dictionary = keyValueStore.getObject(forKey: recipientId, transaction: tx) as! SessionsByDeviceDictionary?
|
|
guard let entry = dictionary?[Int32(bitPattern: deviceId)] else {
|
|
return nil
|
|
}
|
|
return serializedSession(fromDatabaseRepresentation: entry)
|
|
}
|
|
|
|
fileprivate func storeSerializedSession(
|
|
_ sessionData: Data,
|
|
for serviceId: ServiceId,
|
|
deviceId: UInt32,
|
|
tx: DBWriteTransaction
|
|
) throws {
|
|
switch recipientIdFinder.ensureRecipientId(for: serviceId, tx: tx) {
|
|
case .failure(let error):
|
|
switch error {
|
|
case .mustNotUsePniBecauseAciExists:
|
|
throw error
|
|
}
|
|
case .success(let recipientId):
|
|
storeSerializedSession(for: recipientId, deviceId: deviceId, sessionData: sessionData, tx: tx)
|
|
}
|
|
}
|
|
|
|
private func storeSerializedSession(
|
|
for recipientId: String,
|
|
deviceId: UInt32,
|
|
sessionData: Data,
|
|
tx: DBWriteTransaction
|
|
) {
|
|
owsAssertDebug(!recipientId.isEmpty)
|
|
owsAssertDebug(deviceId > 0)
|
|
|
|
var dictionary = (keyValueStore.getObject(forKey: recipientId, transaction: tx) as! SessionsByDeviceDictionary?) ?? [:]
|
|
dictionary[Int32(bitPattern: deviceId)] = sessionData as NSData
|
|
keyValueStore.setObject(dictionary, key: recipientId, transaction: tx)
|
|
}
|
|
|
|
public func mightContainSession(for recipient: SignalRecipient, tx: DBReadTransaction) -> Bool {
|
|
return keyValueStore.hasValue(recipient.uniqueId, transaction: tx)
|
|
}
|
|
|
|
public func mergeRecipient(_ recipient: SignalRecipient, into targetRecipient: SignalRecipient, tx: DBWriteTransaction) {
|
|
let recipientPair = MergePair(fromValue: recipient, intoValue: targetRecipient)
|
|
let sessionBlob = recipientPair.map { keyValueStore.getData($0.uniqueId, transaction: tx) }
|
|
guard let fromValue = sessionBlob.fromValue else {
|
|
return
|
|
}
|
|
if sessionBlob.intoValue == nil {
|
|
keyValueStore.setData(fromValue, key: targetRecipient.uniqueId, transaction: tx)
|
|
}
|
|
keyValueStore.removeValue(forKey: recipient.uniqueId, transaction: tx)
|
|
}
|
|
|
|
public func deleteAllSessions(for serviceId: ServiceId, tx: DBWriteTransaction) {
|
|
Logger.info("deleting all sessions for \(serviceId)")
|
|
switch recipientIdFinder.recipientId(for: serviceId, tx: tx) {
|
|
case .none, .some(.failure(.mustNotUsePniBecauseAciExists)):
|
|
// There can't possibly be any sessions that need to be deleted.
|
|
return
|
|
case .some(.success(let recipientId)):
|
|
owsAssertDebug(!recipientId.isEmpty)
|
|
deleteAllSessions(for: recipientId, tx: tx)
|
|
}
|
|
}
|
|
|
|
public func deleteAllSessions(for recipientId: AccountId, tx: DBWriteTransaction) {
|
|
keyValueStore.removeValue(forKey: recipientId, transaction: tx)
|
|
}
|
|
|
|
public func archiveAllSessions(for serviceId: ServiceId, tx: DBWriteTransaction) {
|
|
Logger.info("archiving all sessions for \(serviceId)")
|
|
switch recipientIdFinder.recipientId(for: serviceId, tx: tx) {
|
|
case .none, .some(.failure(.mustNotUsePniBecauseAciExists)):
|
|
// There can't possibly be any sessions that need to be archived.
|
|
return
|
|
case .some(.success(let recipientId)):
|
|
archiveAllSessions(for: recipientId, tx: tx)
|
|
}
|
|
}
|
|
|
|
public func archiveAllSessions(for address: SignalServiceAddress, tx: DBWriteTransaction) {
|
|
Logger.info("archiving all sessions for \(address)")
|
|
switch recipientIdFinder.recipientId(for: address, tx: tx) {
|
|
case .none, .some(.failure(.mustNotUsePniBecauseAciExists)):
|
|
// There can't possibly be any sessions that need to be archived.
|
|
return
|
|
case .some(.success(let recipientId)):
|
|
archiveAllSessions(for: recipientId, tx: tx)
|
|
}
|
|
}
|
|
|
|
private func archiveAllSessions(for recipientId: AccountId, tx: DBWriteTransaction) {
|
|
owsAssertDebug(!recipientId.isEmpty)
|
|
|
|
guard let dictionary = keyValueStore.getObject(forKey: recipientId, transaction: tx) as! SessionsByDeviceDictionary? else {
|
|
// We never had a session for this account in the first place.
|
|
return
|
|
}
|
|
|
|
let newDictionary: SessionsByDeviceDictionary = dictionary.mapValues { record in
|
|
guard let data = serializedSession(fromDatabaseRepresentation: record) else {
|
|
// We've already logged an error; skip this session.
|
|
return record
|
|
}
|
|
|
|
do {
|
|
let session = try SessionRecord(bytes: data)
|
|
session.archiveCurrentState()
|
|
return Data(session.serialize()) as NSData
|
|
} catch {
|
|
owsFailDebug("\(error)")
|
|
return record
|
|
}
|
|
}
|
|
|
|
keyValueStore.setObject(newDictionary, key: recipientId, transaction: tx)
|
|
}
|
|
|
|
public func resetSessionStore(tx: DBWriteTransaction) {
|
|
Logger.warn("resetting session store")
|
|
keyValueStore.removeAll(transaction: tx)
|
|
}
|
|
|
|
public func printAll(tx: DBReadTransaction) {
|
|
Logger.debug("All Sessions.")
|
|
keyValueStore.enumerateKeysAndObjects(transaction: tx) { key, value, _ in
|
|
guard let deviceSessions = value as? NSDictionary else {
|
|
owsFailDebug("Unexpected type: \(type(of: value)) in collection.")
|
|
return
|
|
}
|
|
|
|
Logger.debug(" Sessions for recipient: \(key)")
|
|
deviceSessions.enumerateKeysAndObjects { key, value, _ in
|
|
guard let data = self.serializedSession(fromDatabaseRepresentation: value) else {
|
|
// We've already logged an error here, just move on.
|
|
return
|
|
}
|
|
do {
|
|
let sessionRecord = try SessionRecord(bytes: data)
|
|
Logger.debug(" Device: \(key) hasCurrentState: \(sessionRecord.hasCurrentState)")
|
|
} catch {
|
|
owsFailDebug("invalid session record: \(error)")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
extension SSKSessionStore {
|
|
public func loadSession(
|
|
for serviceId: ServiceId,
|
|
deviceId: UInt32,
|
|
tx: DBReadTransaction
|
|
) throws -> SessionRecord? {
|
|
guard let serializedData = try loadSerializedSession(for: serviceId, deviceId: deviceId, tx: tx) else {
|
|
return nil
|
|
}
|
|
return try SessionRecord(bytes: serializedData)
|
|
}
|
|
|
|
fileprivate func storeSession(
|
|
_ record: SessionRecord,
|
|
for serviceId: ServiceId,
|
|
deviceId: UInt32,
|
|
tx: DBWriteTransaction
|
|
) throws {
|
|
try storeSerializedSession(Data(record.serialize()), for: serviceId, deviceId: deviceId, tx: tx)
|
|
}
|
|
|
|
public func archiveSession(for serviceId: ServiceId, deviceId: UInt32, tx: DBWriteTransaction) {
|
|
do {
|
|
guard let session = try loadSession(for: serviceId, deviceId: deviceId, tx: tx) else {
|
|
return
|
|
}
|
|
session.archiveCurrentState()
|
|
try storeSession(session, for: serviceId, deviceId: deviceId, tx: tx)
|
|
} catch {
|
|
owsFailDebug("\(error)")
|
|
}
|
|
}
|
|
}
|
|
|
|
extension SSKSessionStore: LibSignalClient.SessionStore {
|
|
public func loadSession(for address: ProtocolAddress, context: StoreContext) throws -> SessionRecord? {
|
|
return try loadSession(for: address.serviceId, deviceId: address.deviceId, tx: context.asTransaction.asV2Read)
|
|
}
|
|
|
|
public func loadExistingSessions(
|
|
for addresses: [ProtocolAddress],
|
|
context: StoreContext
|
|
) throws -> [SessionRecord] {
|
|
|
|
try addresses.compactMap {
|
|
try loadSession(for: $0, context: context)
|
|
}
|
|
}
|
|
|
|
public func storeSession(_ record: SessionRecord, for address: ProtocolAddress, context: StoreContext) throws {
|
|
try storeSession(record, for: address.serviceId, deviceId: address.deviceId, tx: context.asTransaction.asV2Write)
|
|
}
|
|
}
|
|
|
|
#if TESTABLE_BUILD
|
|
|
|
extension SSKSessionStore {
|
|
public func removeAll(tx: DBWriteTransaction) {
|
|
keyValueStore.removeAll(transaction: tx)
|
|
}
|
|
}
|
|
|
|
#endif
|