From 06c99da5e66956506201b2bf24a51d9771dea84a Mon Sep 17 00:00:00 2001
From: Alexander Sack <asac@pantacor.com>
Date: Sun, 18 Mar 2018 13:17:17 +0100
Subject: [PATCH] add support for retrieving total bytes written in response

---
 rest/middleware.go | 6 +++++-
 rest/response.go   | 9 +++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/rest/middleware.go b/rest/middleware.go
index ba03fb8..a1f07a4 100644
--- a/rest/middleware.go
+++ b/rest/middleware.go
@@ -2,6 +2,8 @@ package rest
 
 import (
 	"net/http"
+
+	"github.com/miolini/datacounter"
 )
 
 // HandlerFunc defines the handler function. It is the go-json-rest equivalent of http.HandlerFunc.
@@ -61,8 +63,10 @@ func adapterFunc(handler HandlerFunc) http.HandlerFunc {
 			map[string]interface{}{},
 		}
 
+		countingWriter := datacounter.NewResponseWriterCounter(origWriter)
+
 		writer := &responseWriter{
-			origWriter,
+			countingWriter,
 			false,
 		}
 
diff --git a/rest/response.go b/rest/response.go
index 52529f1..842bc91 100644
--- a/rest/response.go
+++ b/rest/response.go
@@ -5,6 +5,8 @@ import (
 	"encoding/json"
 	"net"
 	"net/http"
+
+	"github.com/miolini/datacounter"
 )
 
 // A ResponseWriter interface dedicated to JSON HTTP response.
@@ -27,6 +29,9 @@ type ResponseWriter interface {
 	// Similar to the http.ResponseWriter interface, with additional JSON related
 	// headers set.
 	WriteHeader(int)
+
+	// Count of bytes written as response
+	Count() uint64
 }
 
 // This allows to customize the field name used in the error response payload.
@@ -125,3 +130,7 @@ func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 	hijacker := w.ResponseWriter.(http.Hijacker)
 	return hijacker.Hijack()
 }
+
+func (w *responseWriter) Count() uint64 {
+	return w.ResponseWriter.(*datacounter.ResponseWriterCounter).Count()
+}