encode.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162
  1. package eventstream
  2. import (
  3. "bytes"
  4. "encoding/binary"
  5. "encoding/hex"
  6. "encoding/json"
  7. "fmt"
  8. "hash"
  9. "hash/crc32"
  10. "io"
  11. "github.com/aws/aws-sdk-go/aws"
  12. )
  13. // Encoder provides EventStream message encoding.
  14. type Encoder struct {
  15. w io.Writer
  16. logger aws.Logger
  17. headersBuf *bytes.Buffer
  18. }
  19. // NewEncoder initializes and returns an Encoder to encode Event Stream
  20. // messages to an io.Writer.
  21. func NewEncoder(w io.Writer, opts ...func(*Encoder)) *Encoder {
  22. e := &Encoder{
  23. w: w,
  24. headersBuf: bytes.NewBuffer(nil),
  25. }
  26. for _, opt := range opts {
  27. opt(e)
  28. }
  29. return e
  30. }
  31. // EncodeWithLogger adds a logger to be used by the encode when decoding
  32. // stream events.
  33. func EncodeWithLogger(logger aws.Logger) func(*Encoder) {
  34. return func(d *Encoder) {
  35. d.logger = logger
  36. }
  37. }
  38. // Encode encodes a single EventStream message to the io.Writer the Encoder
  39. // was created with. An error is returned if writing the message fails.
  40. func (e *Encoder) Encode(msg Message) (err error) {
  41. e.headersBuf.Reset()
  42. writer := e.w
  43. if e.logger != nil {
  44. encodeMsgBuf := bytes.NewBuffer(nil)
  45. writer = io.MultiWriter(writer, encodeMsgBuf)
  46. defer func() {
  47. logMessageEncode(e.logger, encodeMsgBuf, msg, err)
  48. }()
  49. }
  50. if err = EncodeHeaders(e.headersBuf, msg.Headers); err != nil {
  51. return err
  52. }
  53. crc := crc32.New(crc32IEEETable)
  54. hashWriter := io.MultiWriter(writer, crc)
  55. headersLen := uint32(e.headersBuf.Len())
  56. payloadLen := uint32(len(msg.Payload))
  57. if err = encodePrelude(hashWriter, crc, headersLen, payloadLen); err != nil {
  58. return err
  59. }
  60. if headersLen > 0 {
  61. if _, err = io.Copy(hashWriter, e.headersBuf); err != nil {
  62. return err
  63. }
  64. }
  65. if payloadLen > 0 {
  66. if _, err = hashWriter.Write(msg.Payload); err != nil {
  67. return err
  68. }
  69. }
  70. msgCRC := crc.Sum32()
  71. return binary.Write(writer, binary.BigEndian, msgCRC)
  72. }
  73. func logMessageEncode(logger aws.Logger, msgBuf *bytes.Buffer, msg Message, encodeErr error) {
  74. w := bytes.NewBuffer(nil)
  75. defer func() { logger.Log(w.String()) }()
  76. fmt.Fprintf(w, "Message to encode:\n")
  77. encoder := json.NewEncoder(w)
  78. if err := encoder.Encode(msg); err != nil {
  79. fmt.Fprintf(w, "Failed to get encoded message, %v\n", err)
  80. }
  81. if encodeErr != nil {
  82. fmt.Fprintf(w, "Encode error: %v\n", encodeErr)
  83. return
  84. }
  85. fmt.Fprintf(w, "Raw message:\n%s\n", hex.Dump(msgBuf.Bytes()))
  86. }
  87. func encodePrelude(w io.Writer, crc hash.Hash32, headersLen, payloadLen uint32) error {
  88. p := messagePrelude{
  89. Length: minMsgLen + headersLen + payloadLen,
  90. HeadersLen: headersLen,
  91. }
  92. if err := p.ValidateLens(); err != nil {
  93. return err
  94. }
  95. err := binaryWriteFields(w, binary.BigEndian,
  96. p.Length,
  97. p.HeadersLen,
  98. )
  99. if err != nil {
  100. return err
  101. }
  102. p.PreludeCRC = crc.Sum32()
  103. err = binary.Write(w, binary.BigEndian, p.PreludeCRC)
  104. if err != nil {
  105. return err
  106. }
  107. return nil
  108. }
  109. // EncodeHeaders writes the header values to the writer encoded in the event
  110. // stream format. Returns an error if a header fails to encode.
  111. func EncodeHeaders(w io.Writer, headers Headers) error {
  112. for _, h := range headers {
  113. hn := headerName{
  114. Len: uint8(len(h.Name)),
  115. }
  116. copy(hn.Name[:hn.Len], h.Name)
  117. if err := hn.encode(w); err != nil {
  118. return err
  119. }
  120. if err := h.Value.encode(w); err != nil {
  121. return err
  122. }
  123. }
  124. return nil
  125. }
  126. func binaryWriteFields(w io.Writer, order binary.ByteOrder, vs ...interface{}) error {
  127. for _, v := range vs {
  128. if err := binary.Write(w, order, v); err != nil {
  129. return err
  130. }
  131. }
  132. return nil
  133. }