Add stateless dictionary support (#216)

* Add stateless dictionary support

Enable up to 8K dictionary for stateless compression.

README.md

@@ -15,6 +15,8 @@ This package provides various compression algorithms.
 
 # changelog
 
+* Feb 4, 2020: (v1.10.0) Add optional dictionary to [stateless deflate](https://pkg.go.dev/github.com/klauspost/compress/flate?tab=doc#StatelessDeflate). Breaking change, send `nil` for previous behaviour. [#216](https://github.com/klauspost/compress/pull/216)
+* Feb 3, 2020: Fix buffer overflow on repeated small block deflate.  [#218](https://github.com/klauspost/compress/pull/218)
 * Jan 31, 2020: Allow copying content from an existing ZIP file without decompressing+compressing. [#214](https://github.com/klauspost/compress/pull/214)
 * Jan 28, 2020: Added [S2](https://github.com/klauspost/compress/tree/master/s2#s2-compression) AMD64 assembler and various optimizations. Stream speed >10GB/s.  [#186](https://github.com/klauspost/compress/pull/186)
 * Jan 20,2020 (v1.9.8) Optimize gzip/deflate with better size estimates and faster table generation. [#207](https://github.com/klauspost/compress/pull/207) by [luyu6056](https://github.com/luyu6056),  [#206](https://github.com/klauspost/compress/pull/206).

flate/flate_test.go

@@ -138,6 +138,31 @@ func TestRegressions(t *testing.T) {
 				}
 			})
 		}
+		t.Run(tt.Name+"stateless", func(t *testing.T) {
+			// Split into two and use history...
+			buf := new(bytes.Buffer)
+			err = StatelessDeflate(buf, data1[:len(data1)/2], false, nil)
+			if err != nil {
+				t.Error(err)
+			}
+
+			// Use top half as dictionary...
+			dict := data1[:len(data1)/2]
+			err = StatelessDeflate(buf, data1[len(data1)/2:], true, dict)
+			if err != nil {
+				t.Error(err)
+			}
+			t.Log(buf.Len())
+			fr1 := NewReader(buf)
+			data2, err := ioutil.ReadAll(fr1)
+			if err != nil {
+				t.Error(err)
+			}
+			if bytes.Compare(data1, data2) != 0 {
+				fmt.Printf("want:%x\ngot: %x\n", data1, data2)
+				t.Error("not equal")
+			}
+		})
 	}
 }
 

flate/huffman_bit_writer.go

@@ -177,6 +177,11 @@ func (w *huffmanBitWriter) flush() {
 		w.nbits = 0
 		return
 	}
+	if w.lastHeader > 0 {
+		// We owe an EOB
+		w.writeCode(w.literalEncoding.codes[endBlockMarker])
+		w.lastHeader = 0
+	}
 	n := w.nbytes
 	for w.nbits != 0 {
 		w.bytes[n] = byte(w.bits)
@@ -594,8 +599,8 @@ func (w *huffmanBitWriter) writeBlockDynamic(tokens *tokens, eof bool, input []b
 		tokens.AddEOB()
 	}
 
-	// We cannot reuse pure huffman table.
-	if w.lastHuffMan && w.lastHeader > 0 {
+	// We cannot reuse pure huffman table, and must mark as EOF.
+	if (w.lastHuffMan || eof) && w.lastHeader > 0 {
 		// We will not try to reuse.
 		w.writeCode(w.literalEncoding.codes[endBlockMarker])
 		w.lastHeader = 0

flate/stateless.go

@@ -8,6 +8,8 @@ import (
 
 const (
 	maxStatelessBlock = math.MaxInt16
+	// dictionary will be taken from maxStatelessBlock, so limit it.
+	maxStatelessDict = 8 << 10
 
 	slTableBits  = 13
 	slTableSize  = 1 << slTableBits
@@ -25,11 +27,11 @@ func (s *statelessWriter) Close() error {
 	}
 	s.closed = true
 	// Emit EOF block
-	return StatelessDeflate(s.dst, nil, true)
+	return StatelessDeflate(s.dst, nil, true, nil)
 }
 
 func (s *statelessWriter) Write(p []byte) (n int, err error) {
-	err = StatelessDeflate(s.dst, p, false)
+	err = StatelessDeflate(s.dst, p, false, nil)
 	if err != nil {
 		return 0, err
 	}
@@ -59,7 +61,10 @@ var bitWriterPool = sync.Pool{
 
 // StatelessDeflate allows to compress directly to a Writer without retaining state.
 // When returning everything will be flushed.
-func StatelessDeflate(out io.Writer, in []byte, eof bool) error {
+// Up to 8KB of an optional dictionary can be given which is presumed to presumed to precede the block.
+// Longer dictionaries will be truncated and will still produce valid output.
+// Sending nil dictionary is perfectly fine.
+func StatelessDeflate(out io.Writer, in []byte, eof bool, dict []byte) error {
 	var dst tokens
 	bw := bitWriterPool.Get().(*huffmanBitWriter)
 	bw.reset(out)
@@ -76,35 +81,53 @@ func StatelessDeflate(out io.Writer, in []byte, eof bool) error {
 		return bw.err
 	}
 
+	// Truncate dict
+	if len(dict) > maxStatelessDict {
+		dict = dict[len(dict)-maxStatelessDict:]
+	}
+
 	for len(in) > 0 {
 		todo := in
-		if len(todo) > maxStatelessBlock {
-			todo = todo[:maxStatelessBlock]
+		if len(todo) > maxStatelessBlock-len(dict) {
+			todo = todo[:maxStatelessBlock-len(dict)]
 		}
 		in = in[len(todo):]
+		uncompressed := todo
+		if len(dict) > 0 {
+			// combine dict and source
+			bufLen := len(todo) + len(dict)
+			combined := make([]byte, bufLen)
+			copy(combined, dict)
+			copy(combined[len(dict):], todo)
+			todo = combined
+		}
 		// Compress
-		statelessEnc(&dst, todo)
+		statelessEnc(&dst, todo, int16(len(dict)))
 		isEof := eof && len(in) == 0
 
 		if dst.n == 0 {
-			bw.writeStoredHeader(len(todo), isEof)
+			bw.writeStoredHeader(len(uncompressed), isEof)
 			if bw.err != nil {
 				return bw.err
 			}
-			bw.writeBytes(todo)
-		} else if int(dst.n) > len(todo)-len(todo)>>4 {
+			bw.writeBytes(uncompressed)
+		} else if int(dst.n) > len(uncompressed)-len(uncompressed)>>4 {
 			// If we removed less than 1/16th, huffman compress the block.
-			bw.writeBlockHuff(isEof, todo, false)
+			bw.writeBlockHuff(isEof, uncompressed, len(in) == 0)
 		} else {
-			bw.writeBlockDynamic(&dst, isEof, todo, false)
+			bw.writeBlockDynamic(&dst, isEof, uncompressed, len(in) == 0)
+		}
+		if len(in) > 0 {
+			// Retain a dict if we have more
+			dict = todo[len(todo)-maxStatelessDict:]
+			dst.Reset()
 		}
 		if bw.err != nil {
 			return bw.err
 		}
-		dst.Reset()
 	}
 	if !eof {
-		// Align.
+		// Align, only a stored block can do that.
 		bw.writeStoredHeader(0, false)
 	}
 	bw.flush()
@@ -130,7 +153,7 @@ func load6416(b []byte, i int16) uint64 {
 		uint64(b[4])<<32 | uint64(b[5])<<40 | uint64(b[6])<<48 | uint64(b[7])<<56
 }
 
-func statelessEnc(dst *tokens, src []byte) {
+func statelessEnc(dst *tokens, src []byte, startAt int16) {
 	const (
 		inputMargin            = 12 - 1
 		minNonLiteralBlockSize = 1 + 1 + inputMargin
@@ -144,15 +167,23 @@ func statelessEnc(dst *tokens, src []byte) {
 
 	// This check isn't in the Snappy implementation, but there, the caller
 	// instead of the callee handles this case.
-	if len(src) < minNonLiteralBlockSize {
+	if len(src)-int(startAt) < minNonLiteralBlockSize {
 		// We do not fill the token table.
 		// This will be picked up by caller.
-		dst.n = uint16(len(src))
+		dst.n = 0
 		return
 	}
+	// Index until startAt
+	if startAt > 0 {
+		cv := load3232(src, 0)
+		for i := int16(0); i < startAt; i++ {
+			table[hashSL(cv)] = tableEntry{offset: i}
+			cv = (cv >> 8) | (uint32(src[i+4]) << 24)
+		}
+	}
 
-	s := int16(1)
-	nextEmit := int16(0)
+	s := startAt + 1
+	nextEmit := startAt
 	// sLimit is when to stop looking for offset/length copies. The inputMargin
 	// lets us use a fast path for emitLiteral in the main loop, while we are
 	// looking for copies.

flate/testdata/huffman-text.dyn.expect-noinput

@@ -1,3 +1,3 @@
 �`�J�|�ஏb���F��=M/MX�+�K������������ˊ�;��޹���`�.�&;$
 ���A	A
�:��F8T�	h� ͍�˘�P� �"PI&@��	lG p`7�Td�x���D�GA^k�, � �OA�U�!���AV�J��QV�2,��ށ���j(,;]X�`��
-��*xqF_��2>n^��A��Um�� �Œ���2>�T��
�	g�O�� ���U��+�����d��5ʕ�d��6_�i�2
+��*xqF_��2>n^��A��Um�� �Œ���2>�T��
�	g�O�� ���U��+�����d��5ʕ�d��6_�i�2�

flate/testdata/huffman-zero.dyn.expect-noinput

@@ -1 +1 @@
-��@h�m��۶m۶m۶m۶m۶��6�rk
+��@h�m��۶m۶m۶m۶m۶��6�rk�

gzip/gzip.go

@@ -207,7 +207,7 @@ func (z *Writer) Write(p []byte) (int, error) {
 	z.size += uint32(len(p))
 	z.digest = crc32.Update(z.digest, crc32.IEEETable, p)
 	if z.level == StatelessCompression {
-		return len(p), flate.StatelessDeflate(z.w, p, false)
+		return len(p), flate.StatelessDeflate(z.w, p, false, nil)
 	}
 	n, z.err = z.compressor.Write(p)
 	return n, z.err
@@ -255,7 +255,7 @@ func (z *Writer) Close() error {
 		}
 	}
 	if z.level == StatelessCompression {
-		z.err = flate.StatelessDeflate(z.w, nil, true)
+		z.err = flate.StatelessDeflate(z.w, nil, true, nil)
 	} else {
 		z.err = z.compressor.Close()
 	}