From 868729072028bbd5c0ad48e29f3117d83a545c0c Mon Sep 17 00:00:00 2001
From: "Mads R. B. Kristensen" <madsbk@gmail.com>
Date: Thu, 19 Sep 2024 14:50:34 +0200
Subject: [PATCH] overflow_error

---
 cpp/include/kvikio/remote_handle.hpp | 29 +++++++++++++++++++++-------
 python/kvikio/tests/test_http_io.py  |  4 ++--
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/cpp/include/kvikio/remote_handle.hpp b/cpp/include/kvikio/remote_handle.hpp
index e219a06f5c..32362266aa 100644
--- a/cpp/include/kvikio/remote_handle.hpp
+++ b/cpp/include/kvikio/remote_handle.hpp
@@ -162,7 +162,7 @@ class CurlHandle {
     // Make curl_easy_perform() fail when receiving HTTP code errors.
     setopt(CURLOPT_FAILONERROR, 1L);
 
-    // Perform and check for errors.
+    // Perform the curl operation and check for errors.
     CURLcode err = curl_easy_perform(handle());
     if (err != CURLE_OK) {
       std::string msg(_errbuf);
@@ -173,8 +173,6 @@ class CurlHandle {
       } else {
         ss << "(" << msg << ")";
       }
-      if (err == CURLE_WRITE_ERROR) { ss << "[maybe the server doesn't support file ranges?]"; }
-      std::cout << "perform() - error: " << ss.str() << std::endl;
       throw std::runtime_error(ss.str());
     }
   }
@@ -218,6 +216,7 @@ struct CallbackContext {
   char* buf;
   std::size_t size;
   std::size_t offset;
+  bool overflow_error;
 };
 
 inline std::size_t callback_host_memory(char* data,
@@ -227,7 +226,10 @@ inline std::size_t callback_host_memory(char* data,
 {
   auto ctx           = reinterpret_cast<CallbackContext*>(context);
   std::size_t nbytes = size * nmemb;
-  if (ctx->size < ctx->offset + nbytes) { return CURL_WRITEFUNC_ERROR; }
+  if (ctx->size < ctx->offset + nbytes) {
+    ctx->overflow_error = true;
+    return CURL_WRITEFUNC_ERROR;
+  }
 
   // std::cout << "callback_host_memory() - data: " << ((void*)data)
   //           << ", ctx->buf: " << (void*)ctx->buf << ", offset: " << ctx->offset
@@ -245,7 +247,10 @@ inline std::size_t callback_device_memory(char* data,
 {
   auto ctx           = reinterpret_cast<CallbackContext*>(context);
   std::size_t nbytes = size * nmemb;
-  if (ctx->size < ctx->offset + nbytes) { return CURL_WRITEFUNC_ERROR; }
+  if (ctx->size < ctx->offset + nbytes) {
+    ctx->overflow_error = true;
+    return CURL_WRITEFUNC_ERROR;
+  }
 
   CUcontext cuda_ctx = get_context_from_pointer(ctx->buf);
   PushAndPopContext c(cuda_ctx);
@@ -327,11 +332,21 @@ class RemoteHandle {
     } else {
       curl.setopt(CURLOPT_WRITEFUNCTION, detail::callback_device_memory);
     }
-    detail::CallbackContext ctx{.buf = reinterpret_cast<char*>(buf), .size = size, .offset = 0};
+    detail::CallbackContext ctx{
+      .buf = reinterpret_cast<char*>(buf), .size = size, .offset = 0, .overflow_error = false};
     curl.setopt(CURLOPT_WRITEDATA, &ctx);
 
     // std::cout << "read() - buf: " << buf << ", byte_range: " << byte_range << std::endl;
-    curl.perform();
+    try {
+      curl.perform();
+    } catch (std::runtime_error const& e) {
+      if (ctx.overflow_error) {
+        std::stringstream ss;
+        ss << "maybe the server doesn't support file ranges? [" << e.what() << "]";
+        throw std::overflow_error(ss.str());
+      }
+      throw;
+    }
     return size;
   }
 
diff --git a/python/kvikio/tests/test_http_io.py b/python/kvikio/tests/test_http_io.py
index 3edc9377d7..9b510d042e 100644
--- a/python/kvikio/tests/test_http_io.py
+++ b/python/kvikio/tests/test_http_io.py
@@ -108,10 +108,10 @@ def test_no_range_support(http_server, tmpdir, xp):
     with kvikio.RemoteFile(f"{http_server}/a") as f:
         assert f.nbytes() == a.nbytes
         with pytest.raises(
-            RuntimeError, match="maybe the server doesn't support file ranges?"
+            OverflowError, match="maybe the server doesn't support file ranges?"
         ):
             f.read(b, size=10, file_offset=0)
         with pytest.raises(
-            RuntimeError, match="maybe the server doesn't support file ranges?"
+            OverflowError, match="maybe the server doesn't support file ranges?"
         ):
             f.read(b, size=10, file_offset=10)