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