diff --git a/wim/lzx/lzx.go b/wim/lzx/lzx.go index 2e9f851..73f8336 100644 --- a/wim/lzx/lzx.go +++ b/wim/lzx/lzx.go @@ -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< 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<> (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<>(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 } diff --git a/wim/wim.go b/wim/wim.go index 50e82ef..1d02e92 100644 --- a/wim/wim.go +++ b/wim/wim.go @@ -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