session-ios/SessionProtocolKit/Signal/LokiSessionCipher.swift

import Foundation

@objc(SNSessionCipher)
public final class LokiSessionCipher : SessionCipher {
    private let sessionResetImplementation: SessionRestorationProtocol?
    private let sessionStore: SessionStore
    private let preKeyStore: PreKeyStore
    private let recipientID: String
    private let deviceID: Int32
    
    @objc public static let newSessionAdoptedNotification = "LKNewSessionAdoptedNotification"
    @objc public static let contactKey = "LKContactKey"
    
    @objc public init(sessionResetImplementation: SessionRestorationProtocol, sessionStore: SessionStore, preKeyStore: PreKeyStore, signedPreKeyStore: SignedPreKeyStore, identityKeyStore: IdentityKeyStore, recipientID: String, deviceID: Int32) {
        self.sessionResetImplementation = sessionResetImplementation
        self.sessionStore = sessionStore
        self.preKeyStore = preKeyStore
        self.recipientID = recipientID
        self.deviceID = deviceID
        super.init(sessionStore: sessionStore, preKeyStore: preKeyStore, signedPreKeyStore: signedPreKeyStore, identityKeyStore: identityKeyStore, recipientId: recipientID, deviceId: deviceID)
    }
    
    @available(*, unavailable)
    override convenience private init(axolotlStore sessionStore: AxolotlStore, recipientId: String, deviceId: Int32) {
        self.init(sessionStore: sessionStore, preKeyStore: sessionStore, signedPreKeyStore: sessionStore, identityKeyStore: sessionStore, recipientId: recipientId, deviceId: deviceId)
    }
    
    override private init(sessionStore: SessionStore, preKeyStore: PreKeyStore, signedPreKeyStore: SignedPreKeyStore, identityKeyStore: IdentityKeyStore, recipientId: String, deviceId: Int32) {
        self.sessionResetImplementation = nil
        self.sessionStore = sessionStore
        self.preKeyStore = preKeyStore
        self.recipientID = recipientId
        self.deviceID = deviceId
        super.init(sessionStore: sessionStore, preKeyStore: preKeyStore, signedPreKeyStore: signedPreKeyStore, identityKeyStore: identityKeyStore, recipientId: recipientId, deviceId: deviceId)
    }
    
    override public func decrypt(_ whisperMessage: CipherMessage, protocolContext: Any?) throws -> Data {
        // Note that while decrypting our state may change internally
        let currentState = getCurrentState(protocolContext: protocolContext)
        if (currentState == nil && whisperMessage.cipherMessageType == .prekey) {
            try sessionResetImplementation?.validatePreKeyWhisperMessage(for: recipientID, preKeyWhisperMessage: whisperMessage as! PreKeyWhisperMessage, using: protocolContext!)
        }
        let plainText = try super.decrypt(whisperMessage, protocolContext: protocolContext)
        handleSessionReset(for: whisperMessage, previousState: currentState, protocolContext: protocolContext!)
        return plainText
    }
    
    private func getCurrentState(protocolContext: Any?) -> SessionState? {
        let record = sessionStore.loadSession(recipientID, deviceId: deviceID, protocolContext: protocolContext)
        return record.isFresh() ? nil : record.sessionState()
    }
    
    private func handleSessionReset(for whisperMessage: CipherMessage, previousState: SessionState?, protocolContext: Any) {
        // Don't bother doing anything if we didn't have a session before
        guard let previousState = previousState else { return }
        let sessionResetStatus = sessionResetImplementation?.getSessionRestorationStatus(for: recipientID) ?? SessionRestorationStatus.none
        // Bail early if no session reset is in progress
        guard sessionResetStatus != .none else { return }
        let currentState = getCurrentState(protocolContext: protocolContext)
        // Check if our previous state and our current state differ
        if (currentState == nil || currentState!.aliceBaseKey != previousState.aliceBaseKey) {
            if sessionResetStatus == .requestReceived {
                // The other user used an old session to contact us. Wait for them to use a new one
                restoreSession(previousState, protocolContext: protocolContext)
            } else {
                // Our session reset went through successfully.
                // We initiated a session reset and got a different session back from the user.
                deleteAllSessions(except: currentState, protocolContext: protocolContext)
                notifySessionAdopted(protocolContext)
            }
        } else if sessionResetStatus == .requestReceived {
            // Our session reset went through successfully.
            // We got a message with the same session from the other user.
            deleteAllSessions(except: previousState, protocolContext: protocolContext)
            notifySessionAdopted(protocolContext)
        }
    }
    
    private func notifySessionAdopted(_ protocolContext: Any) {
        self.sessionResetImplementation?.handleNewSessionAdopted(for: recipientID, using: protocolContext)
        NotificationCenter.default.post(name: NSNotification.Name(rawValue: LokiSessionCipher.newSessionAdoptedNotification), object: nil, userInfo: [ LokiSessionCipher.contactKey : recipientID ])
    }
    
    private func deleteAllSessions(except state: SessionState?, protocolContext: Any?) {
        let record = sessionStore.loadSession(recipientID, deviceId: deviceID, protocolContext: protocolContext)
        record.removePreviousSessionStates()
        let newState = state ?? SessionState()
        record.setState(newState)
        sessionStore.storeSession(recipientID, deviceId: deviceID, session: record, protocolContext: protocolContext)
    }
    
    private func restoreSession(_ state: SessionState, protocolContext: Any?) {
        let record = sessionStore.loadSession(recipientID, deviceId: deviceID, protocolContext: protocolContext)
        // Remove the state from previous session states
        record.previousSessionStates()?.enumerateObjects(options: .reverse) { obj, index, stop in
            guard let obj = obj as? SessionState, state.aliceBaseKey == obj.aliceBaseKey else { return }
            record.previousSessionStates()?.removeObject(at: index)
            stop.pointee = true
        }
        // Promote so the previous state gets archived
        record.promoteState(state)
        sessionStore.storeSession(recipientID, deviceId: deviceID, session: record, protocolContext: protocolContext)
    }
}