header_rewrite: fix set-body origin replacement across hooks and transforms

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

@@ -1177,8 +1177,24 @@ 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 body to ``<text>``. Can also be used to delete a body with ``""``.
+
+For origin response replacement, ``set-body`` is supported at both
+``READ_RESPONSE_HDR_HOOK`` and ``SEND_RESPONSE_HDR_HOOK``. Prefer
+``READ_RESPONSE_HDR_HOOK`` when possible so body replacement happens before
+response body tunneling starts.
+
+.. note::
+
+   When ``set-body`` replaces an origin response body, ATS emits the replacement
+   through its internal error-body path. ``Content-Type`` defaults to
+   ``text/html`` unless you override it with ``set-header Content-Type``.
+   ``set-body ""`` clears the internal replacement body, but does not suppress an
+   origin response body on this hook; use a non-empty replacement value when
+   sanitizing origin responses.
+   The gold tests cover origin replacement for both hooks with and without a
+   response transform plugin. The no-transform matrix runs with HTTP cache
+   disabled and includes repeated-URL cache-bypass probes.
 
 set-body-from
 ~~~~~~~~~~~~~

include/proxy/http/HttpSM.h

@@ -500,6 +500,12 @@ class HttpSM : public Continuation, public PluginUserArgs<TS_USER_ARGS_TXN>
   int       get_request_method_wksidx() const;
 
 public:
+  TSHttpHookID
+  get_cur_hook_id() const
+  {
+    return cur_hook_id;
+  }
+
   // TODO:  Now that bodies can be empty, should the body counters be set to -1 ? TS-2213
   // Stats & Logging Info
   int     client_request_hdr_bytes        = 0;

include/proxy/http/HttpTransact.h

@@ -761,6 +761,10 @@ class HttpTransact
     char   *internal_msg_buffer_type                = nullptr; // out
     int64_t internal_msg_buffer_size                = 0;       // out
     int64_t internal_msg_buffer_fast_allocator_size = -1;
+    // Hook that was firing when a plugin called TSHttpTxnErrorBodySet().
+    // Used to scope the SERVER_READ/TRANSFORM_READ internal-body divert to
+    // plugin-driven response replacement only.
+    TSHttpHookID internal_msg_buffer_set_on = TS_HTTP_LAST_HOOK;
 
     int  scheme                    = -1;     // out
     int  next_hop_scheme           = scheme; // out
@@ -934,7 +938,8 @@ class HttpTransact
         }
         internal_msg_buffer = nullptr;
       }
-      internal_msg_buffer_size = 0;
+      internal_msg_buffer_size   = 0;
+      internal_msg_buffer_set_on = TS_HTTP_LAST_HOOK;
     }
 
     ProxyProtocol pp_info;

plugins/header_rewrite/operators.cc

@@ -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);
 }
 

src/api/InkAPI.cc

@@ -4948,6 +4948,13 @@ TSHttpTxnErrorBodySet(TSHttpTxn txnp, char *buf, size_t buflength, char *mimetyp
   s->internal_msg_buffer                     = buf;
   s->internal_msg_buffer_size                = buf ? buflength : 0;
   s->internal_msg_buffer_fast_allocator_size = -1;
+  // TSHttpTxnErrorBodySet() and TSHttpTxnServerRequestBodySet() share the same buffer.
+  // Switching to an error/response body override must clear the request-body mode.
+  s->api_server_request_body_set = false;
+  // Record the hook the plugin was running on so HttpSM::handle_api_return() can
+  // scope the SERVER_READ/TRANSFORM_READ internal-body divert to response-stage
+  // body replacement (the canonical TSHttpTxnErrorBodySet() use case).
+  s->internal_msg_buffer_set_on = buf ? sm->get_cur_hook_id() : TS_HTTP_LAST_HOOK;
 
   s->internal_msg_buffer_type = mimetype;
 }

src/proxy/http/HttpSM.cc

@@ -1687,9 +1687,35 @@ 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);
+    // Bypass transform streaming when a plugin replaced the response body.
+    if (t_state.internal_msg_buffer && !t_state.api_server_request_body_set && t_state.hdr_info.server_response.valid()) {
+      SMDbg(dbg_ctl_http, "plugin set internal body, bypassing response transform for internal transfer");
+      t_state.api_info.cache_untransformed = true;
+      if (tunnel.is_tunnel_active()) {
+        tunnel.kill_tunnel();
+      }
+      // Drop transform VC table state because this path no longer drives transform reads.
+      if (transform_info.entry != nullptr) {
+        vc_table.cleanup_entry(transform_info.entry);
+        transform_info.entry = nullptr;
+      }
+      transform_info.vc = nullptr;
+      // Some downstream paths still read client_response; seed it from transform_response when missing.
+      if (t_state.hdr_info.client_response.valid() == 0 && t_state.hdr_info.transform_response.valid()) {
+        t_state.hdr_info.client_response.create(HTTPType::RESPONSE);
+        t_state.hdr_info.client_response.copy(&t_state.hdr_info.transform_response);
+      }
+      // Internal transfer doesn't use the server session.
+      if (server_entry != nullptr && server_entry->in_tunnel == false) {
+        release_server_session();
+      }
+      // Serve the plugin-provided body through the internal tunnel handler.
+      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: {
@@ -1722,6 +1748,16 @@ HttpSM::handle_api_return()
       }
 
       setup_blind_tunnel(true, initial_data);
+    } else if (t_state.internal_msg_buffer && !t_state.api_server_request_body_set && t_state.hdr_info.server_response.valid() &&
+               plugin_tunnel == nullptr &&
+               (t_state.internal_msg_buffer_set_on == TS_HTTP_READ_RESPONSE_HDR_HOOK ||
+                t_state.internal_msg_buffer_set_on == TS_HTTP_SEND_RESPONSE_HDR_HOOK)) {
+      // Plugin replaced the origin response body via TSHttpTxnErrorBodySet(); divert to internal transfer.
+      SMDbg(dbg_ctl_http, "plugin set internal body, using internal transfer instead of server tunnel");
+      if (server_entry != nullptr && server_entry->in_tunnel == false) {
+        release_server_session();
+      }
+      setup_internal_transfer(&HttpSM::tunnel_handler);
     } else {
       HttpTunnelProducer *p = setup_server_transfer();
       perform_cache_write_action();
@@ -7578,12 +7614,18 @@ HttpSM::setup_client_request_plugin_agents(HttpTunnelProducer *p, int num_header
 inline void
 HttpSM::transform_cleanup(TSHttpHookID hook, HttpTransformInfo *info)
 {
+  if (info->entry == nullptr) {
+    return;
+  }
   APIHook *t_hook = api_hooks.get(hook);
   if (t_hook && info->vc == nullptr) {
     do {
-      VConnection *t_vcon = t_hook->m_cont;
-      t_vcon->do_io_close();
-      t_hook = t_hook->m_link.next;
+      APIHook *next = t_hook->m_link.next;
+      // Some transform hooks can already be detached by the time kill_this() runs.
+      if (auto *t_vcon = static_cast<VConnection *>(t_hook->m_cont); t_vcon != nullptr) {
+        t_vcon->do_io_close();
+      }
+      t_hook = next;
     } while (t_hook != nullptr);
   }
 }
@@ -7664,7 +7706,11 @@ HttpSM::kill_this()
     //   In that case, we need to manually close all the
     //   transforms to prevent memory leaks (INKqa06147)
     if (hooks_set) {
-      transform_cleanup(TS_HTTP_RESPONSE_TRANSFORM_HOOK, &transform_info);
+      bool bypassed_response_transform =
+        t_state.api_info.cache_untransformed && t_state.internal_msg_buffer && !t_state.api_server_request_body_set;
+      if (!bypassed_response_transform) {
+        transform_cleanup(TS_HTTP_RESPONSE_TRANSFORM_HOOK, &transform_info);
+      }
       transform_cleanup(TS_HTTP_REQUEST_TRANSFORM_HOOK, &post_transform_info);
       plugin_agents_cleanup();
     }
@@ -8230,6 +8276,21 @@ HttpSM::set_next_state()
   case HttpTransact::StateMachineAction_t::SERVER_READ: {
     t_state.source = HttpTransact::Source_t::HTTP_ORIGIN_SERVER;
 
+    if (transform_info.vc && t_state.internal_msg_buffer && !t_state.api_server_request_body_set &&
+        t_state.hdr_info.server_response.valid()) {
+      SMDbg(dbg_ctl_http, "plugin set internal body, bypassing response transform");
+      t_state.api_info.cache_untransformed = true;
+      if (transform_info.entry != nullptr) {
+        vc_table.cleanup_entry(transform_info.entry);
+        transform_info.entry = nullptr;
+      }
+      transform_info.vc = nullptr;
+      if (t_state.hdr_info.client_response.valid() == 0 && t_state.hdr_info.transform_response.valid()) {
+        t_state.hdr_info.client_response.create(HTTPType::RESPONSE);
+        t_state.hdr_info.client_response.copy(&t_state.hdr_info.transform_response);
+      }
+    }
+
     if (transform_info.vc) {
       ink_assert(t_state.hdr_info.client_response.valid() == 0);
       ink_assert((t_state.hdr_info.transform_response.valid() ? true : false) == true);