ginzip/ginzip.go

package ginzip

import (
	"compress/gzip"
	"fmt"
	"log"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/andybalholm/brotli"
	"github.com/gin-gonic/gin"
)

// Options object is used to configure Ginzip.
type Options struct {
	// Compression levels for gzip/brotli
	// Default setting: ""/"d"/"default"
	// Best speed: "min"/"speed"
	// Best compression: "max"
	// Compression disabled: "false"/"n"/"no"
	// Manual setting: 1-9 (gzip), 0-11 (brotli)
	GzipLevel   string
	BrotliLevel string

	// Filepath extensions where compression should be skipped
	SkipExtensions []string
}

type intOptions struct {
	GzipEn         bool
	GzipLevel      int
	BrotliEn       bool
	BrotliLevel    int
	SkipExtensions map[string]bool
}

// New creates a new Ginzip middleware function.
// Attach to your Gin router like this:
//  router.Use(ginzip.New(ginzip.DefaultOptions()))
//
// Or use groups:
//  ui := r.Group("/", ginzip.New(ginzip.DefaultOptions()))
func New(options Options) gin.HandlerFunc {
	opt := parseOptions(options)

	return func(c *gin.Context) {
		gWriter := getGinzipWriter(c, opt)
		if gWriter == nil {
			return
		}

		c.Writer = gWriter

		c.Header("Content-Encoding", gWriter.Encoding())
		c.Header("Vary", "Accept-Encoding")
		defer func() {
			gWriter.Close()
			c.Header("Content-Length", fmt.Sprint(c.Writer.Size()))
		}()
		c.Next()
	}
}

// DefaultOptions creates a new config object with default values:
//
// Default compression level, ignore most common binary filetypes
//  .png, .gif, .jpeg, .jpg, .mp3, .mp4, .ogg, .zip, .exe
func DefaultOptions() Options {
	return Options{
		GzipLevel:   "",
		BrotliLevel: "",
		SkipExtensions: []string{
			".png", ".gif", ".jpeg", ".jpg", ".mp3", ".mp4", ".ogg",
			".zip", ".exe",
		},
	}
}

// IgnoreExt adds more ignored file extension to the Options object.
func (o Options) IgnoreExt(exts ...string) Options {
	o.SkipExtensions = append(o.SkipExtensions, exts...)
	return o
}

// ginzipWriter is Ginzip's interface for the individual writer types
// handling the different compression algorithms.
type ginzipWriter interface {
	gin.ResponseWriter
	WriteString(s string) (int, error)
	Write(data []byte) (int, error)
	Close() error
	Encoding() string
}

type gzipWriter struct {
	gin.ResponseWriter
	writer *gzip.Writer
}

type brotliWriter struct {
	gin.ResponseWriter
	writer *brotli.Writer
}

func (g *gzipWriter) WriteString(s string) (int, error) {
	return g.writer.Write([]byte(s))
}

func (g *gzipWriter) Write(data []byte) (int, error) {
	return g.writer.Write(data)
}

func (g *gzipWriter) Close() error {
	return g.writer.Close()
}

func (g *gzipWriter) Encoding() string {
	return "gzip"
}

func (b *brotliWriter) WriteString(s string) (int, error) {
	return b.writer.Write([]byte(s))
}

func (b *brotliWriter) Write(data []byte) (int, error) {
	return b.writer.Write(data)
}

func (b *brotliWriter) Close() error {
	return b.writer.Close()
}

func (b *brotliWriter) Encoding() string {
	return "br"
}

func shouldCompress(req *http.Request, skipExtensions map[string]bool) bool {
	// Dont compress websocket connections
	if strings.Contains(req.Header.Get("Connection"), "Upgrade") ||
		strings.Contains(req.Header.Get("Content-Type"), "text/event-stream") {
		return false
	}

	extension := filepath.Ext(req.URL.Path)
	if extension == "" {
		return true
	}

	_, doSkip := skipExtensions[extension]
	return !doSkip
}

func parseLevel(level string, dfault int, min int, max int) (int, bool) {
	level = strings.ToLower(level)

	switch level {
	case "", "default", "d":
		return dfault, true
	case "false", "no", "n":
		return 0, false
	case "max":
		return max, true
	case "min", "speed":
		return min, true
	}

	i, err := strconv.Atoi(level)
	if err == nil && i >= min && i <= max {
		return i, true
	}

	log.Printf("Ginzip level %s is invalid, falling back to default", level)
	return dfault, true
}

func parseOptions(options Options) intOptions {
	gzipLvl, gzipEn := parseLevel(
		options.GzipLevel,
		gzip.DefaultCompression, gzip.BestSpeed, gzip.BestCompression,
	)
	brLvl, brEn := parseLevel(
		options.BrotliLevel,
		brotli.DefaultCompression, brotli.BestSpeed, brotli.BestCompression,
	)

	exts := map[string]bool{}
	for _, ext := range options.SkipExtensions {
		exts[ext] = true
	}

	return intOptions{
		GzipEn:         gzipEn,
		GzipLevel:      gzipLvl,
		BrotliEn:       brEn,
		BrotliLevel:    brLvl,
		SkipExtensions: exts,
	}
}

func getGinzipWriter(c *gin.Context, options intOptions) ginzipWriter {
	if !shouldCompress(c.Request, options.SkipExtensions) {
		return nil
	}

	if strings.Contains(c.Request.Header.Get("Accept-Encoding"), "br") &&
		options.BrotliEn {
		br := brotli.NewWriterLevel(c.Writer, options.BrotliLevel)

		return &brotliWriter{c.Writer, br}
	}
	if strings.Contains(c.Request.Header.Get("Accept-Encoding"), "gzip") &&
		options.GzipEn {
		gz, err := gzip.NewWriterLevel(c.Writer, options.GzipLevel)
		if err != nil {
			return nil
		}

		return &gzipWriter{c.Writer, gz}
	}
	return nil
}