rpc_util.go 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449
  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() (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. // TODO(bradfitz,zhaoq): garbage. reuse buffer after proto decoding instead
  206. // of making it for each message:
  207. msg = make([]byte, int(length))
  208. if _, err := io.ReadFull(p.r, msg); err != nil {
  209. if err == io.EOF {
  210. err = io.ErrUnexpectedEOF
  211. }
  212. return 0, nil, err
  213. }
  214. return pf, msg, nil
  215. }
  216. // encode serializes msg and prepends the message header. If msg is nil, it
  217. // generates the message header of 0 message length.
  218. func encode(c Codec, msg interface{}, cp Compressor, cbuf *bytes.Buffer) ([]byte, error) {
  219. var b []byte
  220. var length uint
  221. if msg != nil {
  222. var err error
  223. // TODO(zhaoq): optimize to reduce memory alloc and copying.
  224. b, err = c.Marshal(msg)
  225. if err != nil {
  226. return nil, err
  227. }
  228. if cp != nil {
  229. if err := cp.Do(cbuf, b); err != nil {
  230. return nil, err
  231. }
  232. b = cbuf.Bytes()
  233. }
  234. length = uint(len(b))
  235. }
  236. if length > math.MaxUint32 {
  237. return nil, Errorf(codes.InvalidArgument, "grpc: message too large (%d bytes)", length)
  238. }
  239. const (
  240. payloadLen = 1
  241. sizeLen = 4
  242. )
  243. var buf = make([]byte, payloadLen+sizeLen+len(b))
  244. // Write payload format
  245. if cp == nil {
  246. buf[0] = byte(compressionNone)
  247. } else {
  248. buf[0] = byte(compressionMade)
  249. }
  250. // Write length of b into buf
  251. binary.BigEndian.PutUint32(buf[1:], uint32(length))
  252. // Copy encoded msg to buf
  253. copy(buf[5:], b)
  254. return buf, nil
  255. }
  256. func checkRecvPayload(pf payloadFormat, recvCompress string, dc Decompressor) error {
  257. switch pf {
  258. case compressionNone:
  259. case compressionMade:
  260. if dc == nil || recvCompress != dc.Type() {
  261. return transport.StreamErrorf(codes.Unimplemented, "grpc: Decompressor is not installed for grpc-encoding %q", recvCompress)
  262. }
  263. default:
  264. return transport.StreamErrorf(codes.Internal, "grpc: received unexpected payload format %d", pf)
  265. }
  266. return nil
  267. }
  268. func recv(p *parser, c Codec, s *transport.Stream, dc Decompressor, m interface{}) error {
  269. pf, d, err := p.recvMsg()
  270. if err != nil {
  271. return err
  272. }
  273. if err := checkRecvPayload(pf, s.RecvCompress(), dc); err != nil {
  274. return err
  275. }
  276. if pf == compressionMade {
  277. d, err = dc.Do(bytes.NewReader(d))
  278. if err != nil {
  279. return transport.StreamErrorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
  280. }
  281. }
  282. if err := c.Unmarshal(d, m); err != nil {
  283. return transport.StreamErrorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
  284. }
  285. return nil
  286. }
  287. // rpcError defines the status from an RPC.
  288. type rpcError struct {
  289. code codes.Code
  290. desc string
  291. }
  292. func (e *rpcError) Error() string {
  293. return fmt.Sprintf("rpc error: code = %d desc = %s", e.code, e.desc)
  294. }
  295. // Code returns the error code for err if it was produced by the rpc system.
  296. // Otherwise, it returns codes.Unknown.
  297. func Code(err error) codes.Code {
  298. if err == nil {
  299. return codes.OK
  300. }
  301. if e, ok := err.(*rpcError); ok {
  302. return e.code
  303. }
  304. return codes.Unknown
  305. }
  306. // ErrorDesc returns the error description of err if it was produced by the rpc system.
  307. // Otherwise, it returns err.Error() or empty string when err is nil.
  308. func ErrorDesc(err error) string {
  309. if err == nil {
  310. return ""
  311. }
  312. if e, ok := err.(*rpcError); ok {
  313. return e.desc
  314. }
  315. return err.Error()
  316. }
  317. // Errorf returns an error containing an error code and a description;
  318. // Errorf returns nil if c is OK.
  319. func Errorf(c codes.Code, format string, a ...interface{}) error {
  320. if c == codes.OK {
  321. return nil
  322. }
  323. return &rpcError{
  324. code: c,
  325. desc: fmt.Sprintf(format, a...),
  326. }
  327. }
  328. // toRPCErr converts an error into a rpcError.
  329. func toRPCErr(err error) error {
  330. switch e := err.(type) {
  331. case *rpcError:
  332. return err
  333. case transport.StreamError:
  334. return &rpcError{
  335. code: e.Code,
  336. desc: e.Desc,
  337. }
  338. case transport.ConnectionError:
  339. return &rpcError{
  340. code: codes.Internal,
  341. desc: e.Desc,
  342. }
  343. default:
  344. switch err {
  345. case context.DeadlineExceeded:
  346. return &rpcError{
  347. code: codes.DeadlineExceeded,
  348. desc: err.Error(),
  349. }
  350. case context.Canceled:
  351. return &rpcError{
  352. code: codes.Canceled,
  353. desc: err.Error(),
  354. }
  355. case ErrClientConnClosing:
  356. return &rpcError{
  357. code: codes.FailedPrecondition,
  358. desc: err.Error(),
  359. }
  360. }
  361. }
  362. return Errorf(codes.Unknown, "%v", err)
  363. }
  364. // convertCode converts a standard Go error into its canonical code. Note that
  365. // this is only used to translate the error returned by the server applications.
  366. func convertCode(err error) codes.Code {
  367. switch err {
  368. case nil:
  369. return codes.OK
  370. case io.EOF:
  371. return codes.OutOfRange
  372. case io.ErrClosedPipe, io.ErrNoProgress, io.ErrShortBuffer, io.ErrShortWrite, io.ErrUnexpectedEOF:
  373. return codes.FailedPrecondition
  374. case os.ErrInvalid:
  375. return codes.InvalidArgument
  376. case context.Canceled:
  377. return codes.Canceled
  378. case context.DeadlineExceeded:
  379. return codes.DeadlineExceeded
  380. }
  381. switch {
  382. case os.IsExist(err):
  383. return codes.AlreadyExists
  384. case os.IsNotExist(err):
  385. return codes.NotFound
  386. case os.IsPermission(err):
  387. return codes.PermissionDenied
  388. }
  389. return codes.Unknown
  390. }
  391. // SupportPackageIsVersion3 is referenced from generated protocol buffer files
  392. // to assert that that code is compatible with this version of the grpc package.
  393. //
  394. // This constant may be renamed in the future if a change in the generated code
  395. // requires a synchronised update of grpc-go and protoc-gen-go. This constant
  396. // should not be referenced from any other code.
  397. const SupportPackageIsVersion3 = true