session-ios/SessionMessagingKit/Sending & Receiving/MessageReceiver.swift

// Copyright © 2022 Rangeproof Pty Ltd. All rights reserved.

import Foundation
import GRDB
import Sodium
import SessionUtilitiesKit

public enum MessageReceiver {
    private static var lastEncryptionKeyPairRequest: [String: Date] = [:]

    public static func parse(
        _ db: Database,
        data: Data,
        openGroupId: String? = nil,
        openGroupMessageServerId: UInt64? = nil,
        isRetry: Bool = false
    ) throws -> (Message, SNProtoContent) {
        let userPublicKey: String = getUserHexEncodedPublicKey()
        let isOpenGroupMessage: Bool = (openGroupMessageServerId != nil)
        
        // Parse the envelope
        let envelope = try SNProtoEnvelope.parseData(data)
        
        // Decrypt the contents
        guard let ciphertext = envelope.content else { throw MessageReceiverError.noData }
        
        var plaintext: Data
        var sender: String
        var groupPublicKey: String? = nil
        
        if isOpenGroupMessage {
            (plaintext, sender) = (envelope.content!, envelope.source!)
        }
        else {
            switch envelope.type {
                case .sessionMessage:
                    guard let userX25519KeyPair: Box.KeyPair = Identity.fetchUserKeyPair() else {
                        throw MessageReceiverError.noUserX25519KeyPair
                    }
                    
                    (plaintext, sender) = try decryptWithSessionProtocol(ciphertext: ciphertext, using: userX25519KeyPair)
                    
                case .closedGroupMessage:
                    guard
                        let hexEncodedGroupPublicKey = envelope.source,
                        let closedGroup: ClosedGroup = try? ClosedGroup.fetchOne(db, id: hexEncodedGroupPublicKey)
                    else {
                        throw MessageReceiverError.invalidGroupPublicKey
                    }
                    guard
                        let encryptionKeyPairs: [ClosedGroupKeyPair] = try? closedGroup.keyPairs.order(ClosedGroupKeyPair.Columns.receivedTimestamp.desc).fetchAll(db),
                        !encryptionKeyPairs.isEmpty
                    else {
                        throw MessageReceiverError.noGroupKeyPair
                    }
                    
                    // Loop through all known group key pairs in reverse order (i.e. try the latest key
                    // pair first (which'll more than likely be the one we want) but try older ones in
                    // case that didn't work)
                    func decrypt(keyPairs: [ClosedGroupKeyPair], lastError: Error? = nil) throws -> (Data, String) {
                        guard let keyPair: ClosedGroupKeyPair = keyPairs.first else {
                            throw (lastError ?? MessageReceiverError.decryptionFailed)
                        }
                        
                        do {
                            return try decryptWithSessionProtocol(
                                ciphertext: ciphertext,
                                using: Box.KeyPair(
                                    publicKey: keyPair.publicKey.bytes,
                                    secretKey: keyPair.secretKey.bytes
                                )
                            )
                        }
                        catch {
                            return try decrypt(keyPairs: Array(keyPairs.suffix(from: 1)), lastError: error)
                        }
                    }
                    
                    groupPublicKey = hexEncodedGroupPublicKey
                    (plaintext, sender) = try decrypt(keyPairs: encryptionKeyPairs)
                    
                    /*
                    do {
                        try decrypt()
                    } catch {
                        do {
                            let now = Date()
                            // Don't spam encryption key pair requests
                            let shouldRequestEncryptionKeyPair = given(lastEncryptionKeyPairRequest[groupPublicKey!]) { now.timeIntervalSince($0) > 30 } ?? true
                            if shouldRequestEncryptionKeyPair {
                                try MessageSender.requestEncryptionKeyPair(for: groupPublicKey!, using: transaction as! YapDatabaseReadWriteTransaction)
                                lastEncryptionKeyPairRequest[groupPublicKey!] = now
                            }
                        }
                        throw error // Throw the * decryption * error and not the error generated by requestEncryptionKeyPair (if it generated one)
                    }
                     */
                    
                default: throw MessageReceiverError.unknownEnvelopeType
            }
        }
        
        // Don't process the envelope any further if the sender is blocked
        guard (try? Contact.fetchOne(db, id: sender))?.isBlocked != true else {
            throw MessageReceiverError.senderBlocked
        }
        
        // Parse the proto
        let proto: SNProtoContent
        
        do {
            proto = try SNProtoContent.parseData((plaintext as NSData).removePadding())
        }
        catch {
            SNLog("Couldn't parse proto due to error: \(error).")
            throw error
        }
        
        // Parse the message
        let message: Message? = {
            if let readReceipt = ReadReceipt.fromProto(proto, sender: sender) { return readReceipt }
            if let typingIndicator = TypingIndicator.fromProto(proto, sender: sender) { return typingIndicator }
            if let closedGroupControlMessage = ClosedGroupControlMessage.fromProto(proto, sender: sender) { return closedGroupControlMessage }
            if let dataExtractionNotification = DataExtractionNotification.fromProto(proto, sender: sender) { return dataExtractionNotification }
            if let expirationTimerUpdate = ExpirationTimerUpdate.fromProto(proto, sender: sender) { return expirationTimerUpdate }
            if let configurationMessage = ConfigurationMessage.fromProto(proto, sender: sender) { return configurationMessage }
            if let unsendRequest = UnsendRequest.fromProto(proto, sender: sender) { return unsendRequest }
            if let messageRequestResponse = MessageRequestResponse.fromProto(proto, sender: sender) { return messageRequestResponse }
            if let visibleMessage = VisibleMessage.fromProto(proto, sender: sender) { return visibleMessage }
            return nil
        }()
        
        if let message = message {
            // Ignore self sends if needed
            if !message.isSelfSendValid {
                guard sender != userPublicKey else { throw MessageReceiverError.selfSend }
            }
            
            // Guard against control messages in open groups
            if isOpenGroupMessage {
                guard message is VisibleMessage else { throw MessageReceiverError.invalidMessage }
            }
            
            // Finish parsing
            message.sender = sender
            message.recipient = userPublicKey
            message.sentTimestamp = envelope.timestamp
            message.receivedTimestamp = UInt64((Date().timeIntervalSince1970) * 1000)
            message.groupPublicKey = groupPublicKey
            message.openGroupServerMessageId = openGroupMessageServerId
            
            // Validate
            var isValid: Bool = message.isValid
            if message is VisibleMessage && !isValid && proto.dataMessage?.attachments.isEmpty == false {
                isValid = true
            }
            
            guard isValid else {
                throw MessageReceiverError.invalidMessage
            }
            
            // If the message failed to process the first time around we retry it later (if the error is retryable). In this case the timestamp
            // will already be in the database but we don't want to treat the message as a duplicate. The isRetry flag is a simple workaround
            // for this issue.
            
            switch (isRetry, message, (message as? ClosedGroupControlMessage)?.kind) {
                // Allow duplicates in this case to avoid the following situation:
                // • The app performed a background poll or received a push notification
                // • This method was invoked and the received message timestamps table was updated
                // • Processing wasn't finished
                // • The user doesn't see theO new closed group
                case (_, _, .new): break
                    
                // All `VisibleMessage` values will have an associated `Interaction` so just let
                // the unique constraints on that table prevent duplicate messages
                case is (Bool, VisibleMessage, ClosedGroupControlMessage.Kind?): break
                
                // If the message failed to process and we are retrying then there will already
                // be a `ControlMessageProcessRecord`, so just allow this through
                case (true, _, _): break
                    
                default:
                    do {
                        try ControlMessageProcessRecord(
                            threadId: {
                                if let openGroupId: String = openGroupId {
                                    return openGroupId
                                }
                                
                                if let groupPublicKey: String = groupPublicKey {
                                    return groupPublicKey
                                }
                                
                                return sender
                            }(),
                            sentTimestampMs: Int64(envelope.timestamp),
                            serverHash: (message.serverHash ?? ""),
                            openGroupMessageServerId: (openGroupMessageServerId.map { Int64($0) } ?? 0)
                        ).insert(db)
                    }
                    catch { throw MessageReceiverError.duplicateMessage }
            }
            
            // Return
            return (message, proto)
        }
        
        throw MessageReceiverError.unknownMessage
    }
}