From b19c900382027552da579bd526f7898a324cad8a Mon Sep 17 00:00:00 2001
From: redraskal <ben@ryben.dev>
Date: Wed, 24 Jan 2024 14:46:28 -0600
Subject: [PATCH] fix: y8s4 zstd reader

---
 dissect/error.go  |  1 -
 dissect/player.go |  4 ---
 dissect/reader.go | 66 +++++++++++++++++++++++------------------------
 dissect/utils.go  | 12 +++++++++
 4 files changed, 45 insertions(+), 38 deletions(-)

diff --git a/dissect/error.go b/dissect/error.go
index 4ede27d..bae3774 100644
--- a/dissect/error.go
+++ b/dissect/error.go
@@ -9,7 +9,6 @@ import (
 
 var ErrInvalidFile = errors.New("dissect: not a dissect file")
 var ErrInvalidFolder = errors.New("dissect: not a match folder")
-var ErrInvalidLength = errors.New("dissect: received an invalid length of bytes")
 var ErrInvalidStringSep = errors.New("dissect: invalid string separator")
 
 // Ok returns true if err only pertains to EOF (read was successful).
diff --git a/dissect/player.go b/dissect/player.go
index fce7e27..8bb5412 100644
--- a/dissect/player.go
+++ b/dissect/player.go
@@ -151,10 +151,6 @@ func readPlayer(r *Reader) error {
 	if !found && len(username) > 0 {
 		r.Header.Players = append(r.Header.Players, p)
 	}
-	//if err := r.seek(unknownIndicator); err != nil {
-	//	return err
-	//}
-	//_, err = r.read(30) // unknown data, see above
 	return err
 }
 
diff --git a/dissect/reader.go b/dissect/reader.go
index 9dc8e85..e631f93 100644
--- a/dissect/reader.go
+++ b/dissect/reader.go
@@ -68,17 +68,13 @@ func NewReader(in io.Reader) (r *Reader, err error) {
 	return r, err
 }
 
-func (r *Reader) readChunkedData(br *bufio.Reader) error {
-	log.Debug().Msg("creating header buffer")
-	r.b = make([]byte, 0)
-	buf := make([]byte, 4096)
-	for i := 0; i < 5; i++ {
-		_, err := br.Read(buf)
-		if err != nil {
-			return err
-		}
-		r.b = append(r.b, buf...)
+func (r *Reader) readChunkedData(genericReader io.Reader) error {
+	log.Debug().Msg("reading data")
+	temp, err := io.ReadAll(genericReader)
+	if err != nil {
+		return err
 	}
+	r.b = temp
 	log.Debug().Msg("reading header magic")
 	if err := r.readHeaderMagic(); err != nil {
 		return err
@@ -89,13 +85,16 @@ func (r *Reader) readChunkedData(br *bufio.Reader) error {
 	if err != nil {
 		return err
 	}
-	pattern := []byte{0x32, 0x30, 0x30, 0x56, 0x52, 0x50, 0x4D, 0x43}
-	i := 0
-	zstdReader, _ := zstd.NewReader(bytes.NewReader([]byte{}))
+	log.Debug().Msg("decompressing data")
+	zstdMagic := []byte{0x28, 0xB5, 0x2F, 0xFD}
+	zstdReader, _ := zstd.NewReader(nil)
+	memoryReader := bytes.NewReader(nil)
+	patternIndex := 0
+	sections := 0
+	data := make([]byte, 0)
 	for !errors.Is(err, io.EOF) {
-		log.Debug().Msg("scanning for dissect frame")
-		for i != 8 {
-			b, scanErr := br.ReadByte()
+		for patternIndex != 4 {
+			b, scanErr := r.Bytes(1)
 			if errors.Is(scanErr, io.EOF) {
 				err = scanErr
 				break
@@ -103,38 +102,39 @@ func (r *Reader) readChunkedData(br *bufio.Reader) error {
 			if scanErr != nil {
 				return scanErr
 			}
-			if b == pattern[i] {
-				i++
+			if b[0] == zstdMagic[patternIndex] {
+				patternIndex++
 			} else {
-				i = 0
+				patternIndex = 0
 			}
 		}
 		if errors.Is(err, io.EOF) {
 			break
 		}
-		i = 0
-		_, err := br.Discard(8)
-		if err != nil {
-			return err
-		}
-		log.Debug().Msg("reading zstd frame")
-		if err = zstdReader.Reset(br); err != nil {
+		sections++
+		patternIndex = 0
+		memoryReader.Reset(r.b[r.offset-4:])
+		tempReader := countedReader{memoryReader, 0}
+		if err = zstdReader.Reset(&tempReader); err != nil {
 			return err
 		}
 		decompressed, err := io.ReadAll(zstdReader)
 		if err != nil && !(len(decompressed) > 0 && errors.Is(err, zstd.ErrMagicMismatch)) {
 			return err
 		}
-		log.Debug().Msg("writing decompressed frame")
 		for _, b := range decompressed {
-			r.b = append(r.b, b)
+			data = append(data, b)
 		}
+		r.offset += tempReader.n
 	}
+	r.b = data
+	r.offset = 0
+	log.Debug().Int("zstd_sections", sections).Send()
 	return nil
 }
 
-func (r *Reader) readUnchunkedData(br *bufio.Reader) error {
-	zstdReader, err := zstd.NewReader(br)
+func (r *Reader) readUnchunkedData(genericReader io.Reader) error {
+	zstdReader, err := zstd.NewReader(genericReader)
 	if err != nil {
 		return err
 	}
@@ -179,7 +179,7 @@ func (r *Reader) worker(start int, end int, wg *sync.WaitGroup, matches chan<- m
 func (r *Reader) Read() (err error) {
 	numWorkers := 5
 	var wg sync.WaitGroup
-	channel := make(chan match)
+	channel := make(chan match, 300)
 	start := r.offset
 	end := len(r.b)
 	if r.readPartial {
@@ -187,8 +187,8 @@ func (r *Reader) Read() (err error) {
 	}
 	blockSize := int(math.Floor(float64(end-start) / float64(numWorkers)))
 	log.Debug().Int("workers", numWorkers).Int("blockSize", blockSize).Send()
+	wg.Add(numWorkers)
 	for i := 0; i < numWorkers; i++ {
-		wg.Add(1)
 		blockStart := r.offset + (i * blockSize)
 		blockEnd := blockStart + blockSize
 		if i > 0 {
@@ -285,7 +285,7 @@ func (r *Reader) Seek(pattern []byte) error {
 func (r *Reader) Skip(n int) error {
 	r.offset += n
 	if r.offset >= len(r.b) {
-		return ErrInvalidLength
+		return io.EOF
 	}
 	return nil
 }
diff --git a/dissect/utils.go b/dissect/utils.go
index 895e626..a011dbf 100644
--- a/dissect/utils.go
+++ b/dissect/utils.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"github.com/rs/zerolog/log"
 	"github.com/xuri/excelize/v2"
+	"io"
 	"strings"
 )
 
@@ -31,6 +32,17 @@ func (r *Reader) PlayerIndexByUsername(username string) int {
 	return -1
 }
 
+type countedReader struct {
+	io.Reader
+	n int
+}
+
+func (r *countedReader) Read(p []byte) (int, error) {
+	n, err := r.Reader.Read(p)
+	r.n += n
+	return n, err
+}
+
 // hexEventComparison - Debugging tool
 type hexEventComparison struct {
 	usernames []string