internal/lz4block: Detect match with no space for offset

This corner case wasn't detected by any of the decoders.

internal/lz4block/decode_amd64.s

@@ -7,8 +7,8 @@
 
 // AX scratch
 // BX scratch
-// CX scratch
-// DX token
+// CX literal and match lengths
+// DX token, match offset
 //
 // DI &dst
 // SI &src
@@ -44,11 +44,9 @@ TEXT ·decodeBlock(SB), NOSPLIT, $48-80
 	MOVQ R9, R13
 	SUBQ $16, R13
 
-loop:
-	// for si < len(src)
-	CMPQ SI, R9
-	JAE end
+	XORL CX, CX
 
+loop:
 	// token := uint32(src[si])
 	MOVBLZX (SI), DX
 	INCQ SI
@@ -70,7 +68,7 @@ loop:
 	// copy shortcut
 
 	// A two-stage shortcut for the most common case:
-	// 1) If the literal length is 1..14, and there is enough space,
+	// 1) If the literal length is 0..14, and there is enough space,
 	// enter the shortcut and copy 16 bytes on behalf of the literals
 	// (in the fast mode, only 8 bytes can be safely copied this way).
 	// 2) Further if the match length is 4..18, copy 18 bytes in a similar
@@ -122,7 +120,7 @@ loop:
 	LEAQ const_minMatch(DI)(CX*1), DI
 
 	// shortcut complete, load next token
-	JMP loop
+	JMP loopcheck
 
 	// Read the rest of the literal length:
 	// do { BX = src[si++]; lit_len += BX } while (BX == 0xFF).
@@ -213,13 +211,13 @@ memmove_lit:
 	SUBQ $16, R13
 
 finish_lit_copy:
-	CMPQ SI, R9
-	JAE end
-
-offset:
 	// CX := mLen
 	// free up DX to use for offset
-	MOVQ DX, CX
+	MOVL DX, CX
+	ANDL $0xF, CX
+
+	CMPQ SI, R9
+	JAE end
 
 	// offset
 	// si += 2
@@ -231,10 +229,8 @@ offset:
 	MOVWQZX -2(SI), DX
 
 	// 0 offset is invalid
-	CMPQ DX, $0
-	JEQ err_corrupt
-
-	ANDB $0xF, CX
+	TESTL DX, DX
+	JEQ   err_corrupt
 
 match_len_loop_pre:
 	// if mlen != 0xF
@@ -301,7 +297,7 @@ copy_match_loop:
 	DECQ CX
 	JNZ copy_match_loop
 
-	JMP loop
+	JMP loopcheck
 
 copy_interior_match:
 	CMPQ CX, $16
@@ -317,7 +313,8 @@ copy_interior_match:
 	MOVOU X0, (DI)
 
 	ADDQ CX, DI
-	JMP loop
+	XORL CX, CX
+	JMP  loopcheck
 
 copy_match_from_dict:
 	// CX = match_len
@@ -420,8 +417,21 @@ memmove_match:
 	SUBQ $16, R13
 	MOVQ dict_base+48(FP), R14
 	MOVQ dict_len+56(FP), R15
+	XORL CX, CX
 
-	JMP loop
+loopcheck:
+	// for si < len(src)
+	CMPQ SI, R9
+	JB   loop
+
+end:
+	// Remaining length must be zero.
+	TESTQ CX, CX
+	JNE   err_corrupt
+
+	SUBQ R11, DI
+	MOVQ DI, ret+72(FP)
+	RET
 
 err_corrupt:
 	MOVQ $-1, ret+72(FP)
@@ -434,8 +444,3 @@ err_short_buf:
 err_short_dict:
 	MOVQ $-3, ret+72(FP)
 	RET
-
-end:
-	SUBQ R11, DI
-	MOVQ DI, ret+72(FP)
-	RET

internal/lz4block/decode_arm.s

@@ -110,13 +110,13 @@ copyLiteralFinish:
 	MOVB.NE   tmp2, -1(dst)
 
 copyLiteralDone:
-	CMP src, srcend
-	BEQ end
-
 	// Initial part of match length.
 	// This frees up the token register for reuse as offset.
 	AND $15, token, len
 
+	CMP src, srcend
+	BEQ end
+
 	// Read offset.
 	ADD.S $2, src
 	BCS   shortSrc
@@ -213,6 +213,8 @@ copyMatchDone:
 	BNE loop
 
 end:
+	CMP  $0, len
+	BNE  corrupt
 	SUB  dstorig, dst, tmp1
 	MOVW tmp1, ret+36(FP)
 	RET

internal/lz4block/decode_arm64.s

@@ -118,6 +118,9 @@ copyLiteralShortEnd:
 	MOVB.P  tmp4, 1(dst)
 
 copyLiteralDone:
+	// Initial part of match length.
+	AND $15, token, len
+
 	CMP src, srcend
 	BEQ end
 
@@ -129,8 +132,7 @@ copyLiteralDone:
 	MOVHU -2(src), offset
 	CBZ   offset, corrupt
 
-	// Read match length.
-	AND $15, token, len
+	// Read rest of match length.
 	CMP $15, len
 	BNE readMatchlenDone
 
@@ -195,6 +197,7 @@ copyMatchLoop8:
 
 	MOVD (match)(len), tmp2 // match+len == match+lenRem-8.
 	ADD  lenRem, dst
+	MOVD $0, len
 	MOVD tmp2, -8(dst)
 	B    copyMatchDone
 
@@ -210,6 +213,7 @@ copyMatchDone:
 	BNE loop
 
 end:
+	CBNZ len, corrupt
 	SUB  dstorig, dst, tmp1
 	MOVD tmp1, ret+72(FP)
 	RET

internal/lz4block/decode_other.go

@@ -73,9 +73,11 @@ func decodeBlock(dst, src, dict []byte) (ret int) {
 				di += lLen
 			}
 		}
-		if si == uint(len(src)) {
+
+		mLen := b & 0xF
+		if si == uint(len(src)) && mLen == 0 {
 			break
-		} else if si > uint(len(src)) {
+		} else if si >= uint(len(src)) {
 			return hasError
 		}
 
@@ -86,7 +88,7 @@ func decodeBlock(dst, src, dict []byte) (ret int) {
 		si += 2
 
 		// Match.
-		mLen := minMatch + b&0xF
+		mLen += minMatch
 		if mLen == minMatch+0xF {
 			for {
 				x := uint(src[si])

internal/lz4block/decode_test.go

@@ -100,6 +100,11 @@ func TestBlockDecode(t *testing.T) {
 			emitSeq(strings.Repeat("A", 15), 0, 0),
 			bytes.Repeat([]byte("A"), 15),
 		},
+		{
+			"literal_only_long_2",
+			emitSeq(strings.Repeat("A", 16), 0, 0),
+			bytes.Repeat([]byte("A"), 16),
+		},
 		{
 			"repeat_match_len",
 			emitSeq("a", 1, 4),
@@ -115,6 +120,13 @@ func TestBlockDecode(t *testing.T) {
 			emitSeq("A", 1, 16),
 			bytes.Repeat([]byte("A"), 17),
 		},
+		{
+			// Triggers a case in the amd64 decoder where its last action
+			// is a call to runtime.memmove.
+			"memmove_last",
+			emitSeq(strings.Repeat("ABCD", 20), 80, 60),
+			bytes.Repeat([]byte("ABCD"), 36)[:140],
+		},
 		{
 			"repeat_match_log_len_2_seq",
 			concat(emitSeq("a", 1, 15), emitSeq("B", 1, 15), emitSeq("end", 0, 0)),
@@ -171,6 +183,11 @@ func TestDecodeBlockInvalid(t *testing.T) {
 			"\x20a", // litlen = 2 but only a single-byte literal
 			100,
 		},
+		{
+			"no_space_for_offset",
+			"\x01", // mLen > 0 but no following offset.
+			10,
+		},
 		{
 			"write_beyond_len_dst",
 			"\x1b0\x01\x00000000000000",