Merge pull request #1 from askrella/compression-caching

Add cache and gzip compression

Dockerfile

@@ -1,14 +1,14 @@
-FROM golang:1.21 as build
+FROM golang:1.22 as build
 
 WORKDIR /go/src/app
 COPY go.mod go.mod
-#COPY go.sum go.sum
+COPY go.sum go.sum
 COPY cmd cmd
 COPY internal internal
 
 RUN go mod download
 
-RUN CGO_ENABLED=0 go build -o /go/bin/app ./cmd/main.go
+RUN CGO_ENABLED=0 go build -ldflags="-w -s" -o /go/bin/app ./cmd/main.go
 
 FROM gcr.io/distroless/static-debian12
 

README.md

@@ -1,7 +1,5 @@
 # Go MockServer
 
-> NOTE: Currently WIP, only forwarding works.
-
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 ![Docker](https://github.com/askrella/whatsapp-chatgpt/actions/workflows/docker.yml/badge.svg)
 ![Docker AMD64](https://img.shields.io/badge/docker-amd64-blue)
@@ -11,21 +9,23 @@
 ![Askrella](https://avatars.githubusercontent.com/u/77694724?s=100)
 
 A tiny (< 5mb) implementation of a request capturing proxy written in Golang. Used internally at [Askrella](https://askrella.de/)
-for mocking external partners in our CI pipelines.
+for mocking external partners in our CI pipelines and sometimes improving compression and loading times for CSV
+datasets.
 
 # Why?
 
 We at [Askrella](https://askrella.de/) sometimes encounter projects with minimal own logic rather than gluing
 multiple APIs together and process their data. In these cases we had troubles with the response times these
-APIs offer (> 3.5s). Since we want to run extensive tests against these providers in our pipeline, we need
+APIs offer (> 3.5s, rare cases even > 1,4min). Since we want to run extensive tests against these providers in our pipeline, we need
 some kind of proxy capturing and mocking these providers in our pipeline.
 
 On our search we encountered [MockServer](https://www.mock-server.com/) written in Java, but the functionality was way to extensive
 and the images too big for fast iterations in our scenarios (possibly dozens of different APIs on a small CI/CD pipeline).
 
 Of course we don't offer the same functionality as the big brothers, but we kept it simple and our images
-minimal (< 5mb)!
-The code itself doesn't need any dependencies and only relies on the Go standard library.
+minimal (< 10mb)!
+The code itself doesn't need many dependencies and only relies on the Go standard library as well as some
+curated compression libraries (e.g. brotli).
 Our final container is based on [distroless](https://github.com/GoogleContainerTools/distroless).
 
 # :gear: Getting Started
@@ -41,6 +41,10 @@ services:
       MOCK_TARGET: https://google.com 
       # MOCK_PORT: 80 # The port MockServer shall bind to
       # MOCK_HOST: 0.0.0.0 # The interface MockServer shall use
+      # CACHE_ENABLED: true # Whether responses to requests shall be cached and re-delivered similar requests
+      # RECOMPRESS: true # When enabled the upstream (external server) response body will be decompressed and compressed
+                         # using gzip level 9 compression. Depending on the upstream, this can lead to extremely reduced
+                         # request times.
     ports:
       - "8080:80/tcp"
 ```
@@ -49,8 +53,8 @@ services:
 
 # :wave: Contributors
 
-<a href="https://github.com/askrella/aws-ses-mock/graphs/contributors">
-  <img src="https://contrib.rocks/image?repo=askrella/aws-ses-mock" />
+<a href="https://github.com/askrella/go-mockserver/graphs/contributors">
+  <img src="https://contrib.rocks/image?repo=askrella/go-mockserver" />
 </a>
 
 * [Askrella Software Agency](askrella.de)

compose.yml

@@ -5,5 +5,7 @@ services:
     build: .
     environment:
       MOCK_TARGET: https://google.com
+      CACHE_ENABLED: true
+      RECOMPRESS: true
     ports:
       - "8080:80/tcp"

go.mod

@@ -1,3 +1,5 @@
 module github.com/askrella/go-mockserver
 
-go 1.21.1
+go 1.22
+
+require github.com/andybalholm/brotli v1.1.0

go.sum

@@ -0,0 +1,2 @@
+github.com/andybalholm/brotli v1.1.0 h1:eLKJA0d02Lf0mVpIDgYnqXcUn0GqVmEFny3VuID1U3M=
+github.com/andybalholm/brotli v1.1.0/go.mod h1:sms7XGricyQI9K10gOSf56VKKWS4oLer58Q+mhRPtnY=

internal/compression/compression.go

@@ -0,0 +1,67 @@
+package compression
+
+import (
+	"bytes"
+	gzip2 "compress/gzip"
+	"github.com/andybalholm/brotli"
+	"io"
+	"log"
+)
+
+func DecompressGzip(data []byte) (decompressed []byte) {
+	reader, err := gzip2.NewReader(bytes.NewBuffer(data))
+	if err != nil {
+		log.Panicln("Cannot create new gzip reader:", err)
+	}
+	defer func(reader *gzip2.Reader) {
+		err := reader.Close()
+		if err != nil {
+			log.Panicln("Error closing gzip reader: ", err)
+		}
+	}(reader)
+
+	decompressed, err = io.ReadAll(reader)
+	if err != nil {
+		log.Panicln("Reading decompressed data from gzip reader caused an error: ", err)
+	}
+
+	return decompressed
+}
+
+func CompressBrotli(data []byte) (compressed []byte, contentEncoding string) {
+	out := bytes.Buffer{}
+	writer := brotli.NewWriterV2(&out, brotli.BestSpeed)
+	_, err := writer.Write(data)
+	if err != nil {
+		log.Panicln("Error writing to brotli writer: ", err)
+	}
+
+	err = writer.Close()
+	if err != nil {
+		log.Panicln("Error closing brotli writer: ", err)
+	}
+
+	return out.Bytes(), "br"
+}
+
+func CompressGzip(data []byte) (compressed []byte, contentEncoding string) {
+	var b bytes.Buffer
+	gzw, err := gzip2.NewWriterLevel(&b, 9)
+	if err != nil {
+		log.Panicln("Error creating gzip writer:", err)
+	}
+	_, err = gzw.Write(data)
+	if err != nil {
+		log.Panicln("gzip write error:", err)
+	}
+	err = gzw.Flush()
+	if err != nil {
+		log.Panicln("gzip flush error:", err)
+	}
+	err = gzw.Close()
+	if err != nil {
+		log.Panicln("gzip close error:", err)
+	}
+
+	return b.Bytes(), "gzip"
+}

internal/config/server.go

@@ -7,6 +7,24 @@ import (
 	"strings"
 )
 
+func Recompress() bool {
+	recompress := os.Getenv("RECOMPRESS")
+	if len(recompress) == 0 {
+		return true
+	}
+
+	return recompress == "true"
+}
+
+func CacheEnabled() bool {
+	cacheEnabled := os.Getenv("CACHE_ENABLED")
+	if cacheEnabled == "" {
+		return true
+	}
+
+	return cacheEnabled == "true"
+}
+
 func GetHost() string {
 	host := os.Getenv("MOCK_HOST")
 	if host == "" {

internal/server/cache_compression_handler.go

@@ -0,0 +1,127 @@
+package server
+
+import (
+	"bytes"
+	"github.com/askrella/go-mockserver/internal/compression"
+	"github.com/askrella/go-mockserver/internal/config"
+	"github.com/askrella/go-mockserver/internal/store"
+	"io"
+	"log"
+	"net/http"
+	"strconv"
+)
+
+// CacheCompressionHandler provides the ability to
+// 1. Store and replay requests and
+// 2. DecompressGzip upstream and compress with stronger algorithms/levels for downstream
+func CacheCompressionHandler(handler http.Handler) http.Handler {
+	return http.HandlerFunc(func(respWriter http.ResponseWriter, request *http.Request) {
+		if config.CacheEnabled() {
+			if cacheUsed := useTransactionResponseCache(respWriter, request); cacheUsed {
+				return
+			}
+		}
+
+		var resultingBody []byte
+
+		// Intercept the response request with a custom request and get the body data.
+		var buffer bytes.Buffer
+		decompressedResponseWriter := &responseWriterInterceptor{ResponseWriter: respWriter, Writer: &buffer}
+		handler.ServeHTTP(decompressedResponseWriter, request)
+		body := buffer.Bytes()
+		resultingBody = body
+
+		// Check if the content encoding is gzip and in case it is not, skip
+		currentEncoding := decompressedResponseWriter.Header().Get("Content-Encoding")
+		if currentEncoding != "gzip" {
+			log.Println("Skip compression, content-encoding: " + currentEncoding)
+			_, err := respWriter.Write(body)
+			if err != nil {
+				respWriter.WriteHeader(http.StatusInternalServerError)
+				log.Println("Error writing original response without compression: ", err)
+				return
+			}
+
+			return
+		}
+
+		if config.Recompress() {
+			// DecompressGzip and re-compress. This seems weird at first, but we want to control how the content
+			// gets compressed and sent to the downstream. For example: when loading a huge CSV dataset, we saw substantial
+			// performance improvements by using level 9 gzip instead of the original gzip compression.
+			decompressed := compression.DecompressGzip(body)
+			compressed, contentEncoding := compression.CompressGzip(decompressed)
+			contentLength := strconv.FormatInt(int64(len(compressed)), 10)
+
+			respWriter.Header().Set("Content-Encoding", contentEncoding)
+			respWriter.Header().Set("Content-Length", contentLength)
+			_, err := respWriter.Write(compressed)
+			if err != nil {
+				respWriter.WriteHeader(http.StatusInternalServerError)
+				log.Println("Failed to write response", err)
+				return
+			}
+
+			resultingBody = compressed
+		}
+
+		if config.CacheEnabled() {
+			storeTransactionResponse(respWriter, request, resultingBody)
+		}
+	})
+}
+
+func useTransactionResponseCache(respWriter http.ResponseWriter, request *http.Request) (cacheUsed bool) {
+	transaction, found := store.FindTransaction(request)
+	// Cache handling: If a transaction is found, use the body from the cached transaction.
+	if !found {
+		return false
+	}
+	log.Println("Found transaction, using cache.")
+	response := transaction.Response
+
+	// Write headers from cached transaction to downstream
+	for key, value := range response.Headers {
+		for _, headerValue := range value {
+			respWriter.Header().Add(key, headerValue)
+		}
+	}
+
+	_, err := respWriter.Write(response.Body)
+	if err != nil {
+		log.Println(err)
+		respWriter.WriteHeader(http.StatusInternalServerError)
+		return true
+	}
+
+	return true
+}
+
+func storeTransactionResponse(respWriter http.ResponseWriter, request *http.Request, compressed []byte) {
+	// Find transaction and store our compressed content
+	transaction, found := store.FindTransaction(request)
+	if !found {
+		respWriter.WriteHeader(http.StatusInternalServerError)
+		log.Panicln("Transaction not found")
+		return
+	}
+
+	transaction.Response = store.Response{
+		Body:    compressed,
+		Headers: respWriter.Header(),
+	}
+}
+
+type responseWriterInterceptor struct {
+	http.ResponseWriter
+	Body   io.Reader
+	Writer io.Writer
+}
+
+func (w *responseWriterInterceptor) Write(b []byte) (int, error) {
+	return w.Writer.Write(b)
+}
+
+func (w *responseWriterInterceptor) WriteHeader(statusCode int) {
+	w.ResponseWriter.WriteHeader(statusCode)
+}

internal/server/server.go

@@ -7,34 +7,69 @@ import (
 	"log"
 	"net/http"
 	"net/http/httputil"
+	_ "net/http/pprof"
 	"net/url"
 	"strconv"
+	"strings"
 )
 
 func InitializeServer() {
+	go func() {
+		err := http.ListenAndServe("0.0.0.0:8080", nil)
+		if err != nil {
+			log.Panicln("Cannot start server for pprof: ", err)
+		}
+	}()
+
 	// Parse the target URL
 	target, err := url.Parse(config.GetTargetURI())
 	if err != nil {
 		log.Fatal("Error parsing target URL: " + err.Error())
 	}
 
+	mux := http.NewServeMux()
+
 	// Create a reverse proxy to forward requests to the target host
 	reverseProxy := httputil.NewSingleHostReverseProxy(target)
 
 	// Create a custom handler function that uses the reverse proxy
-	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-		log.Println("Received request:", r.Method, r.URL)
-
+	reverseProxyHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Capture the request for our cache
 		store.CaptureRequest(r)
+		prepareRequest(r)
 
 		// Forward the request to the target host
 		reverseProxy.ServeHTTP(w, r)
 	})
 
+	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		reverseProxyHandler.ServeHTTP(w, r)
+	})
+
+	handler := CacheCompressionHandler(mux)
+
 	// Start the Go server on port 8080
 	port := config.GetHost() + ":" + strconv.Itoa(config.GetServerPort())
 	fmt.Println("Server listening on port", port)
-	if err := http.ListenAndServe(port, nil); err != nil {
+	if err := http.ListenAndServe(port, handler); err != nil {
 		log.Fatal("Error starting server: " + err.Error())
 	}
 }
+
+func prepareRequest(req *http.Request) {
+	req.Header.Del("Postman-Token")
+
+	// Let's pretend we are chrome
+	req.Header.Set(
+		"User-Agent",
+		"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
+	)
+
+	host, _ := strings.CutPrefix(config.GetTargetURI(), "https://")
+	host, _ = strings.CutPrefix(host, "http://")
+
+	req.Header.Set("Host", host)
+	req.Header.Set("Accept-Encoding", "gzip, deflate, br")
+
+	req.Host = host
+}