481 lines
12 KiB
Go
481 lines
12 KiB
Go
|
package fasthttp
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"fmt"
|
||
|
|
"io"
|
||
|
|
"io/fs"
|
||
|
|
"sync"
|
||
|
|
|
||
|
|
"github.com/klauspost/compress/flate"
|
||
|
|
"github.com/klauspost/compress/gzip"
|
||
|
|
"github.com/klauspost/compress/zlib"
|
||
|
|
"github.com/valyala/bytebufferpool"
|
||
|
|
"github.com/valyala/fasthttp/stackless"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Supported compression levels.
|
||
|
|
const (
|
||
|
|
CompressNoCompression = flate.NoCompression
|
||
|
|
CompressBestSpeed = flate.BestSpeed
|
||
|
|
CompressBestCompression = flate.BestCompression
|
||
|
|
CompressDefaultCompression = 6 // flate.DefaultCompression
|
||
|
|
CompressHuffmanOnly = -2 // flate.HuffmanOnly
|
||
|
|
)
|
||
|
|
|
||
|
|
func acquireGzipReader(r io.Reader) (*gzip.Reader, error) {
|
||
|
|
v := gzipReaderPool.Get()
|
||
|
|
if v == nil {
|
||
|
|
return gzip.NewReader(r)
|
||
|
|
}
|
||
|
|
zr := v.(*gzip.Reader)
|
||
|
|
if err := zr.Reset(r); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return zr, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func releaseGzipReader(zr *gzip.Reader) {
|
||
|
|
zr.Close()
|
||
|
|
gzipReaderPool.Put(zr)
|
||
|
|
}
|
||
|
|
|
||
|
|
var gzipReaderPool sync.Pool
|
||
|
|
|
||
|
|
func acquireFlateReader(r io.Reader) (io.ReadCloser, error) {
|
||
|
|
v := flateReaderPool.Get()
|
||
|
|
if v == nil {
|
||
|
|
zr, err := zlib.NewReader(r)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return zr, nil
|
||
|
|
}
|
||
|
|
zr := v.(io.ReadCloser)
|
||
|
|
if err := resetFlateReader(zr, r); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return zr, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func releaseFlateReader(zr io.ReadCloser) {
|
||
|
|
zr.Close()
|
||
|
|
flateReaderPool.Put(zr)
|
||
|
|
}
|
||
|
|
|
||
|
|
func resetFlateReader(zr io.ReadCloser, r io.Reader) error {
|
||
|
|
zrr, ok := zr.(zlib.Resetter)
|
||
|
|
if !ok {
|
||
|
|
// sanity check. should only be called with a zlib.Reader
|
||
|
|
panic("BUG: zlib.Reader doesn't implement zlib.Resetter???")
|
||
|
|
}
|
||
|
|
return zrr.Reset(r, nil)
|
||
|
|
}
|
||
|
|
|
||
|
|
var flateReaderPool sync.Pool
|
||
|
|
|
||
|
|
func acquireStacklessGzipWriter(w io.Writer, level int) stackless.Writer {
|
||
|
|
nLevel := normalizeCompressLevel(level)
|
||
|
|
p := stacklessGzipWriterPoolMap[nLevel]
|
||
|
|
v := p.Get()
|
||
|
|
if v == nil {
|
||
|
|
return stackless.NewWriter(w, func(w io.Writer) stackless.Writer {
|
||
|
|
return acquireRealGzipWriter(w, level)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
sw := v.(stackless.Writer)
|
||
|
|
sw.Reset(w)
|
||
|
|
return sw
|
||
|
|
}
|
||
|
|
|
||
|
|
func releaseStacklessGzipWriter(sw stackless.Writer, level int) {
|
||
|
|
sw.Close()
|
||
|
|
nLevel := normalizeCompressLevel(level)
|
||
|
|
p := stacklessGzipWriterPoolMap[nLevel]
|
||
|
|
p.Put(sw)
|
||
|
|
}
|
||
|
|
|
||
|
|
func acquireRealGzipWriter(w io.Writer, level int) *gzip.Writer {
|
||
|
|
nLevel := normalizeCompressLevel(level)
|
||
|
|
p := realGzipWriterPoolMap[nLevel]
|
||
|
|
v := p.Get()
|
||
|
|
if v == nil {
|
||
|
|
zw, err := gzip.NewWriterLevel(w, level)
|
||
|
|
if err != nil {
|
||
|
|
// gzip.NewWriterLevel only errors for invalid
|
||
|
|
// compression levels. Clamp it to be min or max.
|
||
|
|
if level < gzip.HuffmanOnly {
|
||
|
|
level = gzip.HuffmanOnly
|
||
|
|
} else {
|
||
|
|
level = gzip.BestCompression
|
||
|
|
}
|
||
|
|
zw, _ = gzip.NewWriterLevel(w, level)
|
||
|
|
}
|
||
|
|
return zw
|
||
|
|
}
|
||
|
|
zw := v.(*gzip.Writer)
|
||
|
|
zw.Reset(w)
|
||
|
|
return zw
|
||
|
|
}
|
||
|
|
|
||
|
|
func releaseRealGzipWriter(zw *gzip.Writer, level int) {
|
||
|
|
zw.Close()
|
||
|
|
nLevel := normalizeCompressLevel(level)
|
||
|
|
p := realGzipWriterPoolMap[nLevel]
|
||
|
|
p.Put(zw)
|
||
|
|
}
|
||
|
|
|
||
|
|
var (
|
||
|
|
stacklessGzipWriterPoolMap = newCompressWriterPoolMap()
|
||
|
|
realGzipWriterPoolMap = newCompressWriterPoolMap()
|
||
|
|
)
|
||
|
|
|
||
|
|
// AppendGzipBytesLevel appends gzipped src to dst using the given
|
||
|
|
// compression level and returns the resulting dst.
|
||
|
|
//
|
||
|
|
// Supported compression levels are:
|
||
|
|
//
|
||
|
|
// - CompressNoCompression
|
||
|
|
// - CompressBestSpeed
|
||
|
|
// - CompressBestCompression
|
||
|
|
// - CompressDefaultCompression
|
||
|
|
// - CompressHuffmanOnly
|
||
|
|
func AppendGzipBytesLevel(dst, src []byte, level int) []byte {
|
||
|
|
w := &byteSliceWriter{dst}
|
||
|
|
WriteGzipLevel(w, src, level) //nolint:errcheck
|
||
|
|
return w.b
|
||
|
|
}
|
||
|
|
|
||
|
|
// WriteGzipLevel writes gzipped p to w using the given compression level
|
||
|
|
// and returns the number of compressed bytes written to w.
|
||
|
|
//
|
||
|
|
// Supported compression levels are:
|
||
|
|
//
|
||
|
|
// - CompressNoCompression
|
||
|
|
// - CompressBestSpeed
|
||
|
|
// - CompressBestCompression
|
||
|
|
// - CompressDefaultCompression
|
||
|
|
// - CompressHuffmanOnly
|
||
|
|
func WriteGzipLevel(w io.Writer, p []byte, level int) (int, error) {
|
||
|
|
switch w.(type) {
|
||
|
|
case *byteSliceWriter,
|
||
|
|
*bytes.Buffer,
|
||
|
|
*bytebufferpool.ByteBuffer:
|
||
|
|
// These writers don't block, so we can just use stacklessWriteGzip
|
||
|
|
ctx := &compressCtx{
|
||
|
|
w: w,
|
||
|
|
p: p,
|
||
|
|
level: level,
|
||
|
|
}
|
||
|
|
stacklessWriteGzip(ctx)
|
||
|
|
return len(p), nil
|
||
|
|
default:
|
||
|
|
zw := acquireStacklessGzipWriter(w, level)
|
||
|
|
n, err := zw.Write(p)
|
||
|
|
releaseStacklessGzipWriter(zw, level)
|
||
|
|
return n, err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
var (
|
||
|
|
stacklessWriteGzipOnce sync.Once
|
||
|
|
stacklessWriteGzipFunc func(ctx interface{}) bool
|
||
|
|
)
|
||
|
|
|
||
|
|
func stacklessWriteGzip(ctx interface{}) {
|
||
|
|
stacklessWriteGzipOnce.Do(func() {
|
||
|
|
stacklessWriteGzipFunc = stackless.NewFunc(nonblockingWriteGzip)
|
||
|
|
})
|
||
|
|
stacklessWriteGzipFunc(ctx)
|
||
|
|
}
|
||
|
|
|
||
|
|
func nonblockingWriteGzip(ctxv interface{}) {
|
||
|
|
ctx := ctxv.(*compressCtx)
|
||
|
|
zw := acquireRealGzipWriter(ctx.w, ctx.level)
|
||
|
|
|
||
|
|
zw.Write(ctx.p) //nolint:errcheck // no way to handle this error anyway
|
||
|
|
|
||
|
|
releaseRealGzipWriter(zw, ctx.level)
|
||
|
|
}
|
||
|
|
|
||
|
|
// WriteGzip writes gzipped p to w and returns the number of compressed
|
||
|
|
// bytes written to w.
|
||
|
|
func WriteGzip(w io.Writer, p []byte) (int, error) {
|
||
|
|
return WriteGzipLevel(w, p, CompressDefaultCompression)
|
||
|
|
}
|
||
|
|
|
||
|
|
// AppendGzipBytes appends gzipped src to dst and returns the resulting dst.
|
||
|
|
func AppendGzipBytes(dst, src []byte) []byte {
|
||
|
|
return AppendGzipBytesLevel(dst, src, CompressDefaultCompression)
|
||
|
|
}
|
||
|
|
|
||
|
|
// WriteGunzip writes ungzipped p to w and returns the number of uncompressed
|
||
|
|
// bytes written to w.
|
||
|
|
func WriteGunzip(w io.Writer, p []byte) (int, error) {
|
||
|
|
r := &byteSliceReader{p}
|
||
|
|
zr, err := acquireGzipReader(r)
|
||
|
|
if err != nil {
|
||
|
|
return 0, err
|
||
|
|
}
|
||
|
|
n, err := copyZeroAlloc(w, zr)
|
||
|
|
releaseGzipReader(zr)
|
||
|
|
nn := int(n)
|
||
|
|
if int64(nn) != n {
|
||
|
|
return 0, fmt.Errorf("too much data gunzipped: %d", n)
|
||
|
|
}
|
||
|
|
return nn, err
|
||
|
|
}
|
||
|
|
|
||
|
|
// AppendGunzipBytes appends gunzipped src to dst and returns the resulting dst.
|
||
|
|
func AppendGunzipBytes(dst, src []byte) ([]byte, error) {
|
||
|
|
w := &byteSliceWriter{dst}
|
||
|
|
_, err := WriteGunzip(w, src)
|
||
|
|
return w.b, err
|
||
|
|
}
|
||
|
|
|
||
|
|
// AppendDeflateBytesLevel appends deflated src to dst using the given
|
||
|
|
// compression level and returns the resulting dst.
|
||
|
|
//
|
||
|
|
// Supported compression levels are:
|
||
|
|
//
|
||
|
|
// - CompressNoCompression
|
||
|
|
// - CompressBestSpeed
|
||
|
|
// - CompressBestCompression
|
||
|
|
// - CompressDefaultCompression
|
||
|
|
// - CompressHuffmanOnly
|
||
|
|
func AppendDeflateBytesLevel(dst, src []byte, level int) []byte {
|
||
|
|
w := &byteSliceWriter{dst}
|
||
|
|
WriteDeflateLevel(w, src, level) //nolint:errcheck
|
||
|
|
return w.b
|
||
|
|
}
|
||
|
|
|
||
|
|
// WriteDeflateLevel writes deflated p to w using the given compression level
|
||
|
|
// and returns the number of compressed bytes written to w.
|
||
|
|
//
|
||
|
|
// Supported compression levels are:
|
||
|
|
//
|
||
|
|
// - CompressNoCompression
|
||
|
|
// - CompressBestSpeed
|
||
|
|
// - CompressBestCompression
|
||
|
|
// - CompressDefaultCompression
|
||
|
|
// - CompressHuffmanOnly
|
||
|
|
func WriteDeflateLevel(w io.Writer, p []byte, level int) (int, error) {
|
||
|
|
switch w.(type) {
|
||
|
|
case *byteSliceWriter,
|
||
|
|
*bytes.Buffer,
|
||
|
|
*bytebufferpool.ByteBuffer:
|
||
|
|
// These writers don't block, so we can just use stacklessWriteDeflate
|
||
|
|
ctx := &compressCtx{
|
||
|
|
w: w,
|
||
|
|
p: p,
|
||
|
|
level: level,
|
||
|
|
}
|
||
|
|
stacklessWriteDeflate(ctx)
|
||
|
|
return len(p), nil
|
||
|
|
default:
|
||
|
|
zw := acquireStacklessDeflateWriter(w, level)
|
||
|
|
n, err := zw.Write(p)
|
||
|
|
releaseStacklessDeflateWriter(zw, level)
|
||
|
|
return n, err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
var (
|
||
|
|
stacklessWriteDeflateOnce sync.Once
|
||
|
|
stacklessWriteDeflateFunc func(ctx interface{}) bool
|
||
|
|
)
|
||
|
|
|
||
|
|
func stacklessWriteDeflate(ctx interface{}) {
|
||
|
|
stacklessWriteDeflateOnce.Do(func() {
|
||
|
|
stacklessWriteDeflateFunc = stackless.NewFunc(nonblockingWriteDeflate)
|
||
|
|
})
|
||
|
|
stacklessWriteDeflateFunc(ctx)
|
||
|
|
}
|
||
|
|
|
||
|
|
func nonblockingWriteDeflate(ctxv interface{}) {
|
||
|
|
ctx := ctxv.(*compressCtx)
|
||
|
|
zw := acquireRealDeflateWriter(ctx.w, ctx.level)
|
||
|
|
|
||
|
|
zw.Write(ctx.p) //nolint:errcheck // no way to handle this error anyway
|
||
|
|
|
||
|
|
releaseRealDeflateWriter(zw, ctx.level)
|
||
|
|
}
|
||
|
|
|
||
|
|
type compressCtx struct {
|
||
|
|
w io.Writer
|
||
|
|
p []byte
|
||
|
|
level int
|
||
|
|
}
|
||
|
|
|
||
|
|
// WriteDeflate writes deflated p to w and returns the number of compressed
|
||
|
|
// bytes written to w.
|
||
|
|
func WriteDeflate(w io.Writer, p []byte) (int, error) {
|
||
|
|
return WriteDeflateLevel(w, p, CompressDefaultCompression)
|
||
|
|
}
|
||
|
|
|
||
|
|
// AppendDeflateBytes appends deflated src to dst and returns the resulting dst.
|
||
|
|
func AppendDeflateBytes(dst, src []byte) []byte {
|
||
|
|
return AppendDeflateBytesLevel(dst, src, CompressDefaultCompression)
|
||
|
|
}
|
||
|
|
|
||
|
|
// WriteInflate writes inflated p to w and returns the number of uncompressed
|
||
|
|
// bytes written to w.
|
||
|
|
func WriteInflate(w io.Writer, p []byte) (int, error) {
|
||
|
|
r := &byteSliceReader{p}
|
||
|
|
zr, err := acquireFlateReader(r)
|
||
|
|
if err != nil {
|
||
|
|
return 0, err
|
||
|
|
}
|
||
|
|
n, err := copyZeroAlloc(w, zr)
|
||
|
|
releaseFlateReader(zr)
|
||
|
|
nn := int(n)
|
||
|
|
if int64(nn) != n {
|
||
|
|
return 0, fmt.Errorf("too much data inflated: %d", n)
|
||
|
|
}
|
||
|
|
return nn, err
|
||
|
|
}
|
||
|
|
|
||
|
|
// AppendInflateBytes appends inflated src to dst and returns the resulting dst.
|
||
|
|
func AppendInflateBytes(dst, src []byte) ([]byte, error) {
|
||
|
|
w := &byteSliceWriter{dst}
|
||
|
|
_, err := WriteInflate(w, src)
|
||
|
|
return w.b, err
|
||
|
|
}
|
||
|
|
|
||
|
|
type byteSliceWriter struct {
|
||
|
|
b []byte
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *byteSliceWriter) Write(p []byte) (int, error) {
|
||
|
|
w.b = append(w.b, p...)
|
||
|
|
return len(p), nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type byteSliceReader struct {
|
||
|
|
b []byte
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *byteSliceReader) Read(p []byte) (int, error) {
|
||
|
|
if len(r.b) == 0 {
|
||
|
|
return 0, io.EOF
|
||
|
|
}
|
||
|
|
n := copy(p, r.b)
|
||
|
|
r.b = r.b[n:]
|
||
|
|
return n, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *byteSliceReader) ReadByte() (byte, error) {
|
||
|
|
if len(r.b) == 0 {
|
||
|
|
return 0, io.EOF
|
||
|
|
}
|
||
|
|
n := r.b[0]
|
||
|
|
r.b = r.b[1:]
|
||
|
|
return n, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func acquireStacklessDeflateWriter(w io.Writer, level int) stackless.Writer {
|
||
|
|
nLevel := normalizeCompressLevel(level)
|
||
|
|
p := stacklessDeflateWriterPoolMap[nLevel]
|
||
|
|
v := p.Get()
|
||
|
|
if v == nil {
|
||
|
|
return stackless.NewWriter(w, func(w io.Writer) stackless.Writer {
|
||
|
|
return acquireRealDeflateWriter(w, level)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
sw := v.(stackless.Writer)
|
||
|
|
sw.Reset(w)
|
||
|
|
return sw
|
||
|
|
}
|
||
|
|
|
||
|
|
func releaseStacklessDeflateWriter(sw stackless.Writer, level int) {
|
||
|
|
sw.Close()
|
||
|
|
nLevel := normalizeCompressLevel(level)
|
||
|
|
p := stacklessDeflateWriterPoolMap[nLevel]
|
||
|
|
p.Put(sw)
|
||
|
|
}
|
||
|
|
|
||
|
|
func acquireRealDeflateWriter(w io.Writer, level int) *zlib.Writer {
|
||
|
|
nLevel := normalizeCompressLevel(level)
|
||
|
|
p := realDeflateWriterPoolMap[nLevel]
|
||
|
|
v := p.Get()
|
||
|
|
if v == nil {
|
||
|
|
zw, err := zlib.NewWriterLevel(w, level)
|
||
|
|
if err != nil {
|
||
|
|
// zlib.NewWriterLevel only errors for invalid
|
||
|
|
// compression levels. Clamp it to be min or max.
|
||
|
|
if level < zlib.HuffmanOnly {
|
||
|
|
level = zlib.HuffmanOnly
|
||
|
|
} else {
|
||
|
|
level = zlib.BestCompression
|
||
|
|
}
|
||
|
|
zw, _ = zlib.NewWriterLevel(w, level)
|
||
|
|
}
|
||
|
|
return zw
|
||
|
|
}
|
||
|
|
zw := v.(*zlib.Writer)
|
||
|
|
zw.Reset(w)
|
||
|
|
return zw
|
||
|
|
}
|
||
|
|
|
||
|
|
func releaseRealDeflateWriter(zw *zlib.Writer, level int) {
|
||
|
|
zw.Close()
|
||
|
|
nLevel := normalizeCompressLevel(level)
|
||
|
|
p := realDeflateWriterPoolMap[nLevel]
|
||
|
|
p.Put(zw)
|
||
|
|
}
|
||
|
|
|
||
|
|
var (
|
||
|
|
stacklessDeflateWriterPoolMap = newCompressWriterPoolMap()
|
||
|
|
realDeflateWriterPoolMap = newCompressWriterPoolMap()
|
||
|
|
)
|
||
|
|
|
||
|
|
func newCompressWriterPoolMap() []*sync.Pool {
|
||
|
|
// Initialize pools for all the compression levels defined
|
||
|
|
// in https://pkg.go.dev/compress/flate#pkg-constants .
|
||
|
|
// Compression levels are normalized with normalizeCompressLevel,
|
||
|
|
// so the fit [0..11].
|
||
|
|
var m []*sync.Pool
|
||
|
|
for i := 0; i < 12; i++ {
|
||
|
|
m = append(m, &sync.Pool{})
|
||
|
|
}
|
||
|
|
return m
|
||
|
|
}
|
||
|
|
|
||
|
|
func isFileCompressible(f fs.File, minCompressRatio float64) bool {
|
||
|
|
// Try compressing the first 4kb of the file
|
||
|
|
// and see if it can be compressed by more than
|
||
|
|
// the given minCompressRatio.
|
||
|
|
b := bytebufferpool.Get()
|
||
|
|
zw := acquireStacklessGzipWriter(b, CompressDefaultCompression)
|
||
|
|
lr := &io.LimitedReader{
|
||
|
|
R: f,
|
||
|
|
N: 4096,
|
||
|
|
}
|
||
|
|
_, err := copyZeroAlloc(zw, lr)
|
||
|
|
releaseStacklessGzipWriter(zw, CompressDefaultCompression)
|
||
|
|
seeker, ok := f.(io.Seeker)
|
||
|
|
if !ok {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
seeker.Seek(0, io.SeekStart) //nolint:errcheck
|
||
|
|
if err != nil {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
|
||
|
|
n := 4096 - lr.N
|
||
|
|
zn := len(b.B)
|
||
|
|
bytebufferpool.Put(b)
|
||
|
|
return float64(zn) < float64(n)*minCompressRatio
|
||
|
|
}
|
||
|
|
|
||
|
|
// normalizes compression level into [0..11], so it could be used as an index
|
||
|
|
// in *PoolMap.
|
||
|
|
func normalizeCompressLevel(level int) int {
|
||
|
|
// -2 is the lowest compression level - CompressHuffmanOnly
|
||
|
|
// 9 is the highest compression level - CompressBestCompression
|
||
|
|
if level < -2 || level > 9 {
|
||
|
|
level = CompressDefaultCompression
|
||
|
|
}
|
||
|
|
return level + 2
|
||
|
|
}
|