lib/async/http/cache/body.rb in async-http-cache-0.1.5 vs lib/async/http/cache/body.rb in async-http-cache-0.2.0

@@ -20,40 +20,51 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
 require 'protocol/http/body/rewindable'
 require 'protocol/http/body/streamable'
+require 'protocol/http/body/digestable'
 
 module Async
 	module HTTP
 		module Cache
 			module Body
-				def self.wrap(message, &block)
-					if body = message.body
+				TRAILERS = 'trailers'
+				ETAG = 'etag'
+				
+				def self.wrap(response, &block)
+					if body = response.body
 						if body.empty?
 							# A body that is empty? at the outset, is immutable. This generally only applies to HEAD requests.
-							yield message, body
+							yield response, body
 						else
-							# Create a rewindable body wrapping the message body:
-							rewindable = ::Protocol::HTTP::Body::Rewindable.new(body)
+							# Insert a rewindable body so that we can cache the response body:
+							rewindable = ::Protocol::HTTP::Body::Rewindable.wrap(response)
 							
-							# Set the message body to the rewindable body:
-							message.body = rewindable
+							unless response.headers.include?(ETAG)
+								# Compute a digest and add it to the response headers:
+								::Protocol::HTTP::Body::Digestable.wrap(response) do |wrapper|
+									response.headers.add(ETAG, wrapper.etag)
+								end
+								
+								# Ensure the etag is listed as a trailer:
+								response.headers.add(TRAILERS, ETAG)
+							end
 							
-							# Wrap the message with the callback:
-							::Protocol::HTTP::Body::Streamable.wrap(message) do |error|
+							# Wrap the response with the callback:
+							::Protocol::HTTP::Body::Streamable.wrap(response) do |error|
 								if error
 									Async.logger.error(self) {error}
 								else
-									yield message, rewindable.buffered
+									yield response, rewindable.buffered
 								end
 							end
 						end
 					else
-						yield message, nil
+						yield response, nil
 					end
 					
-					return message
+					return response
 				end
 			end
 		end
 	end
 end