OpenDiablo2/MPQ/MPQStream.go

package MPQ

import (
	"bytes"
	"compress/zlib"
	"encoding/binary"
	"fmt"
	"log"
	"strings"

	"github.com/JoshVarga/blast"
	"github.com/essial/OpenDiablo2/Common"
	"github.com/essial/OpenDiablo2/Compression"
)

// Stream represents a stream of data in an MPQ archive
type Stream struct {
	MPQData           MPQ
	BlockTableEntry   BlockTableEntry
	FileName          string
	EncryptionSeed    uint32
	BlockPositions    []uint32
	CurrentPosition   uint32
	CurrentData       []byte
	CurrentBlockIndex uint32
	BlockSize         uint32
}

// CreateStream creates an MPQ stream
func CreateStream(mpq MPQ, blockTableEntry BlockTableEntry, fileName string) *Stream {
	result := &Stream{
		MPQData:           mpq,
		BlockTableEntry:   blockTableEntry,
		CurrentBlockIndex: 0xFFFFFFFF,
	}
	fileSegs := strings.Split(fileName, `\`)
	result.EncryptionSeed = hashString(fileSegs[len(fileSegs)-1], 3)
	if result.BlockTableEntry.HasFlag(FileFixKey) {
		result.EncryptionSeed = (result.EncryptionSeed + result.BlockTableEntry.FilePosition) ^ result.BlockTableEntry.UncompressedFileSize
	}
	result.BlockSize = 0x200 << result.MPQData.Data.BlockSize

	if (result.BlockTableEntry.HasFlag(FileCompress) || result.BlockTableEntry.HasFlag(FileImplode)) && !result.BlockTableEntry.HasFlag(FileSingleUnit) {
		result.loadBlockOffsets()
	}
	return result
}

func (v *Stream) loadBlockOffsets() {
	blockPositionCount := ((v.BlockTableEntry.UncompressedFileSize + v.BlockSize - 1) / v.BlockSize) + 1
	v.BlockPositions = make([]uint32, blockPositionCount)
	v.MPQData.File.Seek(int64(v.BlockTableEntry.FilePosition), 0)
	binary.Read(v.MPQData.File, binary.LittleEndian, &v.BlockPositions)
	blockPosSize := blockPositionCount << 2
	if v.BlockTableEntry.HasFlag(FileEncrypted) {
		decrypt(v.BlockPositions, v.EncryptionSeed-1)
		if v.BlockPositions[0] != blockPosSize {
			panic("Decryption of MPQ failed!")
		}
		if v.BlockPositions[1] > v.BlockSize+blockPosSize {
			panic("Decryption of MPQ failed!")
		}
	}
}

func (v *Stream) Read(buffer []byte, offset, count uint32) uint32 {
	if v.BlockTableEntry.HasFlag(FileSingleUnit) {
		return v.readInternalSingleUnit(buffer, offset, count)
	}
	toRead := count
	readTotal := uint32(0)
	for toRead > 0 {
		read := v.readInternal(buffer, offset, count)
		if read == 0 {
			break
		}
		readTotal += read
		offset += read
		toRead -= read
	}
	return readTotal
}

func (v *Stream) readInternalSingleUnit(buffer []byte, offset, count uint32) uint32 {
	if len(v.CurrentData) == 0 {
		v.loadSingleUnit()
	}

	bytesToCopy := Common.Min(uint32(len(v.CurrentData))-v.CurrentPosition, count)
	copy(buffer[offset:offset+bytesToCopy], v.CurrentData[v.CurrentPosition:v.CurrentPosition+bytesToCopy])
	v.CurrentPosition += bytesToCopy
	return bytesToCopy
}

func (v *Stream) readInternal(buffer []byte, offset, count uint32) uint32 {
	v.bufferData()
	localPosition := v.CurrentPosition % v.BlockSize
	bytesToCopy := Common.MinInt32(int32(len(v.CurrentData))-int32(localPosition), int32(count))
	if bytesToCopy <= 0 {
		return 0
	}
	copy(buffer[offset:offset+uint32(bytesToCopy)], v.CurrentData[localPosition:localPosition+uint32(bytesToCopy)])
	v.CurrentPosition += uint32(bytesToCopy)
	return uint32(bytesToCopy)
}

func (v *Stream) bufferData() {
	requiredBlock := uint32(v.CurrentPosition / v.BlockSize)
	if requiredBlock == v.CurrentBlockIndex {
		return
	}
	expectedLength := Common.Min(v.BlockTableEntry.UncompressedFileSize-(requiredBlock*v.BlockSize), v.BlockSize)
	v.CurrentData = v.loadBlock(requiredBlock, expectedLength)
	v.CurrentBlockIndex = requiredBlock
}

func (v *Stream) loadSingleUnit() {
	fileData := make([]byte, v.BlockSize)
	v.MPQData.File.Seek(int64(v.MPQData.Data.HeaderSize), 0)
	binary.Read(v.MPQData.File, binary.LittleEndian, &fileData)
	if v.BlockSize == v.BlockTableEntry.UncompressedFileSize {
		v.CurrentData = fileData
		return
	}
	v.CurrentData = decompressMulti(fileData, v.BlockTableEntry.UncompressedFileSize)
}

func (v *Stream) loadBlock(blockIndex, expectedLength uint32) []byte {
	var (
		offset uint32
		toRead uint32
	)
	if v.BlockTableEntry.HasFlag(FileCompress) || v.BlockTableEntry.HasFlag(FileImplode) {
		offset = v.BlockPositions[blockIndex]
		toRead = v.BlockPositions[blockIndex+1] - offset
	} else {
		offset = blockIndex * v.BlockSize
		toRead = expectedLength
	}
	offset += v.BlockTableEntry.FilePosition
	data := make([]byte, toRead)
	v.MPQData.File.Seek(int64(offset), 0)
	binary.Read(v.MPQData.File, binary.LittleEndian, &data)
	if v.BlockTableEntry.HasFlag(FileEncrypted) && v.BlockTableEntry.UncompressedFileSize > 3 {
		if v.EncryptionSeed == 0 {
			panic("Unable to determine encryption key")
		}

		decryptBytes(data, blockIndex+v.EncryptionSeed)
	}
	if v.BlockTableEntry.HasFlag(FileCompress) && (toRead != expectedLength) {
		if !v.BlockTableEntry.HasFlag(FileSingleUnit) {
			data = decompressMulti(data, expectedLength)
		} else {
			data = pkDecompress(data)
		}
	}
	if v.BlockTableEntry.HasFlag(FileImplode) && (toRead != expectedLength) {
		data = pkDecompress(data)
	}

	return data
}

func decompressMulti(data []byte, expectedLength uint32) []byte {
	copmressionType := data[0]
	switch copmressionType {
	case 1: // Huffman
		panic("huffman decompression not supported")
	case 2: // ZLib/Deflate
		return deflate(data[1:])
	case 8: // PKLib/Impode
		return pkDecompress(data[1:])
	case 0x10: // BZip2
		panic("bzip2 decompression not supported")
	case 0x80: // IMA ADPCM Stereo
		return Compression.WavDecompress(data[1:], 2)
		//return MpqWavCompression.Decompress(sinput, 2);
		//panic("ima adpcm sterio decompression not supported")
	case 0x40: // IMA ADPCM Mono
		//return MpqWavCompression.Decompress(sinput, 1)
		panic("mpq wav decompression not supported")
	case 0x12:
		panic("lzma decompression not supported")
	// Combos
	case 0x22:
		// TODO: sparse then zlib
		panic("sparse decompression + deflate decompression not supported")
	case 0x30:
		// TODO: sparse then bzip2
		panic("sparse decompression + bzip2 decompression not supported")
	case 0x41:
		sinput := Compression.HuffmanDecompress(data[1:])
		sinput = Compression.WavDecompress(sinput, 1)
		tmp := make([]byte, len(sinput))
		copy(tmp, sinput)
		return tmp
	case 0x48:
		//byte[] result = PKDecompress(sinput, outputLength);
		//return MpqWavCompression.Decompress(new MemoryStream(result), 1);
		panic("pk + mpqwav decompression not supported")
	case 0x81:
		sinput := Compression.HuffmanDecompress(data[1:])
		sinput = Compression.WavDecompress(sinput, 2)
		tmp := make([]byte, len(sinput))
		copy(tmp, sinput)
		return tmp
	case 0x88:
		//byte[] result = PKDecompress(sinput, outputLength);
		//return MpqWavCompression.Decompress(new MemoryStream(result), 2);
		panic("pk + wav decompression not supported")
	default:
		panic(fmt.Sprintf("decompression not supported for unknown compression type %X", copmressionType))
	}
}

func deflate(data []byte) []byte {
	b := bytes.NewReader(data)
	r, err := zlib.NewReader(b)
	if err != nil {
		panic(err)
	}
	buffer := new(bytes.Buffer)
	_, err = buffer.ReadFrom(r)
	if err != nil {
		log.Panic(err)
	}
	err = r.Close()
	if err != nil {
		log.Panic(err)
	}
	return buffer.Bytes()
}

func pkDecompress(data []byte) []byte {
	b := bytes.NewReader(data)
	r, err := blast.NewReader(b)
	if err != nil {
		panic(err)
	}
	buffer := new(bytes.Buffer)
	_, err = buffer.ReadFrom(r)
	if err != nil {
		panic(err)
	}
	err = r.Close()
	if err != nil {
		panic(err)
	}
	return buffer.Bytes()
}