2 files changed, 38 insertions, 8 deletions

diff --git a/docs/http.html b/docs/http.html
index c6423ba..15f9c26 100644
--- a/docs/http.html
+++ b/docs/http.html
@@ -136,7 +136,8 @@ http.<b>request{</b><br>
 &nbsp;&nbsp;[proxy = <i>string</i>,]<br>
 &nbsp;&nbsp;[redirect = <i>boolean</i>,]<br>
 &nbsp;&nbsp;[create = <i>function</i>,]<br>
-&nbsp;&nbsp;[maxredirects = <i>number</i>]<br>
+&nbsp;&nbsp;[maxredirects = <i>number</i>,]<br>
+&nbsp;&nbsp;[headers_callback = <i>function</i>]<br>
 <b>}</b>
 </p>
 
@@ -189,6 +190,39 @@ function from  automatically following 301 or 302 server redirect messages;</li>
 <li><tt>maxredirects</tt>: An optional number specifying the maximum number of
     redirects to follow.  Defaults to <tt>5</tt> if not specified. A boolean
     <tt>false</tt> value means no maximum (unlimited).</li>
+<li><tt>headers_callback</tt>: An optional function, called once per request
+    right after the headers have been received (this
+    is after any redirect has already been followed, so it only fires for the
+    final response) and before the response body would be read. It is called
+    as: <tt>ok, new_sink = headers_callback(code, headers, status, sock)</tt>, where
+    <tt>sock</tt> is the raw socket used for the request (the
+    <a href="tcp.html#socket.tcp"><tt>socket.tcp</tt></a>-like object returned
+    by <tt>create</tt>, or the default one). It is not called for an HTTP/0.9
+    reply (no headers at all), for a <tt>408</tt> response, or for a response
+    that is going to be redirected.
+    <br>
+    The callback should return two values, <tt>ok</tt> and a second value whose
+    meaning depends on <tt>ok</tt>:
+    <ul>
+    <li>If <tt>ok</tt> is falsy, the second value is used as an error
+    message: the connection is closed and <tt>request</tt> returns
+    <tt><b>nil</b></tt> followed by that message, just like any other request
+    failure.</li>
+    <li>If <tt>ok</tt> is truthy, the request continues normally. If the
+    second value is also provided, it replaces <tt>sink</tt> for reading the
+    response body; otherwise the original sink is used.</li>
+    </ul>
+    Note that a swapped-in sink is only ever read from when the response
+    actually has a body to receive: for a <tt>HEAD</tt> request or a
+    <tt>204</tt>/<tt>304</tt> response there is nothing to read regardless of
+    which sink is set.
+    <br>
+    This makes it possible to inspect headers before committing to read a
+    body -- for example to reject a response early based on
+    <tt>content-type</tt> or <tt>content-length</tt>, to pick a different
+    sink depending on the headers, or -- combined with the raw
+    <tt>sock</tt> -- to hand the socket off for a protocol upgrade
+    such as WebSockets (<tt>101 Switching Protocols</tt>).</li>
 </ul>
 
 <p class="return">
diff --git a/src/http.lua b/src/http.lua
index 6cdc90a..3e2c789 100644
--- a/src/http.lua
+++ b/src/http.lua
@@ -360,7 +360,7 @@ local trequest, tredirect
         headers = reqt.headers,
         proxy = reqt.proxy,
         maxredirects = reqt.maxredirects,
-        response_headers = reqt.response_headers,
+        headers_callback = reqt.headers_callback,
         nredirects = (reqt.nredirects or 0) + 1,
         create = reqt.create
     }
@@ -405,12 +405,8 @@ end
     end
     -- here we are finally done
     -- provide an opportunity to abort or replace the sink based on the response headers
-    if nreqt.response_headers then
-        local abort, sink = nreqt.response_headers(code, headers, status)
-        if abort then
-            h:close()
-            return 1, code, headers, status
-        end
+    if nreqt.headers_callback then
+        local _, sink = h.try(nreqt.headers_callback(code, headers, status, h.c))
         if sink then
             nreqt.sink = sink
         end