encoding.go raw

   1  /*
   2   *
   3   * Copyright 2017 gRPC authors.
   4   *
   5   * Licensed under the Apache License, Version 2.0 (the "License");
   6   * you may not use this file except in compliance with the License.
   7   * You may obtain a copy of the License at
   8   *
   9   *     http://www.apache.org/licenses/LICENSE-2.0
  10   *
  11   * Unless required by applicable law or agreed to in writing, software
  12   * distributed under the License is distributed on an "AS IS" BASIS,
  13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14   * See the License for the specific language governing permissions and
  15   * limitations under the License.
  16   *
  17   */
  18  
  19  // Package encoding defines the interface for the compressor and codec, and
  20  // functions to register and retrieve compressors and codecs.
  21  //
  22  // # Experimental
  23  //
  24  // Notice: This package is EXPERIMENTAL and may be changed or removed in a
  25  // later release.
  26  package encoding
  27  
  28  import (
  29  	"io"
  30  	"slices"
  31  	"strings"
  32  
  33  	"google.golang.org/grpc/encoding/internal"
  34  	"google.golang.org/grpc/internal/grpcutil"
  35  )
  36  
  37  // Identity specifies the optional encoding for uncompressed streams.
  38  // It is intended for grpc internal use only.
  39  const Identity = "identity"
  40  
  41  func init() {
  42  	internal.RegisterCompressorForTesting = func(c Compressor) func() {
  43  		name := c.Name()
  44  		curCompressor, found := registeredCompressor[name]
  45  		RegisterCompressor(c)
  46  		return func() {
  47  			if found {
  48  				registeredCompressor[name] = curCompressor
  49  				return
  50  			}
  51  			delete(registeredCompressor, name)
  52  			grpcutil.RegisteredCompressorNames = slices.DeleteFunc(grpcutil.RegisteredCompressorNames, func(s string) bool {
  53  				return s == name
  54  			})
  55  		}
  56  	}
  57  }
  58  
  59  // Compressor is used for compressing and decompressing when sending or
  60  // receiving messages.
  61  //
  62  // If a Compressor implements `DecompressedSize(compressedBytes []byte) int`,
  63  // gRPC will invoke it to determine the size of the buffer allocated for the
  64  // result of decompression.  A return value of -1 indicates unknown size.
  65  type Compressor interface {
  66  	// Compress writes the data written to wc to w after compressing it.  If an
  67  	// error occurs while initializing the compressor, that error is returned
  68  	// instead.
  69  	Compress(w io.Writer) (io.WriteCloser, error)
  70  	// Decompress reads data from r, decompresses it, and provides the
  71  	// uncompressed data via the returned io.Reader.  If an error occurs while
  72  	// initializing the decompressor, that error is returned instead.
  73  	Decompress(r io.Reader) (io.Reader, error)
  74  	// Name is the name of the compression codec and is used to set the content
  75  	// coding header.  The result must be static; the result cannot change
  76  	// between calls.
  77  	Name() string
  78  }
  79  
  80  var registeredCompressor = make(map[string]Compressor)
  81  
  82  // RegisterCompressor registers the compressor with gRPC by its name.  It can
  83  // be activated when sending an RPC via grpc.UseCompressor().  It will be
  84  // automatically accessed when receiving a message based on the content coding
  85  // header.  Servers also use it to send a response with the same encoding as
  86  // the request.
  87  //
  88  // NOTE: this function must only be called during initialization time (i.e. in
  89  // an init() function), and is not thread-safe.  If multiple Compressors are
  90  // registered with the same name, the one registered last will take effect.
  91  func RegisterCompressor(c Compressor) {
  92  	registeredCompressor[c.Name()] = c
  93  	if !grpcutil.IsCompressorNameRegistered(c.Name()) {
  94  		grpcutil.RegisteredCompressorNames = append(grpcutil.RegisteredCompressorNames, c.Name())
  95  	}
  96  }
  97  
  98  // GetCompressor returns Compressor for the given compressor name.
  99  func GetCompressor(name string) Compressor {
 100  	return registeredCompressor[name]
 101  }
 102  
 103  // Codec defines the interface gRPC uses to encode and decode messages.  Note
 104  // that implementations of this interface must be thread safe; a Codec's
 105  // methods can be called from concurrent goroutines.
 106  type Codec interface {
 107  	// Marshal returns the wire format of v.
 108  	Marshal(v any) ([]byte, error)
 109  	// Unmarshal parses the wire format into v.
 110  	Unmarshal(data []byte, v any) error
 111  	// Name returns the name of the Codec implementation. The returned string
 112  	// will be used as part of content type in transmission.  The result must be
 113  	// static; the result cannot change between calls.
 114  	Name() string
 115  }
 116  
 117  var registeredCodecs = make(map[string]any)
 118  
 119  // RegisterCodec registers the provided Codec for use with all gRPC clients and
 120  // servers.
 121  //
 122  // The Codec will be stored and looked up by result of its Name() method, which
 123  // should match the content-subtype of the encoding handled by the Codec.  This
 124  // is case-insensitive, and is stored and looked up as lowercase.  If the
 125  // result of calling Name() is an empty string, RegisterCodec will panic. See
 126  // Content-Type on
 127  // https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests for
 128  // more details.
 129  //
 130  // NOTE: this function must only be called during initialization time (i.e. in
 131  // an init() function), and is not thread-safe.  If multiple Codecs are
 132  // registered with the same name, the one registered last will take effect.
 133  func RegisterCodec(codec Codec) {
 134  	if codec == nil {
 135  		panic("cannot register a nil Codec")
 136  	}
 137  	if codec.Name() == "" {
 138  		panic("cannot register Codec with empty string result for Name()")
 139  	}
 140  	contentSubtype := strings.ToLower(codec.Name())
 141  	registeredCodecs[contentSubtype] = codec
 142  }
 143  
 144  // GetCodec gets a registered Codec by content-subtype, or nil if no Codec is
 145  // registered for the content-subtype.
 146  //
 147  // The content-subtype is expected to be lowercase.
 148  func GetCodec(contentSubtype string) Codec {
 149  	c, _ := registeredCodecs[contentSubtype].(Codec)
 150  	return c
 151  }
 152