Allow set-body to replace origin server response bodies

doc/admin-guide/plugins/header_rewrite.en.rst

@@ -1109,35 +1109,76 @@ set-body
 
   set-body <text>
 
-Sets the body to ``<text>``. Can also be used to delete a body with ``""``. This is only useful when overriding the origin status, i.e.
-intercepting/pre-empting a request so that you can override the body from the body-factory with your own.
+Sets the response body to ``<text>``. Can also be used to delete a body with ``""``.
+
+This operator can be used to replace the response body in three scenarios:
+
+1. **Synthetic responses (no origin connection)**: When used at ``REMAP_PSEUDO_HOOK``
+   (before the origin connection is established), ``set-body`` will skip the origin
+   connection entirely and serve a synthetic response directly. Use ``set-status`` to
+   control the response status code (defaults to 200 OK if not specified). This is
+   useful for blocking requests, serving synthetic pages, or returning canned responses
+   without touching the origin.
+
+2. **ATS-generated responses**: When overriding the origin status at remap time (e.g.
+   with ``set-status``), you can use ``set-body`` at ``SEND_RESPONSE_HDR_HOOK`` to
+   override the body from the body-factory with your own.
+
+3. **Origin server responses**: When the origin returns a response with a body,
+   ``set-body`` can replace the origin body with the specified text. This is useful
+   for sanitizing error responses (e.g. replacing a 403 body that contains sensitive
+   information). Use ``READ_RESPONSE_HDR_HOOK`` to inspect the origin response
+   headers and replace the body, or ``SEND_RESPONSE_HDR_HOOK`` to replace the body
+   just before sending to the client.
+
+When replacing an origin response body (scenario 3), ATS will attempt to drain the
+origin body from the server connection buffer so the connection can be reused. If the
+origin body is too large to be fully buffered or uses chunked encoding, the server
+connection will be closed instead.
+
+Example: block a request at remap time without contacting the origin::
+
+   cond %{REMAP_PSEUDO_HOOK} [AND]
+   cond %{CLIENT-URL:PATH} /blocked
+       set-status 403
+       set-body "Access Denied"
+
+Example: sanitize a 403 response from the origin::
+
+   cond %{READ_RESPONSE_HDR_HOOK} [AND]
+   cond %{STATUS} =403
+       set-body "Access Denied"
 
 set-body-from
 ~~~~~~~~~~~~~
 ::
 
   set-body-from <URL>
 
-Will call ``<URL>`` (see URL in `URL Parts`_) to retrieve a custom error response
-and set the body with the result. Triggering this rule on an OK transaction will
-send a 500 status code to the client with the desired response. If this is triggered
-on any error status code, that original status code will be sent to the client.
+Will call ``<URL>`` (see URL in `URL Parts`_) to retrieve a response body from a
+secondary URL and use it to replace the origin's response body. The origin's response
+status code and headers are preserved, and the fetched content replaces only the body.
 
-.. note::
-    This config should only be set using READ_RESPONSE_HDR_HOOK
+This operator can be used at ``READ_RESPONSE_HDR_HOOK`` or ``SEND_RESPONSE_HDR_HOOK``.
+When the origin response body is fully buffered and has a known ``Content-Length``,
+ATS will drain the origin body to allow connection reuse. Otherwise, the origin
+connection is closed.
+
+If the fetch fails or times out, the original origin response is sent unmodified.
 
 An example config would look like::
 
-   cond %{READ_RESPONSE_HDR_HOOK}
-      set-body-from http://www.example.com/second
+   cond %{READ_RESPONSE_HDR_HOOK} [AND]
+   cond %{STATUS} =403
+      set-body-from http://www.example.com/custom-error-page
 
-Where ``http://www.example.com/second`` is the destination to retrieve the custom response from.
-This can be enabled per-mapping or globally.
-Ensure there is a remap rule for the second endpoint as well!
+Where ``http://www.example.com/custom-error-page`` is the destination to retrieve the
+custom response body from. This can be enabled per-mapping or globally.
+Ensure there is a remap rule for the secondary endpoint as well!
 An example remap config would look like::
 
    map /first http://www.example.com/first @plugin=header_rewrite.so @pparam=cond1.conf
-   map /second http://www.example.com/second
+   map /custom-error-page http://www.example.com/custom-error-page
 
 set-config
 ~~~~~~~~~~

include/proxy/http/HttpSM.h

@@ -421,6 +421,7 @@ class HttpSM : public Continuation, public PluginUserArgs<TS_USER_ARGS_TXN>
   void do_redirect();
   void redirect_request(const char *redirect_url, const int redirect_len);
   void do_drain_request_body(HTTPHdr &response);
+  void do_drain_server_response_body();
 
   void wait_for_full_body();
 

plugins/header_rewrite/operators.cc

@@ -68,7 +68,7 @@ handleFetchEvents(TSCont cont, TSEvent event, void *edata)
     } else {
       TSWarning("[%s] Successful set-custom-body fetch did not result in any content", __FUNCTION__);
     }
-    TSHttpTxnReenable(http_txn, TS_EVENT_HTTP_ERROR);
+    TSHttpTxnReenable(http_txn, TS_EVENT_HTTP_CONTINUE);
   } break;
   case OperatorSetBodyFrom::TS_EVENT_FETCHSM_FAILURE: {
     Dbg(pi_dbg_ctl, "OperatorSetBodyFrom: Error getting custom body");
@@ -813,6 +813,7 @@ void
 OperatorSetBody::initialize_hooks()
 {
   add_allowed_hook(TS_REMAP_PSEUDO_HOOK);
+  add_allowed_hook(TS_HTTP_READ_RESPONSE_HDR_HOOK);
   add_allowed_hook(TS_HTTP_SEND_RESPONSE_HDR_HOOK);
 }
 
@@ -1356,6 +1357,7 @@ void
 OperatorSetBodyFrom::initialize_hooks()
 {
   add_allowed_hook(TS_HTTP_READ_RESPONSE_HDR_HOOK);
+  add_allowed_hook(TS_HTTP_SEND_RESPONSE_HDR_HOOK);
 }
 
 bool
@@ -1387,11 +1389,6 @@ OperatorSetBodyFrom::exec(const Resources &res) const
     addr.sin_port        = LOCAL_PORT;
     TSFetchUrl(static_cast<const char *>(req_buf), req_buf_size, reinterpret_cast<struct sockaddr const *>(&addr), fetchCont,
                AFTER_BODY, event_ids);
-
-    // Forces original status code in event TSHttpTxnErrorBodySet changed
-    // the code or another condition was set conflicting with this one.
-    // Set here because res is the only structure that contains the original status code.
-    TSHttpTxnStatusSet(res.state.txnp, res.resp_status, PLUGIN_NAME);
   } else {
     TSError(PLUGIN_NAME, "OperatorSetBodyFrom:exec:: Could not create request");
     return true;

src/proxy/http/HttpSM.cc

@@ -1675,9 +1675,20 @@ HttpSM::handle_api_return()
 
   switch (t_state.next_action) {
   case HttpTransact::StateMachineAction_t::TRANSFORM_READ: {
-    HttpTunnelProducer *p = setup_transfer_from_transform();
-    perform_transform_cache_write_action();
-    tunnel.tunnel_run(p);
+    if (t_state.internal_msg_buffer) {
+      // A plugin replaced the response body via TSHttpTxnErrorBodySet().
+      // Use internal transfer instead of the transform tunnel.
+      SMDbg(dbg_ctl_http, "plugin set internal body, bypassing transform for internal transfer");
+      if (server_txn != nullptr) {
+        do_drain_server_response_body();
+        release_server_session();
+      }
+      setup_internal_transfer(&HttpSM::tunnel_handler);
+    } else {
+      HttpTunnelProducer *p = setup_transfer_from_transform();
+      perform_transform_cache_write_action();
+      tunnel.tunnel_run(p);
+    }
     break;
   }
   case HttpTransact::StateMachineAction_t::SERVER_READ: {
@@ -1709,6 +1720,15 @@ HttpSM::handle_api_return()
       }
 
       setup_blind_tunnel(true, initial_data);
+    } else if (t_state.internal_msg_buffer) {
+      // A plugin replaced the origin response body via TSHttpTxnErrorBodySet().
+      // Drain the origin body if possible, then use internal transfer.
+      SMDbg(dbg_ctl_http, "plugin set internal body, using internal transfer instead of server tunnel");
+      if (server_txn != nullptr) {
+        do_drain_server_response_body();
+        release_server_session();
+      }
+      setup_internal_transfer(&HttpSM::tunnel_handler);
     } else {
       HttpTunnelProducer *p = setup_server_transfer();
       perform_cache_write_action();
@@ -6331,6 +6351,57 @@ HttpSM::do_drain_request_body(HTTPHdr &response)
   _ua.get_txn()->set_close_connection(response);
 }
 
+//
+// void HttpSM::do_drain_server_response_body()
+//
+//  Attempt to synchronously drain the origin server response body
+//  so the connection can be returned to the pool. If the body cannot
+//  be drained (chunked, unknown length, or not fully received),
+//  mark the server connection as no-keepalive so it will be closed.
+//
+void
+HttpSM::do_drain_server_response_body()
+{
+  if (t_state.current.server == nullptr || server_txn == nullptr) {
+    return;
+  }
+
+  int64_t content_length = t_state.hdr_info.response_content_length;
+
+  if (content_length == HTTP_UNDEFINED_CL) {
+    // Chunked or unknown length -- can't drain synchronously
+    SMDbg(dbg_ctl_http, "server response body drain: chunked/unknown length, closing connection");
+    t_state.current.server->keep_alive = HTTPKeepAlive::NO_KEEPALIVE;
+    return;
+  }
+
+  if (content_length == 0) {
+    // No body to drain
+    SMDbg(dbg_ctl_http, "server response body drain: zero length, connection reusable");
+    return;
+  }
+
+  int64_t avail = server_txn->get_remote_reader()->read_avail();
+
+  if (avail >= content_length) {
+    // Entire body is in the buffer -- consume it so the connection can be reused
+    server_txn->get_remote_reader()->consume(content_length);
+    SMDbg(dbg_ctl_http, "server response body drain: consumed %" PRId64 " bytes", content_length);
+
+    // Verify origin didn't send more than Content-Length (protocol violation).
+    // Same check as server_transfer_init().
+    if (server_txn->get_remote_reader()->read_avail() > 0) {
+      SMDbg(dbg_ctl_http, "server response body drain: extra data after Content-Length, closing connection");
+      t_state.current.server->keep_alive = HTTPKeepAlive::NO_KEEPALIVE;
+    }
+  } else {
+    // Body not fully received -- close the connection
+    SMDbg(dbg_ctl_http, "server response body drain: only %" PRId64 " of %" PRId64 " bytes available, closing connection", avail,
+          content_length);
+    t_state.current.server->keep_alive = HTTPKeepAlive::NO_KEEPALIVE;
+  }
+}
+
 void
 HttpSM::do_setup_client_request_body_tunnel(HttpVC_t to_vc_type)
 {
@@ -8487,6 +8558,10 @@ HttpSM::redirect_request(const char *arg_redirect_url, const int arg_redirect_le
     // XXX - doing a destroy() for now, we can do a fileds_clear() if we have performance issue
     t_state.hdr_info.client_response.destroy();
   }
+  // Clear any error body from a previous failed connection attempt (e.g., from
+  // build_error_response) so that how_to_open_connection() does not mistake it
+  // for a plugin-set synthetic body and short-circuit the retry.
+  t_state.free_internal_msg_buffer();
 
   int         scheme          = t_state.next_hop_scheme;
   int         scheme_len      = hdrtoken_index_to_length(scheme);

src/proxy/http/HttpTransact.cc

@@ -786,6 +786,19 @@ how_to_open_connection(HttpTransact::State *s)
 {
   ink_assert((s->pending_work == nullptr) || (s->current.request_to == ResolveInfo::PARENT_PROXY));
 
+  // If a plugin has already set a response body (e.g., set-body at remap),
+  // skip the origin connection and serve the synthetic response directly.
+  if (s->internal_msg_buffer) {
+    HTTPStatus  status = (s->http_return_code != HTTPStatus::NONE) ? s->http_return_code : HTTPStatus::OK;
+    const char *reason = http_hdr_reason_lookup(status);
+    if (s->hdr_info.client_response.valid()) {
+      s->hdr_info.client_response.fields_clear();
+    }
+    HttpTransact::build_response(s, &s->hdr_info.client_response, s->client_info.http_version, status, reason ? reason : "OK");
+    s->source = HttpTransact::Source_t::INTERNAL;
+    return HttpTransact::StateMachineAction_t::INTERNAL_CACHE_NOOP;
+  }
+
   // Originally we returned which type of server to open
   // Now, however, we may want to issue a cache
   // operation first in order to lock the cache
@@ -1220,8 +1233,19 @@ HttpTransact::EndRemapRequest(State *s)
     Metrics::Counter::increment(http_rsb.invalid_client_requests);
     TRANSACT_RETURN(StateMachineAction_t::SEND_ERROR_CACHE_NOOP, nullptr);
   } else {
-    s->hdr_info.client_response.destroy(); // release the underlying memory.
-    s->hdr_info.client_response.clear();   // clear the pointers.
+    // This else branch handles two cases:
+    // 1. Remap succeeded (reverse_proxy == true) - normal request processing
+    // 2. Remap failed but plugin tunnel exists - plugin overrides error
+    //
+    // For case 2, clear the stale error response that build_error_response()
+    // created during remap failure, because the plugin tunnel will provide
+    // its own response. For case 1, preserve any plugin-set internal_msg_buffer
+    // (e.g., from set-body at remap) so how_to_open_connection() can short-circuit.
+    if (s->state_machine->plugin_tunnel_type != HttpPluginTunnel_t::NONE) {
+      s->hdr_info.client_response.destroy(); // release the underlying memory.
+      s->hdr_info.client_response.clear();   // clear the pointers.
+      s->free_internal_msg_buffer();         // clear error body so plugin tunnel is not bypassed.
+    }
     TxnDbg(dbg_ctl_http_trans, "END HttpTransact::EndRemapRequest");
 
     if (s->is_upgrade_request && s->post_remap_upgrade_return_point) {