github.com/aacfactory/fns@v1.2.86-0.20240310083819-80d667fc0a17/transports/middlewares/compress/middleware.go (about)

     1  /*
     2   * Copyright 2023 Wang Min Xiang
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   * 	http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   *
    16   */
    17  
    18  package compress
    19  
    20  import (
    21  	"bytes"
    22  	"github.com/aacfactory/errors"
    23  	"github.com/aacfactory/fns/commons/bytex"
    24  	"github.com/aacfactory/fns/transports"
    25  	"github.com/aacfactory/logs"
    26  	"github.com/valyala/fasthttp"
    27  	"slices"
    28  )
    29  
    30  var (
    31  	strApplicationSlash = []byte("application/")
    32  	strImageSVG         = []byte("image/svg")
    33  	strImageIcon        = []byte("image/x-icon")
    34  	strFontSlash        = []byte("font/")
    35  	strMultipartSlash   = []byte("multipart/")
    36  	strTextSlash        = []byte("text/")
    37  )
    38  
    39  const (
    40  	minCompressLen = 200
    41  )
    42  
    43  func New() transports.Middleware {
    44  	return &Middleware{}
    45  }
    46  
    47  type Middleware struct {
    48  	log        logs.Logger
    49  	enable     bool
    50  	compressor Compressor
    51  	gzip       *GzipCompressor
    52  	deflate    *DeflateCompressor
    53  	brotli     *BrotliCompressor
    54  }
    55  
    56  func (middle *Middleware) Name() string {
    57  	return "compress"
    58  }
    59  
    60  func (middle *Middleware) Construct(options transports.MiddlewareOptions) error {
    61  	middle.log = options.Log
    62  	config := Config{}
    63  	configErr := options.Config.As(&config)
    64  	if configErr != nil {
    65  		return errors.Warning("fns: construct compress middleware failed").WithCause(configErr)
    66  	}
    67  	if !config.Enable {
    68  		return nil
    69  	}
    70  	// gzip
    71  	gzipLevel := config.GzipLevel
    72  	if !slices.Contains([]int{fasthttp.CompressBestSpeed, fasthttp.CompressBestCompression, fasthttp.CompressDefaultCompression, fasthttp.CompressHuffmanOnly}, gzipLevel) {
    73  		gzipLevel = fasthttp.CompressDefaultCompression
    74  	}
    75  	middle.gzip = &GzipCompressor{
    76  		level: gzipLevel,
    77  	}
    78  	// deflate
    79  	deflateLevel := config.DeflateLevel
    80  	if !slices.Contains([]int{fasthttp.CompressBestSpeed, fasthttp.CompressBestCompression, fasthttp.CompressDefaultCompression, fasthttp.CompressHuffmanOnly}, deflateLevel) {
    81  		deflateLevel = fasthttp.CompressDefaultCompression
    82  	}
    83  	middle.deflate = &DeflateCompressor{
    84  		level: deflateLevel,
    85  	}
    86  	// brotli
    87  	brotliLevel := config.BrotliLevel
    88  	if !slices.Contains([]int{fasthttp.CompressBrotliBestSpeed, fasthttp.CompressBrotliBestCompression, fasthttp.CompressBrotliDefaultCompression}, brotliLevel) {
    89  		brotliLevel = fasthttp.CompressBrotliDefaultCompression
    90  	}
    91  	middle.brotli = &BrotliCompressor{
    92  		level: brotliLevel,
    93  	}
    94  	switch config.Default {
    95  	case BrotliName:
    96  		middle.compressor = middle.brotli
    97  		break
    98  	case DeflateName:
    99  		middle.compressor = middle.deflate
   100  		break
   101  	case GzipName:
   102  		middle.compressor = middle.gzip
   103  		break
   104  	default:
   105  		middle.compressor = middle.deflate
   106  		break
   107  	}
   108  	middle.enable = true
   109  	return nil
   110  }
   111  
   112  func (middle *Middleware) Handler(next transports.Handler) transports.Handler {
   113  	if middle.enable {
   114  		return transports.HandlerFunc(func(w transports.ResponseWriter, r transports.Request) {
   115  			next.Handle(w, r)
   116  			if w.BodyLen() < minCompressLen {
   117  				return
   118  			}
   119  			contentType := w.Header().Get(transports.ContentTypeHeaderName)
   120  			canCompress := bytes.HasPrefix(contentType, strTextSlash) ||
   121  				bytes.HasPrefix(contentType, strApplicationSlash) ||
   122  				bytes.HasPrefix(contentType, strImageSVG) ||
   123  				bytes.HasPrefix(contentType, strImageIcon) ||
   124  				bytes.HasPrefix(contentType, strFontSlash) ||
   125  				bytes.HasPrefix(contentType, strMultipartSlash)
   126  			if !canCompress {
   127  				return
   128  			}
   129  			kind := getKind(r)
   130  			var c Compressor = nil
   131  			switch kind {
   132  			case Any, Default:
   133  				c = middle.compressor
   134  				break
   135  			case Gzip:
   136  				c = middle.gzip
   137  				break
   138  			case Deflate:
   139  				c = middle.deflate
   140  				break
   141  			case Brotli:
   142  				c = middle.brotli
   143  				break
   144  			default:
   145  				break
   146  			}
   147  			if c == nil {
   148  				return
   149  			}
   150  			body := w.Body()
   151  			compressed, compressErr := c.Compress(body)
   152  			if compressErr != nil {
   153  				if middle.log.WarnEnabled() {
   154  					middle.log.Warn().Cause(compressErr).With("compress", kind.String()).Message("fns: compress response body failed")
   155  				}
   156  				return
   157  			}
   158  			// header
   159  			w.Header().Set(transports.ContentEncodingHeaderName, bytex.FromString(c.Name()))
   160  			w.Header().Add(transports.VaryHeaderName, transports.AcceptEncodingHeaderName)
   161  			// body
   162  			w.ResetBody()
   163  			_, _ = w.Write(compressed)
   164  		})
   165  	}
   166  	return next
   167  }
   168  
   169  func (middle *Middleware) Close() (err error) {
   170  	return
   171  }