From 0583374c12a0266846458882c5014158545b4751 Mon Sep 17 00:00:00 2001
From: Andrew Gaul <andrew@gaul.org>
Date: Fri, 31 Aug 2018 12:01:20 -0700
Subject: [PATCH] Add /brotli endpoint

Parity with http://httpbin.org/
---
 README.md        |  1 +
 handlers.go      | 21 +++++++++++++++++++++
 handlers_test.go | 25 +++++++++++++++++++++++++
 types.go         |  6 ++++++
 4 files changed, 53 insertions(+)

diff --git a/README.md b/README.md
index 6ef50cf..a47c26b 100644
--- a/README.md
+++ b/README.md
@@ -30,6 +30,7 @@ This way, you can write tests without relying on an external dependency like [ht
 - `/cache/:n` Sets a Cache-Control header for _n_ seconds.
 - `/gzip` Returns gzip-encoded data.
 - `/deflate` Returns deflate-encoded data.
+- `/brotli` Returns brotli-encoded data.
 - `/robots.txt` Returns some robots.txt rules.
 - `/deny` Denied by robots.txt file.
 - `/basic-auth/:user/:passwd` Challenges HTTP Basic Auth.
diff --git a/handlers.go b/handlers.go
index eead2af..f55a14c 100644
--- a/handlers.go
+++ b/handlers.go
@@ -24,6 +24,7 @@ import (
 
 	"github.com/gorilla/mux"
 	"github.com/pkg/errors"
+	brotlienc "gopkg.in/kothar/brotli-go.v0/enc"
 )
 
 var (
@@ -64,6 +65,7 @@ func GetMux() *mux.Router {
 	r.HandleFunc(`/cache`, CacheHandler).Methods(http.MethodGet, http.MethodHead)
 	r.HandleFunc(`/cache/{n:[\d]+}`, SetCacheHandler).Methods(http.MethodGet, http.MethodHead)
 	r.HandleFunc(`/gzip`, GZIPHandler).Methods(http.MethodGet, http.MethodHead)
+	r.HandleFunc(`/brotli`, BrotliHandler).Methods(http.MethodGet, http.MethodHead)
 	r.HandleFunc(`/deflate`, DeflateHandler).Methods(http.MethodGet, http.MethodHead)
 	r.HandleFunc(`/html`, HTMLHandler).Methods(http.MethodGet, http.MethodHead)
 	r.HandleFunc(`/xml`, XMLHandler).Methods(http.MethodGet, http.MethodHead)
@@ -441,6 +443,25 @@ func DeflateHandler(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+// BrotliHandler returns a Brotli-encoded response
+func BrotliHandler(w http.ResponseWriter, r *http.Request) {
+	h, _, _ := net.SplitHostPort(r.RemoteAddr)
+
+	v := brotliResponse{
+		headersResponse: headersResponse{getHeaders(r)},
+		ipResponse:      ipResponse{h},
+		Compressed:      true,
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	w.Header().Add("Content-Encoding", "br")
+	ww := brotlienc.NewBrotliWriter(nil, w)
+	defer ww.Close() // flush
+	if err := writeJSON(ww, v); err != nil {
+		writeErrorJSON(w, errors.Wrap(err, "failed to write json"))
+	}
+}
+
 // RobotsTXTHandler returns a robots.txt response.
 func RobotsTXTHandler(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "text/plain")
diff --git a/handlers_test.go b/handlers_test.go
index 9b838be..a0730b8 100644
--- a/handlers_test.go
+++ b/handlers_test.go
@@ -20,6 +20,7 @@ import (
 
 	"github.com/ahmetb/go-httpbin"
 	"github.com/stretchr/testify/require"
+	brotlidec "gopkg.in/kothar/brotli-go.v0/dec"
 )
 
 var (
@@ -568,6 +569,30 @@ func TestDeflate(t *testing.T) {
 	require.True(t, v.Deflated)
 }
 
+func TestBrotli(t *testing.T) {
+	srv := testServer()
+	defer srv.Close()
+
+	client := new(http.Client)
+	req, err := http.NewRequest("GET", srv.URL+"/brotli", nil)
+	require.Nil(t, err)
+
+	req.Header.Add("Accept-Encoding", "br")
+	resp, err := client.Do(req)
+	require.Nil(t, err)
+	defer resp.Body.Close()
+
+	require.EqualValues(t, "br", resp.Header.Get("Content-Encoding"))
+	require.EqualValues(t, "application/json", resp.Header.Get("Content-Type"))
+	zr := brotlidec.NewBrotliReader(resp.Body)
+
+	var v struct {
+		Compressed bool `json:"compressed"`
+	}
+	require.Nil(t, json.NewDecoder(zr).Decode(&v))
+	require.True(t, v.Compressed)
+}
+
 func TestRobotsTXT(t *testing.T) {
 	srv := testServer()
 	defer srv.Close()
diff --git a/types.go b/types.go
index 3111c5f..807ad0d 100644
--- a/types.go
+++ b/types.go
@@ -54,6 +54,12 @@ type deflateResponse struct {
 	Deflated bool `json:"deflated"`
 }
 
+type brotliResponse struct {
+	headersResponse
+	ipResponse
+	Compressed bool `json:"compressed"`
+}
+
 type basicAuthResponse struct {
 	Authenticated bool   `json:"authenticated"`
 	User          string `json:"string"`