cloudflared-mirror/vendor/github.com/cloudflare/brotli-go/writer.go

170 lines
4.6 KiB
Go

// 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 brotli
/*
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <brotli/encode.h>
struct CompressStreamResult {
size_t bytes_consumed;
const uint8_t* output_data;
size_t output_data_size;
int success;
int has_more;
};
static 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 (
"bytes"
"errors"
"io"
"unsafe"
)
// 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 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)
C.BrotliEncoderSetParameter(
state, C.BROTLI_PARAM_QUALITY, (C.uint32_t)(options.Quality))
if options.LGWin > 0 {
C.BrotliEncoderSetParameter(
state, C.BROTLI_PARAM_LGWIN, (C.uint32_t)(options.LGWin))
}
return &Writer{
dst: dst,
state: state,
}
}
func (w *Writer) SetDictionary(p []byte) {
var data *C.uint8_t
if len(p) != 0 {
data = (*C.uint8_t)(&p[0])
}
C.BrotliEncoderSetCustomDictionary(w.state, C.size_t(len(p)), data)
}
func (w *Writer) writeChunk(p []byte, op C.BrotliEncoderOperation) (n int, err error) {
if w.state == nil {
return 0, errWriterClosed
}
for {
var data *C.uint8_t
if len(p) != 0 {
data = (*C.uint8_t)(&p[0])
}
result := C.CompressStream(w.state, op, data, C.size_t(len(p)))
if result.success == 0 {
return n, errEncode
}
p = p[int(result.bytes_consumed):]
n += int(result.bytes_consumed)
length := int(result.output_data_size)
if length != 0 {
// It is a workaround for non-copying-wrapping of native memory.
// C-encoder never pushes output block longer than ((2 << 25) + 502).
// TODO: use natural wrapper, when it becomes available, see
// https://golang.org/issue/13656.
output := (*[1 << 30]byte)(unsafe.Pointer(result.output_data))[:length:length]
_, err = w.dst.Write(output)
if err != nil {
return n, err
}
}
if len(p) == 0 && result.has_more == 0 {
return n, 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 {
_, err := w.writeChunk(nil, C.BROTLI_OPERATION_FLUSH)
return err
}
// Close flushes remaining data to the decorated writer and frees C resources.
func (w *Writer) Close() error {
// If stream is already closed, it is reported by `writeChunk`.
_, err := w.writeChunk(nil, C.BROTLI_OPERATION_FINISH)
// C-Brotli tolerates `nil` pointer here.
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) {
return w.writeChunk(p, C.BROTLI_OPERATION_PROCESS)
}
// 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 {
err = closeErr
}
return buf.Bytes(), err
}