gotosocial/vendor/google.golang.org/grpc/encoding/gzip/gzip.go
2023-05-09 18:19:48 +01:00

133 lines
3.2 KiB
Go

/*
*
* Copyright 2017 gRPC authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/
// Package gzip implements and registers the gzip compressor
// during the initialization.
//
// # Experimental
//
// Notice: This package is EXPERIMENTAL and may be changed or removed in a
// later release.
package gzip
import (
"compress/gzip"
"encoding/binary"
"fmt"
"io"
"sync"
"google.golang.org/grpc/encoding"
)
// Name is the name registered for the gzip compressor.
const Name = "gzip"
func init() {
c := &compressor{}
c.poolCompressor.New = func() interface{} {
return &writer{Writer: gzip.NewWriter(io.Discard), pool: &c.poolCompressor}
}
encoding.RegisterCompressor(c)
}
type writer struct {
*gzip.Writer
pool *sync.Pool
}
// SetLevel updates the registered gzip compressor to use the compression level specified (gzip.HuffmanOnly is not supported).
// NOTE: this function must only be called during initialization time (i.e. in an init() function),
// and is not thread-safe.
//
// The error returned will be nil if the specified level is valid.
func SetLevel(level int) error {
if level < gzip.DefaultCompression || level > gzip.BestCompression {
return fmt.Errorf("grpc: invalid gzip compression level: %d", level)
}
c := encoding.GetCompressor(Name).(*compressor)
c.poolCompressor.New = func() interface{} {
w, err := gzip.NewWriterLevel(io.Discard, level)
if err != nil {
panic(err)
}
return &writer{Writer: w, pool: &c.poolCompressor}
}
return nil
}
func (c *compressor) Compress(w io.Writer) (io.WriteCloser, error) {
z := c.poolCompressor.Get().(*writer)
z.Writer.Reset(w)
return z, nil
}
func (z *writer) Close() error {
defer z.pool.Put(z)
return z.Writer.Close()
}
type reader struct {
*gzip.Reader
pool *sync.Pool
}
func (c *compressor) Decompress(r io.Reader) (io.Reader, error) {
z, inPool := c.poolDecompressor.Get().(*reader)
if !inPool {
newZ, err := gzip.NewReader(r)
if err != nil {
return nil, err
}
return &reader{Reader: newZ, pool: &c.poolDecompressor}, nil
}
if err := z.Reset(r); err != nil {
c.poolDecompressor.Put(z)
return nil, err
}
return z, nil
}
func (z *reader) Read(p []byte) (n int, err error) {
n, err = z.Reader.Read(p)
if err == io.EOF {
z.pool.Put(z)
}
return n, err
}
// RFC1952 specifies that the last four bytes "contains the size of
// the original (uncompressed) input data modulo 2^32."
// gRPC has a max message size of 2GB so we don't need to worry about wraparound.
func (c *compressor) DecompressedSize(buf []byte) int {
last := len(buf)
if last < 4 {
return -1
}
return int(binary.LittleEndian.Uint32(buf[last-4 : last]))
}
func (c *compressor) Name() string {
return Name
}
type compressor struct {
poolCompressor sync.Pool
poolDecompressor sync.Pool
}