From e66cf65eddbad9b6b6f724f6371e5856686f938b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micka=C3=ABl=20Menu?= <mickael.menu@gmail.com>
Date: Mon, 2 Oct 2023 12:57:15 +0200
Subject: [PATCH] Prevent out of memory errors when sniffing a media type
 (#400)

---
 .../r2/shared/resource/DefaultArchiveFactory.kt |  5 ++++-
 .../readium/r2/shared/resource/MediaTypeExt.kt  | 17 +++++++++++++++--
 .../util/archive/channel/ChannelZipContainer.kt |  3 ++-
 3 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/readium/shared/src/main/java/org/readium/r2/shared/resource/DefaultArchiveFactory.kt b/readium/shared/src/main/java/org/readium/r2/shared/resource/DefaultArchiveFactory.kt
index 326cc428e4..5d38cf3b48 100644
--- a/readium/shared/src/main/java/org/readium/r2/shared/resource/DefaultArchiveFactory.kt
+++ b/readium/shared/src/main/java/org/readium/r2/shared/resource/DefaultArchiveFactory.kt
@@ -15,6 +15,9 @@ import org.readium.r2.shared.util.MessageError
 import org.readium.r2.shared.util.Try
 import org.readium.r2.shared.util.mediatype.MediaTypeRetriever
 
+/**
+ * An [ArchiveFactory] to open local ZIP files with Java's [ZipFile].
+ */
 public class DefaultArchiveFactory(
     private val mediaTypeRetriever: MediaTypeRetriever
 ) : ArchiveFactory {
@@ -28,7 +31,7 @@ public class DefaultArchiveFactory(
             ?.let { open(it) }
             ?: Try.Failure(
                 ArchiveFactory.Error.FormatNotSupported(
-                    MessageError("Resource not supported because file cannot be directly access.")
+                    MessageError("Resource not supported because file cannot be directly accessed.")
                 )
             )
     }
diff --git a/readium/shared/src/main/java/org/readium/r2/shared/resource/MediaTypeExt.kt b/readium/shared/src/main/java/org/readium/r2/shared/resource/MediaTypeExt.kt
index 96fb66bb7b..3d21057b9d 100644
--- a/readium/shared/src/main/java/org/readium/r2/shared/resource/MediaTypeExt.kt
+++ b/readium/shared/src/main/java/org/readium/r2/shared/resource/MediaTypeExt.kt
@@ -1,6 +1,7 @@
 package org.readium.r2.shared.resource
 
 import org.readium.r2.shared.util.Url
+import org.readium.r2.shared.util.getOrDefault
 import org.readium.r2.shared.util.mediatype.ContainerMediaTypeSnifferContent
 import org.readium.r2.shared.util.mediatype.ResourceMediaTypeSnifferContent
 
@@ -9,7 +10,7 @@ public class ResourceMediaTypeSnifferContent(
 ) : ResourceMediaTypeSnifferContent {
 
     override suspend fun read(range: LongRange?): ByteArray? =
-        resource.read(range).getOrNull()
+        resource.safeRead(range)
 }
 
 public class ContainerMediaTypeSnifferContent(
@@ -21,6 +22,18 @@ public class ContainerMediaTypeSnifferContent(
 
     override suspend fun read(path: String, range: LongRange?): ByteArray? =
         Url.fromDecodedPath(path)?.let { url ->
-            container.get(url).read(range).getOrNull()
+            container.get(url).safeRead(range)
         }
 }
+
+private suspend fun Resource.safeRead(range: LongRange?): ByteArray? {
+    try {
+        // We only read files smaller than 5MB to avoid an [OutOfMemoryError].
+        if (range == null && length().getOrDefault(0) > 5 * 1000 * 1000) {
+            return null
+        }
+        return read(range).getOrNull()
+    } catch (e: OutOfMemoryError) {
+        return null
+    }
+}
diff --git a/readium/shared/src/main/java/org/readium/r2/shared/util/archive/channel/ChannelZipContainer.kt b/readium/shared/src/main/java/org/readium/r2/shared/util/archive/channel/ChannelZipContainer.kt
index bff4ac323a..d1b32efaae 100644
--- a/readium/shared/src/main/java/org/readium/r2/shared/util/archive/channel/ChannelZipContainer.kt
+++ b/readium/shared/src/main/java/org/readium/r2/shared/util/archive/channel/ChannelZipContainer.kt
@@ -170,7 +170,8 @@ internal class ChannelZipContainer(
 }
 
 /**
- * An [ArchiveFactory] able to open a ZIP archive served through an HTTP server.
+ * An [ArchiveFactory] able to open a ZIP archive served through a stream (e.g. HTTP server,
+ * content URI, etc.).
  */
 public class ChannelZipArchiveFactory(
     private val mediaTypeRetriever: MediaTypeRetriever