webnet
/
tailscale
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

tsweb: add gzip support to JSONHandlerFunc

Browse Source 
Change-Id: I337e05f92f744bfc7e9d6fb8e67c87c191ba4da8
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
main
Brad Fitzpatrick 4 years ago committed by Brad Fitzpatrick
parent 16652ae52c
commit df8f02db3f
2 changed files with 151 additions and 18 deletions
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Show Stats Download Patch File Download Diff File
  1. 78
      tsweb/jsonhandler.go
  2. 91
      tsweb/jsonhandler_test.go

78
tsweb/jsonhandler.go
Unescape View File

@ -5,9 +5,17 @@
package tsweb
import (
"bytes"
"compress/gzip"
"encoding/json"
"fmt"
"io/ioutil"
"net/http"
"strconv"
"strings"
"sync"
"go4.org/mem"
)
type response struct {
@ -85,7 +93,73 @@ func (fn JSONHandlerFunc) ServeHTTPReturn(w http.ResponseWriter, r *http.Request
return jerr
}
w.WriteHeader(status)
w.Write(b)
if AcceptsEncoding(r, "gzip") {
encb, err := gzipBytes(b)
if err != nil {
return err
}
w.Header().Set("Content-Encoding", "gzip")
w.Header().Set("Content-Length", strconv.Itoa(len(encb)))
w.Write(encb)
} else {
w.Header().Set("Content-Length", strconv.Itoa(len(b)))
w.WriteHeader(status)
w.Write(b)
}
return err
}
var gzWriterPool sync.Pool // of *gzip.Writer
// gzipBytes returns the gzipped encoding of b.
func gzipBytes(b []byte) (zb []byte, err error) {
var buf bytes.Buffer
zw, ok := gzWriterPool.Get().(*gzip.Writer)
if ok {
zw.Reset(&buf)
} else {
zw = gzip.NewWriter(&buf)
}
defer gzWriterPool.Put(zw)
if _, err := zw.Write(b); err != nil {
return nil, err
}
if err := zw.Close(); err != nil {
return nil, err
}
zb = buf.Bytes()
zw.Reset(ioutil.Discard)
return zb, nil
}
// AcceptsEncoding reports whether r accepts the named encoding
// ("gzip", "br", etc).
func AcceptsEncoding(r *http.Request, enc string) bool {
h := r.Header.Get("Accept-Encoding")
if h == "" {
return false
}
if !strings.Contains(h, enc) && !mem.ContainsFold(mem.S(h), mem.S(enc)) {
return false
}
remain := h
for len(remain) > 0 {
comma := strings.Index(remain, ",")
var part string
if comma == -1 {
part = remain
remain = ""
} else {
part = remain[:comma]
remain = remain[comma+1:]
}
part = strings.TrimSpace(part)
if i := strings.Index(part, ";"); i != -1 {
part = part[:i]
}
if part == enc {
return true
}
}
return false
}

91
tsweb/jsonhandler_test.go
Unescape View File

@ -5,8 +5,11 @@
package tsweb
import (
"bytes"
"compress/gzip"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
@ -27,13 +30,25 @@ type Response struct {
}
func TestNewJSONHandler(t *testing.T) {
checkStatus := func(w *httptest.ResponseRecorder, status string, code int) *Response {
checkStatus := func(t *testing.T, w *httptest.ResponseRecorder, status string, code int) *Response {
d := &Response{
Data: &Data{},
}
t.Logf("%s", w.Body.Bytes())
err := json.Unmarshal(w.Body.Bytes(), d)
bodyBytes := w.Body.Bytes()
if w.Result().Header.Get("Content-Encoding") == "gzip" {
zr, err := gzip.NewReader(bytes.NewReader(bodyBytes))
if err != nil {
t.Fatalf("gzip read error at start: %v", err)
}
bodyBytes, err = io.ReadAll(zr)
if err != nil {
t.Fatalf("gzip read error: %v", err)
}
}
t.Logf("%s", bodyBytes)
err := json.Unmarshal(bodyBytes, d)
if err != nil {
t.Logf(err.Error())
return nil
@ -64,7 +79,7 @@ func TestNewJSONHandler(t *testing.T) {
w := httptest.NewRecorder()
r := httptest.NewRequest("GET", "/", nil)
h21.ServeHTTPReturn(w, r)
checkStatus(w, "success", http.StatusOK)
checkStatus(t, w, "success", http.StatusOK)
})
t.Run("403 HTTPError", func(t *testing.T) {
@ -75,7 +90,7 @@ func TestNewJSONHandler(t *testing.T) {
w := httptest.NewRecorder()
r := httptest.NewRequest("GET", "/", nil)
h.ServeHTTPReturn(w, r)
checkStatus(w, "error", http.StatusForbidden)
checkStatus(t, w, "error", http.StatusForbidden)
})
h22 := JSONHandlerFunc(func(r *http.Request) (int, interface{}, error) {
@ -86,7 +101,7 @@ func TestNewJSONHandler(t *testing.T) {
w := httptest.NewRecorder()
r := httptest.NewRequest("GET", "/", nil)
h22.ServeHTTPReturn(w, r)
checkStatus(w, "success", http.StatusOK)
checkStatus(t, w, "success", http.StatusOK)
})
h31 := JSONHandlerFunc(func(r *http.Request) (int, interface{}, error) {
@ -105,21 +120,21 @@ func TestNewJSONHandler(t *testing.T) {
w := httptest.NewRecorder()
r := httptest.NewRequest("POST", "/", strings.NewReader(`{"Name": "tailscale"}`))
h31.ServeHTTPReturn(w, r)
checkStatus(w, "success", http.StatusOK)
checkStatus(t, w, "success", http.StatusOK)
})
t.Run("400 bad json", func(t *testing.T) {
w := httptest.NewRecorder()
r := httptest.NewRequest("POST", "/", strings.NewReader(`{`))
h31.ServeHTTPReturn(w, r)
checkStatus(w, "error", http.StatusBadRequest)
checkStatus(t, w, "error", http.StatusBadRequest)
})
t.Run("400 post data error", func(t *testing.T) {
w := httptest.NewRecorder()
r := httptest.NewRequest("POST", "/", strings.NewReader(`{}`))
h31.ServeHTTPReturn(w, r)
resp := checkStatus(w, "error", http.StatusBadRequest)
resp := checkStatus(t, w, "error", http.StatusBadRequest)
if resp.Error != "name is empty" {
t.Fatalf("wrong error")
}
@ -144,7 +159,23 @@ func TestNewJSONHandler(t *testing.T) {
w := httptest.NewRecorder()
r := httptest.NewRequest("POST", "/", strings.NewReader(`{"Price": 10}`))
h32.ServeHTTPReturn(w, r)
resp := checkStatus(w, "success", http.StatusOK)
resp := checkStatus(t, w, "success", http.StatusOK)
t.Log(resp.Data)
if resp.Data.Price != 20 {
t.Fatalf("wrong price: %d %d", resp.Data.Price, 10)
}
})
t.Run("gzipped", func(t *testing.T) {
w := httptest.NewRecorder()
r := httptest.NewRequest("POST", "/", strings.NewReader(`{"Price": 10}`))
r.Header.Set("Accept-Encoding", "gzip")
h32.ServeHTTPReturn(w, r)
res := w.Result()
if ct := res.Header.Get("Content-Encoding"); ct != "gzip" {
t.Fatalf("encoding = %q; want gzip", ct)
}
resp := checkStatus(t, w, "success", http.StatusOK)
t.Log(resp.Data)
if resp.Data.Price != 20 {
t.Fatalf("wrong price: %d %d", resp.Data.Price, 10)
@ -155,7 +186,7 @@ func TestNewJSONHandler(t *testing.T) {
w := httptest.NewRecorder()
r := httptest.NewRequest("POST", "/", strings.NewReader(`{}`))
h32.ServeHTTPReturn(w, r)
resp := checkStatus(w, "error", http.StatusBadRequest)
resp := checkStatus(t, w, "error", http.StatusBadRequest)
if resp.Error != "price is empty" {
t.Fatalf("wrong error")
}
@ -165,7 +196,7 @@ func TestNewJSONHandler(t *testing.T) {
w := httptest.NewRecorder()
r := httptest.NewRequest("POST", "/", strings.NewReader(`{"Name": "root"}`))
h32.ServeHTTPReturn(w, r)
resp := checkStatus(w, "error", http.StatusInternalServerError)
resp := checkStatus(t, w, "error", http.StatusInternalServerError)
if resp.Error != "internal server error" {
t.Fatalf("wrong error")
}
@ -177,7 +208,7 @@ func TestNewJSONHandler(t *testing.T) {
JSONHandlerFunc(func(r *http.Request) (int, interface{}, error) {
return http.StatusOK, make(chan int), nil
}).ServeHTTPReturn(w, r)
resp := checkStatus(w, "error", http.StatusInternalServerError)
resp := checkStatus(t, w, "error", http.StatusInternalServerError)
if resp.Error != "json marshal error" {
t.Fatalf("wrong error")
}
@ -189,7 +220,7 @@ func TestNewJSONHandler(t *testing.T) {
JSONHandlerFunc(func(r *http.Request) (status int, data interface{}, err error) {
return
}).ServeHTTPReturn(w, r)
checkStatus(w, "error", http.StatusInternalServerError)
checkStatus(t, w, "error", http.StatusInternalServerError)
})
t.Run("403 forbidden, status returned by JSONHandlerFunc and HTTPError agree", func(t *testing.T) {
@ -203,7 +234,7 @@ func TestNewJSONHandler(t *testing.T) {
Data: &Data{},
Error: "403 forbidden",
}
got := checkStatus(w, "error", http.StatusForbidden)
got := checkStatus(t, w, "error", http.StatusForbidden)
if diff := cmp.Diff(want, got); diff != "" {
t.Fatalf(diff)
}
@ -223,9 +254,37 @@ func TestNewJSONHandler(t *testing.T) {
Data: &Data{},
Error: "403 forbidden",
}
got := checkStatus(w, "error", http.StatusForbidden)
got := checkStatus(t, w, "error", http.StatusForbidden)
if diff := cmp.Diff(want, got); diff != "" {
t.Fatalf("(-want,+got):\n%s", diff)
}
})
}
func TestAcceptsEncoding(t *testing.T) {
tests := []struct {
in, enc string
want bool
}{
{"", "gzip", false},
{"gzip", "gzip", true},
{"foo,gzip", "gzip", true},
{"foo, gzip", "gzip", true},
{"foo, gzip ", "gzip", true},
{"gzip, foo ", "gzip", true},
{"gzip, foo ", "br", false},
{"gzip, foo ", "fo", false},
{"gzip;q=1.2, foo ", "gzip", true},
{" gzip;q=1.2, foo ", "gzip", true},
}
for i, tt := range tests {
h := make(http.Header)
if tt.in != "" {
h.Set("Accept-Encoding", tt.in)
}
got := AcceptsEncoding(&http.Request{Header: h}, tt.enc)
if got != tt.want {
t.Errorf("%d. got %v; want %v", i, got, tt.want)
}
}
}

Write Preview
Loading…
Cancel
Save