From 870eab737543387a11d1673654f93610ef3a402e Mon Sep 17 00:00:00 2001
From: da2018 <da2018t@gmail.com>
Date: Wed, 10 Mar 2021 13:42:20 +0800
Subject: [PATCH] Enhance Zstd support

---
 fixtures/sample_skippable.zst | Bin 0 -> 105 bytes
 match_test.go                 |   5 +++++
 matchers/archive.go           |  32 +++++++++++++++++++++++++++++++-
 3 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 fixtures/sample_skippable.zst

diff --git a/fixtures/sample_skippable.zst b/fixtures/sample_skippable.zst
new file mode 100644
index 0000000000000000000000000000000000000000..c8f7d95a8461c9074a319454da4a2be509c4c120
GIT binary patch
literal 105
zcmWI0@|9p=U|@&^VvVi(e_0}0nHWNtOBf7*z|73V+=#)@#K6eNz{JGV%#gvr5C}{d
z6bu+ji%SxV6ciYW6Z5i56EhhMxfnPZezAUVW)|UTWoC$Bn8236Hf!yc-|7vlQo_HB
E0gbU4cK`qY

literal 0
HcmV?d00001

diff --git a/match_test.go b/match_test.go
index 1db8857..e17c37c 100644
--- a/match_test.go
+++ b/match_test.go
@@ -60,6 +60,11 @@ func TestMatchFile(t *testing.T) {
 			t.Fatalf("Invalid type: %s != %s", kind.Extension, test.ext)
 		}
 	}
+	// test zstd with skippable frame
+	kind, _ := MatchFile("./fixtures/sample_skippable.zst")
+	if kind.Extension != "zst" {
+		t.Fatalf("Invalid type: %s != %s", kind.Extension, "zst")
+	}
 }
 
 func TestMatchReader(t *testing.T) {
diff --git a/matchers/archive.go b/matchers/archive.go
index f3a18f0..978796f 100644
--- a/matchers/archive.go
+++ b/matchers/archive.go
@@ -1,5 +1,12 @@
 package matchers
 
+import "encoding/binary"
+
+const (
+	ZstdMagicSkippableStart = 0x184D2A50
+	ZstdMagicSkippableMask  = 0xFFFFFFF0
+)
+
 var (
 	TypeEpub   = newType("epub", "application/epub+zip")
 	TypeZip    = newType("zip", "application/zip")
@@ -40,7 +47,7 @@ var Archive = Map{
 	TypeBz2:    bytePrefixMatcher(bz2Magic),
 	Type7z:     bytePrefixMatcher(sevenzMagic),
 	TypeXz:     bytePrefixMatcher(xzMagic),
-	TypeZstd:   bytePrefixMatcher(zstdMagic),
+	TypeZstd:   Zst,
 	TypePdf:    bytePrefixMatcher(pdfMagic),
 	TypeExe:    bytePrefixMatcher(exeMagic),
 	TypeSwf:    Swf,
@@ -179,3 +186,26 @@ func MachO(buf []byte) bool {
 		(buf[0] == 0xCE && buf[1] == 0xFA && buf[2] == 0xED && buf[3] == 0xFE) ||
 		(buf[0] == 0xCA && buf[1] == 0xFE && buf[2] == 0xBA && buf[3] == 0xBE))
 }
+
+// Zstandard compressed data is made of one or more frames.
+// There are two frame formats defined by Zstandard: Zstandard frames and Skippable frames.
+// See more details from https://tools.ietf.org/id/draft-kucherawy-dispatch-zstd-00.html#rfc.section.2
+func Zst(buf []byte) bool {
+	if compareBytes(buf, zstdMagic, 0) {
+		return true
+	} else {
+		// skippable frames
+    if len(buf) < 8 {
+      return false
+    }
+		if binary.LittleEndian.Uint32(buf[:4]) & ZstdMagicSkippableMask == ZstdMagicSkippableStart {
+			userDataLength := binary.LittleEndian.Uint32(buf[4:8])
+			if len(buf) < 8 + int(userDataLength) {
+			  return false
+      }
+			nextFrame := buf[8+userDataLength:]
+			return Zst(nextFrame)
+		}
+		return false
+	}
+}