|
| 1 | +use core::fmt::Display; |
| 2 | + |
| 3 | +use ngx::core::Status; |
| 4 | +use ngx::http::subrequest::{SubRequestBuilder, SubRequestError}; |
| 5 | +use ngx::http::{ |
| 6 | + HTTPStatus, HttpModule, HttpModuleLocationConf, HttpPhase, HttpRequestHandler, Merge, |
| 7 | + MergeConfigError, Request, add_phase_handler, |
| 8 | +}; |
| 9 | +use ngx::{ngx_log_debug_http, ngx_log_error}; |
| 10 | + |
| 11 | +use nginx_sys::{ |
| 12 | + NGX_CONF_TAKE1, NGX_ERROR, NGX_HTTP_LOC_CONF, NGX_HTTP_LOC_CONF_OFFSET, ngx_command_t, |
| 13 | + ngx_conf_t, ngx_flag_t, ngx_http_complex_value_t, ngx_http_module_t, ngx_http_request_t, |
| 14 | + ngx_http_send_response, ngx_int_t, ngx_module_t, ngx_str_t, ngx_uint_t, |
| 15 | +}; |
| 16 | + |
| 17 | +const NGX_CONF_UNSET_FLAG: ngx_flag_t = nginx_sys::NGX_CONF_UNSET as _; |
| 18 | + |
| 19 | +struct SampleHandler; |
| 20 | + |
| 21 | +enum SampleHandlerError { |
| 22 | + ContextAllocation, |
| 23 | + SubRequestCreation(SubRequestError), |
| 24 | + SubRequest(ngx_int_t), |
| 25 | + Response(ngx_int_t), |
| 26 | +} |
| 27 | + |
| 28 | +impl From<SubRequestError> for SampleHandlerError { |
| 29 | + fn from(e: SubRequestError) -> Self { |
| 30 | + SampleHandlerError::SubRequestCreation(e) |
| 31 | + } |
| 32 | +} |
| 33 | + |
| 34 | +impl Display for SampleHandlerError { |
| 35 | + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { |
| 36 | + match self { |
| 37 | + SampleHandlerError::ContextAllocation => { |
| 38 | + write!(f, "context allocation failed") |
| 39 | + } |
| 40 | + SampleHandlerError::SubRequestCreation(e) => { |
| 41 | + write!(f, "subrequest creation failed: {}", e) |
| 42 | + } |
| 43 | + SampleHandlerError::SubRequest(rc) => { |
| 44 | + write!(f, "subrequest failed with return code: {}", rc) |
| 45 | + } |
| 46 | + SampleHandlerError::Response(rc) => { |
| 47 | + write!(f, "response creation failed with return code: {}", rc) |
| 48 | + } |
| 49 | + } |
| 50 | + } |
| 51 | +} |
| 52 | + |
| 53 | +impl HttpRequestHandler for SampleHandler { |
| 54 | + const PHASE: HttpPhase = HttpPhase::Access; |
| 55 | + type Output = Result<Status, SampleHandlerError>; |
| 56 | + |
| 57 | + fn handler(request: &mut Request) -> Self::Output { |
| 58 | + let co = Module::location_conf(request).expect("module config is none"); |
| 59 | + ngx_log_debug_http!(request, "subrequest module enabled: {}", co.enable); |
| 60 | + |
| 61 | + if co.enable != 1 { |
| 62 | + return Ok(Status::NGX_DECLINED); |
| 63 | + } |
| 64 | + |
| 65 | + let rptr: *mut ngx_http_request_t = request.as_mut(); |
| 66 | + |
| 67 | + match SRCtx::get(request) { |
| 68 | + Some(ctx) => ctx.rc.map_or( |
| 69 | + // `ctx` has been created but not filled yet - subrequest is still in progress |
| 70 | + Ok(Status::NGX_AGAIN), |
| 71 | + // `ctx` has been created and filled - subrequest is completed |
| 72 | + |rc| { |
| 73 | + let status = ctx.status.0; |
| 74 | + let msg = format!("subrequest completed with HTTP status: {status}, rc: {rc}"); |
| 75 | + ngx_log_debug_http!(request, "{msg}"); |
| 76 | + |
| 77 | + if status >= nginx_sys::NGX_HTTP_SPECIAL_RESPONSE as _ { |
| 78 | + Ok(Status::from(ctx.status)) |
| 79 | + } else if rc == nginx_sys::NGX_OK as _ && ctx.out.is_some() { |
| 80 | + let outbuf = unsafe { &*ctx.out.unwrap().buf }; |
| 81 | + let mut ct = ctx.ct; |
| 82 | + let mut cv: ngx_http_complex_value_t = unsafe { core::mem::zeroed() }; |
| 83 | + cv.value = ngx_str_t { |
| 84 | + len: unsafe { outbuf.last.offset_from(outbuf.pos) } as _, |
| 85 | + data: outbuf.pos as _, |
| 86 | + }; |
| 87 | + let resp_rc = unsafe { |
| 88 | + ngx_http_send_response(rptr, status, &raw mut ct, &raw mut cv) |
| 89 | + }; |
| 90 | + if resp_rc == nginx_sys::NGX_OK as _ { |
| 91 | + Ok(Status::from(ctx.status)) |
| 92 | + } else { |
| 93 | + Err(SampleHandlerError::Response(resp_rc)) |
| 94 | + } |
| 95 | + } else if rc == nginx_sys::NGX_OK as _ { |
| 96 | + Ok(Status::from(ctx.status)) |
| 97 | + } else if let Ok(http_status) = HTTPStatus::try_from(rc) { |
| 98 | + Ok(Status::from(http_status)) |
| 99 | + } else { |
| 100 | + Err(SampleHandlerError::SubRequest(rc)) |
| 101 | + } |
| 102 | + }, |
| 103 | + ), |
| 104 | + None => { |
| 105 | + if SRCtx::create(request).is_some() { |
| 106 | + let uri: &str = co.uri.to_str().unwrap_or("/proxy"); |
| 107 | + |
| 108 | + SubRequestBuilder::new(request.pool(), uri)? |
| 109 | + .args("arg1=val1&arg2=val2")? |
| 110 | + .in_memory() |
| 111 | + .waited() |
| 112 | + .build(request, sr_handler)?; |
| 113 | + |
| 114 | + Ok(Status::NGX_AGAIN) |
| 115 | + } else { |
| 116 | + Err(SampleHandlerError::ContextAllocation) |
| 117 | + } |
| 118 | + } |
| 119 | + } |
| 120 | + } |
| 121 | +} |
| 122 | + |
| 123 | +struct SRCtx<'r> { |
| 124 | + rc: Option<ngx_int_t>, |
| 125 | + status: HTTPStatus, |
| 126 | + out: Option<&'r nginx_sys::ngx_chain_t>, |
| 127 | + ct: ngx_str_t, |
| 128 | +} |
| 129 | + |
| 130 | +impl SRCtx<'_> { |
| 131 | + fn create(request: &mut Request) -> Option<&mut Self> { |
| 132 | + let ctx_ref = unsafe { |
| 133 | + request |
| 134 | + .pool() |
| 135 | + .allocate_with_cleanup(Self::default)? |
| 136 | + .as_mut() |
| 137 | + }; |
| 138 | + request.set_module_ctx(ctx_ref as *mut _ as _, Module::module()); |
| 139 | + Some(ctx_ref) |
| 140 | + } |
| 141 | + |
| 142 | + fn get(request: &Request) -> Option<&Self> { |
| 143 | + request.get_module_ctx::<Self>(Module::module()) |
| 144 | + } |
| 145 | + |
| 146 | + fn get_mut(request: &mut Request) -> Option<&mut Self> { |
| 147 | + request.get_module_ctx_mut::<Self>(Module::module()) |
| 148 | + } |
| 149 | +} |
| 150 | + |
| 151 | +impl Default for SRCtx<'_> { |
| 152 | + fn default() -> Self { |
| 153 | + Self { |
| 154 | + rc: None, |
| 155 | + status: HTTPStatus(NGX_ERROR as _), |
| 156 | + out: None, |
| 157 | + ct: ngx_str_t::empty(), |
| 158 | + } |
| 159 | + } |
| 160 | +} |
| 161 | + |
| 162 | +fn sr_handler(r: &mut Request, mut rc: ngx_int_t) -> ngx_int_t { |
| 163 | + let newctx = SRCtx { |
| 164 | + rc: Some(rc), |
| 165 | + status: r.get_status(), |
| 166 | + out: core::ptr::NonNull::new(r.as_ref().out).map(|out| unsafe { out.as_ref() }), |
| 167 | + ct: r.as_ref().headers_out.content_type, |
| 168 | + }; |
| 169 | + if let Some(ctx) = SRCtx::get_mut(r.get_main_mut()) { |
| 170 | + *ctx = newctx; |
| 171 | + } else { |
| 172 | + ngx_log_error!( |
| 173 | + nginx_sys::NGX_LOG_ERR, |
| 174 | + r.log(), |
| 175 | + "subrequest: context not found" |
| 176 | + ); |
| 177 | + rc = NGX_ERROR as _; |
| 178 | + } |
| 179 | + rc |
| 180 | +} |
| 181 | + |
| 182 | +static NGX_HTTP_SUBREQUEST_MODULE_CTX: ngx_http_module_t = ngx_http_module_t { |
| 183 | + preconfiguration: None, |
| 184 | + postconfiguration: Some(Module::postconfiguration), |
| 185 | + create_main_conf: None, |
| 186 | + init_main_conf: None, |
| 187 | + create_srv_conf: None, |
| 188 | + merge_srv_conf: None, |
| 189 | + create_loc_conf: Some(Module::create_loc_conf), |
| 190 | + merge_loc_conf: Some(Module::merge_loc_conf), |
| 191 | +}; |
| 192 | + |
| 193 | +#[cfg(feature = "export-modules")] |
| 194 | +ngx::ngx_modules!(ngx_http_subrequest_module); |
| 195 | + |
| 196 | +#[used] |
| 197 | +#[allow(non_upper_case_globals)] |
| 198 | +#[cfg_attr(not(feature = "export-modules"), unsafe(no_mangle))] |
| 199 | +pub static mut ngx_http_subrequest_module: ngx_module_t = ngx_module_t { |
| 200 | + ctx: &raw const NGX_HTTP_SUBREQUEST_MODULE_CTX as _, |
| 201 | + commands: unsafe { &raw mut NGX_HTTP_SUBREQUEST_COMMANDS[0] }, |
| 202 | + type_: nginx_sys::NGX_HTTP_MODULE as _, |
| 203 | + ..ngx_module_t::default() |
| 204 | +}; |
| 205 | + |
| 206 | +struct Module; |
| 207 | + |
| 208 | +impl HttpModule for Module { |
| 209 | + fn module() -> &'static ngx_module_t { |
| 210 | + unsafe { &*::core::ptr::addr_of!(ngx_http_subrequest_module) } |
| 211 | + } |
| 212 | + |
| 213 | + unsafe extern "C" fn postconfiguration(cf: *mut ngx_conf_t) -> ngx_int_t { |
| 214 | + // SAFETY: this function is called with non-NULL cf always |
| 215 | + let cf = unsafe { &mut *cf }; |
| 216 | + add_phase_handler::<SampleHandler>(cf) |
| 217 | + .map_or(nginx_sys::NGX_ERROR as _, |_| nginx_sys::NGX_OK as _) |
| 218 | + } |
| 219 | +} |
| 220 | + |
| 221 | +#[derive(Debug)] |
| 222 | +struct ModuleConfig { |
| 223 | + enable: ngx_flag_t, |
| 224 | + uri: ngx_str_t, |
| 225 | +} |
| 226 | + |
| 227 | +impl Default for ModuleConfig { |
| 228 | + fn default() -> Self { |
| 229 | + Self { |
| 230 | + enable: NGX_CONF_UNSET_FLAG, |
| 231 | + uri: ngx_str_t::empty(), |
| 232 | + } |
| 233 | + } |
| 234 | +} |
| 235 | + |
| 236 | +impl Merge for ModuleConfig { |
| 237 | + fn merge(&mut self, prev: &ModuleConfig) -> Result<(), MergeConfigError> { |
| 238 | + if self.enable == NGX_CONF_UNSET_FLAG { |
| 239 | + if prev.enable != NGX_CONF_UNSET_FLAG { |
| 240 | + self.enable = prev.enable; |
| 241 | + } else { |
| 242 | + self.enable = 0; |
| 243 | + } |
| 244 | + } |
| 245 | + if self.uri.data.is_null() { |
| 246 | + self.uri = prev.uri; |
| 247 | + } |
| 248 | + if self.enable == 1 && self.uri.data.is_null() { |
| 249 | + self.uri = ngx::ngx_string!("/proxy"); |
| 250 | + } |
| 251 | + Ok(()) |
| 252 | + } |
| 253 | +} |
| 254 | + |
| 255 | +unsafe impl HttpModuleLocationConf for Module { |
| 256 | + type LocationConf = ModuleConfig; |
| 257 | +} |
| 258 | + |
| 259 | +static mut NGX_HTTP_SUBREQUEST_COMMANDS: [ngx_command_t; 3] = [ |
| 260 | + ngx_command_t { |
| 261 | + name: ngx::ngx_string!("subrequest"), |
| 262 | + type_: (NGX_HTTP_LOC_CONF | NGX_CONF_TAKE1) as ngx_uint_t, |
| 263 | + set: Some(nginx_sys::ngx_conf_set_flag_slot), |
| 264 | + conf: NGX_HTTP_LOC_CONF_OFFSET, |
| 265 | + offset: core::mem::offset_of!(ModuleConfig, enable), |
| 266 | + post: core::ptr::null_mut(), |
| 267 | + }, |
| 268 | + ngx_command_t { |
| 269 | + name: ngx::ngx_string!("subrequest_uri"), |
| 270 | + type_: (NGX_HTTP_LOC_CONF | NGX_CONF_TAKE1) as ngx_uint_t, |
| 271 | + set: Some(nginx_sys::ngx_conf_set_str_slot), |
| 272 | + conf: NGX_HTTP_LOC_CONF_OFFSET, |
| 273 | + offset: core::mem::offset_of!(ModuleConfig, uri), |
| 274 | + post: core::ptr::null_mut(), |
| 275 | + }, |
| 276 | + ngx_command_t::empty(), |
| 277 | +]; |
0 commit comments