Support for ZSTD compression/decompression in chunks

go.mod

@@ -5,6 +5,7 @@ go 1.18
 require (
 	github.com/anaminus/parse v0.2.0
 	github.com/bkaradzic/go-lz4 v1.0.0
+	github.com/klauspost/compress v1.18.0
 	golang.org/x/crypto v0.0.0-20210513164829-c07d793c2f9a
 )
 

go.sum

@@ -2,6 +2,8 @@ github.com/anaminus/parse v0.2.0 h1:a8IEzp/INmUHVfZp4zaeggLvgqIjHrBFrBhkM/qdKUM=
 github.com/anaminus/parse v0.2.0/go.mod h1:5EP2T2CqY4EDBxl2S/qhFSSb2VjW+iIqG+8Scy3mZN8=
 github.com/bkaradzic/go-lz4 v1.0.0 h1:RXc4wYsyz985CkXXeX04y4VnZFGG8Rd43pRaHsOXAKk=
 github.com/bkaradzic/go-lz4 v1.0.0/go.mod h1:0YdlkowM3VswSROI7qDxhRvJ3sLhlFrRRwjwegp5jy4=
+github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo=
+github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ=
 golang.org/x/crypto v0.0.0-20210513164829-c07d793c2f9a h1:kr2P4QFmQr29mSLA43kwrOcgcReGTfbE9N577tCTuBc=
 golang.org/x/crypto v0.0.0-20210513164829-c07d793c2f9a/go.mod h1:P+XmwS30IXTQdn5tA2iutPOUgjI07+tq3H3K9MVA1s8=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68 h1:nxC68pudNYkKU6jWhgrqdreuFiOQWj1Fs7T3VrH4Pjw=

rbxl/model.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/anaminus/parse"
 	"github.com/bkaradzic/go-lz4"
+	"github.com/klauspost/compress/zstd"
 )
 
 ////////////////////////////////////////////////////////////////
@@ -21,6 +22,9 @@ const binaryMarker = "!"
 // binaryHeader is the header magic of a binary file.
 const binaryHeader = "\x89\xff\r\n\x1a\n"
 
+// zstdHeader is the header magic of zstd compressed data.
+const zstdHeader = "\x28\xB5\x2F\xFD"
+
 ////////////////////////////////////////////////////////////////
 
 func readString(f *parse.BinaryReader, data *string) (failed bool) {
@@ -143,6 +147,8 @@ type rawChunk struct {
 	signature uint32
 	compressed
 	payload []byte
+	isLZ4   bool
+	isZSTD  bool
 }
 
 func (c rawChunk) Signature() sig {
@@ -180,22 +186,50 @@ func (c *rawChunk) Decode(fr *parse.BinaryReader) bool {
 	} else {
 		c.compressed = true
 
-		// Prepare compressed data for reading by lz4, which requires the
-		// uncompressed length before the compressed data.
-		compressedData := make([]byte, compressedLength+4)
-		binary.LittleEndian.PutUint32(compressedData, decompressedLength)
-
-		if fr.Bytes(compressedData[4:]) {
+		// Prepare compressed data for reading
+		compressedData := make([]byte, compressedLength)
+		if fr.Bytes(compressedData) {
 			return true
 		}
 
-		// ROBLOX ERROR: "Malformed data ([true decompressed length] != [given
-		// decompressed length])". lz4 already does some kind of size
-		// validation, though the error message isn't the same.
+		// Check for ZSTD header
+		if binary.LittleEndian.Uint32(compressedData[:4]) == binary.LittleEndian.Uint32([]byte(zstdHeader)) {
+			c.isZSTD = true
+			// Decompress using ZSTD
+			reader, err := zstd.NewReader(nil)
+			if err != nil {
+				fr.Add(0, fmt.Errorf("zstd: could not create reader: %w", err))
+				return true
+			}
+
+			defer reader.Close()
+			decompressedData, err := reader.DecodeAll(compressedData, make([]byte, 0, decompressedLength))
+
+			if err != nil {
+				fr.Add(0, fmt.Errorf("zstd: %w", err))
+				return true
+			} else if uint32(len(decompressedData)) != decompressedLength {
+				fr.Add(0, fmt.Errorf("zstd: decompressed length mismatch"))
+				return true
+			}
+
+			c.payload = decompressedData
+		} else {
+			c.isLZ4 = true
+			// Prepare compressed data for reading by lz4, which requires the
+			// uncompressed length before the compressed data.
+			lz4Data := make([]byte, compressedLength+4)
+			binary.LittleEndian.PutUint32(lz4Data, decompressedLength)
+			copy(lz4Data[4:], compressedData)
 
-		if _, err := lz4.Decode(c.payload, compressedData); err != nil {
-			fr.Add(0, fmt.Errorf("lz4: %s", err.Error()))
-			return true
+			// ROBLOX ERROR: "Malformed data ([true decompressed length] != [given
+			// decompressed length])". lz4 already does some kind of size
+			// validation, though the error message isn't the same.
+
+			if _, err := lz4.Decode(c.payload, lz4Data); err != nil {
+				fr.Add(0, fmt.Errorf("lz4: %s", err.Error()))
+				return true
+			}
 		}
 	}
 
@@ -210,21 +244,31 @@ func (c *rawChunk) WriteTo(fw *parse.BinaryWriter) bool {
 
 	if c.compressed {
 		var compressedData []byte
-		compressedData, err := lz4.Encode(compressedData, c.payload)
-		if fw.Add(0, err) {
-			return true
-		}
+		var err error
+		if c.isZSTD {
+			encoder, err := zstd.NewWriter(nil)
+			if err != nil {
+				return true
+			}
+			compressedData = encoder.EncodeAll(c.payload, make([]byte, 0, len(c.payload)))
+			err = encoder.Close()
+			if err != nil {
+				return true
+			}
+		} else {
+			compressedData, err = lz4.Encode(compressedData, c.payload)
+			if fw.Add(0, err) {
+				return true
+			}
 
-		// lz4 sanity check
-		if binary.LittleEndian.Uint32(compressedData[:4]) != uint32(len(c.payload)) {
-			panic("lz4 uncompressed length does not match payload length")
+			// lz4 sanity check
+			if binary.LittleEndian.Uint32(compressedData[:4]) != uint32(len(c.payload)) {
+				panic("lz4 uncompressed length does not match payload length")
+			}
+			compressedData = compressedData[4:]
 		}
 
-		// Compressed length; lz4 prepends the length of the uncompressed
-		// payload, so it must be excluded.
-		compressedPayload := compressedData[4:]
-
-		if fw.Number(uint32(len(compressedPayload))) {
+		if fw.Number(uint32(len(compressedData))) {
 			return true
 		}
 
@@ -238,7 +282,7 @@ func (c *rawChunk) WriteTo(fw *parse.BinaryWriter) bool {
 			return true
 		}
 
-		if fw.Bytes(compressedPayload) {
+		if fw.Bytes(compressedData) {
 			return true
 		}
 	} else {