Fix lzxpress decompression on windows. (#5)

There seems to be some problem decompression some files. This change
falls back to manual decompression on if the API is not available (on
older windows or non-windows system).

On windows we call the windows API to decompress the prefetch file in
a more reliable way.

Fixes: #4

Author: scudette

cmd/prefetch.go

@@ -51,7 +51,7 @@ func main() {
 					binparsergen.FatalIfError(err, fmt.Sprintf("Open file: %v", err))
 				}
 
-				decompressed, err := prefetch.LZXpressHuffmanDecompress(
+				decompressed, err := prefetch.LZXpressHuffmanDecompressWithFallback(
 					data[:n], int(header.UncompressedSize()))
 				binparsergen.FatalIfError(err, fmt.Sprintf("Open file: %v", err))
 

go.mod

@@ -0,0 +1,13 @@
+module www.velocidex.com/golang/go-prefetch
+
+go 1.14
+
+require (
+	github.com/davecgh/go-spew v1.1.1
+	github.com/jawher/mow.cli v1.1.0
+	github.com/sebdah/goldie v1.0.0
+	github.com/stretchr/testify v1.3.0
+	golang.org/x/sys v0.0.0-20200720211630-cb9d2d5c5666
+	gopkg.in/yaml.v2 v2.2.2 // indirect
+	www.velocidex.com/golang/binparsergen v0.1.0
+)

go.sum

@@ -0,0 +1,21 @@
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/jawher/mow.cli v1.1.0 h1:NdtHXRc0CwZQ507wMvQ/IS+Q3W3x2fycn973/b8Zuk8=
+github.com/jawher/mow.cli v1.1.0/go.mod h1:aNaQlc7ozF3vw6IJ2dHjp2ZFiA4ozMIYY6PyuRJwlUg=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/sebdah/goldie v1.0.0 h1:9GNhIat69MSlz/ndaBg48vl9dF5fI+NBB6kfOxgfkMc=
+github.com/sebdah/goldie v1.0.0/go.mod h1:jXP4hmWywNEwZzhMuv2ccnqTSFpuq8iyQhtQdkkZBH4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.2.0/go.mod h1:qt09Ya8vawLte6SNmTgCsAVtYtaKzEcn8ATUoHMkEqE=
+github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+golang.org/x/sys v0.0.0-20200720211630-cb9d2d5c5666 h1:gVCS+QOncANNPlmlO1AhlU3oxs4V9z+gTtPwIk3p2N8=
+golang.org/x/sys v0.0.0-20200720211630-cb9d2d5c5666/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
+gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+www.velocidex.com/golang/binparsergen v0.1.0 h1:oNsMHGnlb4jrGwxKxqqmsics6FgYin3HR5UNtLXc8S0=
+www.velocidex.com/golang/binparsergen v0.1.0/go.mod h1:UC43Ecj0mjsidlClTYZ3H4dXdyv7CVI0HsYi4yY3qtc=

lzxpress.go

@@ -97,9 +97,9 @@ Returns the index in treeNodes of the next node to be processed.
 */
 func PrefixCodeTreeAddLeaf(
 	treeNodes []PREFIX_CODE_NODE,
-	leafIndex int,
+	leafIndex uint32,
 	mask uint32,
-	bits uint32) int {
+	bits uint32) uint32 {
 
 	node := &treeNodes[0]
 	i := leafIndex + 1
@@ -126,13 +126,17 @@ func PrefixCodeTreeRebuild(input []byte) *PREFIX_CODE_NODE {
 	symbolInfo := make([]PREFIX_CODE_SYMBOL, 512)
 
 	for i := 0; i < 256; i++ {
+		value := input[i]
+
 		symbolInfo[2*i].id = uint32(2 * i)
 		symbolInfo[2*i].symbol = uint32(2 * i)
-		symbolInfo[2*i].length = uint32(input[i] & 15)
+		symbolInfo[2*i].length = uint32(value & 0xf)
+
+		value >>= 4
 
 		symbolInfo[2*i+1].id = uint32(2*i + 1)
 		symbolInfo[2*i+1].symbol = uint32(2*i + 1)
-		symbolInfo[2*i+1].length = uint32(input[i] >> 4)
+		symbolInfo[2*i+1].length = uint32(value & 0xf)
 	}
 
 	sort.SliceStable(symbolInfo, func(i, j int) bool {
@@ -150,7 +154,8 @@ func PrefixCodeTreeRebuild(input []byte) *PREFIX_CODE_NODE {
 	})
 
 	i := 0
-	for ; i < 512 && symbolInfo[i].length == 0; i++ {
+	for i < 512 && symbolInfo[i].length == 0 {
+		i++
 	}
 
 	mask := uint32(0)
@@ -159,7 +164,7 @@ func PrefixCodeTreeRebuild(input []byte) *PREFIX_CODE_NODE {
 	root := &treeNodes[0]
 	root.leaf = false
 
-	j := 1
+	j := uint32(1)
 	for ; i < 512; i++ {
 		treeNodes[j].id = uint32(j)
 		treeNodes[j].symbol = symbolInfo[i].symbol
@@ -191,7 +196,6 @@ func PrefixCodeTreeDecodeSymbol(bstr *BitStream, root *PREFIX_CODE_NODE) (
 	}
 
 	return node.symbol, nil
-
 }
 
 func LZXpressHuffmanDecompressChunk(
@@ -201,6 +205,12 @@ func LZXpressHuffmanDecompressChunk(
 	output []byte, // The output buffer.
 	chunk_size int, // The required size of uncompressed buffer
 ) (int, int, error) {
+
+	// There must be at least this many bytes available to read.
+	if in_idx+256 > len(input) {
+		return 0, 0, io.EOF
+	}
+
 	root := PrefixCodeTreeRebuild(input[in_idx:])
 	bstr := NewBitStream(input, in_idx+256)
 
@@ -220,10 +230,15 @@ func LZXpressHuffmanDecompressChunk(
 
 		} else {
 			symbol -= 256
-			length := symbol & 15
+			length := uint32(symbol & 15)
 			symbol >>= 4
 
-			offset := (1 << symbol) + bstr.Lookup(symbol)
+			offset := int32(0)
+			if symbol != 0 {
+				offset = int32(bstr.Lookup(symbol))
+			}
+			offset |= 1 << symbol
+			offset = -offset
 
 			if length == 15 {
 				length = uint32(bstr.source[bstr.index]) + 15
@@ -236,16 +251,19 @@ func LZXpressHuffmanDecompressChunk(
 				}
 			}
 
-			bstr.Skip(symbol)
+			err := bstr.Skip(symbol)
+			if err != nil {
+				return int(bstr.index), i, err
+			}
 
 			length = length + 3
 			for {
-				if i-int(offset) < 0 {
+				if i+int(offset) < 0 {
 					return int(bstr.index),
 						i, errors.New("Decompression error")
 				}
 
-				output[i] = output[i-int(offset)]
+				output[i] = output[i+int(offset)]
 				i++
 				length -= 1
 				if length == 0 {

lzxpress_default.go

@@ -0,0 +1,8 @@
+// +build !windows
+
+package prefetch
+
+// Non windows systems fall back to build in decompression.
+func LZXpressHuffmanDecompressWithFallback(input []byte, output_size int) ([]byte, error) {
+	return LZXpressHuffmanDecompress(input, output_size)
+}

lzxpress_windows.go

@@ -0,0 +1,46 @@
+// +build windows
+
+package prefetch
+
+import (
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+)
+
+var (
+	ntoskrnl              = windows.NewLazySystemDLL("ntdll.dll")
+	RtlDecompressBufferEx = ntoskrnl.NewProc("RtlDecompressBufferEx")
+
+	COMPRESSION_FORMAT_XPRESS_HUFF = uint16(0x0004)
+)
+
+func LZXpressHuffmanDecompressWithFallback(input []byte, output_size int) ([]byte, error) {
+
+	// For older windows, we fall back to the build in decompression.
+	err := RtlDecompressBufferEx.Find()
+	if err != nil {
+		return LZXpressHuffmanDecompress(input, output_size)
+	}
+
+	result := make([]byte, output_size)
+	final_size := uint32(output_size)
+
+	workspace := make([]byte, output_size*2)
+
+	ret, _, _ := RtlDecompressBufferEx.Call(
+		uintptr(COMPRESSION_FORMAT_XPRESS_HUFF),
+		uintptr(unsafe.Pointer(&result[0])),
+		uintptr(output_size),
+		uintptr(unsafe.Pointer(&input[0])),
+		uintptr(len(input)),
+		uintptr(unsafe.Pointer(&final_size)),
+		uintptr(unsafe.Pointer(&workspace[0])))
+	if ret != 0 {
+		// Fallback to the built in implementation on error.
+		return LZXpressHuffmanDecompress(input, output_size)
+	}
+
+	// Return the decompressed buffer from the API.
+	return result[:final_size], nil
+}

prefetch.go

@@ -23,7 +23,7 @@ func LoadPrefetch(reader io.ReaderAt) (*PrefetchInfo, error) {
 			return nil, err
 		}
 
-		decompressed, err := LZXpressHuffmanDecompress(
+		decompressed, err := LZXpressHuffmanDecompressWithFallback(
 			data[:n], int(header.UncompressedSize()))
 		if err != nil {
 			return nil, err