1// Copyright 2011 The Go Authors. All rights reserved.
  2// Use of this source code is governed by a BSD-style
  3// license that can be found in the LICENSE file.
  4
  5package packet
  6
  7import (
  8	"compress/bzip2"
  9	"compress/flate"
 10	"compress/zlib"
 11	"io"
 12	"strconv"
 13
 14	"github.com/ProtonMail/go-crypto/openpgp/errors"
 15)
 16
 17// Compressed represents a compressed OpenPGP packet. The decompressed contents
 18// will contain more OpenPGP packets. See RFC 4880, section 5.6.
 19type Compressed struct {
 20	Body io.Reader
 21}
 22
 23const (
 24	NoCompression      = flate.NoCompression
 25	BestSpeed          = flate.BestSpeed
 26	BestCompression    = flate.BestCompression
 27	DefaultCompression = flate.DefaultCompression
 28)
 29
 30// CompressionConfig contains compressor configuration settings.
 31type CompressionConfig struct {
 32	// Level is the compression level to use. It must be set to
 33	// between -1 and 9, with -1 causing the compressor to use the
 34	// default compression level, 0 causing the compressor to use
 35	// no compression and 1 to 9 representing increasing (better,
 36	// slower) compression levels. If Level is less than -1 or
 37	// more then 9, a non-nil error will be returned during
 38	// encryption. See the constants above for convenient common
 39	// settings for Level.
 40	Level int
 41}
 42
 43// decompressionReader ensures that the whole compression packet is read.
 44type decompressionReader struct {
 45	compressed   io.Reader
 46	decompressed io.ReadCloser
 47	readAll      bool
 48}
 49
 50func newDecompressionReader(r io.Reader, decompressor io.ReadCloser) *decompressionReader {
 51	return &decompressionReader{
 52		compressed:   r,
 53		decompressed: decompressor,
 54	}
 55}
 56
 57func (dr *decompressionReader) Read(data []byte) (n int, err error) {
 58	if dr.readAll {
 59		return 0, io.EOF
 60	}
 61	n, err = dr.decompressed.Read(data)
 62	if err == io.EOF {
 63		dr.readAll = true
 64		// Close the decompressor.
 65		if errDec := dr.decompressed.Close(); errDec != nil {
 66			return n, errDec
 67		}
 68		// Consume all remaining data from the compressed packet.
 69		consumeAll(dr.compressed)
 70	}
 71	return n, err
 72}
 73
 74func (c *Compressed) parse(r io.Reader) error {
 75	var buf [1]byte
 76	_, err := readFull(r, buf[:])
 77	if err != nil {
 78		return err
 79	}
 80
 81	switch buf[0] {
 82	case 0:
 83		c.Body = r
 84	case 1:
 85		c.Body = newDecompressionReader(r, flate.NewReader(r))
 86	case 2:
 87		decompressor, err := zlib.NewReader(r)
 88		if err != nil {
 89			return err
 90		}
 91		c.Body = newDecompressionReader(r, decompressor)
 92	case 3:
 93		c.Body = newDecompressionReader(r, io.NopCloser(bzip2.NewReader(r)))
 94	default:
 95		err = errors.UnsupportedError("unknown compression algorithm: " + strconv.Itoa(int(buf[0])))
 96	}
 97
 98	return err
 99}
100
101// compressedWriterCloser represents the serialized compression stream
102// header and the compressor. Its Close() method ensures that both the
103// compressor and serialized stream header are closed. Its Write()
104// method writes to the compressor.
105type compressedWriteCloser struct {
106	sh io.Closer      // Stream Header
107	c  io.WriteCloser // Compressor
108}
109
110func (cwc compressedWriteCloser) Write(p []byte) (int, error) {
111	return cwc.c.Write(p)
112}
113
114func (cwc compressedWriteCloser) Close() (err error) {
115	err = cwc.c.Close()
116	if err != nil {
117		return err
118	}
119
120	return cwc.sh.Close()
121}
122
123// SerializeCompressed serializes a compressed data packet to w and
124// returns a WriteCloser to which the literal data packets themselves
125// can be written and which MUST be closed on completion. If cc is
126// nil, sensible defaults will be used to configure the compression
127// algorithm.
128func SerializeCompressed(w io.WriteCloser, algo CompressionAlgo, cc *CompressionConfig) (literaldata io.WriteCloser, err error) {
129	compressed, err := serializeStreamHeader(w, packetTypeCompressed)
130	if err != nil {
131		return
132	}
133
134	_, err = compressed.Write([]byte{uint8(algo)})
135	if err != nil {
136		return
137	}
138
139	level := DefaultCompression
140	if cc != nil {
141		level = cc.Level
142	}
143
144	var compressor io.WriteCloser
145	switch algo {
146	case CompressionZIP:
147		compressor, err = flate.NewWriter(compressed, level)
148	case CompressionZLIB:
149		compressor, err = zlib.NewWriterLevel(compressed, level)
150	default:
151		s := strconv.Itoa(int(algo))
152		err = errors.UnsupportedError("Unsupported compression algorithm: " + s)
153	}
154	if err != nil {
155		return
156	}
157
158	literaldata = compressedWriteCloser{compressed, compressor}
159
160	return
161}