Downloader: HTTP status checks, 416 resume, disk-space check, Ollama Accept header

app/src/main/java/io/github/abrarul/modularai/model/ModelDownloader.kt

@@ -1,6 +1,7 @@
 package io.github.abrarul.modularai.model
 
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ensureActive
 import kotlinx.coroutines.withContext
 import org.json.JSONObject
 import java.io.File
@@ -8,6 +9,7 @@ import java.io.RandomAccessFile
 import java.net.HttpURLConnection
 import java.net.URL
 import java.security.MessageDigest
+import kotlin.coroutines.coroutineContext
 
 sealed class Source {
     data class HuggingFace(val repo: String, val file: String, val rev: String = "main") : Source()
@@ -25,31 +27,44 @@ class ModelDownloader(private val modelsDir: File) {
         onProgress: (done: Long, total: Long) -> Unit = { _, _ -> }
     ): File = withContext(Dispatchers.IO) {
         val (url, fileName, digest) = resolve(src)
+        require(!fileName.contains('/') && !fileName.contains("..")) { "Unsafe file name: $fileName" }
         val out = File(modelsDir, fileName)
         val part = File(modelsDir, "$fileName.part")
         var have = if (part.exists()) part.length() else 0L
 
         val c = URL(url).openConnection() as HttpURLConnection
         c.instanceFollowRedirects = true
+        c.connectTimeout = 20_000
+        c.readTimeout = 60_000
         if (have > 0) c.setRequestProperty("Range", "bytes=$have-")
         c.connect()
-        val resumed = c.responseCode == 206
-        if (!resumed) have = 0L
-        val total = have + c.contentLengthLong
+        val code = c.responseCode
 
-        RandomAccessFile(part, "rw").use { raf ->
-            if (resumed) raf.seek(have) else raf.setLength(0)
-            c.inputStream.use { inp ->
-                val buf = ByteArray(256 * 1024)
-                while (true) {
-                    val n = inp.read(buf)
-                    if (n < 0) break
-                    raf.write(buf, 0, n)
-                    have += n
-                    onProgress(have, total)
+        if (code != 416) { // 416 = the .part file already holds the whole file
+            if (code !in 200..299) error("HTTP $code from $url")
+            val resumed = code == 206
+            if (!resumed) have = 0L
+            val remaining = c.contentLengthLong
+            val total = if (remaining > 0) have + remaining else -1L
+            if (total > 0 && remaining > modelsDir.usableSpace) {
+                error("Not enough free space: need ${remaining / 1_048_576} MB, have ${modelsDir.usableSpace / 1_048_576} MB")
+            }
+            RandomAccessFile(part, "rw").use { raf ->
+                if (resumed) raf.seek(have) else raf.setLength(0)
+                c.inputStream.use { inp ->
+                    val buf = ByteArray(256 * 1024)
+                    while (true) {
+                        coroutineContext.ensureActive()
+                        val n = inp.read(buf)
+                        if (n < 0) break
+                        raf.write(buf, 0, n)
+                        have += n
+                        onProgress(have, total)
+                    }
                 }
             }
-        }
+            if (total > 0 && part.length() != total) error("Download incomplete (${part.length()} of $total bytes); run the command again to resume")
+        } else c.disconnect()
 
         val expected = sha256 ?: digest
         if (expected != null) {
@@ -60,10 +75,19 @@ class ModelDownloader(private val modelsDir: File) {
             }
         }
         if (out.exists()) out.delete()
-        part.renameTo(out)
+        if (!part.renameTo(out)) { part.copyTo(out, overwrite = true); part.delete() }
         out
     }
 
+    private fun getText(url: String, accept: String): String {
+        val c = URL(url).openConnection() as HttpURLConnection
+        c.setRequestProperty("Accept", accept)
+        c.connectTimeout = 20_000
+        c.readTimeout = 30_000
+        if (c.responseCode !in 200..299) error("HTTP ${c.responseCode} from $url")
+        return c.inputStream.bufferedReader().use { it.readText() }
+    }
+
     private fun resolve(src: Source): Triple<String, String, String?> = when (src) {
         is Source.HuggingFace ->
             Triple("https://huggingface.co/${src.repo}/resolve/${src.rev}/${src.file}", src.file.substringAfterLast('/'), null)
@@ -71,7 +95,7 @@ class ModelDownloader(private val modelsDir: File) {
             Triple(src.url, src.url.substringAfterLast('/').substringBefore('?').ifBlank { "model.gguf" }, null)
         is Source.Ollama -> {
             val base = "https://registry.ollama.ai/v2/library/${src.name}"
-            val m = JSONObject(URL("$base/manifests/${src.tag}").readText())
+            val m = JSONObject(getText("$base/manifests/${src.tag}", "application/vnd.docker.distribution.manifest.v2+json"))
             val layers = m.getJSONArray("layers")
             var digest = ""
             for (i in 0 until layers.length()) {