reader.go 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201
  1. // Copyright 2016 Google Inc. All Rights Reserved.
  2. //
  3. // Distributed under MIT license.
  4. // See file LICENSE for detail or copy at https://opensource.org/licenses/MIT
  5. // Package cbrotli compresses and decompresses data with C-Brotli library.
  6. package cbrotli
  7. /*
  8. #include <stddef.h>
  9. #include <stdint.h>
  10. #include <brotli/decode.h>
  11. static BrotliDecoderResult DecompressStream(BrotliDecoderState* s,
  12. uint8_t* out, size_t out_len,
  13. const uint8_t* in, size_t in_len,
  14. size_t* bytes_written,
  15. size_t* bytes_consumed) {
  16. size_t in_remaining = in_len;
  17. size_t out_remaining = out_len;
  18. BrotliDecoderResult result = BrotliDecoderDecompressStream(
  19. s, &in_remaining, &in, &out_remaining, &out, NULL);
  20. *bytes_written = out_len - out_remaining;
  21. *bytes_consumed = in_len - in_remaining;
  22. return result;
  23. }
  24. */
  25. import "C"
  26. import (
  27. "bytes"
  28. "errors"
  29. "io"
  30. "io/ioutil"
  31. "runtime"
  32. )
  33. type decodeError C.BrotliDecoderErrorCode
  34. func (err decodeError) Error() string {
  35. return "cbrotli: " +
  36. C.GoString(C.BrotliDecoderErrorString(C.BrotliDecoderErrorCode(err)))
  37. }
  38. var errExcessiveInput = errors.New("cbrotli: excessive input")
  39. var errInvalidState = errors.New("cbrotli: invalid state")
  40. var errReaderClosed = errors.New("cbrotli: Reader is closed")
  41. // Reader implements io.ReadCloser by reading Brotli-encoded data from an
  42. // underlying Reader.
  43. type Reader struct {
  44. src io.Reader
  45. state *C.BrotliDecoderState
  46. buf []byte // scratch space for reading from src
  47. in []byte // current chunk to decode; usually aliases buf
  48. pinner *runtime.Pinner // raw dictionary pinner
  49. }
  50. // readBufSize is a "good" buffer size that avoids excessive round-trips
  51. // between C and Go but doesn't waste too much memory on buffering.
  52. // It is arbitrarily chosen to be equal to the constant used in io.Copy.
  53. const readBufSize = 32 * 1024
  54. // NewReader initializes new Reader instance.
  55. // Close MUST be called to free resources.
  56. func NewReader(src io.Reader) *Reader {
  57. return NewReaderWithRawDictionary(src, nil)
  58. }
  59. // NewReaderWithRawDictionary initializes new Reader instance with shared dictionary.
  60. // Close MUST be called to free resources.
  61. func NewReaderWithRawDictionary(src io.Reader, dictionary []byte) *Reader {
  62. s := C.BrotliDecoderCreateInstance(nil, nil, nil)
  63. var p *runtime.Pinner
  64. if dictionary != nil {
  65. p = new(runtime.Pinner)
  66. p.Pin(&dictionary[0])
  67. // TODO(eustas): use return value
  68. C.BrotliDecoderAttachDictionary(s, C.BrotliSharedDictionaryType( /* RAW */ 0),
  69. C.size_t(len(dictionary)), (*C.uint8_t)(&dictionary[0]))
  70. }
  71. return &Reader{
  72. src: src,
  73. state: s,
  74. buf: make([]byte, readBufSize),
  75. pinner: p,
  76. }
  77. }
  78. // Close implements io.Closer. Close MUST be invoked to free native resources.
  79. func (r *Reader) Close() error {
  80. if r.state == nil {
  81. return errReaderClosed
  82. }
  83. // Close despite the state; i.e. there might be some unread decoded data.
  84. C.BrotliDecoderDestroyInstance(r.state)
  85. r.state = nil
  86. if r.pinner != nil {
  87. r.pinner.Unpin()
  88. r.pinner = nil
  89. }
  90. return nil
  91. }
  92. func (r *Reader) Read(p []byte) (n int, err error) {
  93. if r.state == nil {
  94. return 0, errReaderClosed
  95. }
  96. if int(C.BrotliDecoderHasMoreOutput(r.state)) == 0 && len(r.in) == 0 {
  97. m, readErr := r.src.Read(r.buf)
  98. if m == 0 {
  99. // If readErr is `nil`, we just proxy underlying stream behavior.
  100. return 0, readErr
  101. }
  102. r.in = r.buf[:m]
  103. }
  104. if len(p) == 0 {
  105. return 0, nil
  106. }
  107. for {
  108. var written, consumed C.size_t
  109. var data *C.uint8_t
  110. if len(r.in) != 0 {
  111. data = (*C.uint8_t)(&r.in[0])
  112. }
  113. result := C.DecompressStream(r.state,
  114. (*C.uint8_t)(&p[0]), C.size_t(len(p)),
  115. data, C.size_t(len(r.in)),
  116. &written, &consumed)
  117. r.in = r.in[int(consumed):]
  118. n = int(written)
  119. switch result {
  120. case C.BROTLI_DECODER_RESULT_SUCCESS:
  121. if len(r.in) > 0 {
  122. return n, errExcessiveInput
  123. }
  124. return n, nil
  125. case C.BROTLI_DECODER_RESULT_ERROR:
  126. return n, decodeError(C.BrotliDecoderGetErrorCode(r.state))
  127. case C.BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT:
  128. if n == 0 {
  129. return 0, io.ErrShortBuffer
  130. }
  131. return n, nil
  132. case C.BROTLI_DECODER_NEEDS_MORE_INPUT:
  133. }
  134. if len(r.in) != 0 {
  135. return 0, errInvalidState
  136. }
  137. // Calling r.src.Read may block. Don't block if we have data to return.
  138. if n > 0 {
  139. return n, nil
  140. }
  141. // Top off the buffer.
  142. encN, err := r.src.Read(r.buf)
  143. if encN == 0 {
  144. // Not enough data to complete decoding.
  145. if err == io.EOF {
  146. return 0, io.ErrUnexpectedEOF
  147. }
  148. return 0, err
  149. }
  150. r.in = r.buf[:encN]
  151. }
  152. return n, nil
  153. }
  154. // Decode decodes Brotli encoded data.
  155. func Decode(encodedData []byte) ([]byte, error) {
  156. return DecodeWithRawDictionary(encodedData, nil)
  157. }
  158. // DecodeWithRawDictionary decodes Brotli encoded data with shared dictionary.
  159. func DecodeWithRawDictionary(encodedData []byte, dictionary []byte) ([]byte, error) {
  160. s := C.BrotliDecoderCreateInstance(nil, nil, nil)
  161. var p *runtime.Pinner
  162. if dictionary != nil {
  163. p = new(runtime.Pinner)
  164. p.Pin(&dictionary[0])
  165. // TODO(eustas): use return value
  166. C.BrotliDecoderAttachDictionary(s, C.BrotliSharedDictionaryType( /* RAW */ 0),
  167. C.size_t(len(dictionary)), (*C.uint8_t)(&dictionary[0]))
  168. }
  169. r := &Reader{
  170. src: bytes.NewReader(nil),
  171. state: s,
  172. buf: make([]byte, 4), // arbitrarily small but nonzero so that r.src.Read returns io.EOF
  173. in: encodedData,
  174. pinner: p,
  175. }
  176. defer r.Close()
  177. return ioutil.ReadAll(r)
  178. }