Merge pull request #23 from Microsoft/wim_opts

Optimize WIM parsing
This commit is contained in:
John Starks
2016-05-17 13:16:34 -07:00
2 changed files with 255 additions and 143 deletions
+148 -87
View File
@@ -6,7 +6,6 @@
package lzx
import (
"bufio"
"bytes"
"encoding/binary"
"errors"
@@ -17,6 +16,10 @@ const (
maincodecount = 496
maincodesplit = 256
lencodecount = 249
lenshift = 9
codemask = 0x1ff
tablebits = 9
tablesize = 1 << tablebits
maxBlockSize = 32768
windowSize = 32768
@@ -58,7 +61,7 @@ type Reader interface {
}
type decompressor struct {
r Reader
r io.Reader
err error
unaligned bool
nbits byte
@@ -69,28 +72,59 @@ type decompressor struct {
mainlens [maincodecount]byte
lenlens [lencodecount]byte
window [windowSize]byte
b []byte
bv int
bo int
}
//go:noinline
func (f *decompressor) fail(err error) {
if f.err == nil {
f.err = err
}
f.bo = 0
f.bv = 0
}
func (f *decompressor) ensureAtLeast(n int) error {
if f.bv-f.bo >= n {
return nil
}
if f.err != nil {
return f.err
}
if f.bv != f.bo {
copy(f.b[:f.bv-f.bo], f.b[f.bo:f.bv])
}
n, err := io.ReadAtLeast(f.r, f.b[f.bv-f.bo:], n)
if err != nil {
if err == io.EOF {
err = io.ErrUnexpectedEOF
} else {
f.fail(err)
}
return err
}
f.bv = f.bv - f.bo + n
f.bo = 0
return nil
}
// feed retrieves another 16-bit word from the stream and consumes
// it into f.c. It returns false if there are no more bytes available.
// Otherwise, on error, it sets f.err.
func (f *decompressor) feed() bool {
if f.err != nil {
return true
}
var b0, b1 byte
b0, err := f.r.ReadByte()
if err == nil {
b1, err = f.r.ReadByte()
}
err := f.ensureAtLeast(2)
if err != nil {
if err == io.EOF {
if err == io.ErrUnexpectedEOF {
return false
}
f.err = err
}
f.c |= (uint32(b1)<<8 | uint32(b0)) << (16 - f.nbits)
f.c |= (uint32(f.b[f.bo+1])<<8 | uint32(f.b[f.bo])) << (16 - f.nbits)
f.nbits += 16
f.bo += 2
return true
}
@@ -99,7 +133,7 @@ func (f *decompressor) feed() bool {
func (f *decompressor) getBits(n byte) uint16 {
if f.nbits < n {
if !f.feed() {
f.err = io.ErrUnexpectedEOF
f.fail(io.ErrUnexpectedEOF)
}
}
c := uint16(f.c >> (32 - n))
@@ -109,9 +143,9 @@ func (f *decompressor) getBits(n byte) uint16 {
}
type huffman struct {
lens []byte
table []uint16
extra [][]uint16
maxbits byte
table [tablesize]uint16
}
// buildTable builds a huffman decoding table from a slice of code lengths,
@@ -148,54 +182,78 @@ func buildTable(codelens []byte) *huffman {
// Build a table for code lookup. For code sizes < max,
// put all possible suffixes for the code into the table, too.
// Typically a huffman implementation will only do this up to
// a small code length maximum, then fall back to a different
// mechanism; this would probably improve performance.
table := make([]uint16, 1<<max)
for i, cl := range codelens {
if cl != 0 {
code := first[cl]
extendedCode := code << (max - cl)
for j := uint(0); j < 1<<(max-cl); j++ {
table[extendedCode+j] = uint16(i)
}
first[cl]++
// For max > tablebits, split long codes into additional tables
// of suffixes of max-tablebits length.
h := &huffman{maxbits: max}
if max > tablebits {
core := first[tablebits+1] / 2 // Number of codes that fit without extra tables
nextra := 1<<tablebits - core // Number of extra entries
h.extra = make([][]uint16, nextra)
for code := core; code < 1<<tablebits; code++ {
h.table[code] = uint16(code - core)
h.extra[code-core] = make([]uint16, 1<<(max-tablebits))
}
}
return &huffman{
lens: codelens,
table: table,
maxbits: max,
for i, cl := range codelens {
if cl != 0 {
code := first[cl]
first[cl]++
v := uint16(cl)<<lenshift | uint16(i)
if cl <= tablebits {
extendedCode := code << (tablebits - cl)
for j := uint(0); j < 1<<(tablebits-cl); j++ {
h.table[extendedCode+j] = v
}
} else {
prefix := code >> (cl - tablebits)
suffix := code & (1<<(cl-tablebits) - 1)
extendedCode := suffix << (max - cl)
for j := uint(0); j < 1<<(max-cl); j++ {
h.extra[h.table[prefix]][extendedCode+j] = v
}
}
}
}
return h
}
// getCode retrieves the next code using the provided
// huffman tree. It sets f.err on error.
func (f *decompressor) getCode(h *huffman) uint16 {
if h.maxbits == 0 {
// This is an empty tree. It should not be used.
f.err = errCorrupt
if h.maxbits > 0 {
if f.nbits < maxTreePathLen {
f.feed()
}
// For codes with length < tablebits, it doesn't matter
// what the remainder of the bits used for table lookup
// are, since entries with all possible suffixes were
// added to the table.
c := h.table[f.c>>(32-tablebits)]
if c >= 1<<lenshift {
// The code is already in c.
} else {
c = h.extra[c][f.c<<tablebits>>(32-(h.maxbits-tablebits))]
}
n := byte(c >> lenshift)
if f.nbits >= n {
// Only consume the length of the code, not the maximum
// code length.
f.c <<= n
f.nbits -= n
return c & codemask
}
f.fail(io.ErrUnexpectedEOF)
return 0
}
if f.nbits < maxTreePathLen {
f.feed()
}
// For codes with length < h.maxbits, it doesn't matter
// what the remainder of the bits used for table lookup
// are, since entries with all possible suffixes were
// added to the table.
c := h.table[f.c>>(32-h.maxbits)]
n := h.lens[c]
if f.nbits < n {
f.err = io.ErrUnexpectedEOF
return 0
}
// Only consume the length of the code, not the maximum
// code length.
f.c <<= n
f.nbits -= n
return c
// This is an empty tree. It should not be used.
f.fail(errCorrupt)
return 0
}
// mod17 computes the value mod 17.
@@ -279,13 +337,11 @@ func (f *decompressor) readBlockHeader() (byte, uint16, error) {
// If the previous block was an unaligned uncompressed block, restore
// 2-byte alignment.
if f.unaligned {
_, err := f.r.ReadByte()
err := f.ensureAtLeast(1)
if err != nil {
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
return 0, 0, err
}
f.bo++
f.unaligned = false
}
@@ -321,19 +377,17 @@ func (f *decompressor) readBlockHeader() (byte, uint16, error) {
}
f.getBits(n)
if f.err != nil {
return 0, 0, f.err
}
// Read the LRU values for the next block.
var lru [12]byte
_, err := io.ReadFull(f.r, lru[:])
err := f.ensureAtLeast(12)
if err != nil {
return 0, 0, err
}
f.lru[0] = uint16(binary.LittleEndian.Uint32(lru[0:4]))
f.lru[1] = uint16(binary.LittleEndian.Uint32(lru[4:8]))
f.lru[2] = uint16(binary.LittleEndian.Uint32(lru[8:12]))
f.lru[0] = uint16(binary.LittleEndian.Uint32(f.b[f.bo : f.bo+4]))
f.lru[1] = uint16(binary.LittleEndian.Uint32(f.b[f.bo+4 : f.bo+8]))
f.lru[2] = uint16(binary.LittleEndian.Uint32(f.b[f.bo+8 : f.bo+12]))
f.bo += 12
default:
return 0, 0, errCorrupt
@@ -393,10 +447,11 @@ func (f *decompressor) readTrees(readAligned bool) (main *huffman, length *huffm
// readCompressedBlock decodes a compressed block, writing into the window
// starting at start and ending at end, and using the provided huffman trees.
func (f *decompressor) readCompressedBlock(start, end uint16, hmain, hlength, haligned *huffman) (int, error) {
for i := start; i < end; {
i := start
for i < end {
main := f.getCode(hmain)
if f.err != nil {
return int(i - start), f.err
break
}
if main < 256 {
// Literal byte.
@@ -407,16 +462,13 @@ func (f *decompressor) readCompressedBlock(start, end uint16, hmain, hlength, ha
// This is a match backward in the window. Determine
// the offset and dlength.
lenheader := (main - 256) % 8
matchlen := (main - 256) % 8
slot := (main - 256) / 8
// The length is either the low bits of the code,
// or if this is 7, is encoded with the length tree.
var matchlen uint16
if lenheader == 7 {
matchlen = f.getCode(hlength) + 7
} else {
matchlen = lenheader
if matchlen == 7 {
matchlen += f.getCode(hlength)
}
matchlen += 2
@@ -452,16 +504,17 @@ func (f *decompressor) readCompressedBlock(start, end uint16, hmain, hlength, ha
f.lru[0] = matchoffset
}
if matchoffset > i || matchlen > end-i {
return int(i - start), errCorrupt
if matchoffset <= i && matchlen <= end-i {
copyend := i + matchlen
for ; i < copyend; i++ {
f.window[i] = f.window[i-matchoffset]
}
} else {
f.fail(errCorrupt)
break
}
for j := uint16(0); j < matchlen; j++ {
f.window[i+j] = f.window[i+j-matchoffset]
}
i += matchlen
}
return int(end - start), nil
return int(i - start), f.err
}
// readBlock decodes the current block and returns the number of uncompressed bytes.
@@ -476,7 +529,18 @@ func (f *decompressor) readBlock(start uint16) (int, error) {
// Remember to realign the byte stream at the next block.
f.unaligned = true
}
return io.ReadFull(f.r, f.window[start:start+size])
copied := 0
if f.bo < f.bv {
copied = int(size)
s := int(start)
if copied > f.bv-f.bo {
copied = f.bv - f.bo
}
copy(f.window[s:s+copied], f.b[f.bo:f.bo+copied])
f.bo += copied
}
n, err := io.ReadFull(f.r, f.window[start+uint16(copied):start+size])
return copied + n, err
}
hmain, hlength, haligned, err := f.readTrees(blockType == alignedOffsetBlock)
@@ -543,11 +607,8 @@ func NewReader(r io.Reader, uncompressedSize int) (io.ReadCloser, error) {
f := &decompressor{
lru: [3]uint16{1, 1, 1},
uncompressed: uncompressedSize,
}
if br, ok := r.(Reader); ok {
f.r = br
} else {
f.r = bufio.NewReader(r)
b: make([]byte, 4096),
r: r,
}
return f, nil
}
+107 -56
View File
@@ -5,7 +5,6 @@
package wim
import (
"bufio"
"bytes"
"crypto/sha1"
"encoding/binary"
@@ -15,6 +14,7 @@ import (
"io"
"io/ioutil"
"strconv"
"sync"
"time"
"unicode/utf16"
)
@@ -164,7 +164,6 @@ type securityblockDisk struct {
const securityblockDiskSize = 8
type direntry struct {
Length int64
Attributes uint32
SecurityID uint32
SubdirOffset int64
@@ -180,16 +179,15 @@ type direntry struct {
FileNameLength uint16
}
const direntrySize = 102
var direntrySize = int64(binary.Size(direntry{}) + 8) // includes an 8-byte length prefix
type streamentry struct {
Length int64
Unused int64
Hash SHA1Hash
NameLength int16
}
const streamentrySize = 38
var streamentrySize = int64(binary.Size(streamentry{}) + 8) // includes an 8-byte length prefix
// Filetime represents a Windows time.
type Filetime struct {
@@ -299,6 +297,9 @@ type Image struct {
offset resourceDescriptor
sds [][]byte
rootOffset int64
r io.ReadCloser
curOffset int64
m sync.Mutex
ImageInfo
}
@@ -398,6 +399,14 @@ func NewReader(f io.ReaderAt) (*Reader, error) {
return r, nil
}
// Close releases resources associated with the Reader.
func (r *Reader) Close() error {
for _, img := range r.Image {
img.reset()
}
return nil
}
func (r *Reader) resourceReader(hdr *resourceDescriptor) (io.ReadCloser, error) {
return r.resourceReaderWithOffset(hdr, 0)
}
@@ -559,22 +568,23 @@ func (r *Reader) readSecurityDescriptors(rsrc io.Reader) (sds [][]byte, n int64,
// Open parses the image and returns the root directory.
func (img *Image) Open() (*File, error) {
rsrc, err := img.wim.resourceReaderWithOffset(&img.offset, img.rootOffset)
if err != nil {
return nil, err
}
defer rsrc.Close()
if img.sds == nil {
sds, n, err := img.wim.readSecurityDescriptors(rsrc)
rsrc, err := img.wim.resourceReaderWithOffset(&img.offset, img.rootOffset)
if err != nil {
return nil, err
}
sds, n, err := img.wim.readSecurityDescriptors(rsrc)
if err != nil {
rsrc.Close()
return nil, err
}
img.sds = sds
img.r = rsrc
img.rootOffset = n
img.curOffset = n
}
f, err := img.readdir(rsrc)
f, err := img.readdir(img.rootOffset)
if err != nil {
return nil, err
}
@@ -584,16 +594,50 @@ func (img *Image) Open() (*File, error) {
return f[0], err
}
func (img *Image) readdir(rsrc io.Reader) ([]*File, error) {
r := bufio.NewReader(rsrc)
func (img *Image) reset() {
if img.r != nil {
img.r.Close()
img.r = nil
}
img.curOffset = -1
}
func (img *Image) readdir(offset int64) ([]*File, error) {
img.m.Lock()
defer img.m.Unlock()
if offset < img.curOffset || offset > img.curOffset+chunkSize {
// Reset to seek backward or to seek forward very far.
img.reset()
}
if img.r == nil {
rsrc, err := img.wim.resourceReaderWithOffset(&img.offset, offset)
if err != nil {
return nil, err
}
img.r = rsrc
img.curOffset = offset
}
if offset > img.curOffset {
_, err := io.CopyN(ioutil.Discard, img.r, offset-img.curOffset)
if err != nil {
img.reset()
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
return nil, err
}
}
var entries []*File
for {
e, err := img.readNextEntry(r)
e, n, err := img.readNextEntry(img.r)
img.curOffset += n
if err == io.EOF {
break
}
if err != nil {
img.reset()
return nil, err
}
entries = append(entries, e)
@@ -601,38 +645,39 @@ func (img *Image) readdir(rsrc io.Reader) ([]*File, error) {
return entries, nil
}
func (img *Image) readNextEntry(r *bufio.Reader) (*File, error) {
lengthBuf, err := r.Peek(8)
func (img *Image) readNextEntry(r io.Reader) (*File, int64, error) {
var length int64
err := binary.Read(r, binary.LittleEndian, &length)
if err != nil {
return nil, &ParseError{Oper: "directory length check", Err: err}
return nil, 0, &ParseError{Oper: "directory length check", Err: err}
}
left := int(binary.LittleEndian.Uint64(lengthBuf))
if left == 0 {
return nil, io.EOF
if length == 0 {
return nil, 8, io.EOF
}
left := length
if left < direntrySize {
return nil, &ParseError{Oper: "directory entry", Err: errors.New("size too short")}
return nil, 0, &ParseError{Oper: "directory entry", Err: errors.New("size too short")}
}
var dentry direntry
err = binary.Read(r, binary.LittleEndian, &dentry)
if err != nil {
return nil, &ParseError{Oper: "directory entry", Err: err}
return nil, 0, &ParseError{Oper: "directory entry", Err: err}
}
left -= direntrySize
namesLen := int(dentry.FileNameLength + 2 + dentry.ShortNameLength)
namesLen := int64(dentry.FileNameLength + 2 + dentry.ShortNameLength)
if left < namesLen {
return nil, &ParseError{Oper: "directory entry", Err: errors.New("size too short for names")}
return nil, 0, &ParseError{Oper: "directory entry", Err: errors.New("size too short for names")}
}
names := make([]uint16, namesLen/2)
err = binary.Read(r, binary.LittleEndian, names)
if err != nil {
return nil, &ParseError{Oper: "file name", Err: err}
return nil, 0, &ParseError{Oper: "file name", Err: err}
}
left -= namesLen
@@ -652,7 +697,7 @@ func (img *Image) readNextEntry(r *bufio.Reader) (*File, error) {
var ok bool
offset, ok = img.wim.fileData[dentry.Hash]
if !ok {
return nil, &ParseError{Oper: "directory entry", Path: name, Err: fmt.Errorf("could not find file data matching hash %#v", dentry)}
return nil, 0, &ParseError{Oper: "directory entry", Path: name, Err: fmt.Errorf("could not find file data matching hash %#v", dentry)}
}
}
@@ -686,26 +731,30 @@ func (img *Image) readNextEntry(r *bufio.Reader) (*File, error) {
}
if isDir && f.subdirOffset == 0 {
return nil, &ParseError{Oper: "directory entry", Path: name, Err: errors.New("no subdirectory data for directory")}
return nil, 0, &ParseError{Oper: "directory entry", Path: name, Err: errors.New("no subdirectory data for directory")}
} else if !isDir && f.subdirOffset != 0 {
return nil, &ParseError{Oper: "directory entry", Path: name, Err: errors.New("unexpected subdirectory data for non-directory")}
return nil, 0, &ParseError{Oper: "directory entry", Path: name, Err: errors.New("unexpected subdirectory data for non-directory")}
}
if dentry.SecurityID != 0xffffffff {
f.SecurityDescriptor = img.sds[dentry.SecurityID]
}
_, err = r.Discard(left)
_, err = io.CopyN(ioutil.Discard, r, left)
if err != nil {
return nil, err
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
return nil, 0, err
}
if dentry.StreamCount > 0 {
var streams []*Stream
for i := uint16(0); i < dentry.StreamCount; i++ {
s, err := img.readNextStream(r)
s, n, err := img.readNextStream(r)
length += n
if err != nil {
return nil, err
return nil, 0, err
}
// The first unnamed stream should be treated as the file stream.
if i == 0 && s.Name == "" {
@@ -720,42 +769,46 @@ func (img *Image) readNextEntry(r *bufio.Reader) (*File, error) {
}
if dentry.Attributes&FILE_ATTRIBUTE_REPARSE_POINT != 0 && f.Size == 0 {
return nil, &ParseError{Oper: "directory entry", Path: name, Err: errors.New("reparse point is missing reparse stream")}
return nil, 0, &ParseError{Oper: "directory entry", Path: name, Err: errors.New("reparse point is missing reparse stream")}
}
return f, nil
return f, length, nil
}
func (img *Image) readNextStream(r *bufio.Reader) (*Stream, error) {
lengthBuf, err := r.Peek(8)
func (img *Image) readNextStream(r io.Reader) (*Stream, int64, error) {
var length int64
err := binary.Read(r, binary.LittleEndian, &length)
if err != nil {
return nil, &ParseError{Oper: "stream length check", Err: err}
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
return nil, 0, &ParseError{Oper: "stream length check", Err: err}
}
left := int(binary.LittleEndian.Uint64(lengthBuf))
left := length
if left < streamentrySize {
return nil, &ParseError{Oper: "stream entry", Err: errors.New("size too short")}
return nil, 0, &ParseError{Oper: "stream entry", Err: errors.New("size too short")}
}
var sentry streamentry
err = binary.Read(r, binary.LittleEndian, &sentry)
if err != nil {
return nil, &ParseError{Oper: "stream entry", Err: err}
return nil, 0, &ParseError{Oper: "stream entry", Err: err}
}
left -= streamentrySize
if left < int(sentry.NameLength) {
return nil, &ParseError{Oper: "stream entry", Err: errors.New("size too short for name")}
if left < int64(sentry.NameLength) {
return nil, 0, &ParseError{Oper: "stream entry", Err: errors.New("size too short for name")}
}
names := make([]uint16, sentry.NameLength/2)
err = binary.Read(r, binary.LittleEndian, names)
if err != nil {
return nil, &ParseError{Oper: "file name", Err: err}
return nil, 0, &ParseError{Oper: "file name", Err: err}
}
left -= int(sentry.NameLength)
left -= int64(sentry.NameLength)
name := string(utf16.Decode(names))
var offset resourceDescriptor
@@ -763,7 +816,7 @@ func (img *Image) readNextStream(r *bufio.Reader) (*Stream, error) {
var ok bool
offset, ok = img.wim.fileData[sentry.Hash]
if !ok {
return nil, &ParseError{Oper: "stream entry", Path: name, Err: fmt.Errorf("could not find file data matching hash %v", sentry.Hash)}
return nil, 0, &ParseError{Oper: "stream entry", Path: name, Err: fmt.Errorf("could not find file data matching hash %v", sentry.Hash)}
}
}
@@ -777,12 +830,15 @@ func (img *Image) readNextStream(r *bufio.Reader) (*Stream, error) {
offset: offset,
}
_, err = r.Discard(left)
_, err = io.CopyN(ioutil.Discard, r, left)
if err != nil {
return nil, err
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
return nil, 0, err
}
return s, nil
return s, length, nil
}
// Open returns an io.ReadCloser that can be used to read the stream's contents.
@@ -800,12 +856,7 @@ func (f *File) Readdir() ([]*File, error) {
if !f.IsDir() {
return nil, errors.New("not a directory")
}
rsrc, err := f.img.wim.resourceReaderWithOffset(&f.img.offset, f.subdirOffset)
if err != nil {
return nil, err
}
defer rsrc.Close()
return f.img.readdir(rsrc)
return f.img.readdir(f.subdirOffset)
}
// IsDir returns whether the given file is a directory. It returns false when it