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