refactor: overlook pool methods

ensh63 · ensh63 · commit 7f6e4e642a88 · 2025-10-01T14:32:56.000-07:00
diff --git a/examples/httporigdst.rs b/examples/httporigdst.rs
@@ -19,31 +19,13 @@ struct NgxHttpOrigDstCtx {
 }
 
 impl NgxHttpOrigDstCtx {
-    pub fn save(&mut self, addr: &str, port: in_port_t, pool: &core::Pool) -> core::Status {
-        let addr_data = pool.alloc_unaligned(addr.len());
-        if addr_data.is_null() {
-            return core::Status::NGX_ERROR;
-        }
-        unsafe { libc::memcpy(addr_data, addr.as_ptr() as *const c_void, addr.len()) };
-        self.orig_dst_addr.len = addr.len();
-        self.orig_dst_addr.data = addr_data as *mut u8;
+    pub fn save(&mut self, addr: &str, port: in_port_t, pool: &core::Pool) -> core::NgxResult {
+        self.orig_dst_addr = unsafe { ngx_str_t::from_str(pool.as_ptr(), addr) };
 
         let port_str = port.to_string();
-        let port_data = pool.alloc_unaligned(port_str.len());
-        if port_data.is_null() {
-            return core::Status::NGX_ERROR;
-        }
-        unsafe {
-            libc::memcpy(
-                port_data,
-                port_str.as_bytes().as_ptr() as *const c_void,
-                port_str.len(),
-            )
-        };
-        self.orig_dst_port.len = port_str.len();
-        self.orig_dst_port.data = port_data as *mut u8;
+        self.orig_dst_port = unsafe { ngx_str_t::from_str(pool.as_ptr(), &port_str) };
 
-        core::Status::NGX_OK
+        Ok(core::Status::NGX_OK.into())
     }
 
     pub unsafe fn bind_addr(&self, v: *mut ngx_variable_value_t) {
@@ -226,7 +208,7 @@ http_variable_get!(
                     ip,
                     port,
                 );
-                (*new_ctx).save(&ip, port, &request.pool());
+                (*new_ctx).save(&ip, port, &request.pool())?;
                 (*new_ctx).bind_addr(v);
                 request
                     .set_module_ctx(new_ctx as *mut c_void, &*addr_of!(ngx_http_orig_dst_module));
@@ -273,7 +255,7 @@ http_variable_get!(
                     ip,
                     port,
                 );
-                (*new_ctx).save(&ip, port, &request.pool());
+                (*new_ctx).save(&ip, port, &request.pool())?;
                 (*new_ctx).bind_port(v);
                 request
                     .set_module_ctx(new_ctx as *mut c_void, &*addr_of!(ngx_http_orig_dst_module));
diff --git a/examples/shared_dict.rs b/examples/shared_dict.rs
@@ -200,10 +200,10 @@ extern "C" fn ngx_http_shared_dict_add_variable(
     let cf = unsafe { cf.as_mut().unwrap() };
     let pool = unsafe { Pool::from_ngx_pool(cf.pool) };
 
-    let key = pool.calloc_type::<ngx_http_complex_value_t>();
-    if key.is_null() {
-        return NGX_CONF_ERROR;
-    }
+    let key = match pool.allocate_type_zeroed::<ngx_http_complex_value_t>() {
+        Ok(p) => p.as_ptr(),
+        Err(_) => return NGX_CONF_ERROR,
+    };
 
     // SAFETY:
     // - `cf.args` is guaranteed to be a pointer to an array with 3 elements (NGX_CONF_TAKE2).
diff --git a/examples/upstream.rs b/examples/upstream.rs
@@ -120,44 +120,34 @@ http_upstream_init_peer_pt!(
     |request: &mut Request, us: *mut ngx_http_upstream_srv_conf_t| {
         ngx_log_debug_http!(request, "CUSTOM UPSTREAM request peer init");
 
-        let hcpd = request.pool().alloc_type::<UpstreamPeerData>();
-        if hcpd.is_null() {
-            return Status::NGX_ERROR;
-        }
+        let hcpd = request.pool().allocate_type::<UpstreamPeerData>()?.as_ptr();
 
         // SAFETY: this function is called with non-NULL uf always
         let us = unsafe { &mut *us };
-        let hccf = match Module::server_conf(us) {
-            Some(x) => x,
-            None => return Status::NGX_ERROR,
-        };
+        let hccf = Module::server_conf(us).ok_or(ngx::core::NgxError {})?;
 
         let original_init_peer = hccf.original_init_peer.unwrap();
         if unsafe { original_init_peer(request.into(), us) != Status::NGX_OK.into() } {
-            return Status::NGX_ERROR;
+            return Status::NGX_ERROR.into();
         }
 
-        let maybe_upstream = request.upstream();
-        if maybe_upstream.is_none() {
-            return Status::NGX_ERROR;
-        }
-        let upstream_ptr = maybe_upstream.unwrap();
+        let upstream_ptr = request.upstream().ok_or(ngx::core::NgxError {})?;
 
         unsafe {
             (*hcpd).conf = Some(hccf);
-            (*hcpd).upstream = maybe_upstream;
+            (*hcpd).upstream = Some(upstream_ptr);
             (*hcpd).data = (*upstream_ptr).peer.data;
             (*hcpd).client_connection = Some(request.connection());
             (*hcpd).original_get_peer = (*upstream_ptr).peer.get;
             (*hcpd).original_free_peer = (*upstream_ptr).peer.free;
 
-            (*upstream_ptr).peer.data = hcpd as *mut c_void;
+            (*upstream_ptr).peer.data = hcpd as _;
             (*upstream_ptr).peer.get = Some(ngx_http_upstream_get_custom_peer);
             (*upstream_ptr).peer.free = Some(ngx_http_upstream_free_custom_peer);
         }
 
         ngx_log_debug_http!(request, "CUSTOM UPSTREAM end request peer init");
-        Status::NGX_OK
+        Status::NGX_OK.into()
     }
 );
 
@@ -312,24 +302,25 @@ impl HttpModule for Module {
 
     unsafe extern "C" fn create_srv_conf(cf: *mut ngx_conf_t) -> *mut c_void {
         let pool = Pool::from_ngx_pool((*cf).pool);
-        let conf = pool.alloc_type::<SrvConfig>();
-        if conf.is_null() {
+
+        if let Ok(mut conf) = pool.allocate_type::<SrvConfig>() {
+            unsafe {
+                conf.as_mut().max = NGX_CONF_UNSET as u32;
+            }
+            ngx_log_debug_mask!(
+                DebugMask::Http,
+                (*cf).log,
+                "CUSTOM UPSTREAM end create_srv_conf"
+            );
+            conf.as_ptr() as _
+        } else {
             ngx_conf_log_error!(
                 NGX_LOG_EMERG,
                 cf,
                 "CUSTOM UPSTREAM could not allocate memory for config"
             );
-            return std::ptr::null_mut();
+            std::ptr::null_mut()
         }
-
-        (*conf).max = NGX_CONF_UNSET as u32;
-
-        ngx_log_debug_mask!(
-            DebugMask::Http,
-            (*cf).log,
-            "CUSTOM UPSTREAM end create_srv_conf"
-        );
-        conf as *mut c_void
     }
 }
 
diff --git a/src/core/pool.rs b/src/core/pool.rs
@@ -4,7 +4,7 @@ use core::mem;
 use core::ptr::{self, NonNull};
 
 use nginx_sys::{
-    ngx_buf_t, ngx_create_temp_buf, ngx_palloc, ngx_pcalloc, ngx_pfree, ngx_pmemalign, ngx_pnalloc,
+    ngx_buf_t, ngx_create_temp_buf, ngx_palloc, ngx_pfree, ngx_pmemalign, ngx_pnalloc,
     ngx_pool_cleanup_add, ngx_pool_t, NGX_ALIGNMENT,
 };
 
@@ -183,10 +183,7 @@ impl Pool {
     /// Returns `Some(MemoryBuffer)` if the buffer is successfully created, or `None` if allocation
     /// fails.
     pub fn create_buffer_from_static_str(&self, str: &'static str) -> Option<MemoryBuffer> {
-        let buf = self.calloc_type::<ngx_buf_t>();
-        if buf.is_null() {
-            return None;
-        }
+        let buf = self.allocate(Layout::new::<ngx_buf_t>()).ok()?.as_ptr() as *mut ngx_buf_t;
 
         // We cast away const, but buffers with the memory flag are read-only
         let start = str.as_ptr() as *mut u8;
@@ -229,44 +226,6 @@ impl Pool {
         unsafe { ngx_palloc(self.0.as_ptr(), size) }
     }
 
-    /// Allocates memory for a type from the pool.
-    /// The resulting pointer is aligned to a platform word size.
-    ///
-    /// Returns a typed pointer to the allocated memory.
-    pub fn alloc_type<T: Copy>(&self) -> *mut T {
-        self.alloc(mem::size_of::<T>()) as *mut T
-    }
-
-    /// Allocates zeroed memory from the pool of the specified size.
-    /// The resulting pointer is aligned to a platform word size.
-    ///
-    /// Returns a raw pointer to the allocated memory.
-    pub fn calloc(&self, size: usize) -> *mut c_void {
-        unsafe { ngx_pcalloc(self.0.as_ptr(), size) }
-    }
-
-    /// Allocates zeroed memory for a type from the pool.
-    /// The resulting pointer is aligned to a platform word size.
-    ///
-    /// Returns a typed pointer to the allocated memory.
-    pub fn calloc_type<T: Copy>(&self) -> *mut T {
-        self.calloc(mem::size_of::<T>()) as *mut T
-    }
-
-    /// Allocates unaligned memory from the pool of the specified size.
-    ///
-    /// Returns a raw pointer to the allocated memory.
-    pub fn alloc_unaligned(&self, size: usize) -> *mut c_void {
-        unsafe { ngx_pnalloc(self.0.as_ptr(), size) }
-    }
-
-    /// Allocates unaligned memory for a type from the pool.
-    ///
-    /// Returns a typed pointer to the allocated memory.
-    pub fn alloc_type_unaligned<T: Copy>(&self) -> *mut T {
-        self.alloc_unaligned(mem::size_of::<T>()) as *mut T
-    }
-
     /// Allocates memory for a value of a specified type and adds a cleanup handler to the memory
     /// pool.
     ///
@@ -284,6 +243,33 @@ impl Pool {
         }
     }
 
+    /// Allocates unaligned memory from the pool of the specified size.
+    ///
+    /// Returns Result::Ok with a typed pointer to the allocated memory if successful,
+    /// or Result::Err(AllocError) if allocation or cleanup handler addition fails.
+    pub fn allocate_unaligned(&self, size: usize) -> Result<NonNull<[u8]>, AllocError> {
+        self.allocate(unsafe { Layout::from_size_align_unchecked(size, 1) })
+    }
+
+    /// Allocates memory for a type from the pool.
+    /// The resulting pointer is aligned to a platform word size.
+    ///
+    /// Returns Result::Ok with a typed pointer to the allocated memory if successful,
+    /// or Result::Err(AllocError) if allocation fails.
+    pub fn allocate_type<T>(&self) -> Result<NonNull<T>, AllocError> {
+        self.allocate(Layout::new::<T>()).map(|ptr| ptr.cast())
+    }
+
+    /// Allocates zeroed memory for a type from the pool.
+    /// The resulting pointer is aligned to a platform word size.
+    ///
+    /// Returns Result::Ok with a typed pointer to the allocated memory if successful,
+    /// or Result::Err(AllocError) if allocation fails.
+    pub fn allocate_type_zeroed<T>(&self) -> Result<NonNull<T>, AllocError> {
+        self.allocate_zeroed(Layout::new::<T>())
+            .map(|ptr| ptr.cast())
+    }
+
     /// Resizes a memory allocation in place if possible.
     ///
     /// If resizing is requested for the last allocation in the pool, it may be
diff --git a/src/http/upstream.rs b/src/http/upstream.rs
@@ -16,11 +16,11 @@ macro_rules! http_upstream_init_peer_pt {
             r: *mut $crate::ffi::ngx_http_request_t,
             us: *mut $crate::ffi::ngx_http_upstream_srv_conf_t,
         ) -> $crate::ffi::ngx_int_t {
-            let status: $crate::core::Status = $handler(
+            let res: $crate::core::NgxResult = $handler(
                 unsafe { &mut $crate::http::Request::from_ngx_http_request(r) },
                 us,
             );
-            status.0
+            res.map_or_else(|_| $crate::core::Status::NGX_ERROR.into(), |code| code)
         }
     };
 }

Original file line number	Diff line number	Diff line change
`@@ -16,11 +16,11 @@ macro_rules! http_upstream_init_peer_pt {`
`16`	`16`	`r: *mut $crate::ffi::ngx_http_request_t,`
`17`	`17`	`us: *mut $crate::ffi::ngx_http_upstream_srv_conf_t,`
`18`	`18`	`) -> $crate::ffi::ngx_int_t {`
`19`		`- let status: $crate::core::Status = $handler(`
	`19`	`+ let res: $crate::core::NgxResult = $handler(`
`20`	`20`	`unsafe { &mut $crate::http::Request::from_ngx_http_request(r) },`
`21`	`21`	`us,`
`22`	`22`	`);`
`23`		`- status.0`
	`23`	`+ res.map_or_else(\|_\| $crate::core::Status::NGX_ERROR.into(), \|code\| code)`
`24`	`24`	`}`
`25`	`25`	`};`
`26`	`26`	`}`