new package: compression (ported from docker/pkg/archive)

Signed-off-by: Akihiro Suda <suda.akihiro@lab.ntt.co.jp>

archive/compression/compression.go

@@ -0,0 +1,125 @@
+package compression
+
+import (
+	"bufio"
+	"bytes"
+	"compress/gzip"
+	"fmt"
+	"io"
+	"sync"
+)
+
+type (
+	// Compression is the state represents if compressed or not.
+	Compression int
+)
+
+const (
+	// Uncompressed represents the uncompressed.
+	Uncompressed Compression = iota
+	// Gzip is gzip compression algorithm.
+	Gzip
+)
+
+var (
+	bufioReader32KPool = &sync.Pool{
+		New: func() interface{} { return bufio.NewReaderSize(nil, 32*1024) },
+	}
+)
+
+type readCloserWrapper struct {
+	io.Reader
+	closer func() error
+}
+
+func (r *readCloserWrapper) Close() error {
+	if r.closer != nil {
+		return r.closer()
+	}
+	return nil
+}
+
+type writeCloserWrapper struct {
+	io.Writer
+	closer func() error
+}
+
+func (w *writeCloserWrapper) Close() error {
+	if w.closer != nil {
+		w.closer()
+	}
+	return nil
+}
+
+// DetectCompression detects the compression algorithm of the source.
+func DetectCompression(source []byte) Compression {
+	for compression, m := range map[Compression][]byte{
+		Gzip: {0x1F, 0x8B, 0x08},
+	} {
+		if len(source) < len(m) {
+			// Len too short
+			continue
+		}
+		if bytes.Compare(m, source[:len(m)]) == 0 {
+			return compression
+		}
+	}
+	return Uncompressed
+}
+
+// DecompressStream decompresses the archive and returns a ReaderCloser with the decompressed archive.
+func DecompressStream(archive io.Reader) (io.ReadCloser, error) {
+	buf := bufioReader32KPool.Get().(*bufio.Reader)
+	buf.Reset(archive)
+	bs, err := buf.Peek(10)
+	if err != nil && err != io.EOF {
+		// Note: we'll ignore any io.EOF error because there are some odd
+		// cases where the layer.tar file will be empty (zero bytes) and
+		// that results in an io.EOF from the Peek() call. So, in those
+		// cases we'll just treat it as a non-compressed stream and
+		// that means just create an empty layer.
+		// See Issue docker/docker#18170
+		return nil, err
+	}
+
+	closer := func() error {
+		buf.Reset(nil)
+		bufioReader32KPool.Put(buf)
+		return nil
+	}
+	switch compression := DetectCompression(bs); compression {
+	case Uncompressed:
+		readBufWrapper := &readCloserWrapper{buf, closer}
+		return readBufWrapper, nil
+	case Gzip:
+		gzReader, err := gzip.NewReader(buf)
+		if err != nil {
+			return nil, err
+		}
+		readBufWrapper := &readCloserWrapper{gzReader, closer}
+		return readBufWrapper, nil
+	default:
+		return nil, fmt.Errorf("unsupported compression format %s", (&compression).Extension())
+	}
+}
+
+// CompressStream compresseses the dest with specified compression algorithm.
+func CompressStream(dest io.Writer, compression Compression) (io.WriteCloser, error) {
+	switch compression {
+	case Uncompressed:
+		return &writeCloserWrapper{dest, nil}, nil
+	case Gzip:
+		return gzip.NewWriter(dest), nil
+	default:
+		return nil, fmt.Errorf("unsupported compression format %s", (&compression).Extension())
+	}
+}
+
+// Extension returns the extension of a file that uses the specified compression algorithm.
+func (compression *Compression) Extension() string {
+	switch *compression {
+	case Gzip:
+		return "gz"
+	}
+	return ""
+}

archive/compression/compression_test.go

@@ -0,0 +1,67 @@
+package compression
+
+import (
+	"bytes"
+	"crypto/rand"
+	"io/ioutil"
+	"testing"
+)
+
+// generateData generates data that composed of 2 random parts
+// and single zero-filled part within them.
+// Typically, the compression ratio would be about 67%.
+func generateData(t *testing.T, size int) []byte {
+	part0 := size / 3             // random
+	part2 := size / 3             // random
+	part1 := size - part0 - part2 // zero-filled
+	part0Data := make([]byte, part0)
+	if _, err := rand.Read(part0Data); err != nil {
+		t.Fatal(err)
+	}
+	part1Data := make([]byte, part1)
+	part2Data := make([]byte, part2)
+	if _, err := rand.Read(part2Data); err != nil {
+		t.Fatal(err)
+	}
+	return append(part0Data, append(part1Data, part2Data...)...)
+}
+
+func testCompressDecompress(t *testing.T, size int, compression Compression) {
+	orig := generateData(t, size)
+	var b bytes.Buffer
+	compressor, err := CompressStream(&b, compression)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n, err := compressor.Write(orig); err != nil || n != size {
+		t.Fatal(err)
+	}
+	compressor.Close()
+	compressed := b.Bytes()
+	t.Logf("compressed %d bytes to %d bytes (%.2f%%)",
+		len(orig), len(compressed), 100.0*float32(len(compressed))/float32(len(orig)))
+	if compared := bytes.Compare(orig, compressed); (compression == Uncompressed && compared != 0) ||
+		(compression != Uncompressed && compared == 0) {
+		t.Fatal("strange compressed data")
+	}
+
+	decompressor, err := DecompressStream(bytes.NewReader(compressed))
+	if err != nil {
+		t.Fatal(err)
+	}
+	decompressed, err := ioutil.ReadAll(decompressor)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if bytes.Compare(orig, decompressed) != 0 {
+		t.Fatal("strange decompressed data")
+	}
+}
+
+func TestCompressDecompressGzip(t *testing.T) {
+	testCompressDecompress(t, 1024*1024, Gzip)
+}
+
+func TestCompressDecompressUncompressed(t *testing.T) {
+	testCompressDecompress(t, 1024*1024, Uncompressed)
+}