handle encodings

compression.go

@@ -0,0 +1,102 @@
+package traefik_umami_plugin
+
+import (
+	"bytes"
+	"compress/flate"
+	"compress/gzip"
+	"io"
+)
+
+const (
+	// Gzip compression algorithm string.
+	Gzip string = "gzip"
+	// Deflate compression algorithm string.
+	Deflate string = "deflate"
+	// Identity compression algorithm string.
+	Identity string = "identity"
+)
+
+// 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 *Encoding) ([]byte, error) {
+	reader, err := GetRawReader(byteReader, encoding)
+	if err != nil {
+		return nil, &ReaderError{
+			error: err,
+			cause: err,
+		}
+	}
+
+	return io.ReadAll(reader)
+}
+
+func GetRawReader(byteReader *bytes.Buffer, encoding *Encoding) (io.Reader, error) {
+	switch encoding.name {
+	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 *Encoding) ([]byte, error) {
+	switch encoding.name {
+	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
+}

encoding.go

@@ -0,0 +1,74 @@
+package traefik_umami_plugin
+
+import (
+	"strconv"
+	"strings"
+)
+
+type Encoding struct {
+	name string
+	q    float64
+}
+
+func ParseEncoding(encoding string) *Encoding {
+	return &Encoding{
+		name: encoding,
+		q:    1.0,
+	}
+}
+
+type Encodings struct {
+	encodings []Encoding
+}
+
+func ParseEncodings(acceptEncoding string) *Encodings {
+	encodingList := strings.Split(acceptEncoding, ",")
+	result := make([]Encoding, 0, len(encodingList))
+
+	for _, encoding := range encodingList {
+		split := strings.Split(strings.TrimSpace(encoding), ";q=")
+		q := 1.0
+		if len(split) > 1 {
+			q, _ = strconv.ParseFloat(split[1], 64)
+		}
+		result = append(result, Encoding{name: split[0], q: q})
+	}
+
+	return &Encodings{encodings: result}
+}
+
+func (ae *Encodings) String() string {
+	result := make([]string, 0, len(ae.encodings))
+
+	for _, encoding := range ae.encodings {
+		result = append(result, encoding.name)
+	}
+
+	return strings.Join(result, ",")
+}
+
+func (ae *Encodings) FilterSupported() *Encodings {
+	result := make([]Encoding, 0, len(ae.encodings))
+
+	for _, encoding := range ae.encodings {
+		switch encoding.name {
+		case Gzip, Deflate, Identity:
+			result = append(result, encoding)
+		}
+	}
+
+	return &Encodings{encodings: result}
+}
+
+func (ae *Encodings) GetPreferred() *Encoding {
+	maxQ := 0.0
+	var preferred Encoding
+	for _, encoding := range ae.encodings {
+		if encoding.q > maxQ {
+			preferred = encoding
+			maxQ = encoding.q
+		}
+	}
+
+	return &preferred
+}

plugin.go

@@ -142,21 +142,21 @@ func (h *PluginHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 
 	// script injection
 	var injected bool = false
-	if h.config.ScriptInjection {
+	myReq := &Request{Request: *req}
+	if h.config.ScriptInjection && myReq.CouldBeInjectable() {
 		// intercept body
-		myrw := &responseWriter{
-			buffer:         &bytes.Buffer{},
-			ResponseWriter: rw,
-		}
-		myrw.Header().Set("Accept-Encoding", "identity")
-		h.next.ServeHTTP(myrw, req)
+		myRw := NewResponseWriter(rw)
+		myReq.SetSupportedEncoding()
+		h.next.ServeHTTP(myRw, &myReq.Request)
 
-		if myrw.Header().Get("Content-Type") == "text/html" {
+		if myRw.IsInjectable() {
 			// h.log(fmt.Sprintf("Inject %s", req.URL.EscapedPath()))
-			bytes := myrw.buffer.Bytes()
-			newBytes := regexReplaceSingle(bytes, insertBeforeRegex, h.scriptHtml)
-			rw.Write(newBytes)
-			injected = true
+			body, err := myRw.ReadDecoded()
+			if err != nil {
+				newBody := InsertAtBodyEnd(body, h.scriptHtml)
+				myRw.WriteEncoded(newBody, myReq.GetPreferredSupportedEncoding())
+				injected = true
+			}
 		}
 	}
 
@@ -177,8 +177,3 @@ type responseWriter struct {
 	buffer *bytes.Buffer
 	http.ResponseWriter
 }
-
-func (w *responseWriter) Write(p []byte) (int, error) {
-	w.buffer.Reset()
-	return w.buffer.Write(p)
-}

regex.go

@@ -0,0 +1,23 @@
+package traefik_umami_plugin
+
+import (
+	"regexp"
+)
+
+const insertBeforeRegexPattern = `</body>`
+
+var insertBeforeRegex = regexp.MustCompile(insertBeforeRegexPattern)
+
+func RegexReplaceSingle(bytes []byte, match *regexp.Regexp, replace string) []byte {
+	rx := match.FindIndex(bytes)
+	if len(rx) == 0 {
+		return bytes
+	}
+	// insert the script before the head tag
+	newBytes := append(bytes[:rx[0]], append([]byte(replace), bytes[rx[0]:]...)...)
+	return newBytes
+}
+
+func InsertAtBodyEnd(bytes []byte, replace string) []byte {
+	return RegexReplaceSingle(bytes, insertBeforeRegex, replace)
+}

request.go

@@ -0,0 +1,39 @@
+package traefik_umami_plugin
+
+import (
+	"net/http"
+	"strings"
+)
+
+type Request struct {
+	http.Request
+}
+
+func (req *Request) SetSupportedEncoding() {
+	acceptEncoding := ParseEncodings(req.Header.Get("Accept-Encoding"))
+	supported := acceptEncoding.FilterSupported().String()
+	req.Header.Set("Accept-Encoding", supported)
+}
+
+func (req *Request) GetPreferredSupportedEncoding() *Encoding {
+	acceptEncoding := req.Header.Get("Accept-Encoding")
+	return ParseEncodings(acceptEncoding).FilterSupported().GetPreferred()
+}
+
+func (req *Request) CouldBeInjectable() bool {
+	// return false on non-GET requests
+	if req.Method != http.MethodGet {
+		return false
+	}
+
+	// ignore websockets
+	if strings.Contains(req.Header.Get("Upgrade"), "websocket") {
+		return false
+	}
+
+	return true
+}
+
+func (req *Request) IsHtml() bool {
+	return strings.Contains(req.Header.Get("Accept"), "text/html")
+}

response_writer.go

@@ -0,0 +1,65 @@
+package traefik_umami_plugin
+
+import (
+	"bytes"
+	"net/http"
+)
+
+type ResponseWriter struct {
+	buffer *bytes.Buffer
+	http.ResponseWriter
+}
+
+func NewResponseWriter(rw http.ResponseWriter) *ResponseWriter {
+	return &ResponseWriter{
+		buffer:         &bytes.Buffer{},
+		ResponseWriter: rw,
+	}
+}
+
+func (w *ResponseWriter) IsInjectable() bool {
+	return w.Header().Get("Content-Type") == "text/html"
+}
+
+// Body bytes
+// Might be compressed.
+func (w *ResponseWriter) Read() []byte {
+	return w.buffer.Bytes()
+}
+
+// Body bytes
+// Always uncompressed
+// Error if encoding is not supported.
+func (w *ResponseWriter) ReadDecoded() ([]byte, error) {
+	encoding := w.GetContentEncoding()
+	return Decode(w.buffer, encoding)
+}
+
+// Write body bytes.
+func (w *ResponseWriter) Write(p []byte) (int, error) {
+	w.buffer.Reset()
+	return w.buffer.Write(p)
+}
+
+// Write body bytes
+// Compresses the body to the target encoding.
+func (w *ResponseWriter) WriteEncoded(plain []byte, encoding *Encoding) (int, error) {
+	encoded, err := Encode(plain, encoding)
+	if err != nil {
+		return 0, err
+	}
+	w.Write(encoded)
+	w.SetContentEncoding(encoding)
+	return len(plain), nil
+}
+
+// Content-Encoding header.
+func (w *ResponseWriter) GetContentEncoding() *Encoding {
+	str := w.Header().Get("Content-Encoding")
+	return ParseEncoding(str)
+}
+
+// Set Content-Encoding header.
+func (w *ResponseWriter) SetContentEncoding(encoding *Encoding) {
+	w.Header().Set("Content-Encoding", encoding.name)
+}

umami_script.go

@@ -6,24 +6,9 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-	"regexp"
 	"strings"
 )
 
-const insertBeforeRegexPattern = `</body>`
-
-var insertBeforeRegex = regexp.MustCompile(insertBeforeRegexPattern)
-
-// injects the umami script into the response head.
-func regexReplaceSingle(bytes []byte, match *regexp.Regexp, replace string) []byte {
-	rx := match.FindIndex(bytes)
-	if len(rx) == 0 {
-		return bytes
-	}
-	// insert the script before the head tag
-	return append(bytes[:rx[0]], append([]byte(replace), bytes[rx[0]:]...)...)
-}
-
 // builds the umami script.
 func buildUmamiScript(config *Config) (string, error) {
 	// check if the script should be injected