rpc_util.go 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457
  1. /*
  2. *
  3. * Copyright 2014, Google Inc.
  4. * All rights reserved.
  5. *
  6. * Redistribution and use in source and binary forms, with or without
  7. * modification, are permitted provided that the following conditions are
  8. * met:
  9. *
  10. * * Redistributions of source code must retain the above copyright
  11. * notice, this list of conditions and the following disclaimer.
  12. * * Redistributions in binary form must reproduce the above
  13. * copyright notice, this list of conditions and the following disclaimer
  14. * in the documentation and/or other materials provided with the
  15. * distribution.
  16. * * Neither the name of Google Inc. nor the names of its
  17. * contributors may be used to endorse or promote products derived from
  18. * this software without specific prior written permission.
  19. *
  20. * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  21. * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  22. * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  23. * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  24. * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  25. * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  26. * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  27. * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  28. * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  29. * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  30. * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  31. *
  32. */
  33. package grpc
  34. import (
  35. "bytes"
  36. "compress/gzip"
  37. "encoding/binary"
  38. "fmt"
  39. "io"
  40. "io/ioutil"
  41. "math"
  42. "os"
  43. "github.com/golang/protobuf/proto"
  44. "golang.org/x/net/context"
  45. "google.golang.org/grpc/codes"
  46. "google.golang.org/grpc/metadata"
  47. "google.golang.org/grpc/transport"
  48. )
  49. // Codec defines the interface gRPC uses to encode and decode messages.
  50. type Codec interface {
  51. // Marshal returns the wire format of v.
  52. Marshal(v interface{}) ([]byte, error)
  53. // Unmarshal parses the wire format into v.
  54. Unmarshal(data []byte, v interface{}) error
  55. // String returns the name of the Codec implementation. The returned
  56. // string will be used as part of content type in transmission.
  57. String() string
  58. }
  59. // protoCodec is a Codec implementation with protobuf. It is the default codec for gRPC.
  60. type protoCodec struct{}
  61. func (protoCodec) Marshal(v interface{}) ([]byte, error) {
  62. return proto.Marshal(v.(proto.Message))
  63. }
  64. func (protoCodec) Unmarshal(data []byte, v interface{}) error {
  65. return proto.Unmarshal(data, v.(proto.Message))
  66. }
  67. func (protoCodec) String() string {
  68. return "proto"
  69. }
  70. // Compressor defines the interface gRPC uses to compress a message.
  71. type Compressor interface {
  72. // Do compresses p into w.
  73. Do(w io.Writer, p []byte) error
  74. // Type returns the compression algorithm the Compressor uses.
  75. Type() string
  76. }
  77. // NewGZIPCompressor creates a Compressor based on GZIP.
  78. func NewGZIPCompressor() Compressor {
  79. return &gzipCompressor{}
  80. }
  81. type gzipCompressor struct {
  82. }
  83. func (c *gzipCompressor) Do(w io.Writer, p []byte) error {
  84. z := gzip.NewWriter(w)
  85. if _, err := z.Write(p); err != nil {
  86. return err
  87. }
  88. return z.Close()
  89. }
  90. func (c *gzipCompressor) Type() string {
  91. return "gzip"
  92. }
  93. // Decompressor defines the interface gRPC uses to decompress a message.
  94. type Decompressor interface {
  95. // Do reads the data from r and uncompress them.
  96. Do(r io.Reader) ([]byte, error)
  97. // Type returns the compression algorithm the Decompressor uses.
  98. Type() string
  99. }
  100. type gzipDecompressor struct {
  101. }
  102. // NewGZIPDecompressor creates a Decompressor based on GZIP.
  103. func NewGZIPDecompressor() Decompressor {
  104. return &gzipDecompressor{}
  105. }
  106. func (d *gzipDecompressor) Do(r io.Reader) ([]byte, error) {
  107. z, err := gzip.NewReader(r)
  108. if err != nil {
  109. return nil, err
  110. }
  111. defer z.Close()
  112. return ioutil.ReadAll(z)
  113. }
  114. func (d *gzipDecompressor) Type() string {
  115. return "gzip"
  116. }
  117. // callInfo contains all related configuration and information about an RPC.
  118. type callInfo struct {
  119. failFast bool
  120. headerMD metadata.MD
  121. trailerMD metadata.MD
  122. traceInfo traceInfo // in trace.go
  123. }
  124. var defaultCallInfo = callInfo{failFast: true}
  125. // CallOption configures a Call before it starts or extracts information from
  126. // a Call after it completes.
  127. type CallOption interface {
  128. // before is called before the call is sent to any server. If before
  129. // returns a non-nil error, the RPC fails with that error.
  130. before(*callInfo) error
  131. // after is called after the call has completed. after cannot return an
  132. // error, so any failures should be reported via output parameters.
  133. after(*callInfo)
  134. }
  135. type beforeCall func(c *callInfo) error
  136. func (o beforeCall) before(c *callInfo) error { return o(c) }
  137. func (o beforeCall) after(c *callInfo) {}
  138. type afterCall func(c *callInfo)
  139. func (o afterCall) before(c *callInfo) error { return nil }
  140. func (o afterCall) after(c *callInfo) { o(c) }
  141. // Header returns a CallOptions that retrieves the header metadata
  142. // for a unary RPC.
  143. func Header(md *metadata.MD) CallOption {
  144. return afterCall(func(c *callInfo) {
  145. *md = c.headerMD
  146. })
  147. }
  148. // Trailer returns a CallOptions that retrieves the trailer metadata
  149. // for a unary RPC.
  150. func Trailer(md *metadata.MD) CallOption {
  151. return afterCall(func(c *callInfo) {
  152. *md = c.trailerMD
  153. })
  154. }
  155. // FailFast configures the action to take when an RPC is attempted on broken
  156. // connections or unreachable servers. If failfast is true, the RPC will fail
  157. // immediately. Otherwise, the RPC client will block the call until a
  158. // connection is available (or the call is canceled or times out) and will retry
  159. // the call if it fails due to a transient error. Please refer to
  160. // https://github.com/grpc/grpc/blob/master/doc/fail_fast.md
  161. func FailFast(failFast bool) CallOption {
  162. return beforeCall(func(c *callInfo) error {
  163. c.failFast = failFast
  164. return nil
  165. })
  166. }
  167. // The format of the payload: compressed or not?
  168. type payloadFormat uint8
  169. const (
  170. compressionNone payloadFormat = iota // no compression
  171. compressionMade
  172. )
  173. // parser reads complete gRPC messages from the underlying reader.
  174. type parser struct {
  175. // r is the underlying reader.
  176. // See the comment on recvMsg for the permissible
  177. // error types.
  178. r io.Reader
  179. // The header of a gRPC message. Find more detail
  180. // at http://www.grpc.io/docs/guides/wire.html.
  181. header [5]byte
  182. }
  183. // recvMsg reads a complete gRPC message from the stream.
  184. //
  185. // It returns the message and its payload (compression/encoding)
  186. // format. The caller owns the returned msg memory.
  187. //
  188. // If there is an error, possible values are:
  189. // * io.EOF, when no messages remain
  190. // * io.ErrUnexpectedEOF
  191. // * of type transport.ConnectionError
  192. // * of type transport.StreamError
  193. // No other error values or types must be returned, which also means
  194. // that the underlying io.Reader must not return an incompatible
  195. // error.
  196. func (p *parser) recvMsg(maxMsgSize int) (pf payloadFormat, msg []byte, err error) {
  197. if _, err := io.ReadFull(p.r, p.header[:]); err != nil {
  198. return 0, nil, err
  199. }
  200. pf = payloadFormat(p.header[0])
  201. length := binary.BigEndian.Uint32(p.header[1:])
  202. if length == 0 {
  203. return pf, nil, nil
  204. }
  205. if length > uint32(maxMsgSize) {
  206. return 0, nil, Errorf(codes.Internal, "grpc: received message length %d exceeding the max size %d", length, maxMsgSize)
  207. }
  208. // TODO(bradfitz,zhaoq): garbage. reuse buffer after proto decoding instead
  209. // of making it for each message:
  210. msg = make([]byte, int(length))
  211. if _, err := io.ReadFull(p.r, msg); err != nil {
  212. if err == io.EOF {
  213. err = io.ErrUnexpectedEOF
  214. }
  215. return 0, nil, err
  216. }
  217. return pf, msg, nil
  218. }
  219. // encode serializes msg and prepends the message header. If msg is nil, it
  220. // generates the message header of 0 message length.
  221. func encode(c Codec, msg interface{}, cp Compressor, cbuf *bytes.Buffer) ([]byte, error) {
  222. var b []byte
  223. var length uint
  224. if msg != nil {
  225. var err error
  226. // TODO(zhaoq): optimize to reduce memory alloc and copying.
  227. b, err = c.Marshal(msg)
  228. if err != nil {
  229. return nil, err
  230. }
  231. if cp != nil {
  232. if err := cp.Do(cbuf, b); err != nil {
  233. return nil, err
  234. }
  235. b = cbuf.Bytes()
  236. }
  237. length = uint(len(b))
  238. }
  239. if length > math.MaxUint32 {
  240. return nil, Errorf(codes.InvalidArgument, "grpc: message too large (%d bytes)", length)
  241. }
  242. const (
  243. payloadLen = 1
  244. sizeLen = 4
  245. )
  246. var buf = make([]byte, payloadLen+sizeLen+len(b))
  247. // Write payload format
  248. if cp == nil {
  249. buf[0] = byte(compressionNone)
  250. } else {
  251. buf[0] = byte(compressionMade)
  252. }
  253. // Write length of b into buf
  254. binary.BigEndian.PutUint32(buf[1:], uint32(length))
  255. // Copy encoded msg to buf
  256. copy(buf[5:], b)
  257. return buf, nil
  258. }
  259. func checkRecvPayload(pf payloadFormat, recvCompress string, dc Decompressor) error {
  260. switch pf {
  261. case compressionNone:
  262. case compressionMade:
  263. if dc == nil || recvCompress != dc.Type() {
  264. return Errorf(codes.Unimplemented, "grpc: Decompressor is not installed for grpc-encoding %q", recvCompress)
  265. }
  266. default:
  267. return Errorf(codes.Internal, "grpc: received unexpected payload format %d", pf)
  268. }
  269. return nil
  270. }
  271. func recv(p *parser, c Codec, s *transport.Stream, dc Decompressor, m interface{}, maxMsgSize int) error {
  272. pf, d, err := p.recvMsg(maxMsgSize)
  273. if err != nil {
  274. return err
  275. }
  276. if err := checkRecvPayload(pf, s.RecvCompress(), dc); err != nil {
  277. return err
  278. }
  279. if pf == compressionMade {
  280. d, err = dc.Do(bytes.NewReader(d))
  281. if err != nil {
  282. return Errorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
  283. }
  284. }
  285. if len(d) > maxMsgSize {
  286. // TODO: Revisit the error code. Currently keep it consistent with java
  287. // implementation.
  288. return Errorf(codes.Internal, "grpc: received a message of %d bytes exceeding %d limit", len(d), maxMsgSize)
  289. }
  290. if err := c.Unmarshal(d, m); err != nil {
  291. return Errorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
  292. }
  293. return nil
  294. }
  295. // rpcError defines the status from an RPC.
  296. type rpcError struct {
  297. code codes.Code
  298. desc string
  299. }
  300. func (e *rpcError) Error() string {
  301. return fmt.Sprintf("rpc error: code = %d desc = %s", e.code, e.desc)
  302. }
  303. // Code returns the error code for err if it was produced by the rpc system.
  304. // Otherwise, it returns codes.Unknown.
  305. func Code(err error) codes.Code {
  306. if err == nil {
  307. return codes.OK
  308. }
  309. if e, ok := err.(*rpcError); ok {
  310. return e.code
  311. }
  312. return codes.Unknown
  313. }
  314. // ErrorDesc returns the error description of err if it was produced by the rpc system.
  315. // Otherwise, it returns err.Error() or empty string when err is nil.
  316. func ErrorDesc(err error) string {
  317. if err == nil {
  318. return ""
  319. }
  320. if e, ok := err.(*rpcError); ok {
  321. return e.desc
  322. }
  323. return err.Error()
  324. }
  325. // Errorf returns an error containing an error code and a description;
  326. // Errorf returns nil if c is OK.
  327. func Errorf(c codes.Code, format string, a ...interface{}) error {
  328. if c == codes.OK {
  329. return nil
  330. }
  331. return &rpcError{
  332. code: c,
  333. desc: fmt.Sprintf(format, a...),
  334. }
  335. }
  336. // toRPCErr converts an error into a rpcError.
  337. func toRPCErr(err error) error {
  338. switch e := err.(type) {
  339. case *rpcError:
  340. return err
  341. case transport.StreamError:
  342. return &rpcError{
  343. code: e.Code,
  344. desc: e.Desc,
  345. }
  346. case transport.ConnectionError:
  347. return &rpcError{
  348. code: codes.Internal,
  349. desc: e.Desc,
  350. }
  351. default:
  352. switch err {
  353. case context.DeadlineExceeded:
  354. return &rpcError{
  355. code: codes.DeadlineExceeded,
  356. desc: err.Error(),
  357. }
  358. case context.Canceled:
  359. return &rpcError{
  360. code: codes.Canceled,
  361. desc: err.Error(),
  362. }
  363. case ErrClientConnClosing:
  364. return &rpcError{
  365. code: codes.FailedPrecondition,
  366. desc: err.Error(),
  367. }
  368. }
  369. }
  370. return Errorf(codes.Unknown, "%v", err)
  371. }
  372. // convertCode converts a standard Go error into its canonical code. Note that
  373. // this is only used to translate the error returned by the server applications.
  374. func convertCode(err error) codes.Code {
  375. switch err {
  376. case nil:
  377. return codes.OK
  378. case io.EOF:
  379. return codes.OutOfRange
  380. case io.ErrClosedPipe, io.ErrNoProgress, io.ErrShortBuffer, io.ErrShortWrite, io.ErrUnexpectedEOF:
  381. return codes.FailedPrecondition
  382. case os.ErrInvalid:
  383. return codes.InvalidArgument
  384. case context.Canceled:
  385. return codes.Canceled
  386. case context.DeadlineExceeded:
  387. return codes.DeadlineExceeded
  388. }
  389. switch {
  390. case os.IsExist(err):
  391. return codes.AlreadyExists
  392. case os.IsNotExist(err):
  393. return codes.NotFound
  394. case os.IsPermission(err):
  395. return codes.PermissionDenied
  396. }
  397. return codes.Unknown
  398. }
  399. // SupportPackageIsVersion4 is referenced from generated protocol buffer files
  400. // to assert that that code is compatible with this version of the grpc package.
  401. //
  402. // This constant may be renamed in the future if a change in the generated code
  403. // requires a synchronised update of grpc-go and protoc-gen-go. This constant
  404. // should not be referenced from any other code.
  405. const SupportPackageIsVersion4 = true