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