Fix AMD64 (non-bmi) emitting junk values when overreading.

klauspost · Jul 20, 2022 · 5357b67 · 5357b67
1 parent 1ff6043
commit 5357b67
Show file tree

Hide file tree

Showing 5 changed files with 440 additions and 351 deletions.
diff --git a/zstd/_generate/gen.go b/zstd/_generate/gen.go
@@ -594,16 +594,25 @@ func (o options) updateLength(name string, brValue, brBitsRead, state reg.GPVirt
 		MOVQ(state, AX.As64()) // So we can grab high bytes.
 		MOVQ(brBitsRead, CX.As64())
 		MOVQ(brValue, BX)
-		SHLQ(CX, BX)                // BX = br.value << br.bitsRead (part of getBits)
-		MOVB(AX.As8H(), CX.As8L())  // CX = moB  (ofState.addBits(), that is byte #1 of moState)
-		ADDQ(CX.As64(), brBitsRead) // br.bitsRead += n (part of getBits)
-		NEGL(CX.As32())             // CX = 64 - n
-		SHRQ(CX, BX)                // BX = (br.value << br.bitsRead) >> (64 - n) -- getBits() result
-		SHRQ(U8(32), AX)            // AX = mo (ofState.baselineInt(), that's the higher dword of moState)
+		SHLQ(CX, BX)               // BX = br.value << br.bitsRead (part of getBits)
+		MOVB(AX.As8H(), CX.As8L()) // CX = moB  (ofState.addBits(), that is byte #1 of moState)
+		SHRQ(U8(32), AX)           // AX = mo (ofState.baselineInt(), that's the higher dword of moState)
+		// If addBits == 0, skip
 		TESTQ(CX.As64(), CX.As64())
-		CMOVQEQ(CX.As64(), BX) // BX is zero if n is zero
+		JZ(LabelRef(name + "_zero"))
+
+		ADDQ(CX.As64(), brBitsRead) // br.bitsRead += n (part of getBits)
+		// If overread, skip
+		CMPQ(brBitsRead, U8(64))
+		JA(LabelRef(name + "_zero"))
+		CMPQ(CX.As64(), U8(64))
+		JAE(LabelRef(name + "_zero"))
 
-		ADDQ(BX, AX)  // AX - mo + br.getBits(moB)
+		NEGQ(CX.As64()) // CX = 64 - n
+		SHRQ(CX, BX)    // BX = (br.value << br.bitsRead) >> (64 - n) -- getBits() result
+		ADDQ(BX, AX)    // AX - mo + br.getBits(moB)
+
+		Label(name + "_zero")
 		MOVQ(AX, out) // Store result
 	}
 }

diff --git a/zstd/fse_decoder_amd64.go b/zstd/fse_decoder_amd64.go
@@ -34,8 +34,8 @@ const (
 // buildDtable will build the decoding table.
 func (s *fseDecoder) buildDtable() error {
 	ctx := buildDtableAsmContext{
-		stateTable: (*uint16)(&s.stateTable[0]),
-		norm:       (*int16)(&s.norm[0]),
+		stateTable: &s.stateTable[0],
+		norm:       &s.norm[0],
 		dt:         (*uint64)(&s.dt[0]),
 	}
 	code := buildDtable_asm(s, &ctx)

diff --git a/zstd/fuzz_test.go b/zstd/fuzz_test.go
@@ -21,7 +21,6 @@ import (
 )
 
 func FuzzDecodeAll(f *testing.F) {
-	defer cpuinfo.DisableBMI2()()
 	addBytesFromZip(f, "testdata/fuzz/decode-corpus-raw.zip", true)
 	addBytesFromZip(f, "testdata/fuzz/decode-corpus-encoded.zip", false)
 	decLow, err := NewReader(nil, WithDecoderLowmem(true), WithDecoderConcurrency(2), WithDecoderMaxMemory(20<<20), WithDecoderMaxWindow(1<<20), IgnoreChecksum(true))
@@ -57,6 +56,15 @@ func FuzzDecodeAll(f *testing.F) {
 	})
 }
 
+func FuzzDecodeAllNoBMI2(f *testing.F) {
+	if !cpuinfo.HasBMI2() {
+		f.Skip("No BMI, so already tested")
+		return
+	}
+	defer cpuinfo.DisableBMI2()()
+	FuzzDecodeAll(f)
+}
+
 func FuzzDecoder(f *testing.F) {
 	addBytesFromZip(f, "testdata/fuzz/decode-corpus-raw.zip", true)
 	addBytesFromZip(f, "testdata/fuzz/decode-corpus-encoded.zip", false)