From edcfdbcec9a77959cd89c16121ef52283dc24673 Mon Sep 17 00:00:00 2001 From: Aliaksandr Valialkin Date: Fri, 25 Dec 2015 12:26:34 +0200 Subject: [PATCH] Issue #14: added support for response body compression --- http.go | 138 +++++++++++++++++++++++++++++++++++++++++++++++++++ http_test.go | 110 ++++++++++++++++++++++++++++++++++++++++ strings.go | 3 ++ 3 files changed, 251 insertions(+) diff --git a/http.go b/http.go index 64c5f85951..8bf27e664a 100644 --- a/http.go +++ b/http.go @@ -3,9 +3,12 @@ package fasthttp import ( "bufio" "bytes" + "compress/flate" + "compress/gzip" "errors" "fmt" "io" + "io/ioutil" "mime/multipart" "os" "sync" @@ -195,6 +198,45 @@ func (resp *Response) Body() []byte { return resp.body } +// BodyGunzip returns un-gzipped body data. +// +// This method may be used if the response header contains +// 'Content-Encoding: gzip' for reading un-gzipped response body. +// Use Body for reading gzipped response body. +func (resp *Response) BodyGunzip() ([]byte, error) { + // Do not care about memory allocations here, + // since gzip is slow and generates a lot of memory allocations + // by itself. + r, err := gzip.NewReader(bytes.NewBuffer(resp.body)) + if err != nil { + return nil, err + } + defer r.Close() + b, err := ioutil.ReadAll(r) + if err != nil { + return nil, err + } + return b, nil +} + +// BodyInflate returns un-deflated body data. +// +// This method may be used if the response header contains +// 'Content-Encoding: deflate' for reading un-deflated response body. +// Use Body for reading deflated response body. +func (resp *Response) BodyInflate() ([]byte, error) { + // Do not care about memory allocations here, + // since flate is slow and generates a lot of memory allocations + // by itself. + r := flate.NewReader(bytes.NewBuffer(resp.body)) + defer r.Close() + b, err := ioutil.ReadAll(r) + if err != nil { + return nil, err + } + return b, nil +} + // AppendBody appends p to response body. func (resp *Response) AppendBody(p []byte) { resp.closeBodyStream() @@ -520,6 +562,102 @@ func (req *Request) Write(w *bufio.Writer) error { return err } +// WriteGzip writes response with gzipped body to w. +// +// The method sets 'Content-Encoding: gzip' header. +// +// WriteGzip doesn't flush response to w for performance reasons. +func (resp *Response) WriteGzip(w *bufio.Writer) error { + return resp.WriteGzipLevel(w, gzip.DefaultCompression) +} + +// WriteGzipLevel writes response with gzipped body to w. +// +// Level is compression level. See available levels in encoding/gzip package. +// +// The method sets 'Content-Encoding: gzip' header. +// +// WriteGzipLevel doesn't flush response to w for performance reasons. +func (resp *Response) WriteGzipLevel(w *bufio.Writer, level int) error { + // Do not care about memory allocations here, since gzip is slow + // and allocates a lot of memory by itself. + if resp.bodyStream != nil { + bs := resp.bodyStream + resp.bodyStream = NewStreamReader(func(sw *bufio.Writer) { + zw := newGzipWriter(sw, level) + defer zw.Close() + io.Copy(zw, bs) + }) + } else { + var buf bytes.Buffer + zw := newGzipWriter(&buf, level) + if _, err := zw.Write(resp.body); err != nil { + return err + } + zw.Close() + resp.body = buf.Bytes() + } + + resp.Header.SetCanonical(strContentEncoding, strGzip) + return resp.Write(w) +} + +// WriteDeflate writes response with deflated body to w. +// +// The method sets 'Content-Encoding: deflate' header. +// +// WriteDeflate doesn't flush response to w for performance reasons. +func (resp *Response) WriteDeflate(w *bufio.Writer) error { + return resp.WriteDeflateLevel(w, flate.DefaultCompression) +} + +// WriteDeflateLevel writes response with deflated body to w. +// +// Level is compression level. See available levels in encoding/flate package. +// +// The method sets 'Content-Encoding: deflate' header. +// +// WriteDeflateLevel doesn't flush response to w for performance reasons. +func (resp *Response) WriteDeflateLevel(w *bufio.Writer, level int) error { + // Do not care about memory allocations here, since flate is slow + // and allocates a lot of memory by itself. + if resp.bodyStream != nil { + bs := resp.bodyStream + resp.bodyStream = NewStreamReader(func(sw *bufio.Writer) { + zw := newDeflateWriter(sw, level) + defer zw.Close() + io.Copy(zw, bs) + }) + } else { + var buf bytes.Buffer + zw := newDeflateWriter(&buf, level) + if _, err := zw.Write(resp.body); err != nil { + return err + } + zw.Close() + resp.body = buf.Bytes() + } + + resp.Header.SetCanonical(strContentEncoding, strDeflate) + return resp.Write(w) +} + +func newDeflateWriter(w io.Writer, level int) *flate.Writer { + zw, err := flate.NewWriter(w, level) + if err != nil { + panic(fmt.Sprintf("BUG: flate.NewWriter(%d) returns non-nil error: %s", level, err)) + } + return zw +} + +func newGzipWriter(w io.Writer, level int) *gzip.Writer { + zw, err := gzip.NewWriterLevel(w, level) + if err != nil { + panic(fmt.Sprintf("BUG: gzip.NewWriter(%d) returns non-nil error: %s", level, err)) + } + return zw +} + // Write writes response to w. // // Write doesn't flush response to w for performance reasons. diff --git a/http_test.go b/http_test.go index 5ea62b7d73..11284ef768 100644 --- a/http_test.go +++ b/http_test.go @@ -9,6 +9,116 @@ import ( "testing" ) +func TestResponseGzipStream(t *testing.T) { + var r Response + r.SetBodyStreamWriter(func(w *bufio.Writer) { + fmt.Fprintf(w, "foo") + w.Flush() + w.Write([]byte("barbaz")) + w.Flush() + fmt.Fprintf(w, "1234") + if err := w.Flush(); err != nil { + t.Fatalf("unexpected error: %s", err) + } + }) + testResponseGzipExt(t, &r, "foobarbaz1234") +} + +func TestResponseDeflateStream(t *testing.T) { + var r Response + r.SetBodyStreamWriter(func(w *bufio.Writer) { + w.Write([]byte("foo")) + w.Flush() + fmt.Fprintf(w, "barbaz") + w.Flush() + w.Write([]byte("1234")) + if err := w.Flush(); err != nil { + t.Fatalf("unexpected error: %s", err) + } + }) + testResponseDeflateExt(t, &r, "foobarbaz1234") +} + +func TestResponseDeflate(t *testing.T) { + testResponseDeflate(t, "") + testResponseDeflate(t, "abdasdfsdaa") + testResponseDeflate(t, "asoiowqoieroqweiruqwoierqo") +} + +func TestResponseGzip(t *testing.T) { + testResponseGzip(t, "") + testResponseGzip(t, "foobarbaz") + testResponseGzip(t, "abasdwqpweoweporweprowepr") +} + +func testResponseDeflate(t *testing.T, s string) { + var r Response + r.SetBodyString(s) + testResponseDeflateExt(t, &r, s) +} + +func testResponseDeflateExt(t *testing.T, r *Response, s string) { + var buf bytes.Buffer + bw := bufio.NewWriter(&buf) + if err := r.WriteDeflate(bw); err != nil { + t.Fatalf("unexpected error: %s", err) + } + if err := bw.Flush(); err != nil { + t.Fatalf("unexpected error: %s", err) + } + + var r1 Response + br := bufio.NewReader(&buf) + if err := r1.Read(br); err != nil { + t.Fatalf("unexpected error: %s", err) + } + ce := r1.Header.Peek("Content-Encoding") + if string(ce) != "deflate" { + t.Fatalf("unexpected Content-Encoding %q. Expecting %q", ce, "deflate") + } + body, err := r1.BodyInflate() + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + if string(body) != s { + t.Fatalf("unexpected body %q. Expecting %q", body, s) + } +} + +func testResponseGzip(t *testing.T, s string) { + var r Response + r.SetBodyString(s) + testResponseGzipExt(t, &r, s) +} + +func testResponseGzipExt(t *testing.T, r *Response, s string) { + var buf bytes.Buffer + bw := bufio.NewWriter(&buf) + if err := r.WriteGzip(bw); err != nil { + t.Fatalf("unexpected error: %s", err) + } + if err := bw.Flush(); err != nil { + t.Fatalf("unexpected error: %s", err) + } + + var r1 Response + br := bufio.NewReader(&buf) + if err := r1.Read(br); err != nil { + t.Fatalf("unexpected error: %s", err) + } + ce := r1.Header.Peek("Content-Encoding") + if string(ce) != "gzip" { + t.Fatalf("unexpected Content-Encoding %q. Expecting %q", ce, "gzip") + } + body, err := r1.BodyGunzip() + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + if string(body) != s { + t.Fatalf("unexpected body %q. Expecting %q", body, s) + } +} + func TestRequestMultipartForm(t *testing.T) { var w bytes.Buffer mw := multipart.NewWriter(&w) diff --git a/strings.go b/strings.go index ea7943c9f4..6dc42fb5ad 100644 --- a/strings.go +++ b/strings.go @@ -32,6 +32,7 @@ var ( strReferer = []byte("Referer") strServer = []byte("Server") strTransferEncoding = []byte("Transfer-Encoding") + strContentEncoding = []byte("Content-Encoding") strUserAgent = []byte("User-Agent") strCookie = []byte("Cookie") strSetCookie = []byte("Set-Cookie") @@ -44,6 +45,8 @@ var ( strCookiePath = []byte("path") strClose = []byte("close") + strGzip = []byte("gzip") + strDeflate = []byte("deflate") strKeepAlive = []byte("keep-alive") strKeepAliveCamelCase = []byte("Keep-Alive") strUpgrade = []byte("Upgrade")