Skip to content

Commit

Permalink
Fix lzxpress decompression on windows. (#5)
Browse files Browse the repository at this point in the history
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
  • Loading branch information
scudette authored Jul 22, 2020
1 parent 0469fa2 commit 37e4751
Show file tree
Hide file tree
Showing 7 changed files with 120 additions and 14 deletions.
2 changes: 1 addition & 1 deletion cmd/prefetch.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))

Expand Down
13 changes: 13 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
@@ -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
)
21 changes: 21 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
@@ -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=
42 changes: 30 additions & 12 deletions lzxpress.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 {
Expand All @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -191,7 +196,6 @@ func PrefixCodeTreeDecodeSymbol(bstr *BitStream, root *PREFIX_CODE_NODE) (
}

return node.symbol, nil

}

func LZXpressHuffmanDecompressChunk(
Expand All @@ -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)

Expand All @@ -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
Expand All @@ -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 {
Expand Down
8 changes: 8 additions & 0 deletions lzxpress_default.go
Original file line number Diff line number Diff line change
@@ -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)
}
46 changes: 46 additions & 0 deletions lzxpress_windows.go
Original file line number Diff line number Diff line change
@@ -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
}
2 changes: 1 addition & 1 deletion prefetch.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 37e4751

Please sign in to comment.