writer.go 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222
  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
  6. /*
  7. #include <stdbool.h>
  8. #include <stddef.h>
  9. #include <stdint.h>
  10. #include <brotli/encode.h>
  11. struct CompressStreamResult {
  12. size_t bytes_consumed;
  13. const uint8_t* output_data;
  14. size_t output_data_size;
  15. int success;
  16. int has_more;
  17. };
  18. static struct CompressStreamResult CompressStream(
  19. BrotliEncoderState* s, BrotliEncoderOperation op,
  20. const uint8_t* data, size_t data_size) {
  21. struct CompressStreamResult result;
  22. size_t available_in = data_size;
  23. const uint8_t* next_in = data;
  24. size_t available_out = 0;
  25. result.success = BrotliEncoderCompressStream(s, op,
  26. &available_in, &next_in, &available_out, 0, 0) ? 1 : 0;
  27. result.bytes_consumed = data_size - available_in;
  28. result.output_data = 0;
  29. result.output_data_size = 0;
  30. if (result.success) {
  31. result.output_data = BrotliEncoderTakeOutput(s, &result.output_data_size);
  32. }
  33. result.has_more = BrotliEncoderHasMoreOutput(s) ? 1 : 0;
  34. return result;
  35. }
  36. */
  37. import "C"
  38. import (
  39. "bytes"
  40. "errors"
  41. "io"
  42. "runtime"
  43. "unsafe"
  44. )
  45. // PreparedDictionary is a handle to native object.
  46. type PreparedDictionary struct {
  47. opaque *C.BrotliEncoderPreparedDictionary
  48. pinner *runtime.Pinner
  49. }
  50. // DictionaryType is type for shared dictionary
  51. type DictionaryType int
  52. const (
  53. // DtRaw denotes LZ77 prefix dictionary
  54. DtRaw DictionaryType = 0
  55. // DtSerialized denotes serialized format
  56. DtSerialized DictionaryType = 1
  57. )
  58. // NewPreparedDictionary prepares dictionary data for encoder.
  59. // Same instance can be used for multiple encoding sessions.
  60. // Close MUST be called to free resources.
  61. func NewPreparedDictionary(data []byte, dictionaryType DictionaryType, quality int) *PreparedDictionary {
  62. var ptr *C.uint8_t
  63. if len(data) != 0 {
  64. ptr = (*C.uint8_t)(&data[0])
  65. }
  66. p := new(runtime.Pinner)
  67. p.Pin(&data[0])
  68. d := C.BrotliEncoderPrepareDictionary(C.BrotliSharedDictionaryType(dictionaryType), C.size_t(len(data)), ptr, C.int(quality), nil, nil, nil)
  69. return &PreparedDictionary{
  70. opaque: d,
  71. pinner: p,
  72. }
  73. }
  74. // Close frees C resources.
  75. // IMPORTANT: calling Close until all encoders that use that dictionary are closed as well will
  76. // cause crash.
  77. func (p *PreparedDictionary) Close() error {
  78. // C-Brotli tolerates `nil` pointer here.
  79. C.BrotliEncoderDestroyPreparedDictionary(p.opaque)
  80. p.opaque = nil
  81. p.pinner.Unpin()
  82. return nil
  83. }
  84. // WriterOptions configures Writer.
  85. type WriterOptions struct {
  86. // Quality controls the compression-speed vs compression-density trade-offs.
  87. // The higher the quality, the slower the compression. Range is 0 to 11.
  88. Quality int
  89. // LGWin is the base 2 logarithm of the sliding window size.
  90. // Range is 10 to 24. 0 indicates automatic configuration based on Quality.
  91. LGWin int
  92. // Prepared shared dictionary
  93. Dictionary *PreparedDictionary
  94. }
  95. // Writer implements io.WriteCloser by writing Brotli-encoded data to an
  96. // underlying Writer.
  97. type Writer struct {
  98. healthy bool
  99. dst io.Writer
  100. state *C.BrotliEncoderState
  101. buf, encoded []byte
  102. }
  103. var (
  104. errEncode = errors.New("cbrotli: encode error")
  105. errWriterClosed = errors.New("cbrotli: Writer is closed")
  106. errWriterUnhealthy = errors.New("cbrotli: Writer is unhealthy")
  107. )
  108. // NewWriter initializes new Writer instance.
  109. // Close MUST be called to free resources.
  110. func NewWriter(dst io.Writer, options WriterOptions) *Writer {
  111. state := C.BrotliEncoderCreateInstance(nil, nil, nil)
  112. healthy := state != nil
  113. if C.BrotliEncoderSetParameter(
  114. state, C.BROTLI_PARAM_QUALITY, (C.uint32_t)(options.Quality)) == 0 {
  115. healthy = false
  116. }
  117. if options.LGWin > 0 {
  118. if C.BrotliEncoderSetParameter(
  119. state, C.BROTLI_PARAM_LGWIN, (C.uint32_t)(options.LGWin)) == 0 {
  120. healthy = false
  121. }
  122. }
  123. if options.Dictionary != nil {
  124. if C.BrotliEncoderAttachPreparedDictionary(state, options.Dictionary.opaque) == 0 {
  125. healthy = false
  126. }
  127. }
  128. return &Writer{
  129. healthy: healthy,
  130. dst: dst,
  131. state: state,
  132. }
  133. }
  134. func (w *Writer) writeChunk(p []byte, op C.BrotliEncoderOperation) (n int, err error) {
  135. if !w.healthy {
  136. return 0, errWriterUnhealthy
  137. }
  138. if w.state == nil {
  139. return 0, errWriterClosed
  140. }
  141. for {
  142. var data *C.uint8_t
  143. if len(p) != 0 {
  144. data = (*C.uint8_t)(&p[0])
  145. }
  146. result := C.CompressStream(w.state, op, data, C.size_t(len(p)))
  147. if result.success == 0 {
  148. return n, errEncode
  149. }
  150. p = p[int(result.bytes_consumed):]
  151. n += int(result.bytes_consumed)
  152. length := int(result.output_data_size)
  153. if length != 0 {
  154. // It is a workaround for non-copying-wrapping of native memory.
  155. // C-encoder never pushes output block longer than ((2 << 25) + 502).
  156. // TODO(eustas): use natural wrapper, when it becomes available, see
  157. // https://golang.org/issue/13656.
  158. output := (*[1 << 30]byte)(unsafe.Pointer(result.output_data))[:length:length]
  159. _, err = w.dst.Write(output)
  160. if err != nil {
  161. return n, err
  162. }
  163. }
  164. if len(p) == 0 && result.has_more == 0 {
  165. return n, nil
  166. }
  167. }
  168. }
  169. // Flush outputs encoded data for all input provided to Write. The resulting
  170. // output can be decoded to match all input before Flush, but the stream is
  171. // not yet complete until after Close.
  172. // Flush has a negative impact on compression.
  173. func (w *Writer) Flush() error {
  174. _, err := w.writeChunk(nil, C.BROTLI_OPERATION_FLUSH)
  175. return err
  176. }
  177. // Close flushes remaining data to the decorated writer and frees C resources.
  178. func (w *Writer) Close() error {
  179. // If stream is already closed, it is reported by `writeChunk`.
  180. _, err := w.writeChunk(nil, C.BROTLI_OPERATION_FINISH)
  181. // C-Brotli tolerates `nil` pointer here.
  182. C.BrotliEncoderDestroyInstance(w.state)
  183. w.state = nil
  184. return err
  185. }
  186. // Write implements io.Writer. Flush or Close must be called to ensure that the
  187. // encoded bytes are actually flushed to the underlying Writer.
  188. func (w *Writer) Write(p []byte) (n int, err error) {
  189. return w.writeChunk(p, C.BROTLI_OPERATION_PROCESS)
  190. }
  191. // Encode returns content encoded with Brotli.
  192. func Encode(content []byte, options WriterOptions) ([]byte, error) {
  193. var buf bytes.Buffer
  194. writer := NewWriter(&buf, options)
  195. _, err := writer.Write(content)
  196. if closeErr := writer.Close(); err == nil {
  197. err = closeErr
  198. }
  199. return buf.Bytes(), err
  200. }