subrequest

apache · Aug 27, 2024 · 530448f · 530448f
1 parent 2d0a7a1
commit 530448f
Show file tree

Hide file tree

Showing 4 changed files with 88 additions and 3 deletions.
diff --git a/apisix/cli/ngx_tpl.lua b/apisix/cli/ngx_tpl.lua
@@ -809,6 +809,27 @@ http {
             }
         }
 
+        location /subrequest {
+            internal;
+
+            proxy_http_version 1.1;
+            proxy_set_header   Host              $upstream_host;
+            proxy_set_header   Upgrade           $upstream_upgrade;
+            proxy_set_header   Connection        $upstream_connection;
+            proxy_set_header   X-Real-IP         $remote_addr;
+
+            proxy_set_header   X-Forwarded-For      $proxy_add_x_forwarded_for;
+            proxy_set_header   X-Forwarded-Proto    $scheme;
+            proxy_set_header   X-Forwarded-Host     $host;
+            proxy_set_header   X-Forwarded-Port     $server_port;
+
+            proxy_pass_header     Server;
+            proxy_pass_header     Date;
+            proxy_ssl_name        $upstream_host;
+            proxy_ssl_server_name on;
+            proxy_pass      $upstream_scheme://apisix_backend$upstream_uri;
+        }
+
         location @disable_proxy_buffering {
             # http server location configuration snippet starts
             {% if http_server_location_configuration_snippet then %}

diff --git a/apisix/init.lua b/apisix/init.lua
@@ -458,6 +458,45 @@ local function common_phase(phase_name)
     return plugin.run_plugin(phase_name, nil, api_ctx)
 end
 
+local methods_map = {
+    GET       = ngx.HTTP_GET,
+    PUT       = ngx.HTTP_PUT,
+    POST      = ngx.HTTP_POST,
+    PATCH     = ngx.HTTP_PATCH,
+    DELETE    = ngx.HTTP_DELETE,
+    OPTIONS   = ngx.HTTP_OPTIONS,
+    TRACE     = ngx.HTTP_TRACE,
+    HEAD      = ngx.HTTP_HEAD,
+}
+
+local function subrequest(api_ctx)
+    ngx.req.read_body()
+    local options = {
+      always_forward_body = true,
+      share_all_vars = true,
+      method = methods_map[ngx.req.get_method()],
+      ctx = ngx.ctx,
+    }
+
+    local res = ngx.location.capture("/subrequest", options)
+    if not res or res.truncated then
+        return core.response.exit(502)
+    end
+
+    if res.truncated and options.method ~= ngx.HTTP_HEAD then
+        return core.response.exit(503)
+    end
+
+    api_ctx.subreq_status = res.status
+    api_ctx.subreq_headers = res.header
+    api_ctx.subreq_body = res.body
+
+    for key, value in pairs(res.header) do
+        core.response.set_header(key, value)
+    end
+    core.log.info("finishing subrequest")
+    core.response.exit(res.status, res.body)
+end
 
 
 function _M.handle_upstream(api_ctx, route, enable_websocket)
@@ -723,8 +762,10 @@ function _M.http_access_phase()
     end
 
     _M.handle_upstream(api_ctx, route, enable_websocket)
-
-    if ngx.ctx.disable_proxy_buffering then
+    if api_ctx.subrequest then
+        subrequest(api_ctx)
+    end
+    if api_ctx.disable_proxy_buffering then
         stash_ngx_ctx()
         return ngx.exec("@disable_proxy_buffering")
     end

diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua
@@ -64,7 +64,9 @@ function _M.access(conf, ctx)
 
     if conf.model.options and conf.model.options.stream then
         request_table.stream = true
-        ngx.ctx.disable_proxy_buffering = true
+        ctx.disable_proxy_buffering = true
+    else
+        ctx.subrequest = true
     end
 
     if conf.model.name then

diff --git a/t/APISIX.pm b/t/APISIX.pm
@@ -778,6 +778,27 @@ _EOC_
             }
         }
 
+        location /subrequest {
+            internal;
+
+            proxy_http_version 1.1;
+            proxy_set_header   Host              \$upstream_host;
+            proxy_set_header   Upgrade           \$upstream_upgrade;
+            proxy_set_header   Connection        \$upstream_connection;
+            proxy_set_header   X-Real-IP         \$remote_addr;
+
+            proxy_set_header   X-Forwarded-For      \$proxy_add_x_forwarded_for;
+            proxy_set_header   X-Forwarded-Proto    \$scheme;
+            proxy_set_header   X-Forwarded-Host     \$host;
+            proxy_set_header   X-Forwarded-Port     \$server_port;
+
+            proxy_pass_header     Server;
+            proxy_pass_header     Date;
+            proxy_ssl_name        \$upstream_host;
+            proxy_ssl_server_name on;
+            proxy_pass      \$upstream_scheme://apisix_backend\$upstream_uri;
+        }
+
         location / {
             set \$upstream_mirror_host        '';
             set \$upstream_mirror_uri         '';