Restructure system for improved maintenance going forward (#1)

compressutil/compressutil.go

@@ -0,0 +1,92 @@
+// Package compressutil a plugin to handle compression and decompression tasks
+package compressutil
+
+import (
+	"bytes"
+	"compress/flate"
+	"compress/gzip"
+	"io"
+	"log"
+)
+
+// ReaderError for notating that an error occurred while reading compressed data.
+type ReaderError struct {
+	error
+
+	cause error
+}
+
+// Decode data in a bytes.Reader based on supplied encoding.
+func Decode(byteReader *bytes.Buffer, encoding string) (data []byte, err error) {
+	reader, err := getRawReader(byteReader, encoding)
+	if err != nil {
+		return nil, &ReaderError{cause: err}
+	}
+
+	return io.ReadAll(reader)
+}
+
+func getRawReader(byteReader *bytes.Buffer, encoding string) (io.Reader, error) {
+	switch encoding {
+	case "gzip":
+		return gzip.NewReader(byteReader)
+
+	case "deflate":
+		return flate.NewReader(byteReader), nil
+
+	default:
+		return byteReader, nil
+	}
+}
+
+// Encode data in a []byte based on supplied encoding.
+func Encode(data []byte, encoding string) ([]byte, error) {
+	switch encoding {
+	case "gzip":
+		return compressWithGzip(data)
+
+	case "deflate":
+		return compressWithZlib(data)
+
+	default:
+		return data, nil
+	}
+}
+
+func compressWithGzip(bodyBytes []byte) ([]byte, error) {
+	var buf bytes.Buffer
+	gzipWriter := gzip.NewWriter(&buf)
+
+	if _, err := gzipWriter.Write(bodyBytes); err != nil {
+		log.Printf("unable to recompress rewrited body: %v", err)
+
+		return nil, err
+	}
+
+	if err := gzipWriter.Close(); err != nil {
+		log.Printf("unable to close gzip writer: %v", err)
+
+		return nil, err
+	}
+
+	return buf.Bytes(), nil
+}
+
+func compressWithZlib(bodyBytes []byte) ([]byte, error) {
+	var buf bytes.Buffer
+	zlibWriter, _ := flate.NewWriter(&buf, flate.DefaultCompression)
+
+	if _, err := zlibWriter.Write(bodyBytes); err != nil {
+		log.Printf("unable to recompress rewrited body: %v", err)
+
+		return nil, err
+	}
+
+	if err := zlibWriter.Close(); err != nil {
+		log.Printf("unable to close zlib writer: %v", err)
+
+		return nil, err
+	}
+
+	return buf.Bytes(), nil
+}

compressutil/compressutil_test.go

@@ -0,0 +1,154 @@
+package compressutil_test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/packruler/rewrite-body/compressutil"
+)
+
+type TestStruct struct {
+	desc        string
+	input       []byte
+	expected    []byte
+	encoding    string
+	shouldMatch bool
+}
+
+func TestEncode(t *testing.T) {
+	var (
+		deflatedBytes = []byte{
+			74, 203, 207, 87, 200, 44, 86, 40, 201, 72, 85,
+			200, 75, 45, 87, 72, 74, 44, 2, 4, 0, 0, 255, 255,
+		}
+		gzippedBytes = []byte{
+			31, 139, 8, 0, 0, 0, 0, 0, 0, 255, 74, 203, 207, 87, 200, 44, 86, 40, 201, 72, 85,
+			200, 75, 45, 87, 72, 74, 44, 2, 4, 0, 0, 255, 255, 251, 28, 166, 187, 18, 0, 0, 0,
+		}
+		normalBytes = []byte("foo is the new bar")
+	)
+
+	tests := []TestStruct{
+		{
+			desc:        "should support identity",
+			input:       normalBytes,
+			expected:    normalBytes,
+			encoding:    "identity",
+			shouldMatch: true,
+		},
+		{
+			desc:        "should support gzip",
+			input:       normalBytes,
+			expected:    gzippedBytes,
+			encoding:    "gzip",
+			shouldMatch: false,
+		},
+		{
+			desc:        "should support deflate",
+			input:       normalBytes,
+			expected:    deflatedBytes,
+			encoding:    "deflate",
+			shouldMatch: false,
+		},
+		{
+			desc:        "should NOT support brotli",
+			input:       normalBytes,
+			expected:    normalBytes,
+			encoding:    "br",
+			shouldMatch: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			output, err := compressutil.Encode(test.input, test.encoding)
+			if err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+
+			isBad := !bytes.Equal(test.expected, output)
+
+			if isBad {
+				t.Errorf("expected error got body: %v\n wanted: %v", output, test.expected)
+			}
+
+			if test.shouldMatch {
+				isBad = !bytes.Equal(test.input, output)
+			} else {
+				isBad = bytes.Equal(test.input, output)
+			}
+			if isBad {
+				t.Errorf("match error got body: %v\n wanted: %v", output, test.input)
+			}
+		})
+	}
+}
+
+func TestDecode(t *testing.T) {
+	var (
+		deflatedBytes = []byte{
+			74, 203, 207, 87, 200, 44, 86, 40, 201, 72, 85,
+			200, 75, 45, 87, 72, 74, 44, 2, 4, 0, 0, 255, 255,
+		}
+		gzippedBytes = []byte{
+			31, 139, 8, 0, 0, 0, 0, 0, 0, 255, 74, 203, 207, 87, 200, 44, 86, 40, 201, 72, 85,
+			200, 75, 45, 87, 72, 74, 44, 2, 4, 0, 0, 255, 255, 251, 28, 166, 187, 18, 0, 0, 0,
+		}
+		normalBytes = []byte("foo is the new bar")
+	)
+
+	tests := []TestStruct{
+		{
+			desc:        "should support identity",
+			input:       normalBytes,
+			expected:    normalBytes,
+			encoding:    "identity",
+			shouldMatch: true,
+		},
+		{
+			desc:        "should support gzip",
+			input:       gzippedBytes,
+			expected:    normalBytes,
+			encoding:    "gzip",
+			shouldMatch: false,
+		},
+		{
+			desc:        "should support deflate",
+			input:       deflatedBytes,
+			expected:    normalBytes,
+			encoding:    "deflate",
+			shouldMatch: false,
+		},
+		{
+			desc:        "should NOT support brotli",
+			input:       normalBytes,
+			expected:    normalBytes,
+			encoding:    "br",
+			shouldMatch: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			output, err := compressutil.Decode(bytes.NewBuffer(test.input), test.encoding)
+			if err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+
+			isBad := !bytes.Equal(test.expected, output)
+
+			if isBad {
+				t.Errorf("expected error got body: %v\n wanted: %v", output, test.expected)
+			}
+
+			if test.shouldMatch {
+				isBad = !bytes.Equal(test.input, output)
+			} else {
+				isBad = bytes.Equal(test.input, output)
+			}
+			if isBad {
+				t.Errorf("match error got body: %s\n wanted: %s", output, test.input)
+			}
+		})
+	}
+}

httputil/response_writer.go

@@ -0,0 +1,129 @@
+// Package httputil a package for handling http data tasks
+package httputil
+
+import (
+	"bytes"
+	"log"
+	"net/http"
+	"strings"
+
+	"github.com/packruler/rewrite-body/compressutil"
+)
+
+// ResponseWrapper a wrapper used to simplify ResponseWriter data access and manipulation.
+type ResponseWrapper struct {
+	buffer       bytes.Buffer
+	lastModified bool
+	wroteHeader  bool
+
+	http.ResponseWriter
+}
+
+// WriteHeader into wrapped ResponseWriter.
+func (wrapper *ResponseWrapper) WriteHeader(statusCode int) {
+	if !wrapper.lastModified {
+		wrapper.ResponseWriter.Header().Del("Last-Modified")
+	}
+
+	wrapper.wroteHeader = true
+
+	// Delegates the Content-Length Header creation to the final body write.
+	wrapper.ResponseWriter.Header().Del("Content-Length")
+
+	wrapper.ResponseWriter.WriteHeader(statusCode)
+}
+
+// Write data to internal buffer and mark the status code as http.StatusOK.
+func (wrapper *ResponseWrapper) Write(data []byte) (int, error) {
+	if !wrapper.wroteHeader {
+		wrapper.WriteHeader(http.StatusOK)
+	}
+
+	return wrapper.buffer.Write(data)
+}
+
+// GetBuffer get a pointer to the ResponseWriter buffer.
+func (wrapper *ResponseWrapper) GetBuffer() *bytes.Buffer {
+	return &wrapper.buffer
+}
+
+// GetContent load the content currently in the internal buffer
+// accounting for possible encoding.
+func (wrapper *ResponseWrapper) GetContent() ([]byte, error) {
+	encoding := wrapper.GetContentEncoding()
+
+	return compressutil.Decode(wrapper.GetBuffer(), encoding)
+}
+
+// SetContent write data to the internal ResponseWriter buffer
+// and match initial encoding.
+func (wrapper *ResponseWrapper) SetContent(data []byte) {
+	encoding := wrapper.GetContentEncoding()
+
+	bodyBytes, _ := compressutil.Encode(data, encoding)
+
+	if !wrapper.wroteHeader {
+		wrapper.WriteHeader(http.StatusOK)
+	}
+
+	if _, err := wrapper.ResponseWriter.Write(bodyBytes); err != nil {
+		log.Printf("unable to write rewrited body: %v", err)
+	}
+}
+
+// SupportsProcessing determine if http.Request is supported by this plugin.
+func SupportsProcessing(request *http.Request) bool {
+	// Ignore non GET requests
+	if request.Method != "GET" {
+		return false
+	}
+
+	if strings.Contains(request.Header.Get("Upgrade"), "websocket") {
+		log.Printf("Ignoring websocket request for %s", request.RequestURI)
+
+		return false
+	}
+
+	return true
+}
+
+// GetContentEncoding get the Content-Encoding header value.
+func (wrapper *ResponseWrapper) GetContentEncoding() string {
+	return wrapper.Header().Get("Content-Encoding")
+}
+
+// GetContentType get the Content-Encoding header value.
+func (wrapper *ResponseWrapper) GetContentType() string {
+	return wrapper.Header().Get("Content-Type")
+}
+
+// SupportsProcessing determine if HttpWrapper is supported by this plugin based on encoding.
+func (wrapper *ResponseWrapper) SupportsProcessing() bool {
+	contentType := wrapper.GetContentType()
+
+	// If content type does not match return values with false
+	if contentType != "" && !strings.Contains(contentType, "text") {
+		return false
+	}
+
+	encoding := wrapper.GetContentEncoding()
+
+	// If content type is supported validate encoding as well
+	switch encoding {
+	case "gzip":
+		fallthrough
+	case "deflate":
+		fallthrough
+	case "identity":
+		fallthrough
+	case "":
+		return true
+	default:
+		return false
+	}
+}
+
+// SetLastModified update the local lastModified variable from non-package-based users.
+func (wrapper *ResponseWrapper) SetLastModified(value bool) {
+	wrapper.lastModified = value
+}