Expose CURL handle to write function via Handler

Author: bdbai

src/easy/context.rs

@@ -0,0 +1,50 @@
+use std::fmt;
+
+use super::handle::EasyData;
+use super::handler::Inner;
+#[cfg(doc)]
+use super::{Easy, Handler};
+
+/// Provides access to the handle inside [`Handler::write`] callback.
+pub struct WriteContext2<H: ?Sized> {
+    inner: *mut Inner<H>,
+}
+
+/// Provides access to the handle inside [`Easy::write_function`] callback.
+#[repr(transparent)]
+pub struct WriteContext(WriteContext2<EasyData>);
+
+impl<H: ?Sized> WriteContext2<H> {
+    /// Returns the raw Easy pointer.
+    #[inline]
+    pub fn raw(&self) -> *mut curl_sys::CURL {
+        // Safety: make sure not to borrow `inner` that would be an alias to the inner handle
+        // activated in a Handler callback.
+        unsafe { *std::ptr::addr_of!((*self.inner).handle) }
+    }
+}
+
+impl WriteContext {
+    /// Returns the raw Easy pointer.
+    #[inline]
+    pub fn raw(&self) -> *mut curl_sys::CURL {
+        self.0.raw()
+    }
+
+    pub(super) fn from_mut(inner: &mut WriteContext2<EasyData>) -> &mut Self {
+        // Safety: `inner` has repr transparent over WriteContext2<EasyData>.
+        unsafe { std::mem::transmute::<&mut WriteContext2<EasyData>, &mut Self>(inner) }
+    }
+}
+
+impl<H: ?Sized> fmt::Debug for WriteContext2<H> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("WriteContext2").finish()
+    }
+}
+
+impl fmt::Debug for WriteContext {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("WriteContext").finish()
+    }
+}

src/easy/handle.rs

@@ -13,6 +13,7 @@ use crate::easy::handler::{Auth, NetRc, PostRedirections, ProxyType, SslOpt};
 use crate::easy::handler::{HttpVersion, IpResolve, SslVersion, TimeCondition};
 use crate::easy::{Easy2, Handler};
 use crate::easy::{Form, List};
+use crate::easy::{WriteContext, WriteContext2};
 use crate::Error;
 
 /// Raw bindings to a libcurl "easy session".
@@ -105,6 +106,8 @@ unsafe impl Send for EasyData {}
 #[derive(Default)]
 struct Callbacks<'a> {
     write: Option<Box<dyn FnMut(&[u8]) -> Result<usize, WriteError> + 'a>>,
+    write_context:
+        Option<Box<dyn FnMut(&[u8], &mut WriteContext) -> Result<usize, WriteError> + 'a>>,
     read: Option<Box<dyn FnMut(&mut [u8]) -> Result<usize, ReadError> + 'a>>,
     seek: Option<Box<dyn FnMut(SeekFrom) -> SeekResult + 'a>>,
     debug: Option<Box<dyn FnMut(InfoType, &[u8]) + 'a>>,
@@ -251,6 +254,15 @@ impl Easy {
         Ok(())
     }
 
+    /// Same as [`Easy::write_function`] but with access to the [`WriteContext`].
+    pub fn write_function_with_context<F>(&mut self, f: F) -> Result<(), Error>
+    where
+        F: FnMut(&[u8], &mut WriteContext) -> Result<usize, WriteError> + Send + 'static,
+    {
+        self.inner.get_mut().owned.write_context = Some(Box::new(f));
+        Ok(())
+    }
+
     /// Read callback for data uploads.
     ///
     /// This callback function gets called by libcurl as soon as it needs to
@@ -1515,6 +1527,22 @@ impl Handler for EasyData {
         }
     }
 
+    fn write_context(
+        &mut self,
+        data: &[u8],
+        ctx: &mut WriteContext2<Self>,
+    ) -> Result<usize, WriteError> {
+        unsafe {
+            match self.callback(|s| &mut s.write_context) {
+                Some(write) => write(data, WriteContext::from_mut(ctx)),
+                None => match self.callback(|s| &mut s.write) {
+                    Some(write) => write(data),
+                    None => Ok(data.len()),
+                },
+            }
+        }
+    }
+
     fn read(&mut self, data: &mut [u8]) -> Result<usize, ReadError> {
         unsafe {
             match self.callback(|s| &mut s.read) {
@@ -1587,6 +1615,16 @@ impl<'easy, 'data> Transfer<'easy, 'data> {
         Ok(())
     }
 
+    /// Same as `Easy::write_context_function`, just takes a non `'static`
+    /// lifetime corresponding to the lifetime of this transfer.
+    pub fn write_context_function<F>(&mut self, f: F) -> Result<(), Error>
+    where
+        F: FnMut(&[u8], &mut WriteContext) -> Result<usize, WriteError> + 'data,
+    {
+        self.data.write_context = Some(Box::new(f));
+        Ok(())
+    }
+
     /// Same as `Easy::read_function`, just takes a non `'static` lifetime
     /// corresponding to the lifetime of this transfer.
     pub fn read_function<F>(&mut self, f: F) -> Result<(), Error>

src/easy/handler.rs

@@ -15,6 +15,7 @@ use socket2::Socket;
 use crate::easy::form;
 use crate::easy::list;
 use crate::easy::windows;
+use crate::easy::WriteContext2;
 use crate::easy::{Form, List};
 use crate::panic;
 use crate::Error;
@@ -81,6 +82,16 @@ pub trait Handler {
         Ok(data.len())
     }
 
+    /// Same as [`write`], but with a context allowing access to some functionalities within
+    /// the callback and the raw handle.
+    fn write_context(
+        &mut self,
+        data: &[u8],
+        _ctx: &mut WriteContext2<Self>,
+    ) -> Result<usize, WriteError> {
+        self.write(data)
+    }
+
     /// Read callback for data uploads.
     ///
     /// This callback function gets called by libcurl as soon as it needs to
@@ -379,8 +390,8 @@ pub struct Easy2<H> {
     inner: Box<Inner<H>>,
 }
 
-struct Inner<H> {
-    handle: *mut curl_sys::CURL,
+pub(super) struct Inner<H: ?Sized> {
+    pub(super) handle: *mut curl_sys::CURL,
     header_list: Option<List>,
     resolve_list: Option<List>,
     connect_to_list: Option<List>,

src/easy/mod.rs

@@ -7,12 +7,14 @@
 //! Most simple usage of libcurl will likely use the `Easy` structure here, and
 //! you can find more docs about its usage on that struct.
 
+mod context;
 mod form;
 mod handle;
 mod handler;
 mod list;
 mod windows;
 
+pub use self::context::{WriteContext, WriteContext2};
 pub use self::form::{Form, Part};
 pub use self::handle::{Easy, Transfer};
 pub use self::handler::{Auth, NetRc, PostRedirections, ProxyType, SslOpt};