vendor/google.golang.org/grpc/rpc_util.go

   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         "encoding/binary"
  25         "io"
  26         "io/ioutil"
  27         "math"
  28         "sync"
  29         "time"
  30
  31         "golang.org/x/net/context"
  32         "google.golang.org/grpc/codes"
  33         "google.golang.org/grpc/credentials"
  34         "google.golang.org/grpc/metadata"
  35         "google.golang.org/grpc/peer"
  36         "google.golang.org/grpc/stats"
  37         "google.golang.org/grpc/status"
  38         "google.golang.org/grpc/transport"
  39 )
  40
  41 // Compressor defines the interface gRPC uses to compress a message.
  42 type Compressor interface {
  43         // Do compresses p into w.
  44         Do(w io.Writer, p []byte) error
  45         // Type returns the compression algorithm the Compressor uses.
  46         Type() string
  47 }
  48
  49 type gzipCompressor struct {
  50         pool sync.Pool
  51 }
  52
  53 // NewGZIPCompressor creates a Compressor based on GZIP.
  54 func NewGZIPCompressor() Compressor {
  55         return &gzipCompressor{
  56                 pool: sync.Pool{
  57                         New: func() interface{} {
  58                                 return gzip.NewWriter(ioutil.Discard)
  59                         },
  60                 },
  61         }
  62 }
  63
  64 func (c *gzipCompressor) Do(w io.Writer, p []byte) error {
  65         z := c.pool.Get().(*gzip.Writer)
  66         defer c.pool.Put(z)
  67         z.Reset(w)
  68         if _, err := z.Write(p); err != nil {
  69                 return err
  70         }
  71         return z.Close()
  72 }
  73
  74 func (c *gzipCompressor) Type() string {
  75         return "gzip"
  76 }
  77
  78 // Decompressor defines the interface gRPC uses to decompress a message.
  79 type Decompressor interface {
  80         // Do reads the data from r and uncompress them.
  81         Do(r io.Reader) ([]byte, error)
  82         // Type returns the compression algorithm the Decompressor uses.
  83         Type() string
  84 }
  85
  86 type gzipDecompressor struct {
  87         pool sync.Pool
  88 }
  89
  90 // NewGZIPDecompressor creates a Decompressor based on GZIP.
  91 func NewGZIPDecompressor() Decompressor {
  92         return &gzipDecompressor{}
  93 }
  94
  95 func (d *gzipDecompressor) Do(r io.Reader) ([]byte, error) {
  96         var z *gzip.Reader
  97         switch maybeZ := d.pool.Get().(type) {
  98         case nil:
  99                 newZ, err := gzip.NewReader(r)
 100                 if err != nil {
 101                         return nil, err
 102                 }
 103                 z = newZ
 104         case *gzip.Reader:
 105                 z = maybeZ
 106                 if err := z.Reset(r); err != nil {
 107                         d.pool.Put(z)
 108                         return nil, err
 109                 }
 110         }
 111
 112         defer func() {
 113                 z.Close()
 114                 d.pool.Put(z)
 115         }()
 116         return ioutil.ReadAll(z)
 117 }
 118
 119 func (d *gzipDecompressor) Type() string {
 120         return "gzip"
 121 }
 122
 123 // callInfo contains all related configuration and information about an RPC.
 124 type callInfo struct {
 125         failFast              bool
 126         headerMD              metadata.MD
 127         trailerMD             metadata.MD
 128         peer                  *peer.Peer
 129         traceInfo             traceInfo // in trace.go
 130         maxReceiveMessageSize *int
 131         maxSendMessageSize    *int
 132         creds                 credentials.PerRPCCredentials
 133 }
 134
 135 func defaultCallInfo() *callInfo {
 136         return &callInfo{failFast: true}
 137 }
 138
 139 // CallOption configures a Call before it starts or extracts information from
 140 // a Call after it completes.
 141 type CallOption interface {
 142         // before is called before the call is sent to any server.  If before
 143         // returns a non-nil error, the RPC fails with that error.
 144         before(*callInfo) error
 145
 146         // after is called after the call has completed.  after cannot return an
 147         // error, so any failures should be reported via output parameters.
 148         after(*callInfo)
 149 }
 150
 151 // EmptyCallOption does not alter the Call configuration.
 152 // It can be embedded in another structure to carry satellite data for use
 153 // by interceptors.
 154 type EmptyCallOption struct{}
 155
 156 func (EmptyCallOption) before(*callInfo) error { return nil }
 157 func (EmptyCallOption) after(*callInfo)        {}
 158
 159 type beforeCall func(c *callInfo) error
 160
 161 func (o beforeCall) before(c *callInfo) error { return o(c) }
 162 func (o beforeCall) after(c *callInfo)        {}
 163
 164 type afterCall func(c *callInfo)
 165
 166 func (o afterCall) before(c *callInfo) error { return nil }
 167 func (o afterCall) after(c *callInfo)        { o(c) }
 168
 169 // Header returns a CallOptions that retrieves the header metadata
 170 // for a unary RPC.
 171 func Header(md *metadata.MD) CallOption {
 172         return afterCall(func(c *callInfo) {
 173                 *md = c.headerMD
 174         })
 175 }
 176
 177 // Trailer returns a CallOptions that retrieves the trailer metadata
 178 // for a unary RPC.
 179 func Trailer(md *metadata.MD) CallOption {
 180         return afterCall(func(c *callInfo) {
 181                 *md = c.trailerMD
 182         })
 183 }
 184
 185 // Peer returns a CallOption that retrieves peer information for a
 186 // unary RPC.
 187 func Peer(peer *peer.Peer) CallOption {
 188         return afterCall(func(c *callInfo) {
 189                 if c.peer != nil {
 190                         *peer = *c.peer
 191                 }
 192         })
 193 }
 194
 195 // FailFast configures the action to take when an RPC is attempted on broken
 196 // connections or unreachable servers. If failFast is true, the RPC will fail
 197 // immediately. Otherwise, the RPC client will block the call until a
 198 // connection is available (or the call is canceled or times out). Please refer
 199 // to https://github.com/grpc/grpc/blob/master/doc/wait-for-ready.md.
 200 // The default behavior of RPCs is to fail fast.
 201 func FailFast(failFast bool) CallOption {
 202         return beforeCall(func(c *callInfo) error {
 203                 c.failFast = failFast
 204                 return nil
 205         })
 206 }
 207
 208 // MaxCallRecvMsgSize returns a CallOption which sets the maximum message size the client can receive.
 209 func MaxCallRecvMsgSize(s int) CallOption {
 210         return beforeCall(func(o *callInfo) error {
 211                 o.maxReceiveMessageSize = &s
 212                 return nil
 213         })
 214 }
 215
 216 // MaxCallSendMsgSize returns a CallOption which sets the maximum message size the client can send.
 217 func MaxCallSendMsgSize(s int) CallOption {
 218         return beforeCall(func(o *callInfo) error {
 219                 o.maxSendMessageSize = &s
 220                 return nil
 221         })
 222 }
 223
 224 // PerRPCCredentials returns a CallOption that sets credentials.PerRPCCredentials
 225 // for a call.
 226 func PerRPCCredentials(creds credentials.PerRPCCredentials) CallOption {
 227         return beforeCall(func(c *callInfo) error {
 228                 c.creds = creds
 229                 return nil
 230         })
 231 }
 232
 233 // The format of the payload: compressed or not?
 234 type payloadFormat uint8
 235
 236 const (
 237         compressionNone payloadFormat = iota // no compression
 238         compressionMade
 239 )
 240
 241 // parser reads complete gRPC messages from the underlying reader.
 242 type parser struct {
 243         // r is the underlying reader.
 244         // See the comment on recvMsg for the permissible
 245         // error types.
 246         r io.Reader
 247
 248         // The header of a gRPC message. Find more detail
 249         // at https://grpc.io/docs/guides/wire.html.
 250         header [5]byte
 251 }
 252
 253 // recvMsg reads a complete gRPC message from the stream.
 254 //
 255 // It returns the message and its payload (compression/encoding)
 256 // format. The caller owns the returned msg memory.
 257 //
 258 // If there is an error, possible values are:
 259 //   * io.EOF, when no messages remain
 260 //   * io.ErrUnexpectedEOF
 261 //   * of type transport.ConnectionError
 262 //   * of type transport.StreamError
 263 // No other error values or types must be returned, which also means
 264 // that the underlying io.Reader must not return an incompatible
 265 // error.
 266 func (p *parser) recvMsg(maxReceiveMessageSize int) (pf payloadFormat, msg []byte, err error) {
 267         if _, err := p.r.Read(p.header[:]); err != nil {
 268                 return 0, nil, err
 269         }
 270
 271         pf = payloadFormat(p.header[0])
 272         length := binary.BigEndian.Uint32(p.header[1:])
 273
 274         if length == 0 {
 275                 return pf, nil, nil
 276         }
 277         if int64(length) > int64(maxInt) {
 278                 return 0, nil, Errorf(codes.ResourceExhausted, "grpc: received message larger than max length allowed on current machine (%d vs. %d)", length, maxInt)
 279         }
 280         if int(length) > maxReceiveMessageSize {
 281                 return 0, nil, Errorf(codes.ResourceExhausted, "grpc: received message larger than max (%d vs. %d)", length, maxReceiveMessageSize)
 282         }
 283         // TODO(bradfitz,zhaoq): garbage. reuse buffer after proto decoding instead
 284         // of making it for each message:
 285         msg = make([]byte, int(length))
 286         if _, err := p.r.Read(msg); err != nil {
 287                 if err == io.EOF {
 288                         err = io.ErrUnexpectedEOF
 289                 }
 290                 return 0, nil, err
 291         }
 292         return pf, msg, nil
 293 }
 294
 295 // encode serializes msg and returns a buffer of message header and a buffer of msg.
 296 // If msg is nil, it generates the message header and an empty msg buffer.
 297 func encode(c Codec, msg interface{}, cp Compressor, cbuf *bytes.Buffer, outPayload *stats.OutPayload) ([]byte, []byte, error) {
 298         var b []byte
 299         const (
 300                 payloadLen = 1
 301                 sizeLen    = 4
 302         )
 303
 304         if msg != nil {
 305                 var err error
 306                 b, err = c.Marshal(msg)
 307                 if err != nil {
 308                         return nil, nil, Errorf(codes.Internal, "grpc: error while marshaling: %v", err.Error())
 309                 }
 310                 if outPayload != nil {
 311                         outPayload.Payload = msg
 312                         // TODO truncate large payload.
 313                         outPayload.Data = b
 314                         outPayload.Length = len(b)
 315                 }
 316                 if cp != nil {
 317                         if err := cp.Do(cbuf, b); err != nil {
 318                                 return nil, nil, Errorf(codes.Internal, "grpc: error while compressing: %v", err.Error())
 319                         }
 320                         b = cbuf.Bytes()
 321                 }
 322         }
 323
 324         if uint(len(b)) > math.MaxUint32 {
 325                 return nil, nil, Errorf(codes.ResourceExhausted, "grpc: message too large (%d bytes)", len(b))
 326         }
 327
 328         bufHeader := make([]byte, payloadLen+sizeLen)
 329         if cp == nil {
 330                 bufHeader[0] = byte(compressionNone)
 331         } else {
 332                 bufHeader[0] = byte(compressionMade)
 333         }
 334         // Write length of b into buf
 335         binary.BigEndian.PutUint32(bufHeader[payloadLen:], uint32(len(b)))
 336         if outPayload != nil {
 337                 outPayload.WireLength = payloadLen + sizeLen + len(b)
 338         }
 339         return bufHeader, b, nil
 340 }
 341
 342 func checkRecvPayload(pf payloadFormat, recvCompress string, dc Decompressor) error {
 343         switch pf {
 344         case compressionNone:
 345         case compressionMade:
 346                 if dc == nil || recvCompress != dc.Type() {
 347                         return Errorf(codes.Unimplemented, "grpc: Decompressor is not installed for grpc-encoding %q", recvCompress)
 348                 }
 349         default:
 350                 return Errorf(codes.Internal, "grpc: received unexpected payload format %d", pf)
 351         }
 352         return nil
 353 }
 354
 355 func recv(p *parser, c Codec, s *transport.Stream, dc Decompressor, m interface{}, maxReceiveMessageSize int, inPayload *stats.InPayload) error {
 356         pf, d, err := p.recvMsg(maxReceiveMessageSize)
 357         if err != nil {
 358                 return err
 359         }
 360         if inPayload != nil {
 361                 inPayload.WireLength = len(d)
 362         }
 363         if err := checkRecvPayload(pf, s.RecvCompress(), dc); err != nil {
 364                 return err
 365         }
 366         if pf == compressionMade {
 367                 d, err = dc.Do(bytes.NewReader(d))
 368                 if err != nil {
 369                         return Errorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
 370                 }
 371         }
 372         if len(d) > maxReceiveMessageSize {
 373                 // TODO: Revisit the error code. Currently keep it consistent with java
 374                 // implementation.
 375                 return Errorf(codes.ResourceExhausted, "grpc: received message larger than max (%d vs. %d)", len(d), maxReceiveMessageSize)
 376         }
 377         if err := c.Unmarshal(d, m); err != nil {
 378                 return Errorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
 379         }
 380         if inPayload != nil {
 381                 inPayload.RecvTime = time.Now()
 382                 inPayload.Payload = m
 383                 // TODO truncate large payload.
 384                 inPayload.Data = d
 385                 inPayload.Length = len(d)
 386         }
 387         return nil
 388 }
 389
 390 type rpcInfo struct {
 391         failfast      bool
 392         bytesSent     bool
 393         bytesReceived bool
 394 }
 395
 396 type rpcInfoContextKey struct{}
 397
 398 func newContextWithRPCInfo(ctx context.Context, failfast bool) context.Context {
 399         return context.WithValue(ctx, rpcInfoContextKey{}, &rpcInfo{failfast: failfast})
 400 }
 401
 402 func rpcInfoFromContext(ctx context.Context) (s *rpcInfo, ok bool) {
 403         s, ok = ctx.Value(rpcInfoContextKey{}).(*rpcInfo)
 404         return
 405 }
 406
 407 func updateRPCInfoInContext(ctx context.Context, s rpcInfo) {
 408         if ss, ok := rpcInfoFromContext(ctx); ok {
 409                 ss.bytesReceived = s.bytesReceived
 410                 ss.bytesSent = s.bytesSent
 411         }
 412         return
 413 }
 414
 415 // Code returns the error code for err if it was produced by the rpc system.
 416 // Otherwise, it returns codes.Unknown.
 417 //
 418 // Deprecated; use status.FromError and Code method instead.
 419 func Code(err error) codes.Code {
 420         if s, ok := status.FromError(err); ok {
 421                 return s.Code()
 422         }
 423         return codes.Unknown
 424 }
 425
 426 // ErrorDesc returns the error description of err if it was produced by the rpc system.
 427 // Otherwise, it returns err.Error() or empty string when err is nil.
 428 //
 429 // Deprecated; use status.FromError and Message method instead.
 430 func ErrorDesc(err error) string {
 431         if s, ok := status.FromError(err); ok {
 432                 return s.Message()
 433         }
 434         return err.Error()
 435 }
 436
 437 // Errorf returns an error containing an error code and a description;
 438 // Errorf returns nil if c is OK.
 439 //
 440 // Deprecated; use status.Errorf instead.
 441 func Errorf(c codes.Code, format string, a ...interface{}) error {
 442         return status.Errorf(c, format, a...)
 443 }
 444
 445 // The SupportPackageIsVersion variables are referenced from generated protocol
 446 // buffer files to ensure compatibility with the gRPC version used.  The latest
 447 // support package version is 5.
 448 //
 449 // Older versions are kept for compatibility. They may be removed if
 450 // compatibility cannot be maintained.
 451 //
 452 // These constants should not be referenced from any other code.
 453 const (
 454         SupportPackageIsVersion3 = true
 455         SupportPackageIsVersion4 = true
 456         SupportPackageIsVersion5 = true
 457 )
 458
 459 // Version is the current grpc version.
 460 const Version = "1.8.0-dev"
 461
 462 const grpcUA = "grpc-go/" + Version