Compress Middleware for Fiber v3

compress.go

package utils

import (
	"compress/flate"
	"io"
	"lawang/bytes"
	"lawang/configs"

	"github.com/andybalholm/brotli"
	"github.com/gofiber/fiber/v3"
	"github.com/klauspost/compress/gzip"
	"github.com/klauspost/compress/zlib"
	"github.com/klauspost/compress/zstd"
)

func Compress(ctx fiber.Ctx) error {
	acceptEncoding := ctx.Get(configs.LAWANG_ACCEPT_ENCODING_HEADER)
	if acceptEncoding == "" {
		return ctx.Next()
	}

	if err := ctx.Next(); err != nil {
		return err
	}

	body := ctx.Response().Body()
	if len(body) < 1024 {
		return nil
	}

	buffer := bytes.BufferPool.Get()
	defer bytes.BufferPool.Put(buffer)

	var writer io.WriteCloser
	var err error

	switch acceptEncoding {
	case configs.LAWANG_CONTENT_ENCODING_GZIP:
		writer, err = gzip.NewWriterLevel(buffer, gzip.BestCompression)
	case configs.LAWANG_CONTENT_ENCODING_BROTLI:
		writer = brotli.NewWriterLevel(buffer, brotli.BestCompression)
	case configs.LAWANG_CONTENT_ENCODING_ZSTD:
		writer, err = zstd.NewWriter(buffer)
	case configs.LAWANG_CONTENT_ENCODING_DEFLATE:
		writer, err = flate.NewWriter(buffer, flate.BestCompression)
	case configs.LAWANG_CONTENT_ENCODING_ZLIB:
		writer = zlib.NewWriter(buffer)
	default:
		return nil
	}

	if err != nil {
		return err
	}
	defer writer.Close()

	_, err = writer.Write(body)
	if err != nil {
		return err
	}

	if flusher, ok := writer.(interface{ Flush() error }); ok {
		flusher.Flush()
	}

	ctx.Response().Header.Set(configs.LAWANG_CONTENT_ENCODING_HEADER, acceptEncoding)
	ctx.Response().Header.Del(fiber.HeaderContentLength)
	ctx.Response().SetBodyRaw(buffer.Bytes())

	return nil
}