...

Source file src/google.golang.org/grpc/rpc_util.go

Documentation: google.golang.org/grpc

     1  /*
     2   *
     3   * Copyright 2014 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 grpc
    20  
    21  import (
    22  	"bytes"
    23  	"compress/gzip"
    24  	"context"
    25  	"encoding/binary"
    26  	"fmt"
    27  	"io"
    28  	"math"
    29  	"strings"
    30  	"sync"
    31  	"time"
    32  
    33  	"google.golang.org/grpc/codes"
    34  	"google.golang.org/grpc/credentials"
    35  	"google.golang.org/grpc/encoding"
    36  	"google.golang.org/grpc/encoding/proto"
    37  	"google.golang.org/grpc/internal/transport"
    38  	"google.golang.org/grpc/metadata"
    39  	"google.golang.org/grpc/peer"
    40  	"google.golang.org/grpc/stats"
    41  	"google.golang.org/grpc/status"
    42  )
    43  
    44  // Compressor defines the interface gRPC uses to compress a message.
    45  //
    46  // Deprecated: use package encoding.
    47  type Compressor interface {
    48  	// Do compresses p into w.
    49  	Do(w io.Writer, p []byte) error
    50  	// Type returns the compression algorithm the Compressor uses.
    51  	Type() string
    52  }
    53  
    54  type gzipCompressor struct {
    55  	pool sync.Pool
    56  }
    57  
    58  // NewGZIPCompressor creates a Compressor based on GZIP.
    59  //
    60  // Deprecated: use package encoding/gzip.
    61  func NewGZIPCompressor() Compressor {
    62  	c, _ := NewGZIPCompressorWithLevel(gzip.DefaultCompression)
    63  	return c
    64  }
    65  
    66  // NewGZIPCompressorWithLevel is like NewGZIPCompressor but specifies the gzip compression level instead
    67  // of assuming DefaultCompression.
    68  //
    69  // The error returned will be nil if the level is valid.
    70  //
    71  // Deprecated: use package encoding/gzip.
    72  func NewGZIPCompressorWithLevel(level int) (Compressor, error) {
    73  	if level < gzip.DefaultCompression || level > gzip.BestCompression {
    74  		return nil, fmt.Errorf("grpc: invalid compression level: %d", level)
    75  	}
    76  	return &gzipCompressor{
    77  		pool: sync.Pool{
    78  			New: func() any {
    79  				w, err := gzip.NewWriterLevel(io.Discard, level)
    80  				if err != nil {
    81  					panic(err)
    82  				}
    83  				return w
    84  			},
    85  		},
    86  	}, nil
    87  }
    88  
    89  func (c *gzipCompressor) Do(w io.Writer, p []byte) error {
    90  	z := c.pool.Get().(*gzip.Writer)
    91  	defer c.pool.Put(z)
    92  	z.Reset(w)
    93  	if _, err := z.Write(p); err != nil {
    94  		return err
    95  	}
    96  	return z.Close()
    97  }
    98  
    99  func (c *gzipCompressor) Type() string {
   100  	return "gzip"
   101  }
   102  
   103  // Decompressor defines the interface gRPC uses to decompress a message.
   104  //
   105  // Deprecated: use package encoding.
   106  type Decompressor interface {
   107  	// Do reads the data from r and uncompress them.
   108  	Do(r io.Reader) ([]byte, error)
   109  	// Type returns the compression algorithm the Decompressor uses.
   110  	Type() string
   111  }
   112  
   113  type gzipDecompressor struct {
   114  	pool sync.Pool
   115  }
   116  
   117  // NewGZIPDecompressor creates a Decompressor based on GZIP.
   118  //
   119  // Deprecated: use package encoding/gzip.
   120  func NewGZIPDecompressor() Decompressor {
   121  	return &gzipDecompressor{}
   122  }
   123  
   124  func (d *gzipDecompressor) Do(r io.Reader) ([]byte, error) {
   125  	var z *gzip.Reader
   126  	switch maybeZ := d.pool.Get().(type) {
   127  	case nil:
   128  		newZ, err := gzip.NewReader(r)
   129  		if err != nil {
   130  			return nil, err
   131  		}
   132  		z = newZ
   133  	case *gzip.Reader:
   134  		z = maybeZ
   135  		if err := z.Reset(r); err != nil {
   136  			d.pool.Put(z)
   137  			return nil, err
   138  		}
   139  	}
   140  
   141  	defer func() {
   142  		z.Close()
   143  		d.pool.Put(z)
   144  	}()
   145  	return io.ReadAll(z)
   146  }
   147  
   148  func (d *gzipDecompressor) Type() string {
   149  	return "gzip"
   150  }
   151  
   152  // callInfo contains all related configuration and information about an RPC.
   153  type callInfo struct {
   154  	compressorType        string
   155  	failFast              bool
   156  	maxReceiveMessageSize *int
   157  	maxSendMessageSize    *int
   158  	creds                 credentials.PerRPCCredentials
   159  	contentSubtype        string
   160  	codec                 baseCodec
   161  	maxRetryRPCBufferSize int
   162  	onFinish              []func(err error)
   163  }
   164  
   165  func defaultCallInfo() *callInfo {
   166  	return &callInfo{
   167  		failFast:              true,
   168  		maxRetryRPCBufferSize: 256 * 1024, // 256KB
   169  	}
   170  }
   171  
   172  // CallOption configures a Call before it starts or extracts information from
   173  // a Call after it completes.
   174  type CallOption interface {
   175  	// before is called before the call is sent to any server.  If before
   176  	// returns a non-nil error, the RPC fails with that error.
   177  	before(*callInfo) error
   178  
   179  	// after is called after the call has completed.  after cannot return an
   180  	// error, so any failures should be reported via output parameters.
   181  	after(*callInfo, *csAttempt)
   182  }
   183  
   184  // EmptyCallOption does not alter the Call configuration.
   185  // It can be embedded in another structure to carry satellite data for use
   186  // by interceptors.
   187  type EmptyCallOption struct{}
   188  
   189  func (EmptyCallOption) before(*callInfo) error      { return nil }
   190  func (EmptyCallOption) after(*callInfo, *csAttempt) {}
   191  
   192  // StaticMethod returns a CallOption which specifies that a call is being made
   193  // to a method that is static, which means the method is known at compile time
   194  // and doesn't change at runtime. This can be used as a signal to stats plugins
   195  // that this method is safe to include as a key to a measurement.
   196  func StaticMethod() CallOption {
   197  	return StaticMethodCallOption{}
   198  }
   199  
   200  // StaticMethodCallOption is a CallOption that specifies that a call comes
   201  // from a static method.
   202  type StaticMethodCallOption struct {
   203  	EmptyCallOption
   204  }
   205  
   206  // Header returns a CallOptions that retrieves the header metadata
   207  // for a unary RPC.
   208  func Header(md *metadata.MD) CallOption {
   209  	return HeaderCallOption{HeaderAddr: md}
   210  }
   211  
   212  // HeaderCallOption is a CallOption for collecting response header metadata.
   213  // The metadata field will be populated *after* the RPC completes.
   214  //
   215  // # Experimental
   216  //
   217  // Notice: This type is EXPERIMENTAL and may be changed or removed in a
   218  // later release.
   219  type HeaderCallOption struct {
   220  	HeaderAddr *metadata.MD
   221  }
   222  
   223  func (o HeaderCallOption) before(c *callInfo) error { return nil }
   224  func (o HeaderCallOption) after(c *callInfo, attempt *csAttempt) {
   225  	*o.HeaderAddr, _ = attempt.s.Header()
   226  }
   227  
   228  // Trailer returns a CallOptions that retrieves the trailer metadata
   229  // for a unary RPC.
   230  func Trailer(md *metadata.MD) CallOption {
   231  	return TrailerCallOption{TrailerAddr: md}
   232  }
   233  
   234  // TrailerCallOption is a CallOption for collecting response trailer metadata.
   235  // The metadata field will be populated *after* the RPC completes.
   236  //
   237  // # Experimental
   238  //
   239  // Notice: This type is EXPERIMENTAL and may be changed or removed in a
   240  // later release.
   241  type TrailerCallOption struct {
   242  	TrailerAddr *metadata.MD
   243  }
   244  
   245  func (o TrailerCallOption) before(c *callInfo) error { return nil }
   246  func (o TrailerCallOption) after(c *callInfo, attempt *csAttempt) {
   247  	*o.TrailerAddr = attempt.s.Trailer()
   248  }
   249  
   250  // Peer returns a CallOption that retrieves peer information for a unary RPC.
   251  // The peer field will be populated *after* the RPC completes.
   252  func Peer(p *peer.Peer) CallOption {
   253  	return PeerCallOption{PeerAddr: p}
   254  }
   255  
   256  // PeerCallOption is a CallOption for collecting the identity of the remote
   257  // peer. The peer field will be populated *after* the RPC completes.
   258  //
   259  // # Experimental
   260  //
   261  // Notice: This type is EXPERIMENTAL and may be changed or removed in a
   262  // later release.
   263  type PeerCallOption struct {
   264  	PeerAddr *peer.Peer
   265  }
   266  
   267  func (o PeerCallOption) before(c *callInfo) error { return nil }
   268  func (o PeerCallOption) after(c *callInfo, attempt *csAttempt) {
   269  	if x, ok := peer.FromContext(attempt.s.Context()); ok {
   270  		*o.PeerAddr = *x
   271  	}
   272  }
   273  
   274  // WaitForReady configures the action to take when an RPC is attempted on broken
   275  // connections or unreachable servers. If waitForReady is false and the
   276  // connection is in the TRANSIENT_FAILURE state, the RPC will fail
   277  // immediately. Otherwise, the RPC client will block the call until a
   278  // connection is available (or the call is canceled or times out) and will
   279  // retry the call if it fails due to a transient error.  gRPC will not retry if
   280  // data was written to the wire unless the server indicates it did not process
   281  // the data.  Please refer to
   282  // https://github.com/grpc/grpc/blob/master/doc/wait-for-ready.md.
   283  //
   284  // By default, RPCs don't "wait for ready".
   285  func WaitForReady(waitForReady bool) CallOption {
   286  	return FailFastCallOption{FailFast: !waitForReady}
   287  }
   288  
   289  // FailFast is the opposite of WaitForReady.
   290  //
   291  // Deprecated: use WaitForReady.
   292  func FailFast(failFast bool) CallOption {
   293  	return FailFastCallOption{FailFast: failFast}
   294  }
   295  
   296  // FailFastCallOption is a CallOption for indicating whether an RPC should fail
   297  // fast or not.
   298  //
   299  // # Experimental
   300  //
   301  // Notice: This type is EXPERIMENTAL and may be changed or removed in a
   302  // later release.
   303  type FailFastCallOption struct {
   304  	FailFast bool
   305  }
   306  
   307  func (o FailFastCallOption) before(c *callInfo) error {
   308  	c.failFast = o.FailFast
   309  	return nil
   310  }
   311  func (o FailFastCallOption) after(c *callInfo, attempt *csAttempt) {}
   312  
   313  // OnFinish returns a CallOption that configures a callback to be called when
   314  // the call completes. The error passed to the callback is the status of the
   315  // RPC, and may be nil. The onFinish callback provided will only be called once
   316  // by gRPC. This is mainly used to be used by streaming interceptors, to be
   317  // notified when the RPC completes along with information about the status of
   318  // the RPC.
   319  //
   320  // # Experimental
   321  //
   322  // Notice: This API is EXPERIMENTAL and may be changed or removed in a
   323  // later release.
   324  func OnFinish(onFinish func(err error)) CallOption {
   325  	return OnFinishCallOption{
   326  		OnFinish: onFinish,
   327  	}
   328  }
   329  
   330  // OnFinishCallOption is CallOption that indicates a callback to be called when
   331  // the call completes.
   332  //
   333  // # Experimental
   334  //
   335  // Notice: This type is EXPERIMENTAL and may be changed or removed in a
   336  // later release.
   337  type OnFinishCallOption struct {
   338  	OnFinish func(error)
   339  }
   340  
   341  func (o OnFinishCallOption) before(c *callInfo) error {
   342  	c.onFinish = append(c.onFinish, o.OnFinish)
   343  	return nil
   344  }
   345  
   346  func (o OnFinishCallOption) after(c *callInfo, attempt *csAttempt) {}
   347  
   348  // MaxCallRecvMsgSize returns a CallOption which sets the maximum message size
   349  // in bytes the client can receive. If this is not set, gRPC uses the default
   350  // 4MB.
   351  func MaxCallRecvMsgSize(bytes int) CallOption {
   352  	return MaxRecvMsgSizeCallOption{MaxRecvMsgSize: bytes}
   353  }
   354  
   355  // MaxRecvMsgSizeCallOption is a CallOption that indicates the maximum message
   356  // size in bytes the client can receive.
   357  //
   358  // # Experimental
   359  //
   360  // Notice: This type is EXPERIMENTAL and may be changed or removed in a
   361  // later release.
   362  type MaxRecvMsgSizeCallOption struct {
   363  	MaxRecvMsgSize int
   364  }
   365  
   366  func (o MaxRecvMsgSizeCallOption) before(c *callInfo) error {
   367  	c.maxReceiveMessageSize = &o.MaxRecvMsgSize
   368  	return nil
   369  }
   370  func (o MaxRecvMsgSizeCallOption) after(c *callInfo, attempt *csAttempt) {}
   371  
   372  // MaxCallSendMsgSize returns a CallOption which sets the maximum message size
   373  // in bytes the client can send. If this is not set, gRPC uses the default
   374  // `math.MaxInt32`.
   375  func MaxCallSendMsgSize(bytes int) CallOption {
   376  	return MaxSendMsgSizeCallOption{MaxSendMsgSize: bytes}
   377  }
   378  
   379  // MaxSendMsgSizeCallOption is a CallOption that indicates the maximum message
   380  // size in bytes the client can send.
   381  //
   382  // # Experimental
   383  //
   384  // Notice: This type is EXPERIMENTAL and may be changed or removed in a
   385  // later release.
   386  type MaxSendMsgSizeCallOption struct {
   387  	MaxSendMsgSize int
   388  }
   389  
   390  func (o MaxSendMsgSizeCallOption) before(c *callInfo) error {
   391  	c.maxSendMessageSize = &o.MaxSendMsgSize
   392  	return nil
   393  }
   394  func (o MaxSendMsgSizeCallOption) after(c *callInfo, attempt *csAttempt) {}
   395  
   396  // PerRPCCredentials returns a CallOption that sets credentials.PerRPCCredentials
   397  // for a call.
   398  func PerRPCCredentials(creds credentials.PerRPCCredentials) CallOption {
   399  	return PerRPCCredsCallOption{Creds: creds}
   400  }
   401  
   402  // PerRPCCredsCallOption is a CallOption that indicates the per-RPC
   403  // credentials to use for the call.
   404  //
   405  // # Experimental
   406  //
   407  // Notice: This type is EXPERIMENTAL and may be changed or removed in a
   408  // later release.
   409  type PerRPCCredsCallOption struct {
   410  	Creds credentials.PerRPCCredentials
   411  }
   412  
   413  func (o PerRPCCredsCallOption) before(c *callInfo) error {
   414  	c.creds = o.Creds
   415  	return nil
   416  }
   417  func (o PerRPCCredsCallOption) after(c *callInfo, attempt *csAttempt) {}
   418  
   419  // UseCompressor returns a CallOption which sets the compressor used when
   420  // sending the request.  If WithCompressor is also set, UseCompressor has
   421  // higher priority.
   422  //
   423  // # Experimental
   424  //
   425  // Notice: This API is EXPERIMENTAL and may be changed or removed in a
   426  // later release.
   427  func UseCompressor(name string) CallOption {
   428  	return CompressorCallOption{CompressorType: name}
   429  }
   430  
   431  // CompressorCallOption is a CallOption that indicates the compressor to use.
   432  //
   433  // # Experimental
   434  //
   435  // Notice: This type is EXPERIMENTAL and may be changed or removed in a
   436  // later release.
   437  type CompressorCallOption struct {
   438  	CompressorType string
   439  }
   440  
   441  func (o CompressorCallOption) before(c *callInfo) error {
   442  	c.compressorType = o.CompressorType
   443  	return nil
   444  }
   445  func (o CompressorCallOption) after(c *callInfo, attempt *csAttempt) {}
   446  
   447  // CallContentSubtype returns a CallOption that will set the content-subtype
   448  // for a call. For example, if content-subtype is "json", the Content-Type over
   449  // the wire will be "application/grpc+json". The content-subtype is converted
   450  // to lowercase before being included in Content-Type. See Content-Type on
   451  // https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests for
   452  // more details.
   453  //
   454  // If ForceCodec is not also used, the content-subtype will be used to look up
   455  // the Codec to use in the registry controlled by RegisterCodec. See the
   456  // documentation on RegisterCodec for details on registration. The lookup of
   457  // content-subtype is case-insensitive. If no such Codec is found, the call
   458  // will result in an error with code codes.Internal.
   459  //
   460  // If ForceCodec is also used, that Codec will be used for all request and
   461  // response messages, with the content-subtype set to the given contentSubtype
   462  // here for requests.
   463  func CallContentSubtype(contentSubtype string) CallOption {
   464  	return ContentSubtypeCallOption{ContentSubtype: strings.ToLower(contentSubtype)}
   465  }
   466  
   467  // ContentSubtypeCallOption is a CallOption that indicates the content-subtype
   468  // used for marshaling messages.
   469  //
   470  // # Experimental
   471  //
   472  // Notice: This type is EXPERIMENTAL and may be changed or removed in a
   473  // later release.
   474  type ContentSubtypeCallOption struct {
   475  	ContentSubtype string
   476  }
   477  
   478  func (o ContentSubtypeCallOption) before(c *callInfo) error {
   479  	c.contentSubtype = o.ContentSubtype
   480  	return nil
   481  }
   482  func (o ContentSubtypeCallOption) after(c *callInfo, attempt *csAttempt) {}
   483  
   484  // ForceCodec returns a CallOption that will set codec to be used for all
   485  // request and response messages for a call. The result of calling Name() will
   486  // be used as the content-subtype after converting to lowercase, unless
   487  // CallContentSubtype is also used.
   488  //
   489  // See Content-Type on
   490  // https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests for
   491  // more details. Also see the documentation on RegisterCodec and
   492  // CallContentSubtype for more details on the interaction between Codec and
   493  // content-subtype.
   494  //
   495  // This function is provided for advanced users; prefer to use only
   496  // CallContentSubtype to select a registered codec instead.
   497  //
   498  // # Experimental
   499  //
   500  // Notice: This API is EXPERIMENTAL and may be changed or removed in a
   501  // later release.
   502  func ForceCodec(codec encoding.Codec) CallOption {
   503  	return ForceCodecCallOption{Codec: codec}
   504  }
   505  
   506  // ForceCodecCallOption is a CallOption that indicates the codec used for
   507  // marshaling messages.
   508  //
   509  // # Experimental
   510  //
   511  // Notice: This type is EXPERIMENTAL and may be changed or removed in a
   512  // later release.
   513  type ForceCodecCallOption struct {
   514  	Codec encoding.Codec
   515  }
   516  
   517  func (o ForceCodecCallOption) before(c *callInfo) error {
   518  	c.codec = o.Codec
   519  	return nil
   520  }
   521  func (o ForceCodecCallOption) after(c *callInfo, attempt *csAttempt) {}
   522  
   523  // CallCustomCodec behaves like ForceCodec, but accepts a grpc.Codec instead of
   524  // an encoding.Codec.
   525  //
   526  // Deprecated: use ForceCodec instead.
   527  func CallCustomCodec(codec Codec) CallOption {
   528  	return CustomCodecCallOption{Codec: codec}
   529  }
   530  
   531  // CustomCodecCallOption is a CallOption that indicates the codec used for
   532  // marshaling messages.
   533  //
   534  // # Experimental
   535  //
   536  // Notice: This type is EXPERIMENTAL and may be changed or removed in a
   537  // later release.
   538  type CustomCodecCallOption struct {
   539  	Codec Codec
   540  }
   541  
   542  func (o CustomCodecCallOption) before(c *callInfo) error {
   543  	c.codec = o.Codec
   544  	return nil
   545  }
   546  func (o CustomCodecCallOption) after(c *callInfo, attempt *csAttempt) {}
   547  
   548  // MaxRetryRPCBufferSize returns a CallOption that limits the amount of memory
   549  // used for buffering this RPC's requests for retry purposes.
   550  //
   551  // # Experimental
   552  //
   553  // Notice: This API is EXPERIMENTAL and may be changed or removed in a
   554  // later release.
   555  func MaxRetryRPCBufferSize(bytes int) CallOption {
   556  	return MaxRetryRPCBufferSizeCallOption{bytes}
   557  }
   558  
   559  // MaxRetryRPCBufferSizeCallOption is a CallOption indicating the amount of
   560  // memory to be used for caching this RPC for retry purposes.
   561  //
   562  // # Experimental
   563  //
   564  // Notice: This type is EXPERIMENTAL and may be changed or removed in a
   565  // later release.
   566  type MaxRetryRPCBufferSizeCallOption struct {
   567  	MaxRetryRPCBufferSize int
   568  }
   569  
   570  func (o MaxRetryRPCBufferSizeCallOption) before(c *callInfo) error {
   571  	c.maxRetryRPCBufferSize = o.MaxRetryRPCBufferSize
   572  	return nil
   573  }
   574  func (o MaxRetryRPCBufferSizeCallOption) after(c *callInfo, attempt *csAttempt) {}
   575  
   576  // The format of the payload: compressed or not?
   577  type payloadFormat uint8
   578  
   579  const (
   580  	compressionNone payloadFormat = 0 // no compression
   581  	compressionMade payloadFormat = 1 // compressed
   582  )
   583  
   584  // parser reads complete gRPC messages from the underlying reader.
   585  type parser struct {
   586  	// r is the underlying reader.
   587  	// See the comment on recvMsg for the permissible
   588  	// error types.
   589  	r io.Reader
   590  
   591  	// The header of a gRPC message. Find more detail at
   592  	// https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md
   593  	header [5]byte
   594  
   595  	// recvBufferPool is the pool of shared receive buffers.
   596  	recvBufferPool SharedBufferPool
   597  }
   598  
   599  // recvMsg reads a complete gRPC message from the stream.
   600  //
   601  // It returns the message and its payload (compression/encoding)
   602  // format. The caller owns the returned msg memory.
   603  //
   604  // If there is an error, possible values are:
   605  //   - io.EOF, when no messages remain
   606  //   - io.ErrUnexpectedEOF
   607  //   - of type transport.ConnectionError
   608  //   - an error from the status package
   609  //
   610  // No other error values or types must be returned, which also means
   611  // that the underlying io.Reader must not return an incompatible
   612  // error.
   613  func (p *parser) recvMsg(maxReceiveMessageSize int) (pf payloadFormat, msg []byte, err error) {
   614  	if _, err := p.r.Read(p.header[:]); err != nil {
   615  		return 0, nil, err
   616  	}
   617  
   618  	pf = payloadFormat(p.header[0])
   619  	length := binary.BigEndian.Uint32(p.header[1:])
   620  
   621  	if length == 0 {
   622  		return pf, nil, nil
   623  	}
   624  	if int64(length) > int64(maxInt) {
   625  		return 0, nil, status.Errorf(codes.ResourceExhausted, "grpc: received message larger than max length allowed on current machine (%d vs. %d)", length, maxInt)
   626  	}
   627  	if int(length) > maxReceiveMessageSize {
   628  		return 0, nil, status.Errorf(codes.ResourceExhausted, "grpc: received message larger than max (%d vs. %d)", length, maxReceiveMessageSize)
   629  	}
   630  	msg = p.recvBufferPool.Get(int(length))
   631  	if _, err := p.r.Read(msg); err != nil {
   632  		if err == io.EOF {
   633  			err = io.ErrUnexpectedEOF
   634  		}
   635  		return 0, nil, err
   636  	}
   637  	return pf, msg, nil
   638  }
   639  
   640  // encode serializes msg and returns a buffer containing the message, or an
   641  // error if it is too large to be transmitted by grpc.  If msg is nil, it
   642  // generates an empty message.
   643  func encode(c baseCodec, msg any) ([]byte, error) {
   644  	if msg == nil { // NOTE: typed nils will not be caught by this check
   645  		return nil, nil
   646  	}
   647  	b, err := c.Marshal(msg)
   648  	if err != nil {
   649  		return nil, status.Errorf(codes.Internal, "grpc: error while marshaling: %v", err.Error())
   650  	}
   651  	if uint(len(b)) > math.MaxUint32 {
   652  		return nil, status.Errorf(codes.ResourceExhausted, "grpc: message too large (%d bytes)", len(b))
   653  	}
   654  	return b, nil
   655  }
   656  
   657  // compress returns the input bytes compressed by compressor or cp.
   658  // If both compressors are nil, or if the message has zero length, returns nil,
   659  // indicating no compression was done.
   660  //
   661  // TODO(dfawley): eliminate cp parameter by wrapping Compressor in an encoding.Compressor.
   662  func compress(in []byte, cp Compressor, compressor encoding.Compressor) ([]byte, error) {
   663  	if compressor == nil && cp == nil {
   664  		return nil, nil
   665  	}
   666  	if len(in) == 0 {
   667  		return nil, nil
   668  	}
   669  	wrapErr := func(err error) error {
   670  		return status.Errorf(codes.Internal, "grpc: error while compressing: %v", err.Error())
   671  	}
   672  	cbuf := &bytes.Buffer{}
   673  	if compressor != nil {
   674  		z, err := compressor.Compress(cbuf)
   675  		if err != nil {
   676  			return nil, wrapErr(err)
   677  		}
   678  		if _, err := z.Write(in); err != nil {
   679  			return nil, wrapErr(err)
   680  		}
   681  		if err := z.Close(); err != nil {
   682  			return nil, wrapErr(err)
   683  		}
   684  	} else {
   685  		if err := cp.Do(cbuf, in); err != nil {
   686  			return nil, wrapErr(err)
   687  		}
   688  	}
   689  	return cbuf.Bytes(), nil
   690  }
   691  
   692  const (
   693  	payloadLen = 1
   694  	sizeLen    = 4
   695  	headerLen  = payloadLen + sizeLen
   696  )
   697  
   698  // msgHeader returns a 5-byte header for the message being transmitted and the
   699  // payload, which is compData if non-nil or data otherwise.
   700  func msgHeader(data, compData []byte) (hdr []byte, payload []byte) {
   701  	hdr = make([]byte, headerLen)
   702  	if compData != nil {
   703  		hdr[0] = byte(compressionMade)
   704  		data = compData
   705  	} else {
   706  		hdr[0] = byte(compressionNone)
   707  	}
   708  
   709  	// Write length of payload into buf
   710  	binary.BigEndian.PutUint32(hdr[payloadLen:], uint32(len(data)))
   711  	return hdr, data
   712  }
   713  
   714  func outPayload(client bool, msg any, data, payload []byte, t time.Time) *stats.OutPayload {
   715  	return &stats.OutPayload{
   716  		Client:           client,
   717  		Payload:          msg,
   718  		Data:             data,
   719  		Length:           len(data),
   720  		WireLength:       len(payload) + headerLen,
   721  		CompressedLength: len(payload),
   722  		SentTime:         t,
   723  	}
   724  }
   725  
   726  func checkRecvPayload(pf payloadFormat, recvCompress string, haveCompressor bool) *status.Status {
   727  	switch pf {
   728  	case compressionNone:
   729  	case compressionMade:
   730  		if recvCompress == "" || recvCompress == encoding.Identity {
   731  			return status.New(codes.Internal, "grpc: compressed flag set with identity or empty encoding")
   732  		}
   733  		if !haveCompressor {
   734  			return status.Newf(codes.Unimplemented, "grpc: Decompressor is not installed for grpc-encoding %q", recvCompress)
   735  		}
   736  	default:
   737  		return status.Newf(codes.Internal, "grpc: received unexpected payload format %d", pf)
   738  	}
   739  	return nil
   740  }
   741  
   742  type payloadInfo struct {
   743  	compressedLength  int // The compressed length got from wire.
   744  	uncompressedBytes []byte
   745  }
   746  
   747  // recvAndDecompress reads a message from the stream, decompressing it if necessary.
   748  //
   749  // Cancelling the returned cancel function releases the buffer back to the pool. So the caller should cancel as soon as
   750  // the buffer is no longer needed.
   751  func recvAndDecompress(p *parser, s *transport.Stream, dc Decompressor, maxReceiveMessageSize int, payInfo *payloadInfo, compressor encoding.Compressor,
   752  ) (uncompressedBuf []byte, cancel func(), err error) {
   753  	pf, compressedBuf, err := p.recvMsg(maxReceiveMessageSize)
   754  	if err != nil {
   755  		return nil, nil, err
   756  	}
   757  
   758  	if st := checkRecvPayload(pf, s.RecvCompress(), compressor != nil || dc != nil); st != nil {
   759  		return nil, nil, st.Err()
   760  	}
   761  
   762  	var size int
   763  	if pf == compressionMade {
   764  		// To match legacy behavior, if the decompressor is set by WithDecompressor or RPCDecompressor,
   765  		// use this decompressor as the default.
   766  		if dc != nil {
   767  			uncompressedBuf, err = dc.Do(bytes.NewReader(compressedBuf))
   768  			size = len(uncompressedBuf)
   769  		} else {
   770  			uncompressedBuf, size, err = decompress(compressor, compressedBuf, maxReceiveMessageSize)
   771  		}
   772  		if err != nil {
   773  			return nil, nil, status.Errorf(codes.Internal, "grpc: failed to decompress the received message: %v", err)
   774  		}
   775  		if size > maxReceiveMessageSize {
   776  			// TODO: Revisit the error code. Currently keep it consistent with java
   777  			// implementation.
   778  			return nil, nil, status.Errorf(codes.ResourceExhausted, "grpc: received message after decompression larger than max (%d vs. %d)", size, maxReceiveMessageSize)
   779  		}
   780  	} else {
   781  		uncompressedBuf = compressedBuf
   782  	}
   783  
   784  	if payInfo != nil {
   785  		payInfo.compressedLength = len(compressedBuf)
   786  		payInfo.uncompressedBytes = uncompressedBuf
   787  
   788  		cancel = func() {}
   789  	} else {
   790  		cancel = func() {
   791  			p.recvBufferPool.Put(&compressedBuf)
   792  		}
   793  	}
   794  
   795  	return uncompressedBuf, cancel, nil
   796  }
   797  
   798  // Using compressor, decompress d, returning data and size.
   799  // Optionally, if data will be over maxReceiveMessageSize, just return the size.
   800  func decompress(compressor encoding.Compressor, d []byte, maxReceiveMessageSize int) ([]byte, int, error) {
   801  	dcReader, err := compressor.Decompress(bytes.NewReader(d))
   802  	if err != nil {
   803  		return nil, 0, err
   804  	}
   805  	if sizer, ok := compressor.(interface {
   806  		DecompressedSize(compressedBytes []byte) int
   807  	}); ok {
   808  		if size := sizer.DecompressedSize(d); size >= 0 {
   809  			if size > maxReceiveMessageSize {
   810  				return nil, size, nil
   811  			}
   812  			// size is used as an estimate to size the buffer, but we
   813  			// will read more data if available.
   814  			// +MinRead so ReadFrom will not reallocate if size is correct.
   815  			//
   816  			// TODO: If we ensure that the buffer size is the same as the DecompressedSize,
   817  			// we can also utilize the recv buffer pool here.
   818  			buf := bytes.NewBuffer(make([]byte, 0, size+bytes.MinRead))
   819  			bytesRead, err := buf.ReadFrom(io.LimitReader(dcReader, int64(maxReceiveMessageSize)+1))
   820  			return buf.Bytes(), int(bytesRead), err
   821  		}
   822  	}
   823  	// Read from LimitReader with limit max+1. So if the underlying
   824  	// reader is over limit, the result will be bigger than max.
   825  	d, err = io.ReadAll(io.LimitReader(dcReader, int64(maxReceiveMessageSize)+1))
   826  	return d, len(d), err
   827  }
   828  
   829  // For the two compressor parameters, both should not be set, but if they are,
   830  // dc takes precedence over compressor.
   831  // TODO(dfawley): wrap the old compressor/decompressor using the new API?
   832  func recv(p *parser, c baseCodec, s *transport.Stream, dc Decompressor, m any, maxReceiveMessageSize int, payInfo *payloadInfo, compressor encoding.Compressor) error {
   833  	buf, cancel, err := recvAndDecompress(p, s, dc, maxReceiveMessageSize, payInfo, compressor)
   834  	if err != nil {
   835  		return err
   836  	}
   837  	defer cancel()
   838  
   839  	if err := c.Unmarshal(buf, m); err != nil {
   840  		return status.Errorf(codes.Internal, "grpc: failed to unmarshal the received message: %v", err)
   841  	}
   842  	return nil
   843  }
   844  
   845  // Information about RPC
   846  type rpcInfo struct {
   847  	failfast      bool
   848  	preloaderInfo *compressorInfo
   849  }
   850  
   851  // Information about Preloader
   852  // Responsible for storing codec, and compressors
   853  // If stream (s) has  context s.Context which stores rpcInfo that has non nil
   854  // pointers to codec, and compressors, then we can use preparedMsg for Async message prep
   855  // and reuse marshalled bytes
   856  type compressorInfo struct {
   857  	codec baseCodec
   858  	cp    Compressor
   859  	comp  encoding.Compressor
   860  }
   861  
   862  type rpcInfoContextKey struct{}
   863  
   864  func newContextWithRPCInfo(ctx context.Context, failfast bool, codec baseCodec, cp Compressor, comp encoding.Compressor) context.Context {
   865  	return context.WithValue(ctx, rpcInfoContextKey{}, &rpcInfo{
   866  		failfast: failfast,
   867  		preloaderInfo: &compressorInfo{
   868  			codec: codec,
   869  			cp:    cp,
   870  			comp:  comp,
   871  		},
   872  	})
   873  }
   874  
   875  func rpcInfoFromContext(ctx context.Context) (s *rpcInfo, ok bool) {
   876  	s, ok = ctx.Value(rpcInfoContextKey{}).(*rpcInfo)
   877  	return
   878  }
   879  
   880  // Code returns the error code for err if it was produced by the rpc system.
   881  // Otherwise, it returns codes.Unknown.
   882  //
   883  // Deprecated: use status.Code instead.
   884  func Code(err error) codes.Code {
   885  	return status.Code(err)
   886  }
   887  
   888  // ErrorDesc returns the error description of err if it was produced by the rpc system.
   889  // Otherwise, it returns err.Error() or empty string when err is nil.
   890  //
   891  // Deprecated: use status.Convert and Message method instead.
   892  func ErrorDesc(err error) string {
   893  	return status.Convert(err).Message()
   894  }
   895  
   896  // Errorf returns an error containing an error code and a description;
   897  // Errorf returns nil if c is OK.
   898  //
   899  // Deprecated: use status.Errorf instead.
   900  func Errorf(c codes.Code, format string, a ...any) error {
   901  	return status.Errorf(c, format, a...)
   902  }
   903  
   904  var errContextCanceled = status.Error(codes.Canceled, context.Canceled.Error())
   905  var errContextDeadline = status.Error(codes.DeadlineExceeded, context.DeadlineExceeded.Error())
   906  
   907  // toRPCErr converts an error into an error from the status package.
   908  func toRPCErr(err error) error {
   909  	switch err {
   910  	case nil, io.EOF:
   911  		return err
   912  	case context.DeadlineExceeded:
   913  		return errContextDeadline
   914  	case context.Canceled:
   915  		return errContextCanceled
   916  	case io.ErrUnexpectedEOF:
   917  		return status.Error(codes.Internal, err.Error())
   918  	}
   919  
   920  	switch e := err.(type) {
   921  	case transport.ConnectionError:
   922  		return status.Error(codes.Unavailable, e.Desc)
   923  	case *transport.NewStreamError:
   924  		return toRPCErr(e.Err)
   925  	}
   926  
   927  	if _, ok := status.FromError(err); ok {
   928  		return err
   929  	}
   930  
   931  	return status.Error(codes.Unknown, err.Error())
   932  }
   933  
   934  // setCallInfoCodec should only be called after CallOptions have been applied.
   935  func setCallInfoCodec(c *callInfo) error {
   936  	if c.codec != nil {
   937  		// codec was already set by a CallOption; use it, but set the content
   938  		// subtype if it is not set.
   939  		if c.contentSubtype == "" {
   940  			// c.codec is a baseCodec to hide the difference between grpc.Codec and
   941  			// encoding.Codec (Name vs. String method name).  We only support
   942  			// setting content subtype from encoding.Codec to avoid a behavior
   943  			// change with the deprecated version.
   944  			if ec, ok := c.codec.(encoding.Codec); ok {
   945  				c.contentSubtype = strings.ToLower(ec.Name())
   946  			}
   947  		}
   948  		return nil
   949  	}
   950  
   951  	if c.contentSubtype == "" {
   952  		// No codec specified in CallOptions; use proto by default.
   953  		c.codec = encoding.GetCodec(proto.Name)
   954  		return nil
   955  	}
   956  
   957  	// c.contentSubtype is already lowercased in CallContentSubtype
   958  	c.codec = encoding.GetCodec(c.contentSubtype)
   959  	if c.codec == nil {
   960  		return status.Errorf(codes.Internal, "no codec registered for content-subtype %s", c.contentSubtype)
   961  	}
   962  	return nil
   963  }
   964  
   965  // The SupportPackageIsVersion variables are referenced from generated protocol
   966  // buffer files to ensure compatibility with the gRPC version used.  The latest
   967  // support package version is 9.
   968  //
   969  // Older versions are kept for compatibility.
   970  //
   971  // These constants should not be referenced from any other code.
   972  const (
   973  	SupportPackageIsVersion3 = true
   974  	SupportPackageIsVersion4 = true
   975  	SupportPackageIsVersion5 = true
   976  	SupportPackageIsVersion6 = true
   977  	SupportPackageIsVersion7 = true
   978  	SupportPackageIsVersion8 = true
   979  	SupportPackageIsVersion9 = true
   980  )
   981  
   982  const grpcUA = "grpc-go/" + Version
   983  

View as plain text