From 0882f353000126c0699b0f84bb30ffa28dd0e245 Mon Sep 17 00:00:00 2001
From: zach <zach@dylibso.com>
Date: Tue, 21 May 2024 20:11:00 -0700
Subject: [PATCH] fix: respect overall timeout when using http requests (#717)

Currently HTTP requests can extend beyond the configured timeout for a
plugin - this PR sets the timeout of the HTTP request to the remaining
time left if a timeout is set in the manifest.
---
 runtime/src/current_plugin.rs | 22 ++++++++++++++--------
 runtime/src/internal.rs       |  4 ----
 runtime/src/pdk.rs            | 11 +++++++++++
 runtime/src/plugin.rs         |  9 +--------
 runtime/src/tests/runtime.rs  | 28 ++++++++++++++++++++++++++++
 5 files changed, 54 insertions(+), 20 deletions(-)

diff --git a/runtime/src/current_plugin.rs b/runtime/src/current_plugin.rs
index bc6e4a7..8e6190b 100644
--- a/runtime/src/current_plugin.rs
+++ b/runtime/src/current_plugin.rs
@@ -15,6 +15,7 @@ pub struct CurrentPlugin {
     pub(crate) available_pages: Option<u32>,
     pub(crate) memory_limiter: Option<MemoryLimiter>,
     pub(crate) id: uuid::Uuid,
+    pub(crate) start_time: std::time::Instant,
 }
 
 unsafe impl Send for CurrentPlugin {}
@@ -359,6 +360,7 @@ impl CurrentPlugin {
             available_pages,
             memory_limiter,
             id,
+            start_time: std::time::Instant::now(),
         })
     }
 
@@ -469,6 +471,18 @@ impl CurrentPlugin {
         let length = self.memory_length(offs).unwrap_or_default();
         (offs, length)
     }
+
+    /// Returns the remaining time before a plugin will timeout, or
+    /// `None` if no timeout is configured in the manifest
+    pub fn time_remaining(&self) -> Option<std::time::Duration> {
+        if let Some(x) = &self.manifest.timeout_ms {
+            let elapsed = &self.start_time.elapsed().as_millis();
+            let ms_left = x.saturating_sub(*elapsed as u64);
+            return Some(std::time::Duration::from_millis(ms_left));
+        }
+
+        None
+    }
 }
 
 impl Internal for CurrentPlugin {
@@ -480,14 +494,6 @@ impl Internal for CurrentPlugin {
         unsafe { &mut *self.store }
     }
 
-    fn linker(&self) -> &Linker<CurrentPlugin> {
-        unsafe { &*self.linker }
-    }
-
-    fn linker_mut(&mut self) -> &mut Linker<CurrentPlugin> {
-        unsafe { &mut *self.linker }
-    }
-
     fn linker_and_store(&mut self) -> (&mut Linker<CurrentPlugin>, &mut Store<CurrentPlugin>) {
         unsafe { (&mut *self.linker, &mut *self.store) }
     }
diff --git a/runtime/src/internal.rs b/runtime/src/internal.rs
index b0e4aa1..7808530 100644
--- a/runtime/src/internal.rs
+++ b/runtime/src/internal.rs
@@ -12,10 +12,6 @@ pub(crate) trait Internal {
 
     fn store_mut(&mut self) -> &mut Store<CurrentPlugin>;
 
-    fn linker(&self) -> &Linker<CurrentPlugin>;
-
-    fn linker_mut(&mut self) -> &mut Linker<CurrentPlugin>;
-
     fn linker_and_store(&mut self) -> (&mut Linker<CurrentPlugin>, &mut Store<CurrentPlugin>);
 
     fn current_plugin(&self) -> &CurrentPlugin {
diff --git a/runtime/src/pdk.rs b/runtime/src/pdk.rs
index 772f8f6..9857aa5 100644
--- a/runtime/src/pdk.rs
+++ b/runtime/src/pdk.rs
@@ -217,6 +217,11 @@ pub(crate) fn http_request(
             r = r.set(k, v);
         }
 
+        // Set HTTP timeout to respect the manifest timeout
+        if let Some(remaining) = data.time_remaining() {
+            r = r.timeout(remaining);
+        }
+
         let res = if body_offset > 0 {
             let handle = match data.memory_handle(body_offset) {
                 Some(h) => h,
@@ -236,6 +241,12 @@ pub(crate) fn http_request(
                 Some(res.into_reader())
             }
             Err(e) => {
+                // Catch timeout and return
+                if let Some(d) = data.time_remaining() {
+                    if e.kind() == ureq::ErrorKind::Io && d.as_nanos() == 0 {
+                        anyhow::bail!("timeout");
+                    }
+                }
                 let msg = e.to_string();
                 if let Some(res) = e.into_response() {
                     data.http_status = res.status();
diff --git a/runtime/src/plugin.rs b/runtime/src/plugin.rs
index c33b856..be320ab 100644
--- a/runtime/src/plugin.rs
+++ b/runtime/src/plugin.rs
@@ -101,14 +101,6 @@ impl Internal for Plugin {
         &mut self.store
     }
 
-    fn linker(&self) -> &Linker<CurrentPlugin> {
-        &self.linker
-    }
-
-    fn linker_mut(&mut self) -> &mut Linker<CurrentPlugin> {
-        &mut self.linker
-    }
-
     fn linker_and_store(&mut self) -> (&mut Linker<CurrentPlugin>, &mut Store<CurrentPlugin>) {
         (&mut self.linker, &mut self.store)
     }
@@ -737,6 +729,7 @@ impl Plugin {
             .expect("Timer should start");
         self.store.epoch_deadline_trap();
         self.store.set_epoch_deadline(1);
+        self.current_plugin_mut().start_time = std::time::Instant::now();
 
         // Call the function
         let mut results = vec![wasmtime::Val::null(); n_results];
diff --git a/runtime/src/tests/runtime.rs b/runtime/src/tests/runtime.rs
index 42a1772..7dd6ca6 100644
--- a/runtime/src/tests/runtime.rs
+++ b/runtime/src/tests/runtime.rs
@@ -240,6 +240,34 @@ fn test_timeout() {
     assert!(err == "timeout");
 }
 
+#[test]
+fn test_http_timeout() {
+    let f = Function::new(
+        "hello_world",
+        [PTR],
+        [PTR],
+        UserData::default(),
+        hello_world,
+    );
+
+    let manifest = Manifest::new([extism_manifest::Wasm::data(WASM_HTTP)])
+        .with_timeout(std::time::Duration::from_millis(1))
+        .with_allowed_host("www.extism.org");
+    let mut plugin = Plugin::new(manifest, [f], true).unwrap();
+
+    let start = std::time::Instant::now();
+    let output: Result<&[u8], Error> =
+        plugin.call("http_request", r#"{"url": "https://www.extism.org"}"#);
+    let end = std::time::Instant::now();
+    let time = end - start;
+    let err = output.unwrap_err().root_cause().to_string();
+    println!(
+        "Timed out plugin ran for {:?}, with error: {:?}",
+        time, &err
+    );
+    assert!(err == "timeout");
+}
+
 typed_plugin!(pub TestTypedPluginGenerics {
     count_vowels<T: FromBytes<'a>>(&str) -> T
 });