wip: add custom http resolver that prioritizes ipv4 addresses

runtime/extism.h

@@ -132,6 +132,12 @@ ExtismSize extism_current_plugin_memory_length(ExtismCurrentPlugin *plugin, Exti
  */
 void extism_current_plugin_memory_free(ExtismCurrentPlugin *plugin, ExtismMemoryHandle ptr);
 
+/**
+ * Add milliseconds to a plug-in's timeout
+ * NOTE: this should only be called from host functions.
+ */
+bool extism_current_plugin_timeout_add_ms(ExtismCurrentPlugin *plugin, uint64_t ms);
+
 /**
  * Create a new host function
  *

runtime/src/current_plugin.rs

@@ -421,6 +421,12 @@ impl CurrentPlugin {
         let length = self.memory_length(offs).unwrap_or_default();
         (offs, length)
     }
+
+    /// Create a new `TimeoutManager` that can be used to adjust timeouts for the
+    /// current plugin. Returns `None` when no timeout has been configured.
+    pub fn timeout_manager(&self) -> Option<TimeoutManager> {
+        TimeoutManager::new(self)
+    }
 }
 
 impl Internal for CurrentPlugin {

runtime/src/lib.rs

@@ -14,6 +14,7 @@ pub(crate) mod manifest;
 pub(crate) mod pdk;
 mod plugin;
 mod plugin_builder;
+mod timeout_manager;
 mod timer;
 
 /// Extism C API
@@ -25,6 +26,7 @@ pub use extism_manifest::{Manifest, Wasm, WasmMetadata};
 pub use function::{Function, UserData, Val, ValType, PTR};
 pub use plugin::{CancelHandle, Plugin, WasmInput, EXTISM_ENV_MODULE, EXTISM_USER_MODULE};
 pub use plugin_builder::{DebugOptions, PluginBuilder};
+pub use timeout_manager::TimeoutManager;
 
 pub(crate) use internal::{Internal, Wasi};
 pub(crate) use timer::{Timer, TimerAction};

runtime/src/pdk.rs

@@ -1,3 +1,5 @@
+use std::cmp::Ordering;
+
 /// All the functions in the file are exposed from inside WASM plugins
 use crate::*;
 
@@ -140,6 +142,26 @@ pub(crate) fn var_set(
     Ok(())
 }
 
+#[derive(Default, Debug)]
+struct Resolver;
+
+impl ureq::Resolver for Resolver {
+    fn resolve(&self, netloc: &str) -> std::io::Result<Vec<std::net::SocketAddr>> {
+        let addrs = std::net::ToSocketAddrs::to_socket_addrs(netloc)?.into_iter();
+        let mut addrs: Vec<_> = addrs.collect();
+        addrs.sort_by(|a, b| {
+            if a.is_ipv4() && b.is_ipv6() {
+                Ordering::Less
+            } else if a.is_ipv6() && b.is_ipv4() {
+                Ordering::Greater
+            } else {
+                Ordering::Equal
+            }
+        });
+        Ok(addrs)
+    }
+}
+
 /// Make an HTTP request
 /// Params: i64 (offset to JSON encoded HttpRequest), i64 (offset to body or 0)
 /// Returns: i64 (offset)
@@ -150,6 +172,7 @@ pub(crate) fn http_request(
 ) -> Result<(), Error> {
     let data: &mut CurrentPlugin = caller.data_mut();
     let http_req_offset = args!(input, 0, i64) as u64;
+
     #[cfg(not(feature = "http"))]
     {
         let handle = match data.memory_handle(http_req_offset) {
@@ -201,7 +224,8 @@ pub(crate) fn http_request(
             )));
         }
 
-        let mut r = ureq::request(req.method.as_deref().unwrap_or("GET"), &req.url);
+        let agent = ureq::builder().resolver(Resolver).build();
+        let mut r = agent.request(req.method.as_deref().unwrap_or("GET"), &req.url);
 
         for (k, v) in req.headers.iter() {
             r = r.set(k, v);

runtime/src/plugin.rs

@@ -238,6 +238,20 @@ impl Plugin {
             CurrentPlugin::new(manifest, with_wasi, available_pages, id)?,
         );
         store.set_epoch_deadline(1);
+        store.call_hook(|data, hook| {
+            if hook.entering_host() {
+                let tx = Timer::tx();
+                tx.send(TimerAction::EnterHost {
+                    id: data.id.clone(),
+                })?;
+            } else if hook.exiting_host() {
+                let tx = Timer::tx();
+                tx.send(TimerAction::ExitHost {
+                    id: data.id.clone(),
+                })?;
+            }
+            Ok(())
+        });
 
         let mut linker = Linker::new(&engine);
         linker.allow_shadowing(true);

runtime/src/sdk.rs

@@ -143,6 +143,25 @@ pub unsafe extern "C" fn extism_current_plugin_memory_free(
     }
 }
 
+/// Add milliseconds to a plug-in's timeout
+/// NOTE: this should only be called from host functions.
+#[no_mangle]
+pub unsafe extern "C" fn extism_current_plugin_timeout_add_ms(
+    plugin: *mut CurrentPlugin,
+    ms: u64,
+) -> bool {
+    if plugin.is_null() {
+        return false;
+    }
+
+    let plugin = &mut *plugin;
+    if let Some(mgr) = plugin.timeout_manager() {
+        return mgr.add(std::time::Duration::from_millis(ms)).is_ok();
+    }
+
+    false
+}
+
 /// Create a new host function
 ///
 /// Arguments

runtime/src/tests/runtime.rs

@@ -235,6 +235,43 @@ fn test_timeout() {
     assert!(err == "timeout");
 }
 
+fn hello_world_timeout_manager(
+    plugin: &mut CurrentPlugin,
+    inputs: &[Val],
+    outputs: &mut [Val],
+    _user_data: UserData<()>,
+) -> Result<(), Error> {
+    let mgr = plugin.timeout_manager().map(|x| x.add_on_drop());
+    std::thread::sleep(std::time::Duration::from_secs(3));
+    outputs[0] = inputs[0].clone();
+    drop(mgr);
+    Ok(())
+}
+
+#[test]
+fn test_timeout_manager() {
+    let f = Function::new(
+        "hello_world",
+        [PTR],
+        [PTR],
+        UserData::default(),
+        hello_world_timeout_manager,
+    );
+
+    let manifest = Manifest::new([extism_manifest::Wasm::data(WASM)])
+        .with_timeout(std::time::Duration::from_secs(1));
+    let mut plugin = Plugin::new(manifest, [f], true).unwrap();
+
+    let start = std::time::Instant::now();
+    let output: Result<&[u8], Error> = plugin.call("count_vowels", "testing");
+    println!("Result {:?}", output);
+    assert!(output.is_ok());
+    let end = std::time::Instant::now();
+    let time = end - start;
+    println!("Plugin ran for {:?}", time);
+    assert!(time.as_secs() >= 3);
+}
+
 typed_plugin!(pub TestTypedPluginGenerics {
     count_vowels<T: FromBytes<'a>>(&str) -> T
 });

runtime/src/timeout_manager.rs

@@ -0,0 +1,92 @@
+use crate::*;
+
+enum DropBehavior {
+    Add,
+    Sub,
+}
+
+/// `TimeoutManager` is used to control `Plugin` timeouts from within host functions
+///
+/// It can be used to add or subtract time from a plug-in's timeout. If a plugin is not
+/// configured to have a timeout then this will have no effect.
+pub struct TimeoutManager {
+    start_time: std::time::Instant,
+    id: uuid::Uuid,
+    tx: std::sync::mpsc::Sender<TimerAction>,
+    cost: f64,
+    drop_behavior: Option<DropBehavior>,
+}
+
+impl TimeoutManager {
+    /// Create a new `TimeoutManager` from the `CurrentPlugin`, this will return `None` if no timeout
+    /// is configured
+    pub fn new(plugin: &CurrentPlugin) -> Option<TimeoutManager> {
+        plugin.manifest.timeout_ms.map(|_| TimeoutManager {
+            start_time: std::time::Instant::now(),
+            id: plugin.id.clone(),
+            tx: Timer::tx(),
+            cost: 1.0,
+            drop_behavior: None,
+        })
+    }
+
+    /// Add to the configured timeout
+    pub fn add(&self, duration: std::time::Duration) -> Result<(), Error> {
+        let d = duration.mul_f64(self.cost);
+        self.tx.send(TimerAction::Extend {
+            id: self.id.clone(),
+            duration: d.into(),
+        })?;
+        Ok(())
+    }
+
+    /// Subtract from the configured timeout
+    pub fn sub(&self, duration: std::time::Duration) -> Result<(), Error> {
+        let d: timer::ExtendTimeout = duration.mul_f64(self.cost).into();
+        self.tx.send(TimerAction::Extend {
+            id: self.id.clone(),
+            duration: -d,
+        })?;
+        Ok(())
+    }
+
+    /// When the `TimeoutManager` is dropped the elapsed duration since it was created will be
+    /// added to the plug-in's timeout
+    pub fn add_on_drop(mut self) -> Self {
+        self.drop_behavior = Some(DropBehavior::Add);
+        self
+    }
+
+    /// When the `TimeoutManager` is dropped the elapsed duration since it was created will be
+    /// added to the plug-in's timeout
+    pub fn sub_on_drop(mut self) -> Self {
+        self.drop_behavior = Some(DropBehavior::Sub);
+        self
+    }
+
+    /// Adjust the cost of added/subtracted values, this will scale all durations
+    /// submitted to this manager by the provided factor.
+    pub fn cost(mut self, cost: f64) -> Self {
+        self.cost = cost;
+        self
+    }
+}
+
+impl Drop for TimeoutManager {
+    fn drop(&mut self) {
+        if let Some(b) = &self.drop_behavior {
+            let duration = self.start_time.elapsed();
+            let x = match b {
+                DropBehavior::Add => self.add(duration),
+                DropBehavior::Sub => self.sub(duration),
+            };
+            if let Err(e) = x {
+                error!(
+                    plugin = self.id.to_string(),
+                    "unable to extend timeout: {}",
+                    e.to_string()
+                );
+            }
+        }
+    }
+}

runtime/src/timer.rs

@@ -1,17 +1,51 @@
 use crate::*;
 
+#[derive(Debug)]
+pub(crate) struct ExtendTimeout {
+    duration: std::time::Duration,
+    negative: bool,
+}
+
+impl From<std::time::Duration> for ExtendTimeout {
+    fn from(value: std::time::Duration) -> Self {
+        ExtendTimeout {
+            duration: value,
+            negative: false,
+        }
+    }
+}
+
+impl std::ops::Neg for ExtendTimeout {
+    type Output = ExtendTimeout;
+
+    fn neg(mut self) -> Self::Output {
+        self.negative = !self.negative;
+        self
+    }
+}
+
 pub(crate) enum TimerAction {
     Start {
         id: uuid::Uuid,
         engine: Engine,
         duration: Option<std::time::Duration>,
     },
+    Extend {
+        id: uuid::Uuid,
+        duration: ExtendTimeout,
+    },
     Stop {
         id: uuid::Uuid,
     },
     Cancel {
         id: uuid::Uuid,
     },
+    EnterHost {
+        id: uuid::Uuid,
+    },
+    ExitHost {
+        id: uuid::Uuid,
+    },
     Shutdown,
 }
 
@@ -31,6 +65,8 @@ extern "C" fn cleanup_timer() {
 
 static mut TIMER: std::sync::Mutex<Option<Timer>> = std::sync::Mutex::new(None);
 
+type TimerMap = std::collections::BTreeMap<uuid::Uuid, (Engine, Option<std::time::Instant>)>;
+
 impl Timer {
     pub(crate) fn tx() -> std::sync::mpsc::Sender<TimerAction> {
         let mut timer = match unsafe { TIMER.lock() } {
@@ -49,7 +85,8 @@ impl Timer {
     pub fn init(timer: &mut Option<Timer>) -> std::sync::mpsc::Sender<TimerAction> {
         let (tx, rx) = std::sync::mpsc::channel();
         let thread = std::thread::spawn(move || {
-            let mut plugins = std::collections::BTreeMap::new();
+            let mut plugins = TimerMap::new();
+            let mut in_host = TimerMap::new();
 
             macro_rules! handle {
                 ($x:expr) => {
@@ -70,12 +107,16 @@ impl Timer {
                         TimerAction::Stop { id } => {
                             trace!(plugin = id.to_string(), "handling stop event");
                             plugins.remove(&id);
+                            in_host.remove(&id);
                         }
                         TimerAction::Cancel { id } => {
                             trace!(plugin = id.to_string(), "handling cancel event");
                             if let Some((engine, _)) = plugins.remove(&id) {
                                 engine.increment_epoch();
                             }
+                            if let Some((engine, _)) = in_host.remove(&id) {
+                                engine.increment_epoch();
+                            }
                         }
                         TimerAction::Shutdown => {
                             trace!("Shutting down timer");
@@ -83,8 +124,42 @@ impl Timer {
                                 trace!(plugin = id.to_string(), "handling shutdown event");
                                 engine.increment_epoch();
                             }
+
+                            for (id, (engine, _)) in in_host.iter() {
+                                trace!(plugin = id.to_string(), "handling shutdown event");
+                                engine.increment_epoch();
+                            }
                             return;
                         }
+                        TimerAction::Extend { id, duration } => {
+                            if let Some((_engine, Some(timeout))) = plugins.get_mut(&id) {
+                                let x = if duration.negative {
+                                    timeout.checked_sub(duration.duration)
+                                } else {
+                                    timeout.checked_add(duration.duration)
+                                };
+                                if let Some(t) = x {
+                                    *timeout = t;
+                                } else {
+                                    error!(
+                                        plugin = id.to_string(),
+                                        "unable to extend timeout by {:?}", duration.duration
+                                    );
+                                }
+                            }
+                        }
+                        TimerAction::EnterHost { id } => {
+                            trace!(plugin = id.to_string(), "enter host function");
+                            if let Some(x) = plugins.remove(&id) {
+                                in_host.insert(id, x);
+                            }
+                        }
+                        TimerAction::ExitHost { id } => {
+                            trace!(plugin = id.to_string(), "exit host function");
+                            if let Some(x) = in_host.remove(&id) {
+                                plugins.insert(id, x);
+                            }
+                        }
                     }
                 };
             }
@@ -96,6 +171,10 @@ impl Timer {
                     }
                 }
 
+                for x in rx.try_iter() {
+                    handle!(x)
+                }
+
                 plugins = plugins
                     .into_iter()
                     .filter(|(_k, (engine, end))| {
@@ -109,10 +188,6 @@ impl Timer {
                         true
                     })
                     .collect();
-
-                for x in rx.try_iter() {
-                    handle!(x)
-                }
             }
         });
         *timer = Some(Timer {