session-ios/SessionMessagingKit/Utilities/ProfileManager.swift

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

import UIKit
import GRDB
import PromiseKit
import SignalCoreKit
import SessionUtilitiesKit

public struct ProfileManager {
    public enum Error: LocalizedError {
        case avatarImageTooLarge
        case avatarWriteFailed
        case avatarEncryptionFailed
        case avatarUploadFailed
        case avatarUploadMaxFileSizeExceeded
        
        var localizedDescription: String {
            switch self {
                case .avatarImageTooLarge: return "Avatar image too large."
                case .avatarWriteFailed: return "Avatar write failed."
                case .avatarEncryptionFailed: return "Avatar encryption failed."
                case .avatarUploadFailed: return "Avatar upload failed."
                case .avatarUploadMaxFileSizeExceeded: return "Maximum file size exceeded."
            }
        }
    }
    
    // The max bytes for a user's profile name, encoded in UTF8.
    // Before encrypting and submitting we NULL pad the name data to this length.
    private static let nameDataLength: UInt = 26
    public static let maxAvatarDiameter: CGFloat = 640
    
    private static var profileAvatarCache: Atomic<[String: UIImage]> = Atomic([:])
    private static var currentAvatarDownloads: Atomic<Set<String>> = Atomic([])
    
    // MARK: - Functions
    
    public static func isToLong(profileName: String) -> Bool {
        return ((profileName.data(using: .utf8)?.count ?? 0) > nameDataLength)
    }
    
    public static func profileAvatar(for id: String) -> UIImage? {
        guard let profile: Profile = GRDBStorage.shared.read({ db in try Profile.fetchOne(db, id: id) }) else {
            return nil
        }
        
        if let profileFileName: String = profile.profilePictureFileName, !profileFileName.isEmpty {
            return loadProfileAvatar(for: profileFileName)
        }
        
        if let profilePictureUrl: String = profile.profilePictureUrl, !profilePictureUrl.isEmpty {
            downloadAvatar(for: profile)
        }
        
        return nil
    }
    
    private static func loadProfileAvatar(for fileName: String) -> UIImage? {
        if let cachedImage: UIImage = profileAvatarCache.wrappedValue[fileName] {
            return cachedImage
        }
        
        guard
            !fileName.isEmpty,
            let data: Data = loadProfileData(with: fileName),
            data.isValidImage,
            let image: UIImage = UIImage(data: data)
        else {
            return nil
        }
    
        profileAvatarCache.mutate { $0[fileName] = image }
        return image
    }
    
    private static func loadProfileData(with fileName: String) -> Data? {
        let filePath: String = OWSUserProfile.profileAvatarFilepath(withFilename: fileName)
        
        return try? Data(contentsOf: URL(fileURLWithPath: filePath))
    }
    
    // MARK: - Profile Encryption
    
    private static func encryptProfileData(data: Data, key: OWSAES256Key) -> Data? {
        guard key.keyData.count == kAES256_KeyByteLength else { return nil }
        
        return Cryptography.encryptAESGCMProfileData(plainTextData: data, key: key)
    }
    
    private static func decryptProfileData(data: Data, key: OWSAES256Key) -> Data? {
        guard key.keyData.count == kAES256_KeyByteLength else { return nil }
        
        return Cryptography.decryptAESGCMProfileData(encryptedData: data, key: key)
    }
    
    // MARK: - Other Users' Profiles
    
    public static func downloadAvatar(for profile: Profile, funcName: String = #function) {
        guard !currentAvatarDownloads.wrappedValue.contains(profile.id) else {
            // Download already in flight; ignore
            return
        }
        guard
            let profileUrlStringAtStart: String = profile.profilePictureUrl,
            let profileUrlAtStart: URL = URL(string: profileUrlStringAtStart)
        else {
            SNLog("Skipping downloading avatar for \(profile.id) because url is not set")
            return
        }
        guard
            let fileId: UInt64 = UInt64(profileUrlAtStart.lastPathComponent),
            let profileKeyAtStart: OWSAES256Key = profile.profileEncryptionKey,
            profileKeyAtStart.keyData.count > 0
        else {
            return
        }
        
        let fileName: String = UUID().uuidString.appendingFileExtension("jpg")
        let filePath: String = OWSUserProfile.profileAvatarFilepath(withFilename: fileName)
        var backgroundTask: OWSBackgroundTask? = OWSBackgroundTask(label: funcName)
        
        DispatchQueue.global(qos: .default).async {
            OWSLogger.verbose("downloading profile avatar: \(profile.id)")
            currentAvatarDownloads.mutate { $0.insert(profile.id) }
            
            let useOldServer: Bool = (profileUrlStringAtStart.contains(FileServerAPIV2.oldServer))
            
            FileServerAPIV2
                .download(fileId, useOldServer: useOldServer)
                .done { data in
                    currentAvatarDownloads.mutate { $0.remove(profile.id) }
                    
                    GRDBStorage.shared.write { db in
                        guard let latestProfile: Profile = try Profile.fetchOne(db, id: profile.id) else {
                            return
                        }
                        
                        guard
                            let latestProfileKey: OWSAES256Key = latestProfile.profileEncryptionKey,
                            !latestProfileKey.keyData.isEmpty,
                            latestProfileKey == profileKeyAtStart
                        else {
                            OWSLogger.warn("Ignoring avatar download for obsolete user profile.")
                            return
                        }
                        
                        guard profileUrlStringAtStart == latestProfile.profilePictureUrl else {
                            OWSLogger.warn("Avatar url has changed during download.")
                            
                            if latestProfile.profilePictureUrl?.isEmpty == false {
                                self.downloadAvatar(for: latestProfile)
                            }
                            return
                        }
                        
                        guard let decryptedData: Data = decryptProfileData(data: data, key: profileKeyAtStart) else {
                            OWSLogger.warn("Avatar data for \(profile.id) could not be decrypted.")
                            return
                        }
                        
                        try? decryptedData.write(to: URL(fileURLWithPath: filePath), options: [.atomic])
                        
                        guard let image: UIImage = UIImage(contentsOfFile: filePath) else {
                            OWSLogger.warn("Avatar image for \(profile.id) could not be loaded.")
                            return
                        }
                        
                        try? latestProfile
                            .with(profilePictureFileName: .update(fileName))
                            .update(db)
                        profileAvatarCache.mutate { $0[fileName] = image }
                    }
                    
                    // Redundant but without reading 'backgroundTask' it will warn that the variable
                    // isn't used
                    if backgroundTask != nil { backgroundTask = nil }
                }
                .retainUntilComplete()
        }
    }
    
    // MARK: - Current User Profile
    
    public static func updateLocal(
        profileName: String,
        avatarImage: UIImage?,
        requiredSync: Bool,
        success: (() -> ())? = nil,
        failure: ((Error) -> ())? = nil
    ) {
        DispatchQueue.global(qos: .default).async {
            // If the profile avatar was updated or removed then encrypt with a new profile key
            // to ensure that other users know that our profile picture was updated
            let newProfileKey: OWSAES256Key = OWSAES256Key.generateRandom()
            
            guard let avatarImage: UIImage = avatarImage else {
                // If we have no image then we need to make sure to remove it from the profile
                GRDBStorage.shared.writeAsync(
                    updates: { db in
                        let existingProfile: Profile = Profile.fetchOrCreateCurrentUser(db)
                        
                        OWSLogger.verbose(existingProfile.profilePictureUrl != nil ?
                            "Updating local profile on service with cleared avatar." :
                            "Updating local profile on service with no avatar."
                        )
                        
                        try? existingProfile
                            .with(
                                name: profileName,
                                profilePictureUrl: nil,
                                profilePictureFileName: nil,
                                profileEncryptionKey: (existingProfile.profilePictureUrl != nil ?
                                    .update(newProfileKey) :
                                    .existing
                                )
                            )
                            .save(db)
                        
                        // Remove any cached avatar image value
                        if let fileName: String = existingProfile.profilePictureFileName {
                            profileAvatarCache.mutate { $0[fileName] = nil }
                        }
                    },
                    completion: { _, _ in
                        SNLog("Successfully updated service with profile.")
                        
                        DispatchQueue.main.async {
                            success?()
                        }
                    }
                )
                return
            }

            // If we have a new avatar image, we must first:
            //
            // * Encode it to JPEG.
            // * Write it to disk.
            // * Encrypt it
            // * Upload it to asset service
            // * Send asset service info to Signal Service
            OWSLogger.verbose("Updating local profile on service with new avatar.")
            let maxAvatarBytes: UInt = (5 * 1000 * 1000)
            var image: UIImage = avatarImage

            if image.size.width != maxAvatarDiameter || image.size.height != maxAvatarDiameter {
                // To help ensure the user is being shown the same cropping of their avatar as
                // everyone else will see, we want to be sure that the image was resized before this point.
                SNLog("Avatar image should have been resized before trying to upload")
                image = image.resizedImage(toFillPixelSize: CGSize(width: maxAvatarDiameter, height: maxAvatarDiameter))
            }

            guard let data: Data = image.jpegData(compressionQuality: 0.95) else {
                DispatchQueue.main.async {
                    SNLog("Updating service with profile failed.")
                    failure?(.avatarWriteFailed)
                }
                return
            }
            
            guard data.count <= maxAvatarBytes else {
                // Our avatar dimensions are so small that it's incredibly unlikely we wouldn't
                // be able to fit our profile photo (eg. generating pure noise at our resolution
                // compresses to ~200k)
                DispatchQueue.main.async {
                    SNLog("Suprised to find profile avatar was too large. Was it scaled properly? image: \(image)")
                    SNLog("Updating service with profile failed.")
                    failure?(.avatarImageTooLarge)
                }
                return
            }
            
            let fileName: String = UUID().uuidString.appendingFileExtension("jpg")
            let filePath: String = OWSUserProfile.profileAvatarFilepath(withFilename: fileName)
            
            // Write the avatar to disk
            do { try data.write(to: URL(fileURLWithPath: filePath), options: [.atomic]) }
            catch {
                DispatchQueue.main.async {
                    SNLog("Updating service with profile failed.")
                    failure?(.avatarWriteFailed)
                }
                return
            }
            
            // Encrypt the avatar for upload
            guard let encryptedAvatarData: Data = encryptProfileData(data: data, key: newProfileKey) else {
                DispatchQueue.main.async {
                    SNLog("Updating service with profile failed.")
                    failure?(.avatarEncryptionFailed)
                }
                return
            }
            
            // Upload the avatar to the FileServer
            FileServerAPIV2
                .upload(encryptedAvatarData)
                .done { fileId in
                    let downloadUrl: String = "\(FileServerAPIV2.server)/files/\(fileId)"
                    UserDefaults.standard[.lastProfilePictureUpload] = Date()
                    
                    GRDBStorage.shared.writeAsync(
                        updates: { db in
                            try? Profile
                                .fetchOrCreateCurrentUser(db)
                                .with(
                                    name: profileName,
                                    profilePictureUrl: .update(downloadUrl),
                                    profilePictureFileName: .update(fileName),
                                    profileEncryptionKey: .update(newProfileKey)
                                )
                                .save(db)
                        },
                        completion: { _, _ in
                            // Update the cached avatar image value
                            profileAvatarCache.mutate { $0[fileName] = avatarImage }
                            
                            DispatchQueue.main.async {
                                SNLog("Successfully updated service with profile.")
                                success?()
                            }
                        }
                    )
                }
                .recover { error in
                    DispatchQueue.main.async {
                        SNLog("Updating service with profile failed.")
                        
                        let isMaxFileSizeExceeded: Bool = ((error as? FileServerAPIV2.Error) == FileServerAPIV2.Error.maxFileSizeExceeded)
                        failure?(isMaxFileSizeExceeded ?
                            .avatarUploadMaxFileSizeExceeded :
                            .avatarUploadFailed
                        )
                    }
                }
                .retainUntilComplete()
        }
    }
}

// MARK: - Objective-C Support
@objc(SMKProfileManager)
public class SMKProfileManager: NSObject {
    @objc public static func profileAvatar(recipientId: String) -> UIImage? {
        return ProfileManager.profileAvatar(for: recipientId)
    }
    
    @objc public static func updateLocal(profileName: String, avatarImage: UIImage?, requiresSync: Bool) {
        ProfileManager.updateLocal(profileName: profileName, avatarImage: avatarImage, requiredSync: requiresSync)
    }
}