1 // 2 // Copyright 2020-2021 Signal Messenger, LLC 3 // SPDX-License-Identifier: AGPL-3.0-only 4 // 5 6 import Foundation 7 8 /// A dummy StoreContext usable with InMemorySignalProtocolStore. 9 public struct NullContext: StoreContext { 10 public init() {} 11 } 12 13 private struct SenderKeyName: Hashable { 14 var sender: ProtocolAddress 15 var distributionId: UUID 16 } 17 18 public class InMemorySignalProtocolStore: IdentityKeyStore, PreKeyStore, SignedPreKeyStore, SessionStore, SenderKeyStore { 19 private var publicKeys: [ProtocolAddress: IdentityKey] = [:] 20 private var privateKey: IdentityKeyPair 21 private var registrationId: UInt32 22 private var prekeyMap: [UInt32: PreKeyRecord] = [:] 23 private var signedPrekeyMap: [UInt32: SignedPreKeyRecord] = [:] 24 private var sessionMap: [ProtocolAddress: SessionRecord] = [:] 25 private var senderKeyMap: [SenderKeyName: SenderKeyRecord] = [:] 26 27 public init() { 28 privateKey = IdentityKeyPair.generate() 29 registrationId = UInt32.random(in: 0...0x3FFF) 30 } 31 32 public init(identity: IdentityKeyPair, registrationId: UInt32) { 33 self.privateKey = identity 34 self.registrationId = registrationId 35 } 36 identityKeyPairnull37 public func identityKeyPair(context: StoreContext) throws -> IdentityKeyPair { 38 return privateKey 39 } 40 localRegistrationIdnull41 public func localRegistrationId(context: StoreContext) throws -> UInt32 { 42 return registrationId 43 } 44 saveIdentitynull45 public func saveIdentity(_ identity: IdentityKey, for address: ProtocolAddress, context: StoreContext) throws -> Bool { 46 if publicKeys.updateValue(identity, forKey: address) == nil { 47 return false; // newly created 48 } else { 49 return true 50 } 51 } 52 isTrustedIdentitynull53 public func isTrustedIdentity(_ identity: IdentityKey, for address: ProtocolAddress, direction: Direction, context: StoreContext) throws -> Bool { 54 if let pk = publicKeys[address] { 55 return pk == identity 56 } else { 57 return true // tofu 58 } 59 } 60 identitynull61 public func identity(for address: ProtocolAddress, context: StoreContext) throws -> IdentityKey? { 62 return publicKeys[address] 63 } 64 loadPreKeynull65 public func loadPreKey(id: UInt32, context: StoreContext) throws -> PreKeyRecord { 66 if let record = prekeyMap[id] { 67 return record 68 } else { 69 throw SignalError.invalidKeyIdentifier("no prekey with this identifier") 70 } 71 } 72 storePreKeynull73 public func storePreKey(_ record: PreKeyRecord, id: UInt32, context: StoreContext) throws { 74 prekeyMap[id] = record 75 } 76 removePreKeynull77 public func removePreKey(id: UInt32, context: StoreContext) throws { 78 prekeyMap.removeValue(forKey: id) 79 } 80 loadSignedPreKeynull81 public func loadSignedPreKey(id: UInt32, context: StoreContext) throws -> SignedPreKeyRecord { 82 if let record = signedPrekeyMap[id] { 83 return record 84 } else { 85 throw SignalError.invalidKeyIdentifier("no signed prekey with this identifier") 86 } 87 } 88 storeSignedPreKeynull89 public func storeSignedPreKey(_ record: SignedPreKeyRecord, id: UInt32, context: StoreContext) throws { 90 signedPrekeyMap[id] = record 91 } 92 loadSessionnull93 public func loadSession(for address: ProtocolAddress, context: StoreContext) throws -> SessionRecord? { 94 return sessionMap[address] 95 } 96 loadExistingSessionsnull97 public func loadExistingSessions(for addresses: [ProtocolAddress], context: StoreContext) throws -> [SessionRecord] { 98 return try addresses.map { address in 99 if let session = sessionMap[address] { 100 return session 101 } 102 throw SignalError.sessionNotFound("\(address)") 103 } 104 } 105 storeSessionnull106 public func storeSession(_ record: SessionRecord, for address: ProtocolAddress, context: StoreContext) throws { 107 sessionMap[address] = record 108 } 109 storeSenderKeynull110 public func storeSenderKey(from sender: ProtocolAddress, distributionId: UUID, record: SenderKeyRecord, context: StoreContext) throws { 111 senderKeyMap[SenderKeyName(sender: sender, distributionId: distributionId)] = record 112 } 113 loadSenderKeynull114 public func loadSenderKey(from sender: ProtocolAddress, distributionId: UUID, context: StoreContext) throws -> SenderKeyRecord? { 115 return senderKeyMap[SenderKeyName(sender: sender, distributionId: distributionId)] 116 } 117 } 118