From 0e1e9fec2613da97ad1cae9fbac1208dbfb3352c Mon Sep 17 00:00:00 2001 From: John Starks Date: Mon, 4 Apr 2016 13:24:28 -0700 Subject: [PATCH] Add WIM parser This new code supports parsing WIM files. Currently, only WIM files that were created with no compression or LZX compression are supported. Xpress compression is not. --- reparse.go | 10 +- wim/decompress.go | 138 ++++++++++ wim/lzx/lzx.go | 546 +++++++++++++++++++++++++++++++++++++++ wim/wim.go | 646 ++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 1337 insertions(+), 3 deletions(-) create mode 100644 wim/decompress.go create mode 100644 wim/lzx/lzx.go create mode 100644 wim/wim.go diff --git a/reparse.go b/reparse.go index 96d7b9a..9425711 100644 --- a/reparse.go +++ b/reparse.go @@ -43,8 +43,12 @@ func (e *UnsupportedReparsePointError) Error() string { // DecodeReparsePoint decodes a Win32 REPARSE_DATA_BUFFER structure containing either a symlink // or a mount point. func DecodeReparsePoint(b []byte) (*ReparsePoint, error) { - isMountPoint := false tag := binary.LittleEndian.Uint32(b[0:4]) + return DecodeReparsePointData(tag, b[8:]) +} + +func DecodeReparsePointData(tag uint32, b []byte) (*ReparsePoint, error) { + isMountPoint := false switch tag { case reparseTagMountPoint: isMountPoint = true @@ -52,11 +56,11 @@ func DecodeReparsePoint(b []byte) (*ReparsePoint, error) { default: return nil, &UnsupportedReparsePointError{tag} } - nameOffset := 16 + binary.LittleEndian.Uint16(b[12:14]) + nameOffset := 8 + binary.LittleEndian.Uint16(b[4:6]) if !isMountPoint { nameOffset += 4 } - nameLength := binary.LittleEndian.Uint16(b[14:16]) + nameLength := binary.LittleEndian.Uint16(b[6:8]) name := make([]uint16, nameLength/2) err := binary.Read(bytes.NewReader(b[nameOffset:nameOffset+nameLength]), binary.LittleEndian, &name) if err != nil { diff --git a/wim/decompress.go b/wim/decompress.go new file mode 100644 index 0000000..f4e67f8 --- /dev/null +++ b/wim/decompress.go @@ -0,0 +1,138 @@ +package wim + +import ( + "encoding/binary" + "io" + "io/ioutil" + + "github.com/Microsoft/go-winio/wim/lzx" +) + +const chunkSize = 32768 // Compressed resource chunk size + +type compressedReader struct { + r *io.SectionReader + d io.ReadCloser + chunks []int64 + curChunk int + originalSize int64 +} + +func newCompressedReader(r *io.SectionReader, originalSize int64, offset int64) (*compressedReader, error) { + nchunks := (originalSize + chunkSize - 1) / chunkSize + var base int64 + chunks := make([]int64, nchunks) + if originalSize <= 0xffffffff { + // 32-bit chunk offsets + base = (nchunks - 1) * 4 + chunks32 := make([]uint32, nchunks-1) + err := binary.Read(r, binary.LittleEndian, chunks32) + if err != nil { + return nil, err + } + for i, n := range chunks32 { + chunks[i+1] = int64(n) + } + + } else { + // 64-bit chunk offsets + base = (nchunks - 1) * 8 + err := binary.Read(r, binary.LittleEndian, chunks[1:]) + if err != nil { + return nil, err + } + } + + for i, c := range chunks { + chunks[i] = c + base + } + + cr := &compressedReader{ + r: r, + chunks: chunks, + originalSize: originalSize, + } + + err := cr.reset(int(offset / chunkSize)) + if err != nil { + return nil, err + } + + suboff := offset % chunkSize + if suboff != 0 { + _, err := io.CopyN(ioutil.Discard, cr.d, suboff) + if err != nil { + return nil, err + } + } + return cr, nil +} + +func (r *compressedReader) chunkOffset(n int) int64 { + if n == len(r.chunks) { + return r.r.Size() + } + return r.chunks[n] +} + +func (r *compressedReader) chunkSize(n int) int { + return int(r.chunkOffset(n+1) - r.chunkOffset(n)) +} + +func (r *compressedReader) uncompressedSize(n int) int { + if n < len(r.chunks)-1 { + return chunkSize + } + size := int(r.originalSize % chunkSize) + if size == 0 { + size = chunkSize + } + return size +} + +func (r *compressedReader) reset(n int) error { + if n >= len(r.chunks) { + return io.EOF + } + if r.d != nil { + r.d.Close() + } + r.curChunk = n + size := r.chunkSize(n) + uncompressedSize := r.uncompressedSize(n) + section := io.NewSectionReader(r.r, r.chunkOffset(n), int64(size)) + if size != uncompressedSize { + d, err := lzx.NewReader(section, uncompressedSize) + if err != nil { + return err + } + r.d = d + } else { + r.d = ioutil.NopCloser(section) + } + + return nil +} + +func (r *compressedReader) Read(b []byte) (int, error) { + for { + n, err := r.d.Read(b) + if err != io.EOF { + return n, err + } + + err = r.reset(r.curChunk + 1) + if err != nil { + return n, err + } + } +} + +func (r *compressedReader) Close() error { + var err error + if r.d != nil { + err = r.d.Close() + r.d = nil + } + return err +} diff --git a/wim/lzx/lzx.go b/wim/lzx/lzx.go new file mode 100644 index 0000000..7addf45 --- /dev/null +++ b/wim/lzx/lzx.go @@ -0,0 +1,546 @@ +// Package lzx implements a decompressor for the the WIM variant of the +// LZX compression algorithm. +// +// The LZX algorithm is an earlier variant of LZX DELTA, which is documented +// at https://msdn.microsoft.com/en-us/library/cc483133(v=exchg.80).aspx. +package lzx + +import ( + "bufio" + "bytes" + "encoding/binary" + "errors" + "io" +) + +const ( + maincodecount = 496 + maincodesplit = 256 + lencodecount = 249 + + maxBlockSize = 32768 + windowSize = 32768 + + treePathLenCount = 17 + + e8filesize = 12000000 + maxe8offset = 0x3fffffff + + verbatimBlock = 1 + alignedOffsetBlock = 2 + uncompressedBlock = 3 +) + +var footerBits = [...]byte{ + 0, 0, 0, 0, 1, 1, 2, 2, + 3, 3, 4, 4, 5, 5, 6, 6, + 7, 7, 8, 8, 9, 9, 10, 10, + 11, 11, 12, 12, 13, 13, 14, +} + +var basePosition = [...]uint16{ + 0, 1, 2, 3, 4, 6, 8, 12, + 16, 24, 32, 48, 64, 96, 128, 192, + 256, 384, 512, 768, 1024, 1536, 2048, 3072, + 4096, 6144, 8192, 12288, 16384, 24576, 32768, +} + +var ( + errCorrupt = errors.New("LZX data corrupt") +) + +// Reader is an interface used by the decompressor to access +// the input stream. If the provided io.Reader does not implement +// Reader, then a bufio.Reader is used. +type Reader interface { + io.Reader + io.ByteReader +} + +type decompressor struct { + r Reader + err error + unaligned bool + nbits byte + c uint32 + lru [3]uint16 + uncompressed int + windowReader *bytes.Reader + mainlens [maincodecount]byte + lenlens [lencodecount]byte + window [windowSize]byte +} + +// 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() + } + if err != nil { + if err == io.EOF { + return false + } + f.err = err + } + f.c |= (uint32(b1)<<8 | uint32(b0)) << (16 - f.nbits) + f.nbits += 16 + return true +} + +// getBits retrieves the next n bits from the byte stream. n +// must be <= 16. It sets f.err on error. +func (f *decompressor) getBits(n byte) uint16 { + if f.nbits < 16 { + if !f.feed() && n > f.nbits { + f.err = io.ErrUnexpectedEOF + } + } + c := uint16(f.c >> (32 - n)) + f.c <<= n + f.nbits -= n + return c +} + +type huffman struct { + lens []byte + table []uint16 + maxbits byte +} + +// buildTable builds a huffman decoding table from a slice of code lengths, +// one per code, in order. Each code length must be less than treePathLenCount. +// See https://en.wikipedia.org/wiki/Canonical_Huffman_code. +func buildTable(codelens []byte) *huffman { + // Determine the number of codes of each length, and the + // maximum length. + var count [treePathLenCount]uint + var max byte + for _, cl := range codelens { + count[cl]++ + if max < cl { + max = cl + } + } + + if max == 0 { + return &huffman{} + } + + // Determine the first code of each length. + var first [treePathLenCount]uint + code := uint(0) + for i := byte(1); i <= max; i++ { + code <<= 1 + first[i] = code + code += count[i] + } + + if code != 1<>(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 +} + +// mod17 computes the value mod 17. +func mod17(b byte) byte { + for b >= 17 { + b -= 17 + } + return b +} + +// readTree updates the huffman tree path lengths in lens by +// reading and decoding lengths from the byte stream. lens +// should be prepopulated with the previous block's tree's path +// lengths. For the first block, lens should be zero. +func (f *decompressor) readTree(lens []byte) error { + // Get the pre-tree for the main tree. + var pretreeLen [20]byte + for i := range pretreeLen { + pretreeLen[i] = byte(f.getBits(4)) + } + if f.err != nil { + return f.err + } + h := buildTable(pretreeLen[:]) + + // The lengths are encoded as a series of huffman codes + // encoded by the pre-tree. + for i := 0; i < len(lens); { + c := byte(f.getCode(h)) + if f.err != nil { + return f.err + } + switch { + case c <= 16: // length is delta from previous length + lens[i] = mod17(lens[i] + 17 - c) + i++ + case c == 17: // next n + 4 lengths are zero + zeroes := int(f.getBits(4)) + 4 + if i+zeroes > len(lens) { + return errCorrupt + } + for j := 0; j < zeroes; j++ { + lens[i+j] = 0 + } + i += zeroes + case c == 18: // next n + 20 lengths are zero + zeroes := int(f.getBits(5)) + 20 + if i+zeroes > len(lens) { + return errCorrupt + } + for j := 0; j < zeroes; j++ { + lens[i+j] = 0 + } + i += zeroes + case c == 19: // next n + 4 lengths all have the same value + same := int(f.getBits(1)) + 4 + if i+same > len(lens) { + return errCorrupt + } + c = byte(f.getCode(h)) + if c > 16 { + return errCorrupt + } + l := mod17(lens[i] + 17 - c) + for j := 0; j < same; j++ { + lens[i+j] = l + } + i += same + default: + return errCorrupt + } + } + + if f.err != nil { + return f.err + } + return nil +} + +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() + if err != nil { + if err == io.EOF { + err = io.ErrUnexpectedEOF + } + return 0, 0, err + } + f.unaligned = false + } + + blockType := f.getBits(3) + full := f.getBits(1) + var blockSize uint16 + if full != 0 { + blockSize = maxBlockSize + } else { + blockSize = f.getBits(16) + if blockSize > maxBlockSize { + return 0, 0, errCorrupt + } + } + + if f.err != nil { + return 0, 0, f.err + } + + switch blockType { + case verbatimBlock, alignedOffsetBlock: + // The caller will read the huffman trees. + case uncompressedBlock: + // Not sure if this can happen... + if f.nbits > 16 || f.nbits == 0 { + return 0, 0, errCorrupt + } + + // Drop the remaining bits in the current 16-bit word. + f.nbits = 0 + f.c = 0 + + // Read the LRU values for the next block. + var lru [12]byte + _, err := io.ReadFull(f.r, lru[:]) + 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])) + + default: + return 0, 0, errCorrupt + } + + return byte(blockType), blockSize, nil +} + +// readTrees reads the two or three huffman trees for the current block. +// readAligned specifies whether to read the aligned offset tree. +func (f *decompressor) readTrees(readAligned bool) (main *huffman, length *huffman, aligned *huffman, err error) { + // Aligned offset blocks start with a small aligned offset tree. + if readAligned { + var alignedLen [8]byte + for i := range alignedLen { + alignedLen[i] = byte(f.getBits(3)) + } + aligned = buildTable(alignedLen[:]) + if aligned == nil { + err = errors.New("corrupt") + return + } + } + + // The main tree is encoded in two parts. + err = f.readTree(f.mainlens[:maincodesplit]) + if err != nil { + return + } + err = f.readTree(f.mainlens[maincodesplit:]) + if err != nil { + return + } + + main = buildTable(f.mainlens[:]) + if main == nil { + err = errors.New("corrupt") + return + } + + // The length tree is encoding in a single part. + err = f.readTree(f.lenlens[:]) + if err != nil { + return + } + + length = buildTable(f.lenlens[:]) + if length == nil { + err = errors.New("corrupt") + return + } + + err = f.err + return +} + +// 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; { + main := f.getCode(hmain) + if f.err != nil { + return int(i - start), f.err + } + if main < 256 { + // Literal byte. + f.window[i] = byte(main) + i++ + continue + } + + // This is a match backward in the window. Determine + // the offset and dlength. + lenheader := (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 + } + matchlen += 2 + + var matchoffset uint16 + if slot < 3 { + // The offset is one of the LRU values. + matchoffset = f.lru[slot] + f.lru[slot] = f.lru[0] + f.lru[0] = matchoffset + } else { + // The offset is encoded as a combination of the + // slot and more bits from the bit stream. + offsetbits := footerBits[slot] + var verbatimbits, alignedbits uint16 + if offsetbits > 0 { + if haligned != nil && offsetbits >= 3 { + // This is an aligned offset block. Combine + // the bits written verbatim with the aligned + // offset tree code. + verbatimbits = f.getBits(offsetbits-3) * 8 + alignedbits = f.getCode(haligned) + } else { + // There are no aligned offset bits to read, + // only verbatim bits. + verbatimbits = f.getBits(offsetbits) + alignedbits = 0 + } + } + matchoffset = basePosition[slot] + verbatimbits + alignedbits - 2 + // Update the LRU cache. + f.lru[2] = f.lru[1] + f.lru[1] = f.lru[0] + f.lru[0] = matchoffset + } + + if matchoffset > i || matchlen > end-i { + return int(i - start), errCorrupt + } + + for j := uint16(0); j < matchlen; j++ { + f.window[i+j] = f.window[i+j-matchoffset] + } + i += matchlen + } + return int(end - start), nil +} + +// readBlock decodes the current block and returns the number of uncompressed bytes. +func (f *decompressor) readBlock(start uint16) (int, error) { + blockType, size, err := f.readBlockHeader() + if err != nil { + return 0, err + } + + if blockType == uncompressedBlock { + if size%2 == 1 { + // Remember to realign the byte stream at the next block. + f.unaligned = true + } + return io.ReadFull(f.r, f.window[start:start+size]) + } + + hmain, hlength, haligned, err := f.readTrees(blockType == alignedOffsetBlock) + if err != nil { + return 0, err + } + + return f.readCompressedBlock(start, start+size, hmain, hlength, haligned) +} + +// decodeE8 reverses the 0xe8 x86 instruction encoding that was performed +// to the uncompressed data before it was compressed. +func decodeE8(b []byte, off int64) { + if off > maxe8offset || len(b) < 10 { + return + } + for i := 0; i < len(b)-10; i++ { + if b[i] == 0xe8 { + currentPtr := int32(off) + int32(i) + abs := int32(binary.LittleEndian.Uint32(b[i+1 : i+5])) + if abs >= -currentPtr && abs < e8filesize { + var rel int32 + if abs >= 0 { + rel = abs - currentPtr + } else { + rel = abs + e8filesize + } + binary.LittleEndian.PutUint32(b[i+1:i+5], uint32(rel)) + } + i += 4 + } + } +} + +func (f *decompressor) Read(b []byte) (int, error) { + // Read and uncompress everything. + if f.windowReader == nil { + n := 0 + for n < f.uncompressed { + k, err := f.readBlock(uint16(n)) + if err != nil { + return 0, err + } + n += k + } + decodeE8(f.window[:f.uncompressed], 0) + f.windowReader = bytes.NewReader(f.window[:f.uncompressed]) + } + + // Just read directly from the window. + return f.windowReader.Read(b) +} + +func (f *decompressor) Close() error { + return nil +} + +// NewReader returns a new io.ReadCloser that decompresses a +// WIM LZX stream until uncompressedSize bytes have been returned. +func NewReader(r io.Reader, uncompressedSize int) (io.ReadCloser, error) { + if uncompressedSize > windowSize { + return nil, errors.New("uncompressed size is limited to 32KB") + } + 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) + } + return f, nil +} diff --git a/wim/wim.go b/wim/wim.go new file mode 100644 index 0000000..ffd2ee5 --- /dev/null +++ b/wim/wim.go @@ -0,0 +1,646 @@ +// Package wim implements a WIM file parser. +// +// WIM files are used to distribute Windows file system and container images. +// They are documented at https://msdn.microsoft.com/en-us/library/windows/desktop/dd861280.aspx. +package wim + +import ( + "bufio" + "bytes" + "crypto/sha1" + "encoding/binary" + "errors" + "fmt" + "io" + "io/ioutil" + "syscall" + "unicode/utf16" +) + +var wimImageTag = [...]byte{'M', 'S', 'W', 'I', 'M', 0, 0, 0} + +type guid struct { + Data1 uint32 + Data2 uint16 + Data3 uint16 + Data4 [8]byte +} + +func (g guid) String() string { + return fmt.Sprintf("%08x-%04x-%04x-%02x%02x-%02x%02x%02x%02x%02x%02x", g.Data1, g.Data2, g.Data3, g.Data4[0], g.Data4[1], g.Data4[2], g.Data4[3], g.Data4[4], g.Data4[5], g.Data4[6], g.Data4[7]) +} + +type resourceDescriptor struct { + FlagsAndCompressedSize uint64 + Offset int64 + OriginalSize int64 +} + +type resFlag byte + +const ( + resFlagFree resFlag = 1 << iota + resFlagMetadata + resFlagCompressed + resFlagSpanned +) + +const validate = false + +const supportedResFlags = resFlagMetadata | resFlagCompressed + +func (r *resourceDescriptor) Flags() resFlag { + return resFlag(r.FlagsAndCompressedSize >> 56) +} + +func (r *resourceDescriptor) CompressedSize() int64 { + return int64(r.FlagsAndCompressedSize & 0xffffffffffffff) +} + +func (r *resourceDescriptor) String() string { + s := fmt.Sprintf("%d bytes at %d", r.CompressedSize(), r.Offset) + if r.Flags()&4 != 0 { + s += fmt.Sprintf(" (uncompresses to %d)", r.OriginalSize) + } + return s +} + +// SHA1Hash contains the SHA1 hash of a file or stream. +type SHA1Hash [20]byte + +type streamDescriptor struct { + resourceDescriptor + PartNumber uint16 + RefCount uint32 + Hash SHA1Hash +} + +type hdrFlag uint32 + +const ( + hdrFlagReserved hdrFlag = 1 << iota + hdrFlagCompressed + hdrFlagReadOnly + hdrFlagSpanned + hdrFlagResourceOnly + hdrFlagMetadataOnly + hdrFlagWriteInProgress + hdrFlagRpFix +) + +const ( + hdrFlagCompressReserved hdrFlag = 1 << (iota + 16) + hdrFlagCompressXpress + hdrFlagCompressLzx +) + +const supportedHdrFlags = hdrFlagRpFix | hdrFlagReadOnly | hdrFlagCompressed | hdrFlagCompressLzx + +type wimHeader struct { + ImageTag [8]byte + Size uint32 + Version uint32 + Flags hdrFlag + CompressionSize uint32 + WIMGuid guid + PartNumber uint16 + TotalParts uint16 + ImageCount uint32 + OffsetTable resourceDescriptor + XMLData resourceDescriptor + BootMetadata resourceDescriptor + BootIndex uint32 + Padding uint32 + Integrity resourceDescriptor + Unused [60]byte +} + +type securityblockDisk struct { + TotalLength uint32 + NumEntries uint32 +} + +const securityblockDiskSize = 8 + +type direntry struct { + Length int64 + Attributes uint32 + SecurityID uint32 + SubdirOffset int64 + Unused1, Unused2 int64 + CreationTime syscall.Filetime + LastAccessTime syscall.Filetime + LastWriteTime syscall.Filetime + Hash SHA1Hash + Padding uint32 + ReparseHardLink int64 + StreamCount uint16 + ShortNameLength uint16 + FileNameLength uint16 +} + +const direntrySize = 102 + +type streamentry struct { + Length int64 + Unused int64 + Hash SHA1Hash + NameLength int16 +} + +const streamentrySize = 38 + +// ParseError is returned when the WIM cannot be parsed. +type ParseError struct { + Oper string + Err error +} + +func (e *ParseError) Error() string { + return "WIM parse error at " + e.Oper + ": " + e.Err.Error() +} + +// Reader provides functions to read a WIM file. +type Reader struct { + hdr wimHeader + r io.ReaderAt + fileData map[SHA1Hash]resourceDescriptor + + Image []*Image // The WIM's images. +} + +// Image represents an image within a WIM file. +type Image struct { + wim *Reader + offset resourceDescriptor + sds [][]byte + rootOffset int64 +} + +// StreamHeader contains alternate data stream metadata. +type StreamHeader struct { + Name string + Hash SHA1Hash + Size int64 +} + +// Stream represents an alternate data stream or reparse point data stream. +type Stream struct { + StreamHeader + wim *Reader + offset resourceDescriptor +} + +// FileHeader contains file metadata. +type FileHeader struct { + Name string + ShortName string + Attributes uint32 + SecurityDescriptor []byte + CreationTime syscall.Filetime + LastAccessTime syscall.Filetime + LastWriteTime syscall.Filetime + Hash SHA1Hash + Size int64 + LinkID int64 + ReparseTag uint32 + ReparseReserved uint32 + ReparseStream *Stream +} + +// File represents a file or directory in a WIM image. +type File struct { + FileHeader + Streams []*Stream + offset resourceDescriptor + img *Image + subdirOffset int64 +} + +// NewReader returns a Reader that can be used to read WIM file data. +func NewReader(f io.ReaderAt) (*Reader, error) { + r := &Reader{r: f} + section := io.NewSectionReader(f, 0, 0xffff) + err := binary.Read(section, binary.LittleEndian, &r.hdr) + if err != nil { + return nil, err + } + + if r.hdr.ImageTag != wimImageTag { + return nil, &ParseError{"image tag", errors.New("not a WIM file")} + } + + if r.hdr.Flags&^supportedHdrFlags != 0 { + return nil, fmt.Errorf("unsupported WIM flags %x", r.hdr.Flags&^supportedHdrFlags) + } + + if r.hdr.CompressionSize != 0x8000 { + return nil, fmt.Errorf("unsupported compression size %d", r.hdr.CompressionSize) + } + + if r.hdr.TotalParts != 1 { + return nil, errors.New("multi-part WIM not supported") + } + + fileData, images, err := r.readOffsetTable(&r.hdr.OffsetTable) + if err != nil { + return nil, err + } + r.fileData = fileData + r.Image = images + return r, nil +} + +func (r *Reader) resourceReader(hdr *resourceDescriptor) (io.ReadCloser, error) { + return r.resourceReaderWithOffset(hdr, 0) +} + +func (r *Reader) resourceReaderWithOffset(hdr *resourceDescriptor, offset int64) (io.ReadCloser, error) { + var sr io.ReadCloser + section := io.NewSectionReader(r.r, hdr.Offset, hdr.CompressedSize()) + if hdr.Flags()&resFlagCompressed == 0 { + section.Seek(offset, 0) + sr = ioutil.NopCloser(section) + } else { + cr, err := newCompressedReader(section, hdr.OriginalSize, offset) + if err != nil { + return nil, err + } + sr = cr + } + + return sr, nil +} + +func (r *Reader) readResource(hdr *resourceDescriptor) ([]byte, error) { + rsrc, err := r.resourceReader(hdr) + if err != nil { + return nil, err + } + defer rsrc.Close() + return ioutil.ReadAll(rsrc) +} + +// ReadXML reads the XML metadata from a WIM. +func (r *Reader) ReadXML() (string, error) { + if r.hdr.XMLData.CompressedSize() == 0 { + return "", nil + } + rsrc, err := r.resourceReader(&r.hdr.XMLData) + if err != nil { + return "", err + } + defer rsrc.Close() + + XMLData := make([]uint16, r.hdr.XMLData.OriginalSize/2) + err = binary.Read(rsrc, binary.LittleEndian, XMLData) + if err != nil { + return "", &ParseError{"XML data", err} + } + + // The BOM will always indicate little-endian UTF-16. + if XMLData[0] != 0xfeff { + return "", &ParseError{"XML data", errors.New("invalid BOM")} + } + return string(utf16.Decode(XMLData[1:])), nil +} + +func (r *Reader) readOffsetTable(res *resourceDescriptor) (map[SHA1Hash]resourceDescriptor, []*Image, error) { + fileData := make(map[SHA1Hash]resourceDescriptor) + var images []*Image + + offsetTable, err := r.readResource(res) + if err != nil { + return nil, nil, &ParseError{"offset table", err} + } + + br := bytes.NewReader(offsetTable) + for { + var res streamDescriptor + err := binary.Read(br, binary.LittleEndian, &res) + if err == io.EOF { + break + } + if err != nil { + return nil, nil, &ParseError{"offset table", err} + } + if res.Flags()&^supportedResFlags != 0 { + return nil, nil, &ParseError{"offset table", errors.New("unsupported resource flag")} + } + + // Validation for ad-hoc testing + if validate { + sec, err := r.resourceReader(&res.resourceDescriptor) + if err != nil { + return nil, nil, err + } + hash := sha1.New() + _, err = io.Copy(hash, sec) + sec.Close() + if err != nil { + return nil, nil, err + } + var cmphash SHA1Hash + copy(cmphash[:], hash.Sum(nil)) + if cmphash != res.Hash { + return nil, nil, errors.New("hash mismatch") + } + } + + if res.Flags()&resFlagMetadata != 0 { + image := &Image{ + wim: r, + offset: res.resourceDescriptor, + } + images = append(images, image) + } else { + fileData[res.Hash] = res.resourceDescriptor + } + } + + if len(images) != int(r.hdr.ImageCount) { + return nil, nil, &ParseError{"offset table", errors.New("mismatched image count")} + } + + return fileData, images, nil +} + +func (r *Reader) readSecurityDescriptors(rsrc io.Reader) (sds [][]byte, n int64, err error) { + var secBlock securityblockDisk + err = binary.Read(rsrc, binary.LittleEndian, &secBlock) + if err != nil { + err = &ParseError{"security table", err} + return + } + + n += securityblockDiskSize + + secSizes := make([]int64, secBlock.NumEntries) + err = binary.Read(rsrc, binary.LittleEndian, &secSizes) + if err != nil { + err = &ParseError{"security table sizes", err} + return + } + + n += int64(secBlock.NumEntries * 8) + + sds = make([][]byte, secBlock.NumEntries) + for i, size := range secSizes { + sd := make([]byte, size&0xffffffff) + _, err = io.ReadFull(rsrc, sd) + if err != nil { + err = &ParseError{"security descriptor", err} + return + } + n += int64(len(sd)) + sds[i] = sd + } + + secsize := int64((secBlock.TotalLength + 7) &^ 7) + if n > secsize { + err = &ParseError{"security descriptor", errors.New("security descriptor table too small")} + return + } + + _, err = io.CopyN(ioutil.Discard, rsrc, secsize-n) + if err != nil { + return + } + + return +} + +// 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) + if err != nil { + return nil, err + } + img.sds = sds + img.rootOffset = n + } + + f, err := img.readdir(rsrc) + if err != nil { + return nil, err + } + if len(f) != 1 { + return nil, &ParseError{"root directory", errors.New("expected exactly 1 root directory entry")} + } + return f[0], err +} + +func (img *Image) readdir(rsrc io.Reader) ([]*File, error) { + r := bufio.NewReader(rsrc) + + var entries []*File + for { + e, err := img.readNextEntry(r) + if err == io.EOF { + break + } + if err != nil { + return nil, err + } + entries = append(entries, e) + } + return entries, nil +} + +func (img *Image) readNextEntry(r *bufio.Reader) (*File, error) { + lengthBuf, err := r.Peek(8) + if err != nil { + return nil, &ParseError{"directory length check", err} + } + + left := int(binary.LittleEndian.Uint64(lengthBuf)) + if left == 0 { + return nil, io.EOF + } + + if left < direntrySize { + return nil, &ParseError{"directory entry", errors.New("size too short")} + } + + var dentry direntry + err = binary.Read(r, binary.LittleEndian, &dentry) + if err != nil { + return nil, &ParseError{"directory entry", err} + } + + left -= direntrySize + + var offset resourceDescriptor + zerohash := SHA1Hash{} + if dentry.Hash != zerohash { + var ok bool + offset, ok = img.wim.fileData[dentry.Hash] + if !ok { + return nil, &ParseError{"directory entry", fmt.Errorf("could not find file data matching hash %v", dentry.Hash)} + } + } + + f := &File{ + FileHeader: FileHeader{ + Attributes: dentry.Attributes, + CreationTime: dentry.CreationTime, + LastAccessTime: dentry.LastAccessTime, + LastWriteTime: dentry.LastWriteTime, + Hash: dentry.Hash, + Size: offset.OriginalSize, + }, + + offset: offset, + img: img, + subdirOffset: dentry.SubdirOffset, + } + + if dentry.Attributes&syscall.FILE_ATTRIBUTE_REPARSE_POINT == 0 { + f.LinkID = dentry.ReparseHardLink + } else { + f.ReparseTag = uint32(dentry.ReparseHardLink) + f.ReparseReserved = uint32(dentry.ReparseHardLink >> 32) + } + + if dentry.SecurityID != 0xffffffff { + f.SecurityDescriptor = img.sds[dentry.SecurityID] + } + + namesLen := int(dentry.FileNameLength + 2 + dentry.ShortNameLength) + if left < namesLen { + return nil, &ParseError{"directory entry", 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{"file name", err} + } + + left -= namesLen + + if dentry.FileNameLength > 0 { + f.Name = string(utf16.Decode(names[:dentry.FileNameLength/2])) + } + + if dentry.ShortNameLength > 0 { + f.ShortName = string(utf16.Decode(names[dentry.FileNameLength/2+1:])) + } + + _, err = r.Discard(left) + if err != nil { + return nil, err + } + + if dentry.StreamCount > 0 { + var streams []*Stream + for i := uint16(0); i < dentry.StreamCount; i++ { + s, err := img.readNextStream(r) + if err != nil { + return nil, err + } + if !(s.Name == "" && s.Size == 0) { + if dentry.Attributes&syscall.FILE_ATTRIBUTE_REPARSE_POINT != 0 && s.Name == "" { + f.ReparseStream = s + } else { + streams = append(streams, s) + } + } + } + f.Streams = streams + } + + if dentry.Attributes&syscall.FILE_ATTRIBUTE_REPARSE_POINT != 0 && f.ReparseStream == nil { + return nil, &ParseError{"directory entry", errors.New("reparse point is missing reparse stream")} + } + + return f, nil +} + +func (img *Image) readNextStream(r *bufio.Reader) (*Stream, error) { + lengthBuf, err := r.Peek(8) + if err != nil { + return nil, &ParseError{"stream length check", err} + } + + left := int(binary.LittleEndian.Uint64(lengthBuf)) + if left < streamentrySize { + return nil, &ParseError{"stream entry", errors.New("size too short")} + } + + var sentry streamentry + err = binary.Read(r, binary.LittleEndian, &sentry) + if err != nil { + return nil, &ParseError{"stream entry", err} + } + + left -= streamentrySize + + var offset resourceDescriptor + if sentry.Hash != (SHA1Hash{}) { + var ok bool + offset, ok = img.wim.fileData[sentry.Hash] + if !ok { + return nil, &ParseError{"stream entry", fmt.Errorf("could not find file data matching hash %v", sentry.Hash)} + } + } + + s := &Stream{ + StreamHeader: StreamHeader{ + Hash: sentry.Hash, + Size: offset.OriginalSize, + }, + wim: img.wim, + offset: offset, + } + + if left < int(sentry.NameLength) { + return nil, &ParseError{"stream entry", 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{"file name", err} + } + + left -= int(sentry.NameLength) + s.Name = string(utf16.Decode(names)) + + _, err = r.Discard(left) + if err != nil { + return nil, err + } + + return s, nil +} + +// Open returns an io.ReadCloser that can be used to read the stream's contents. +func (s *Stream) Open() (io.ReadCloser, error) { + return s.wim.resourceReader(&s.offset) +} + +// Open returns an io.ReadCloser that can be used to read the file's contents. +func (f *File) Open() (io.ReadCloser, error) { + return f.img.wim.resourceReader(&f.offset) +} + +// Readdir reads the directory entries. +func (f *File) Readdir() ([]*File, error) { + if f.Attributes&syscall.FILE_ATTRIBUTE_DIRECTORY == 0 { + 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) +}