Issue #14: added support for response body compression

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.

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)

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")