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         z.Reset(w)
  67         if _, err := z.Write(p); err != nil {
  68                 return err
  69         }
  70         return z.Close()
  71 }
  72
  73 func (c *gzipCompressor) Type() string {
  74         return "gzip"
  75 }
  76
  77 // Decompressor defines the interface gRPC uses to decompress a message.
  78 type Decompressor interface {
  79         // Do reads the data from r and uncompress them.
  80         Do(r io.Reader) ([]byte, error)
  81         // Type returns the compression algorithm the Decompressor uses.
  82         Type() string
  83 }
  84
  85 type gzipDecompressor struct {
  86         pool sync.Pool
  87 }
  88
  89 // NewGZIPDecompressor creates a Decompressor based on GZIP.
  90 func NewGZIPDecompressor() Decompressor {
  91         return &gzipDecompressor{}
  92 }
  93
  94 func (d *gzipDecompressor) Do(r io.Reader) ([]byte, error) {
  95         var z *gzip.Reader
  96         switch maybeZ := d.pool.Get().(type) {
  97         case nil:
  98                 newZ, err := gzip.NewReader(r)
  99                 if err != nil {
 100                         return nil, err
 101                 }
 102                 z = newZ
 103         case *gzip.Reader:
 104                 z = maybeZ
 105                 if err := z.Reset(r); err != nil {
 106                         d.pool.Put(z)
 107                         return nil, err
 108                 }
 109         }
 110
 111         defer func() {
 112                 z.Close()
 113                 d.pool.Put(z)
 114         }()
 115         return ioutil.ReadAll(z)
 116 }
 117
 118 func (d *gzipDecompressor) Type() string {
 119         return "gzip"
 120 }
 121
 122 // callInfo contains all related configuration and information about an RPC.
 123 type callInfo struct {
 124         failFast              bool
 125         headerMD              metadata.MD
 126         trailerMD             metadata.MD
 127         peer                  *peer.Peer
 128         traceInfo             traceInfo // in trace.go
 129         maxReceiveMessageSize *int
 130         maxSendMessageSize    *int
 131         creds                 credentials.PerRPCCredentials
 132 }
 133
 134 var defaultCallInfo = callInfo{failFast: true}
 135
 136 // CallOption configures a Call before it starts or extracts information from
 137 // a Call after it completes.
 138 type CallOption interface {
 139         // before is called before the call is sent to any server.  If before
 140         // returns a non-nil error, the RPC fails with that error.
 141         before(*callInfo) error
 142
 143         // after is called after the call has completed.  after cannot return an
 144         // error, so any failures should be reported via output parameters.
 145         after(*callInfo)
 146 }
 147
 148 // EmptyCallOption does not alter the Call configuration.
 149 // It can be embedded in another structure to carry satellite data for use
 150 // by interceptors.
 151 type EmptyCallOption struct{}
 152
 153 func (EmptyCallOption) before(*callInfo) error { return nil }
 154 func (EmptyCallOption) after(*callInfo)        {}
 155
 156 type beforeCall func(c *callInfo) error
 157
 158 func (o beforeCall) before(c *callInfo) error { return o(c) }
 159 func (o beforeCall) after(c *callInfo)        {}
 160
 161 type afterCall func(c *callInfo)
 162
 163 func (o afterCall) before(c *callInfo) error { return nil }
 164 func (o afterCall) after(c *callInfo)        { o(c) }
 165
 166 // Header returns a CallOptions that retrieves the header metadata
 167 // for a unary RPC.
 168 func Header(md *metadata.MD) CallOption {
 169         return afterCall(func(c *callInfo) {
 170                 *md = c.headerMD
 171         })
 172 }
 173
 174 // Trailer returns a CallOptions that retrieves the trailer metadata
 175 // for a unary RPC.
 176 func Trailer(md *metadata.MD) CallOption {
 177         return afterCall(func(c *callInfo) {
 178                 *md = c.trailerMD
 179         })
 180 }
 181
 182 // Peer returns a CallOption that retrieves peer information for a
 183 // unary RPC.
 184 func Peer(peer *peer.Peer) CallOption {
 185         return afterCall(func(c *callInfo) {
 186                 if c.peer != nil {
 187                         *peer = *c.peer
 188                 }
 189         })
 190 }
 191
 192 // FailFast configures the action to take when an RPC is attempted on broken
 193 // connections or unreachable servers. If failfast is true, the RPC will fail
 194 // immediately. Otherwise, the RPC client will block the call until a
 195 // connection is available (or the call is canceled or times out) and will retry
 196 // the call if it fails due to a transient error. Please refer to
 197 // https://github.com/grpc/grpc/blob/master/doc/wait-for-ready.md.
 198 // Note: failFast is default to true.
 199 func FailFast(failFast bool) CallOption {
 200         return beforeCall(func(c *callInfo) error {
 201                 c.failFast = failFast
 202                 return nil
 203         })
 204 }
 205
 206 // MaxCallRecvMsgSize returns a CallOption which sets the maximum message size the client can receive.
 207 func MaxCallRecvMsgSize(s int) CallOption {
 208         return beforeCall(func(o *callInfo) error {
 209                 o.maxReceiveMessageSize = &s
 210                 return nil
 211         })
 212 }
 213
 214 // MaxCallSendMsgSize returns a CallOption which sets the maximum message size the client can send.
 215 func MaxCallSendMsgSize(s int) CallOption {
 216         return beforeCall(func(o *callInfo) error {
 217                 o.maxSendMessageSize = &s
 218                 return nil
 219         })
 220 }
 221
 222 // PerRPCCredentials returns a CallOption that sets credentials.PerRPCCredentials
 223 // for a call.
 224 func PerRPCCredentials(creds credentials.PerRPCCredentials) CallOption {
 225         return beforeCall(func(c *callInfo) error {
 226                 c.creds = creds
 227                 return nil
 228         })
 229 }
 230
 231 // The format of the payload: compressed or not?
 232 type payloadFormat uint8
 233
 234 const (
 235         compressionNone payloadFormat = iota // no compression
 236         compressionMade
 237 )
 238
 239 // parser reads complete gRPC messages from the underlying reader.
 240 type parser struct {
 241         // r is the underlying reader.
 242         // See the comment on recvMsg for the permissible
 243         // error types.
 244         r io.Reader
 245
 246         // The header of a gRPC message. Find more detail
 247         // at https://grpc.io/docs/guides/wire.html.
 248         header [5]byte
 249 }
 250
 251 // recvMsg reads a complete gRPC message from the stream.
 252 //
 253 // It returns the message and its payload (compression/encoding)
 254 // format. The caller owns the returned msg memory.
 255 //
 256 // If there is an error, possible values are:
 257 //   * io.EOF, when no messages remain
 258 //   * io.ErrUnexpectedEOF
 259 //   * of type transport.ConnectionError
 260 //   * of type transport.StreamError
 261 // No other error values or types must be returned, which also means
 262 // that the underlying io.Reader must not return an incompatible
 263 // error.
 264 func (p *parser) recvMsg(maxReceiveMessageSize int) (pf payloadFormat, msg []byte, err error) {
 265         if _, err := p.r.Read(p.header[:]); err != nil {
 266                 return 0, nil, err
 267         }
 268
 269         pf = payloadFormat(p.header[0])
 270         length := binary.BigEndian.Uint32(p.header[1:])
 271
 272         if length == 0 {
 273                 return pf, nil, nil
 274         }
 275         if length > uint32(maxReceiveMessageSize) {
 276                 return 0, nil, Errorf(codes.ResourceExhausted, "grpc: received message larger than max (%d vs. %d)", length, maxReceiveMessageSize)
 277         }
 278         // TODO(bradfitz,zhaoq): garbage. reuse buffer after proto decoding instead
 279         // of making it for each message:
 280         msg = make([]byte, int(length))
 281         if _, err := p.r.Read(msg); err != nil {
 282                 if err == io.EOF {
 283                         err = io.ErrUnexpectedEOF
 284                 }
 285                 return 0, nil, err
 286         }
 287         return pf, msg, nil
 288 }
 289
 290 // encode serializes msg and prepends the message header. If msg is nil, it
 291 // generates the message header of 0 message length.
 292 func encode(c Codec, msg interface{}, cp Compressor, cbuf *bytes.Buffer, outPayload *stats.OutPayload) ([]byte, error) {
 293         var (
 294                 b      []byte
 295                 length uint
 296         )
 297         if msg != nil {
 298                 var err error
 299                 // TODO(zhaoq): optimize to reduce memory alloc and copying.
 300                 b, err = c.Marshal(msg)
 301                 if err != nil {
 302                         return nil, Errorf(codes.Internal, "grpc: error while marshaling: %v", err.Error())
 303                 }
 304                 if outPayload != nil {
 305                         outPayload.Payload = msg
 306                         // TODO truncate large payload.
 307                         outPayload.Data = b
 308                         outPayload.Length = len(b)
 309                 }
 310                 if cp != nil {
 311                         if err := cp.Do(cbuf, b); err != nil {
 312                                 return nil, Errorf(codes.Internal, "grpc: error while compressing: %v", err.Error())
 313                         }
 314                         b = cbuf.Bytes()
 315                 }
 316                 length = uint(len(b))
 317         }
 318         if length > math.MaxUint32 {
 319                 return nil, Errorf(codes.ResourceExhausted, "grpc: message too large (%d bytes)", length)
 320         }
 321
 322         const (
 323                 payloadLen = 1
 324                 sizeLen    = 4
 325         )
 326
 327         var buf = make([]byte, payloadLen+sizeLen+len(b))
 328
 329         // Write payload format
 330         if cp == nil {
 331                 buf[0] = byte(compressionNone)
 332         } else {
 333                 buf[0] = byte(compressionMade)
 334         }
 335         // Write length of b into buf
 336         binary.BigEndian.PutUint32(buf[1:], uint32(length))
 337         // Copy encoded msg to buf
 338         copy(buf[5:], b)
 339
 340         if outPayload != nil {
 341                 outPayload.WireLength = len(buf)
 342         }
 343
 344         return buf, nil
 345 }
 346
 347 func checkRecvPayload(pf payloadFormat, recvCompress string, dc Decompressor) error {
 348         switch pf {
 349         case compressionNone:
 350         case compressionMade:
 351                 if dc == nil || recvCompress != dc.Type() {
 352                         return Errorf(codes.Unimplemented, "grpc: Decompressor is not installed for grpc-encoding %q", recvCompress)
 353                 }
 354         default:
 355                 return Errorf(codes.Internal, "grpc: received unexpected payload format %d", pf)
 356         }
 357         return nil
 358 }
 359
 360 func recv(p *parser, c Codec, s *transport.Stream, dc Decompressor, m interface{}, maxReceiveMessageSize int, inPayload *stats.InPayload) error {
 361         pf, d, err := p.recvMsg(maxReceiveMessageSize)
 362         if err != nil {
 363                 return err
 364         }
 365         if inPayload != nil {
 366                 inPayload.WireLength = len(d)
 367         }
 368         if err := checkRecvPayload(pf, s.RecvCompress(), dc); err != nil {
 369                 return err
 370         }
 371         if pf == compressionMade {
 372                 d, err = dc.Do(bytes.NewReader(d))
 373                 if err != nil {
 374                         return Errorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
 375                 }
 376         }
 377         if len(d) > maxReceiveMessageSize {
 378                 // TODO: Revisit the error code. Currently keep it consistent with java
 379                 // implementation.
 380                 return Errorf(codes.ResourceExhausted, "grpc: received message larger than max (%d vs. %d)", len(d), maxReceiveMessageSize)
 381         }
 382         if err := c.Unmarshal(d, m); err != nil {
 383                 return Errorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
 384         }
 385         if inPayload != nil {
 386                 inPayload.RecvTime = time.Now()
 387                 inPayload.Payload = m
 388                 // TODO truncate large payload.
 389                 inPayload.Data = d
 390                 inPayload.Length = len(d)
 391         }
 392         return nil
 393 }
 394
 395 type rpcInfo struct {
 396         bytesSent     bool
 397         bytesReceived bool
 398 }
 399
 400 type rpcInfoContextKey struct{}
 401
 402 func newContextWithRPCInfo(ctx context.Context) context.Context {
 403         return context.WithValue(ctx, rpcInfoContextKey{}, &rpcInfo{})
 404 }
 405
 406 func rpcInfoFromContext(ctx context.Context) (s *rpcInfo, ok bool) {
 407         s, ok = ctx.Value(rpcInfoContextKey{}).(*rpcInfo)
 408         return
 409 }
 410
 411 func updateRPCInfoInContext(ctx context.Context, s rpcInfo) {
 412         if ss, ok := rpcInfoFromContext(ctx); ok {
 413                 *ss = s
 414         }
 415         return
 416 }
 417
 418 // Code returns the error code for err if it was produced by the rpc system.
 419 // Otherwise, it returns codes.Unknown.
 420 //
 421 // Deprecated; use status.FromError and Code method instead.
 422 func Code(err error) codes.Code {
 423         if s, ok := status.FromError(err); ok {
 424                 return s.Code()
 425         }
 426         return codes.Unknown
 427 }
 428
 429 // ErrorDesc returns the error description of err if it was produced by the rpc system.
 430 // Otherwise, it returns err.Error() or empty string when err is nil.
 431 //
 432 // Deprecated; use status.FromError and Message method instead.
 433 func ErrorDesc(err error) string {
 434         if s, ok := status.FromError(err); ok {
 435                 return s.Message()
 436         }
 437         return err.Error()
 438 }
 439
 440 // Errorf returns an error containing an error code and a description;
 441 // Errorf returns nil if c is OK.
 442 //
 443 // Deprecated; use status.Errorf instead.
 444 func Errorf(c codes.Code, format string, a ...interface{}) error {
 445         return status.Errorf(c, format, a...)
 446 }
 447
 448 // MethodConfig defines the configuration recommended by the service providers for a
 449 // particular method.
 450 // This is EXPERIMENTAL and subject to change.
 451 type MethodConfig struct {
 452         // WaitForReady indicates whether RPCs sent to this method should wait until
 453         // the connection is ready by default (!failfast). The value specified via the
 454         // gRPC client API will override the value set here.
 455         WaitForReady *bool
 456         // Timeout is the default timeout for RPCs sent to this method. The actual
 457         // deadline used will be the minimum of the value specified here and the value
 458         // set by the application via the gRPC client API.  If either one is not set,
 459         // then the other will be used.  If neither is set, then the RPC has no deadline.
 460         Timeout *time.Duration
 461         // MaxReqSize is the maximum allowed payload size for an individual request in a
 462         // stream (client->server) in bytes. The size which is measured is the serialized
 463         // payload after per-message compression (but before stream compression) in bytes.
 464         // The actual value used is the minumum of the value specified here and the value set
 465         // by the application via the gRPC client API. If either one is not set, then the other
 466         // will be used.  If neither is set, then the built-in default is used.
 467         MaxReqSize *int
 468         // MaxRespSize is the maximum allowed payload size for an individual response in a
 469         // stream (server->client) in bytes.
 470         MaxRespSize *int
 471 }
 472
 473 // ServiceConfig is provided by the service provider and contains parameters for how
 474 // clients that connect to the service should behave.
 475 // This is EXPERIMENTAL and subject to change.
 476 type ServiceConfig struct {
 477         // LB is the load balancer the service providers recommends. The balancer specified
 478         // via grpc.WithBalancer will override this.
 479         LB Balancer
 480         // Methods contains a map for the methods in this service.
 481         // If there is an exact match for a method (i.e. /service/method) in the map, use the corresponding MethodConfig.
 482         // If there's no exact match, look for the default config for the service (/service/) and use the corresponding MethodConfig if it exists.
 483         // Otherwise, the method has no MethodConfig to use.
 484         Methods map[string]MethodConfig
 485 }
 486
 487 func min(a, b *int) *int {
 488         if *a < *b {
 489                 return a
 490         }
 491         return b
 492 }
 493
 494 func getMaxSize(mcMax, doptMax *int, defaultVal int) *int {
 495         if mcMax == nil && doptMax == nil {
 496                 return &defaultVal
 497         }
 498         if mcMax != nil && doptMax != nil {
 499                 return min(mcMax, doptMax)
 500         }
 501         if mcMax != nil {
 502                 return mcMax
 503         }
 504         return doptMax
 505 }
 506
 507 // SupportPackageIsVersion3 is referenced from generated protocol buffer files.
 508 // The latest support package version is 4.
 509 // SupportPackageIsVersion3 is kept for compability. It will be removed in the
 510 // next support package version update.
 511 const SupportPackageIsVersion3 = true
 512
 513 // SupportPackageIsVersion4 is referenced from generated protocol buffer files
 514 // to assert that that code is compatible with this version of the grpc package.
 515 //
 516 // This constant may be renamed in the future if a change in the generated code
 517 // requires a synchronised update of grpc-go and protoc-gen-go. This constant
 518 // should not be referenced from any other code.
 519 const SupportPackageIsVersion4 = true
 520
 521 // Version is the current grpc version.
 522 const Version = "1.6.0-dev"
 523
 524 const grpcUA = "grpc-go/" + Version