diff --git a/go/cbrotli/BUILD b/go/cbrotli/BUILD index 4bea8a0..e045b61 100755 --- a/go/cbrotli/BUILD +++ b/go/cbrotli/BUILD @@ -2,16 +2,19 @@ package(default_visibility = ["//visibility:public"]) licenses(["notice"]) # MIT -load("@io_bazel_rules_go//go:def.bzl", "go_prefix", "go_library", "go_test") +load("@io_bazel_rules_go//go:def.bzl", "go_prefix", "cgo_library", "go_test") go_prefix("github.com/google/brotli") -go_library( +cgo_library( name = "cbrotli", - srcs = ["cbrotli.go"], - deps = [ - "//go/cbrotli/internal:decoder", - "//go/cbrotli/internal:encoder", + srcs = [ + "reader.go", + "writer.go", + ], + cdeps = [ + "//:brotlidec", + "//:brotlienc", ], ) diff --git a/go/cbrotli/cbrotli.go b/go/cbrotli/cbrotli.go deleted file mode 100755 index a7009d0..0000000 --- a/go/cbrotli/cbrotli.go +++ /dev/null @@ -1,254 +0,0 @@ -// Copyright 2016 Google Inc. All Rights Reserved. -// -// Distributed under MIT license. -// See file LICENSE for detail or copy at https://opensource.org/licenses/MIT - -// Package cbrotli compresses and decompresses data with C-Brotli library. -package cbrotli - -import ( - "bytes" - "io" - - "github.com/google/brotli/go/cbrotli/internal/decoder" - "github.com/google/brotli/go/cbrotli/internal/encoder" -) - -// An internalError reports an error in the (c-)brotli code itself. -type internalError string - -func (e internalError) Error() string { - return "cbrotli: internal error: " + string(e) -} - -//------------------------------------------------------------------------------ -// Encoder -//------------------------------------------------------------------------------ - -// WriterOptions configures Writer. -type WriterOptions struct { - // Quality controls the compression-speed vs compression-density trade-offs. - // The higher the quality, the slower the compression. Range is 0 to 11. - Quality int - // LGWin is the base 2 logarithm of the sliding window size. - // Range is 10 to 24. 0 indicates automatic configuration based on Quality. - LGWin int -} - -// Writer implements io.WriteCloser, an io.Writer decorator that produces -// Brotli-encoded data. -type Writer struct { - dst io.Writer - encoder encoder.Encoder - closed bool -} - -// NewWriter initializes new Writer instance. -// Close MUST be called to free resources. -func NewWriter(dst io.Writer, options WriterOptions) *Writer { - return &Writer{ - dst: dst, - encoder: encoder.New(options.Quality, options.LGWin), - } -} - -// Close implements io.Closer. Close MUST be invoked to free native resources. -// Also Close implicitly flushes remaining data to the decorated writer. -func (z *Writer) Close() error { - if z.closed { - return nil - } - defer z.encoder.Close() - _, err := z.writeChunk(nil, encoder.Finish) - z.closed = true - return err -} - -func (z *Writer) writeChunk(p []byte, op encoder.Operation) (int, error) { - if z.closed { - return 0, internalError("write after close") - } - var totalBytesConsumed int - var err error - for { - bytesConsumed, output, status := z.encoder.CompressStream(p, op) - if status == encoder.Error { - err = internalError("encoder failure") - break - } - p = p[bytesConsumed:] - totalBytesConsumed += bytesConsumed - _, err = z.dst.Write(output) - if err != nil { - break - } - if len(p) == 0 && status == encoder.Done { - break - } - } - return totalBytesConsumed, err -} - -// Write implements io.Writer. -func (z *Writer) Write(p []byte) (int, error) { - return z.writeChunk(p, encoder.Process) -} - -// Flush outputs encoded data for all input provided to Write. The resulting -// output can be decoded to match all input before Flush, but the stream is -// not yet complete until after Close. -// Flush has a negative impact on compression. -func (z *Writer) Flush() error { - _, err := z.writeChunk(nil, encoder.Finish) - return err -} - -// Encode returns content encoded with Brotli. -func Encode(content []byte, options WriterOptions) ([]byte, error) { - var buf bytes.Buffer - writer := NewWriter(&buf, options) - defer writer.Close() - _, err := writer.Write(content) - if err != nil { - return nil, err - } - if err := writer.Close(); err != nil { - return nil, err - } - return buf.Bytes(), nil -} - -//------------------------------------------------------------------------------ -// Decoder -//------------------------------------------------------------------------------ - -// Reader implements io.ReadCloser, an io.Reader decorator that decodes -// Brotli-encoded data. -type Reader struct { - src io.Reader - decoder decoder.Decoder - buf []byte // intermediate read buffer pointed to by next - eof bool // true if all compressed stream is decoded - next []byte // buffered data to be passed to decoder - output []byte // data produced by decoder, but not yet consumed - srcErr error // last source reader error - err error // reader state; nil if it is OK to read further - closed bool // true is stream is already closed -} - -// NewReader initializes new Reader instance. -// Close MUST be called to free resources. -func NewReader(src io.Reader) *Reader { - return &Reader{ - src: src, - decoder: decoder.New(), - buf: make([]byte, 32*1024), - eof: false, - } -} - -// Close implements io.Closer. Close MUST be invoked to free native resources. -func (z *Reader) Close() error { - if z.closed { - return nil - } - z.decoder.Close() - z.err = internalError("read after close") - z.closed = true - return nil -} - -func isEOF(src io.Reader) bool { - n, err := src.Read(make([]byte, 1)) - return n == 0 && err == io.EOF -} - -// Read implements io.Reader. -func (z *Reader) Read(p []byte) (int, error) { - // Any error state is unrecoverable. - if z.err != nil { - return 0, z.err - } - // See io.Reader documentation. - if len(p) == 0 { - return 0, nil - } - - var totalOutBytes int - - // There is no practical limit for amount of bytes being consumed by decoder - // before producing any output. Continue feeding decoder until some data is - // produced - for { - // Push already produced output first. - if outBytes := len(z.output); outBytes != 0 { - outBytes = copy(p, z.output) - p = p[outBytes:] - z.output = z.output[outBytes:] - totalOutBytes += outBytes - // Output buffer is full. - if len(p) == 0 { - break - } - continue - } - // No more produced output left. - // If no more output is expected, then we are finished. - if z.eof { - z.err = io.EOF - break - } - // Replenish buffer (might cause blocking read), only if necessary. - if len(z.next) == 0 && totalOutBytes == 0 && z.srcErr != io.EOF { - var n int - n, z.srcErr = z.src.Read(z.buf) - z.next = z.buf[:n] - if z.srcErr != nil && z.srcErr != io.EOF { - z.err = z.srcErr - break - } - } - // Do decoding. - consumed, output, status := z.decoder.DecompressStream(z.next) - z.output = output - z.next = z.next[consumed:] - if status == decoder.Error { - // When error happens, the remaining output does not matter. - z.err = internalError("decoder failure") - break - } else if status == decoder.Done { - // Decoder stream is closed; no further input is expected. - if len(z.next) != 0 || (z.srcErr != io.EOF && !isEOF(z.src)) { - z.err = internalError("excessive input") - break - } - // No more output is expected; keep pushing output. - z.eof = true - continue - } else { - // If can not move any further... - if consumed == 0 && len(z.output) == 0 { - // Unexpected end of input. - if z.srcErr == io.EOF || totalOutBytes == 0 { - z.err = io.ErrUnexpectedEOF - } - // Postpone blocking reads for the next invocation. - break - } - // Continue pushing output. - } - } - return totalOutBytes, z.err -} - -// Decode decodes Brotli encoded data. -func Decode(encodedData []byte) ([]byte, error) { - var buf bytes.Buffer - reader := NewReader(bytes.NewReader(encodedData)) - defer reader.Close() - _, err := io.Copy(&buf, reader) - if err != nil { - return nil, err - } - return buf.Bytes(), nil -} diff --git a/go/cbrotli/cbrotli_test.go b/go/cbrotli/cbrotli_test.go index 4fe0960..0a9368b 100755 --- a/go/cbrotli/cbrotli_test.go +++ b/go/cbrotli/cbrotli_test.go @@ -91,7 +91,7 @@ func TestEncoderStreams(t *testing.T) { // to fill the window. const lgWin = 16 windowSize := int(math.Pow(2, lgWin)) - input := make([]byte, 2*windowSize) + input := make([]byte, 8*windowSize) rand.Read(input) out := bytes.Buffer{} e := NewWriter(&out, WriterOptions{Quality: 11, LGWin: lgWin}) @@ -106,7 +106,7 @@ func TestEncoderStreams(t *testing.T) { // We've fed more data than the sliding window size. Check that some // compressed data has been output. if out.Len() == 0 { - t.Errorf("Output length is after %d bytes written", n) + t.Errorf("Output length is 0 after %d bytes written", n) } if err := e.Close(); err != nil { t.Errorf("Close Error after copied %d bytes: %v", n, err) @@ -151,6 +151,9 @@ func TestEncoderFlush(t *testing.T) { if err := e.Flush(); err != nil { t.Fatalf("Flush(): %v", err) } + if out.Len() == 0 { + t.Fatalf("0 bytes written after Flush()") + } decompressed := make([]byte, 1000) reader := NewReader(bytes.NewReader(out.Bytes())) n, err := reader.Read(decompressed) @@ -187,8 +190,8 @@ func TestReader(t *testing.T) { "Reader output:\n"+ "%q\n"+ "want:\n"+ - "%q", - got, content) + "<%d bytes>", + got, len(content)) } } @@ -204,8 +207,8 @@ func TestDecode(t *testing.T) { "Decode content:\n"+ "%q\n"+ "want:\n"+ - "%q", - decoded, content) + "<%d bytes>", + decoded, len(content)) } } @@ -213,7 +216,13 @@ func TestDecodeFuzz(t *testing.T) { // Test that the decoder terminates with corrupted input. content := bytes.Repeat([]byte("hello world!"), 100) src := rand.NewSource(0) - encoded, _ := Encode(content, WriterOptions{Quality: 5}) + encoded, err := Encode(content, WriterOptions{Quality: 5}) + if err != nil { + t.Fatalf("Encode(<%d bytes>, _) = _, %s", len(content), err) + } + if len(encoded) == 0 { + t.Fatalf("Encode(<%d bytes>, _) produced empty output", len(content)) + } for i := 0; i < 100; i++ { enc := append([]byte{}, encoded...) for j := 0; j < 5; j++ { @@ -260,12 +269,18 @@ func TestEncodeDecode(t *testing.T) { t.Errorf("Decode: %v", err) } if !bytes.Equal(decoded, input) { + var want string + if len(input) > 320 { + want = fmt.Sprintf("<%d bytes>", len(input)) + } else { + want = fmt.Sprintf("%q", input) + } t.Errorf(""+ "Decode content:\n"+ "%q\n"+ "want:\n"+ - "%q", - decoded, input) + "%s", + decoded, want) } } } diff --git a/go/cbrotli/internal/BUILD b/go/cbrotli/internal/BUILD deleted file mode 100755 index a3e3b87..0000000 --- a/go/cbrotli/internal/BUILD +++ /dev/null @@ -1,19 +0,0 @@ -package(default_visibility = ["//visibility:public"]) - -licenses(["notice"]) # MIT - -load("@io_bazel_rules_go//go:def.bzl", "cgo_library") - -cgo_library( - name = "decoder", - srcs = ["decoder.go"], - visibility = ["//go/cbrotli:__subpackages__"], - cdeps = ["//:brotlidec"], -) - -cgo_library( - name = "encoder", - srcs = ["encoder.go"], - visibility = ["//go/cbrotli:__subpackages__"], - cdeps = ["//:brotlienc"], -) diff --git a/go/cbrotli/internal/decoder.go b/go/cbrotli/internal/decoder.go deleted file mode 100755 index 74d4c21..0000000 --- a/go/cbrotli/internal/decoder.go +++ /dev/null @@ -1,110 +0,0 @@ -// Copyright 2016 Google Inc. All Rights Reserved. -// -// Distributed under MIT license. -// See file LICENSE for detail or copy at https://opensource.org/licenses/MIT - -// Package decoder wraps the brotli decoder C API used by package brotli. -package decoder - -/* -#include - -// Wrap BrotliDecoderDecompressStream so that it doesn't take variable (in-out) -// pointers. Instead of updated pointer, deltas are saved in auxiliary struct. - -struct DecompressStreamResult { - size_t bytes_consumed; - const uint8_t* output_data; - size_t output_data_size; - BrotliDecoderResult status; -}; - -struct DecompressStreamResult DecompressStream(BrotliDecoderState* s, - const uint8_t* encoded_data, size_t encoded_data_size) { - struct DecompressStreamResult result; - size_t available_in = encoded_data_size; - const uint8_t* next_in = encoded_data; - size_t available_out = 0; - result.status = BrotliDecoderDecompressStream(s, - &available_in, &next_in, &available_out, 0, 0); - result.bytes_consumed = encoded_data_size - available_in; - result.output_data = 0; - result.output_data_size = 0; - if (result.status != BROTLI_DECODER_RESULT_ERROR) { - result.output_data = BrotliDecoderTakeOutput(s, &result.output_data_size); - if (BrotliDecoderIsFinished(s)) { - result.status = BROTLI_DECODER_RESULT_SUCCESS; - } - } - return result; -} - -*/ -import "C" -import ( - "unsafe" -) - -// Status represents internal state after DecompressStream invokation -type Status int - -const ( - // Error happened - Error Status = iota - // Done means that no more output will be produced - Done - // Ok means that more output might be produced with no additional input - Ok -) - -// Decoder is the Brotli c-decoder handle. -type Decoder struct { - state *C.BrotliDecoderState -} - -// New returns a new Brotli c-decoder handle. -// Close MUST be called to free resources. -func New() Decoder { - return Decoder{state: C.BrotliDecoderCreateInstance(nil, nil, nil)} -} - -// Close frees resources used by decoder. -func (z *Decoder) Close() { - C.BrotliDecoderDestroyInstance(z.state) - z.state = nil -} - -func goStatus(cStatus C.BrotliDecoderResult) (status Status) { - switch cStatus { - case C.BROTLI_DECODER_RESULT_SUCCESS: - return Done - case C.BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT: - return Ok - case C.BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT: - return Ok - } - return Error -} - -// cBytes casts a Go []byte into a C uint8_t*. We pass &buf[0] directly to C, -// which is legal because C doesn't save the pointer longer than the call and -// the byte array itself doesn't contain any pointers. -func cBytes(buf []byte) (*C.uint8_t, C.size_t) { - if len(buf) == 0 { - return (*C.uint8_t)(nil), 0 - } - return (*C.uint8_t)(unsafe.Pointer(&buf[0])), C.size_t(len(buf)) -} - -// DecompressStream reads Brotli-encoded bytes from in, and returns produced -// bytes. Output contents should not be modified. Liveness of output is -// hard-limited by Decoder liveness; slice becomes invalid when any Decoder -// method is invoked. -func (z *Decoder) DecompressStream(in []byte) ( - bytesConsumed int, output []byte, status Status) { - cin, cinSize := cBytes(in) - result := C.DecompressStream(z.state, cin, cinSize) - output = C.GoBytes( - unsafe.Pointer(result.output_data), C.int(result.output_data_size)) - return int(result.bytes_consumed), output, goStatus(result.status) -} diff --git a/go/cbrotli/internal/encoder.go b/go/cbrotli/internal/encoder.go deleted file mode 100755 index b8dd21c..0000000 --- a/go/cbrotli/internal/encoder.go +++ /dev/null @@ -1,135 +0,0 @@ -// Copyright 2016 Google Inc. All Rights Reserved. -// -// Distributed under MIT license. -// See file LICENSE for detail or copy at https://opensource.org/licenses/MIT - -// Package encoder wraps the brotli encoder C API used by package brotli. -package encoder - -/* -#include - -// Wrap BrotliEncoderCompressStream so that it doesn't take variable (in-out) -// pointers. Instead of updated pointer, deltas are saved in auxiliary struct. -struct CompressStreamResult { - size_t bytes_consumed; - const uint8_t* output_data; - size_t output_data_size; - int success; - int has_more; -}; - -struct CompressStreamResult CompressStream( - BrotliEncoderState* s, BrotliEncoderOperation op, - const uint8_t* data, size_t data_size) { - struct CompressStreamResult result; - size_t available_in = data_size; - const uint8_t* next_in = data; - size_t available_out = 0; - result.success = BrotliEncoderCompressStream(s, op, - &available_in, &next_in, &available_out, 0, 0) ? 1 : 0; - result.bytes_consumed = data_size - available_in; - result.output_data = 0; - result.output_data_size = 0; - if (result.success) { - result.output_data = BrotliEncoderTakeOutput(s, &result.output_data_size); - } - result.has_more = BrotliEncoderHasMoreOutput(s) ? 1 : 0; - return result; -} -*/ -import "C" -import ( - "unsafe" -) - -// Operation represents type of request to CompressStream -type Operation int - -const ( - // Process input - Process Operation = iota - // Flush input processed so far - Flush - // Finish stream - Finish -) - -// Status represents internal state after CompressStream invocation -type Status int - -const ( - // Error happened - Error Status = iota - // Done means that no more output will be produced - Done - // Ok means that more output might be produced with no additional input - Ok -) - -// Encoder is the Brotli c-encoder handle. -type Encoder struct { - state *C.BrotliEncoderState -} - -// New returns a new Brotli c-encoder handle. -// quality and lgWin are described in third_party/Brotli/enc/encode.h. -// Close MUST be called to free resources. -func New(quality, lgWin int) Encoder { - state := C.BrotliEncoderCreateInstance(nil, nil, nil) - // TODO(b/18187008): Check if LGBLOCK or MODE are useful to Flywheel. - C.BrotliEncoderSetParameter( - state, C.BROTLI_PARAM_QUALITY, (C.uint32_t)(quality)) - C.BrotliEncoderSetParameter( - state, C.BROTLI_PARAM_LGWIN, (C.uint32_t)(lgWin)) - return Encoder{state} -} - -// Close frees resources used by encoder. -func (z *Encoder) Close() { - C.BrotliEncoderDestroyInstance(z.state) - z.state = nil -} - -// cBytes casts a Go []byte into a C uint8_t*. We pass &buf[0] directly to C, -// which is legal because C doesn't save the pointer longer than the call and -// the byte array itself doesn't contain any pointers. -func cBytes(buf []byte) (*C.uint8_t, C.size_t) { - if len(buf) == 0 { - return (*C.uint8_t)(nil), 0 - } - return (*C.uint8_t)(unsafe.Pointer(&buf[0])), C.size_t(len(buf)) -} - -func cOperation(op Operation) (cOp C.BrotliEncoderOperation) { - switch op { - case Flush: - return C.BROTLI_OPERATION_FLUSH - case Finish: - return C.BROTLI_OPERATION_FINISH - } - return C.BROTLI_OPERATION_PROCESS -} - -// CompressStream processes data and produces Brotli-encoded bytes. Encoder may -// consume considerable amount of input before the first output bytes come out. -// Flush and Finish operations force Encoder to produce output that corresponds -// to input consumed so far. Output contents should not be modified. Liveness of -// output is hard-limited by Encoder liveness; slice becomes invalid when any -// Encoder method is invoked. -func (z *Encoder) CompressStream(in []byte, op Operation) ( - bytesConsumed int, output []byte, status Status) { - cin, cinSize := cBytes(in) - result := C.CompressStream(z.state, cOperation(op), cin, cinSize) - output = C.GoBytes( - unsafe.Pointer(result.output_data), C.int(result.output_data_size)) - var outcome Status - if result.success == 0 { - outcome = Error - } else if result.has_more != 0 { - outcome = Ok - } else { - outcome = Done - } - return int(result.bytes_consumed), output, outcome -} diff --git a/go/cbrotli/reader.go b/go/cbrotli/reader.go new file mode 100644 index 0000000..bb5930f --- /dev/null +++ b/go/cbrotli/reader.go @@ -0,0 +1,154 @@ +// Copyright 2016 Google Inc. All Rights Reserved. +// +// Distributed under MIT license. +// See file LICENSE for detail or copy at https://opensource.org/licenses/MIT + +// Package cbrotli compresses and decompresses data with C-Brotli library. +package cbrotli + +/* +#include +#include + +#include + +static BrotliDecoderResult DecompressStream(BrotliDecoderState* s, + uint8_t* out, size_t out_len, + const uint8_t* in, size_t in_len, + size_t* out_written, + size_t* in_consumed) { + size_t in_remaining = in_len; + size_t out_remaining = out_len; + BrotliDecoderResult result = BrotliDecoderDecompressStream( + s, &in_remaining, &in, &out_remaining, &out, NULL); + *out_written = out_len - out_remaining; + *in_consumed = in_len - in_remaining; + return result; +} + +*/ +import "C" + +import ( + "bytes" + "errors" + "io" +) + +type decodeError C.BrotliDecoderErrorCode + +func (err decodeError) Error() string { + return "cbrotli: " + C.GoString(C.BrotliDecoderErrorString(C.BrotliDecoderErrorCode(err))) +} + +var errExcessiveInput = errors.New("cbrotli: excessive input") + +// Reader implements io.ReadCloser by reading Brotli-encoded data from an +// underlying Reader. +type Reader struct { + src io.Reader + state *C.BrotliDecoderState + buf []byte // scratch space for reading from src + in []byte // current chunk to decode; usually aliases buf +} + +// initialBufSize is a "good" buffer size that avoids excessive round-trips +// between C and Go but doesn't waste too much memory on buffering. +// It is arbitrarily chosen to be equal to the constant used in io.Copy. +// +// TODO(bcmills): This constant should be based on empirical measurements. +const initialReadBufSize = 32 * 1024 + +// NewReader initializes new Reader instance. +// Close MUST be called to free resources. +func NewReader(src io.Reader) *Reader { + return &Reader{ + src: src, + state: C.BrotliDecoderCreateInstance(nil, nil, nil), + buf: make([]byte, initialReadBufSize), + } +} + +func (r *Reader) Close() error { + if r.state != nil { + C.BrotliDecoderDestroyInstance(r.state) + r.state = nil + } + return nil +} + +func (r *Reader) Read(p []byte) (n int, err error) { + if len(r.in) == 0 { + n, err := r.src.Read(r.buf) + if n == 0 && err != nil { + return 0, err + } + r.in = r.buf[:n] + } + + if len(p) == 0 { + return 0, nil + } + + for n == 0 { + var written, consumed C.size_t + result := C.DecompressStream(r.state, + (*C.uint8_t)(&p[0]), C.size_t(len(p)), + (*C.uint8_t)(&r.in[0]), C.size_t(len(r.in)), + &written, &consumed) + r.in = r.in[int(consumed):] + n = int(written) + + switch result { + case C.BROTLI_DECODER_RESULT_SUCCESS: + if len(r.in) > 0 { + return n, errExcessiveInput + } + return n, nil + case C.BROTLI_DECODER_RESULT_ERROR: + return n, decodeError(C.BrotliDecoderGetErrorCode(r.state)) + case C.BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT: + if n == 0 { + return 0, io.ErrShortBuffer + } + return n, nil + case C.BROTLI_DECODER_NEEDS_MORE_INPUT: + } + + if len(r.in) >= len(r.buf)/2 { + // Too much buffer was left over: expand it to make faster progress. + // + // TODO(bcmills): Is there some way to get better feedback about buffer + // sizes from the C API? + r.buf = make([]byte, len(r.in)*2) + } + + // Move unread bytes to the start of the buffer. + r.in = r.buf[:copy(r.buf, r.in)] + + // Top off the buffer. + encN, err := r.src.Read(r.in[len(r.in):cap(r.in)]) + if encN == 0 && n == 0 { + // We need more input to make progress, but there isn't any available. + if err == io.EOF { + return 0, io.ErrUnexpectedEOF + } + return 0, err + } + r.in = r.in[:len(r.in)+encN] + } + + return n, nil +} + +// Decode decodes Brotli encoded data. +func Decode(encodedData []byte) ([]byte, error) { + r := NewReader(bytes.NewReader(encodedData)) + defer r.Close() + var buf bytes.Buffer + _, err := io.Copy(&buf, r) + if err != nil { + return nil, err + } + return buf.Bytes(), nil +} diff --git a/go/cbrotli/writer.go b/go/cbrotli/writer.go new file mode 100644 index 0000000..aefcb94 --- /dev/null +++ b/go/cbrotli/writer.go @@ -0,0 +1,188 @@ +// Copyright 2016 Google Inc. All Rights Reserved. +// +// Distributed under MIT license. +// See file LICENSE for detail or copy at https://opensource.org/licenses/MIT + +package cbrotli + +/* +#include +#include +#include + +#include + +static bool CompressStream(BrotliEncoderState* s, BrotliEncoderOperation op, + uint8_t* out, size_t out_len, + const uint8_t* in, size_t in_len, + size_t* out_written, size_t* in_consumed) { + size_t in_remaining = in_len; + size_t out_remaining = out_len; + bool ok = !!BrotliEncoderCompressStream(s, op, &in_remaining, &in, + &out_remaining, &out, NULL); + *out_written = out_len - out_remaining; + if (in_consumed != NULL) { + *in_consumed = in_len - in_remaining; + } + return ok; +} +*/ +import "C" + +import ( + "bytes" + "errors" + "io" +) + +// WriterOptions configures Writer. +type WriterOptions struct { + // Quality controls the compression-speed vs compression-density trade-offs. + // The higher the quality, the slower the compression. Range is 0 to 11. + Quality int + // LGWin is the base 2 logarithm of the sliding window size. + // Range is 10 to 24. 0 indicates automatic configuration based on Quality. + LGWin int + // BufferSize is the number of bytes to use to buffer encoded output. + // 0 indicates an implementation-defined default. + BufferSize int +} + +// Writer implements io.WriteCloser by writing Brotli-encoded data to an +// underlying Writer. +type Writer struct { + dst io.Writer + state *C.BrotliEncoderState + buf, encoded []byte +} + +var ( + errEncode = errors.New("cbrotli: encode error") + errWriterClosed = errors.New("cbrotli: Writer is closed") +) + +// NewWriter initializes new Writer instance. +// Close MUST be called to free resources. +func NewWriter(dst io.Writer, options WriterOptions) *Writer { + state := C.BrotliEncoderCreateInstance(nil, nil, nil) + // TODO(b/18187008): Check if LGBLOCK or MODE are useful to Flywheel. + C.BrotliEncoderSetParameter( + state, C.BROTLI_PARAM_QUALITY, (C.uint32_t)(options.Quality)) + C.BrotliEncoderSetParameter( + state, C.BROTLI_PARAM_LGWIN, (C.uint32_t)(options.LGWin)) + + // TODO(bcmills): If the underlying io.Writer implements io.ReaderFrom, use + // that instead of copying through a secondary buffer. + bufSize := options.BufferSize + if bufSize <= 0 { + // TODO(bcmills): Are there better default buffer sizes to use here? + // + // Ideally it would be nice to dynamically resize the buffer based on + // feedback from the encoder. + if options.LGWin > 0 { + bufSize = 1 << uint(options.LGWin) + } else { + bufSize = initialReadBufSize + } + } + return &Writer{ + dst: dst, + state: state, + buf: make([]byte, bufSize), + } +} + +func (w *Writer) flushBuf() error { + if len(w.encoded) == 0 { + return nil + } + n, err := w.dst.Write(w.encoded) + w.encoded = w.encoded[n:] + return err +} + +func (w *Writer) untilEmpty(op C.BrotliEncoderOperation) error { + if w.state == nil { + return errWriterClosed + } + if err := w.flushBuf(); err != nil { + return err + } + for { + var written C.size_t + ok := C.CompressStream(w.state, op, + (*C.uint8_t)(&w.buf[0]), C.size_t(len(w.buf)), nil, 0, &written, nil) + w.encoded = w.buf[:int(written)] + if err := w.flushBuf(); err != nil { + return err + } + if written == 0 { + if !ok { + return errEncode + } + return nil + } + } +} + +// Flush outputs encoded data for all input provided to Write. The resulting +// output can be decoded to match all input before Flush, but the stream is +// not yet complete until after Close. +// Flush has a negative impact on compression. +func (w *Writer) Flush() error { + return w.untilEmpty(C.BROTLI_OPERATION_FLUSH) +} + +// Close flushes remaining data to the decorated writer +// and frees C resources. +func (w *Writer) Close() error { + err := w.Flush() + if err == nil { + err = w.untilEmpty(C.BROTLI_OPERATION_FINISH) + } + C.BrotliEncoderDestroyInstance(w.state) + w.state = nil + return err +} + +// Write implements io.Writer. Flush or Close must be called to ensure that the +// encoded bytes are actually flushed to the underlying Writer. +func (w *Writer) Write(p []byte) (n int, err error) { + if w.state == nil { + return 0, errWriterClosed + } + + for len(p) > 0 { + if err := w.flushBuf(); err != nil { + return n, err + } + + var written, consumed C.size_t + ok := C.CompressStream(w.state, C.BROTLI_OPERATION_PROCESS, + (*C.uint8_t)(&w.buf[0]), C.size_t(len(w.buf)), + (*C.uint8_t)(&p[0]), C.size_t(len(p)), + &written, &consumed) + w.encoded = w.buf[:int(written)] + n += int(consumed) + p = p[int(consumed):] + + if !ok { + return n, errEncode + } + } + return n, nil +} + +// Encode returns content encoded with Brotli. +func Encode(content []byte, options WriterOptions) ([]byte, error) { + var buf bytes.Buffer + writer := NewWriter(&buf, options) + _, err := writer.Write(content) + if closeErr := writer.Close(); err == nil && closeErr != nil { + err = closeErr + } + if err != nil { + return nil, err + } + return buf.Bytes(), nil +}