Support blob range requests

cli/internal/worker/mime.go

@@ -0,0 +1,22 @@
+package worker
+
+import "mime"
+
+// Register MIME types for extensions that aren't in Go's built-in table (and
+// that can't be relied on to be in the system table - slim container images
+// typically don't ship /etc/mime.types), so that asset entries get consistent
+// types regardless of where the worker runs.
+func init() {
+	types := map[string]string{
+		".parquet": "application/vnd.apache.parquet",
+		".csv":     "text/csv",
+		".tsv":     "text/tab-separated-values",
+		".jsonl":   "application/x-ndjson",
+		".ndjson":  "application/x-ndjson",
+		".md":      "text/markdown",
+	}
+
+	for extension, mimeType := range types {
+		_ = mime.AddExtensionType(extension, mimeType)
+	}
+}

server/lib/coflux/handlers/blobs.ex

@@ -4,6 +4,7 @@ defmodule Coflux.Handlers.Blobs do
   alias Coflux.{Auth, Utils}
 
   @key_regex ~r/\A[0-9a-f]{64}\z/
+  @range_regex ~r/\Abytes=(\d*)-(\d*)\z/
 
   def init(req, opts) do
     bindings = :cowboy_req.bindings(req)
@@ -41,27 +42,51 @@ defmodule Coflux.Handlers.Blobs do
   end
 
   defp handle(req, "HEAD", key, opts) do
-    status =
-      cond do
-        not valid_key?(key) -> 404
-        File.exists?(blob_path(key)) -> 200
-        true -> 404
+    req =
+      case blob_stat(key) do
+        {:ok, path, size} ->
+          :cowboy_req.reply(200, %{"accept-ranges" => "bytes"}, {:sendfile, 0, size, path}, req)
+
+        :error ->
+          :cowboy_req.reply(404, %{}, req)
       end
 
-    req = :cowboy_req.reply(status, %{}, req)
     {:ok, req, opts}
   end
 
   defp handle(req, "GET", key, opts) do
-    with true <- valid_key?(key),
-         {:ok, content} <- File.read(blob_path(key)) do
-      req = :cowboy_req.reply(200, %{}, content, req)
-      {:ok, req, opts}
-    else
-      _ ->
-        req = :cowboy_req.reply(404, %{}, "Not found", req)
-        {:ok, req, opts}
-    end
+    req =
+      case blob_stat(key) do
+        {:ok, path, size} ->
+          case parse_range(:cowboy_req.header("range", req), size) do
+            :none ->
+              :cowboy_req.reply(
+                200,
+                %{"accept-ranges" => "bytes"},
+                {:sendfile, 0, size, path},
+                req
+              )
+
+            {:range, first, last} ->
+              :cowboy_req.reply(
+                206,
+                %{
+                  "accept-ranges" => "bytes",
+                  "content-range" => "bytes #{first}-#{last}/#{size}"
+                },
+                {:sendfile, first, last - first + 1, path},
+                req
+              )
+
+            :unsatisfiable ->
+              :cowboy_req.reply(416, %{"content-range" => "bytes */#{size}"}, "", req)
+          end
+
+        :error ->
+          :cowboy_req.reply(404, %{}, "Not found", req)
+      end
+
+    {:ok, req, opts}
   end
 
   defp handle(req, "PUT", key, opts) do
@@ -92,6 +117,62 @@ defmodule Coflux.Handlers.Blobs do
     Utils.data_path("blobs/#{a}/#{b}/#{c}")
   end
 
+  defp blob_stat(key) do
+    if valid_key?(key) do
+      path = blob_path(key)
+
+      case File.stat(path) do
+        {:ok, %File.Stat{type: :regular, size: size}} -> {:ok, path, size}
+        _ -> :error
+      end
+    else
+      :error
+    end
+  end
+
+  # Parses a single byte range. Anything unsupported (multiple ranges, other
+  # units, unparsable values) is treated as if no range were requested.
+  defp parse_range(:undefined, _size), do: :none
+
+  defp parse_range(header, size) do
+    case Regex.run(@range_regex, header, capture: :all_but_first) do
+      [first, last] -> resolve_range(first, last, size)
+      nil -> :none
+    end
+  end
+
+  defp resolve_range("", "", _size), do: :none
+
+  defp resolve_range("", suffix, size) do
+    case String.to_integer(suffix) do
+      0 -> :unsatisfiable
+      count -> satisfiable(max(size - count, 0), size - 1, size)
+    end
+  end
+
+  defp resolve_range(first, "", size) do
+    satisfiable(String.to_integer(first), size - 1, size)
+  end
+
+  defp resolve_range(first, last, size) do
+    first = String.to_integer(first)
+    last = String.to_integer(last)
+
+    if first > last do
+      :none
+    else
+      satisfiable(first, min(last, size - 1), size)
+    end
+  end
+
+  defp satisfiable(first, last, size) do
+    if first >= size do
+      :unsatisfiable
+    else
+      {:range, first, last}
+    end
+  end
+
   defp valid_key?(key) when is_binary(key), do: Regex.match?(@key_regex, key)
   defp valid_key?(_), do: false
 

server/lib/coflux/handlers/utils.ex

@@ -115,7 +115,9 @@ defmodule Coflux.Handlers.Utils do
 
     headers = %{
       "access-control-allow-methods" => "OPTIONS, GET, POST, PUT, PATCH, DELETE",
-      "access-control-allow-headers" => "content-type,authorization,x-api-version,x-project",
+      "access-control-allow-headers" =>
+        "content-type,authorization,x-api-version,x-project,range",
+      "access-control-expose-headers" => "content-length,content-range,accept-ranges",
       "access-control-max-age" => "86400"
     }
 

tests/test_blobs.py

@@ -26,6 +26,21 @@ def _request(server, project_id, method, raw_path, body=None):
         conn.close()
 
 
+def _request_headers(server, project_id, method, raw_path, headers=None):
+    """Like ``_request``, but also returns the response headers (lowercased)."""
+    conn, request_headers = _conn(server, project_id)
+    if headers:
+        request_headers = {**request_headers, **headers}
+    try:
+        conn.request(method, raw_path, headers=request_headers)
+        resp = conn.getresponse()
+        data = resp.read()
+        response_headers = {k.lower(): v for k, v in resp.getheaders()}
+        return resp.status, response_headers, data
+    finally:
+        conn.close()
+
+
 def _hex_sha256(b):
     return hashlib.sha256(b).hexdigest()
 
@@ -116,3 +131,140 @@ def test_no_files_created_outside_blobs_dir(isolated_server):
             assert len(entry) == 2 and all(c in "0123456789abcdef" for c in entry), (
                 f"unexpected blob shard: {entry!r}"
             )
+
+
+def _put_blob(server, project_id, body):
+    """PUT a blob and return its key."""
+    key = _hex_sha256(body)
+    status, _ = _request(server, project_id, "PUT", f"/blobs/{key}", body=body)
+    assert status == 204
+    return key
+
+
+@pytest.fixture
+def sample_blob(server, project_id):
+    """A ~1000 byte blob, returning ``(key, body)``."""
+    body = (bytes(range(256)) * 4)[:1000]
+    return _put_blob(server, project_id, body), body
+
+
+def test_get_full_advertises_ranges(server, project_id, sample_blob):
+    """A rangeless GET returns the whole blob and advertises range support."""
+    key, body = sample_blob
+    status, headers, data = _request_headers(server, project_id, "GET", f"/blobs/{key}")
+    assert status == 200
+    assert headers["accept-ranges"] == "bytes"
+    assert headers["content-length"] == str(len(body))
+    assert data == body
+
+
+def test_head_returns_size(server, project_id, sample_blob):
+    """HEAD reports the size without a body."""
+    key, body = sample_blob
+    status, headers, data = _request_headers(
+        server, project_id, "HEAD", f"/blobs/{key}"
+    )
+    assert status == 200
+    assert headers["content-length"] == str(len(body))
+    assert headers["accept-ranges"] == "bytes"
+    assert data == b""
+
+
+def test_get_range_prefix(server, project_id, sample_blob):
+    key, body = sample_blob
+    status, headers, data = _request_headers(
+        server, project_id, "GET", f"/blobs/{key}", {"Range": "bytes=0-9"}
+    )
+    assert status == 206
+    assert data == body[:10]
+    assert headers["content-range"] == f"bytes 0-9/{len(body)}"
+    assert headers["content-length"] == "10"
+
+
+def test_get_range_open_ended(server, project_id, sample_blob):
+    key, body = sample_blob
+    status, headers, data = _request_headers(
+        server, project_id, "GET", f"/blobs/{key}", {"Range": "bytes=990-"}
+    )
+    assert status == 206
+    assert data == body[990:]
+    assert headers["content-range"] == f"bytes 990-999/{len(body)}"
+
+
+def test_get_range_suffix(server, project_id, sample_blob):
+    key, body = sample_blob
+    status, headers, data = _request_headers(
+        server, project_id, "GET", f"/blobs/{key}", {"Range": "bytes=-10"}
+    )
+    assert status == 206
+    assert data == body[-10:]
+    assert headers["content-range"] == f"bytes 990-999/{len(body)}"
+
+
+def test_get_range_clamped_to_size(server, project_id, sample_blob):
+    """A last-byte beyond the end is clamped rather than rejected."""
+    key, body = sample_blob
+    status, headers, data = _request_headers(
+        server, project_id, "GET", f"/blobs/{key}", {"Range": "bytes=990-5000"}
+    )
+    assert status == 206
+    assert data == body[990:]
+    assert headers["content-range"] == f"bytes 990-999/{len(body)}"
+
+
+def test_get_range_beyond_end_unsatisfiable(server, project_id, sample_blob):
+    key, body = sample_blob
+    status, headers, data = _request_headers(
+        server, project_id, "GET", f"/blobs/{key}", {"Range": "bytes=1000-"}
+    )
+    assert status == 416
+    assert headers["content-range"] == f"bytes */{len(body)}"
+    assert data == b""
+
+
+@pytest.mark.parametrize(
+    "range_header",
+    [
+        "bytes=5-2",  # inverted
+        "items=0-1",  # unsupported unit
+        "bytes=0-1,5-6",  # multiple ranges
+        "bytes=abc",  # unparsable
+    ],
+)
+def test_get_ignores_unsupported_ranges(server, project_id, sample_blob, range_header):
+    """Ranges we don't support fall back to serving the whole blob."""
+    key, body = sample_blob
+    status, headers, data = _request_headers(
+        server, project_id, "GET", f"/blobs/{key}", {"Range": range_header}
+    )
+    assert status == 200
+    assert data == body
+    assert "content-range" not in headers
+
+
+def test_empty_blob_ranges(server, project_id):
+    """An empty blob serves a zero-length 200, and any range is unsatisfiable."""
+    key = _put_blob(server, project_id, b"")
+
+    status, headers, data = _request_headers(server, project_id, "GET", f"/blobs/{key}")
+    assert status == 200
+    assert headers["content-length"] == "0"
+    assert data == b""
+
+    status, headers, _ = _request_headers(
+        server, project_id, "GET", f"/blobs/{key}", {"Range": "bytes=0-"}
+    )
+    assert status == 416
+    assert headers["content-range"] == "bytes */0"
+
+
+def test_options_advertises_range_cors(server, project_id):
+    """Preflight allows the Range header and exposes the range response headers."""
+    key = "0" * 64
+    status, headers, _ = _request_headers(
+        server, project_id, "OPTIONS", f"/blobs/{key}"
+    )
+    assert status == 204
+    assert "range" in headers["access-control-allow-headers"].split(",")
+    exposed = headers["access-control-expose-headers"].split(",")
+    assert "content-range" in exposed