From 18670545f91031ca2390b4922a74438f6125e692 Mon Sep 17 00:00:00 2001
From: Niels Andriesse <andriesseniels@gmail.com>
Date: Mon, 10 Feb 2020 15:47:15 +1100
Subject: [PATCH] Implement file size limit

---
 SignalServiceKit/src/Loki/API/LokiAPI.swift       |  2 +-
 SignalServiceKit/src/Loki/API/LokiDotNetAPI.swift | 15 +++++++++++++--
 .../src/Loki/API/LokiFileServerAPI.swift          |  1 +
 .../API/Public Chat/LokiPublicChatPoller.swift    |  2 +-
 4 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/SignalServiceKit/src/Loki/API/LokiAPI.swift b/SignalServiceKit/src/Loki/API/LokiAPI.swift
index 42533db9f..b16adf68d 100644
--- a/SignalServiceKit/src/Loki/API/LokiAPI.swift
+++ b/SignalServiceKit/src/Loki/API/LokiAPI.swift
@@ -136,7 +136,7 @@ public final class LokiAPI : NSObject {
                     getDestinations()
                     lastDeviceLinkUpdate[hexEncodedPublicKey] = Date()
                 }.catch(on: DispatchQueue.global()) { error in
-                    if case LokiDotNetAPI.Error.parsingFailed = error {
+                    if (error as? LokiDotNetAPI.Error) == LokiDotNetAPI.Error.parsingFailed {
                         // Don't immediately re-fetch in case of failure due to a parsing error
                         lastDeviceLinkUpdate[hexEncodedPublicKey] = Date()
                         getDestinations()
diff --git a/SignalServiceKit/src/Loki/API/LokiDotNetAPI.swift b/SignalServiceKit/src/Loki/API/LokiDotNetAPI.swift
index c2846a275..470386a9a 100644
--- a/SignalServiceKit/src/Loki/API/LokiDotNetAPI.swift
+++ b/SignalServiceKit/src/Loki/API/LokiDotNetAPI.swift
@@ -12,8 +12,14 @@ public class LokiDotNetAPI : NSObject {
     private static let attachmentType = "network.loki"
     
     // MARK: Error
-    public enum Error : Swift.Error {
-        case generic, parsingFailed, encryptionFailed, decryptionFailed, signingFailed
+    @objc public class Error : NSError {
+        
+        @objc public static let generic = Error(domain: "com.loki-project.loki-messenger", code: 1, userInfo: [ NSLocalizedDescriptionKey : "An error occurred." ])
+        @objc public static let parsingFailed = Error(domain: "com.loki-project.loki-messenger", code: 2, userInfo: [ NSLocalizedDescriptionKey : "Invalid file server response." ])
+        @objc public static let signingFailed = Error(domain: "com.loki-project.loki-messenger", code: 3, userInfo: [ NSLocalizedDescriptionKey : "Couldn't sign message." ])
+        @objc public static let encryptionFailed = Error(domain: "com.loki-project.loki-messenger", code: 4, userInfo: [ NSLocalizedDescriptionKey : "Couldn't encrypt file." ])
+        @objc public static let decryptionFailed = Error(domain: "com.loki-project.loki-messenger", code: 5, userInfo: [ NSLocalizedDescriptionKey : "Couldn't decrypt file." ])
+        @objc public static let maxFileSizeExceeded = Error(domain: "com.loki-project.loki-messenger", code: 6, userInfo: [ NSLocalizedDescriptionKey : "Maximum file size exceeded." ])
     }
 
     // MARK: Database
@@ -62,6 +68,11 @@ public class LokiDotNetAPI : NSObject {
                 } else {
                     data = unencryptedAttachmentData
                 }
+                // Check the file size if needed
+                let isLokiFileServer = (server == LokiFileServerAPI.server)
+                if isLokiFileServer && data.count > LokiFileServerAPI.maxFileSize {
+                    return seal.reject(Error.maxFileSizeExceeded)
+                }
                 // Create the request
                 let url = "\(server)/files"
                 let parameters: JSON = [ "type" : attachmentType, "Content-Type" : "application/binary" ]
diff --git a/SignalServiceKit/src/Loki/API/LokiFileServerAPI.swift b/SignalServiceKit/src/Loki/API/LokiFileServerAPI.swift
index e19e880db..6a0a9face 100644
--- a/SignalServiceKit/src/Loki/API/LokiFileServerAPI.swift
+++ b/SignalServiceKit/src/Loki/API/LokiFileServerAPI.swift
@@ -9,6 +9,7 @@ public final class LokiFileServerAPI : LokiDotNetAPI {
     #else
     @objc public static let server = "https://file.getsession.org"
     #endif
+    public static let maxFileSize = 10_000_000 // 10 MB
     private static let deviceLinkType = "network.loki.messenger.devicemapping"
     private static let attachmentType = "net.app.core.oembed"
 
diff --git a/SignalServiceKit/src/Loki/API/Public Chat/LokiPublicChatPoller.swift b/SignalServiceKit/src/Loki/API/Public Chat/LokiPublicChatPoller.swift
index 466558199..a334e8554 100644
--- a/SignalServiceKit/src/Loki/API/Public Chat/LokiPublicChatPoller.swift	
+++ b/SignalServiceKit/src/Loki/API/Public Chat/LokiPublicChatPoller.swift	
@@ -198,7 +198,7 @@ public final class LokiPublicChatPoller : NSObject {
                             LokiAPI.lastDeviceLinkUpdate[$0] = Date()
                         }
                     }.catch(on: DispatchQueue.global()) { error in
-                        if case LokiDotNetAPI.Error.parsingFailed = error {
+                        if (error as? LokiDotNetAPI.Error) == LokiDotNetAPI.Error.parsingFailed {
                             // Don't immediately re-fetch in case of failure due to a parsing error
                             hexEncodedPublicKeysToUpdate.forEach {
                                 LokiAPI.lastDeviceLinkUpdate[$0] = Date()