1 files changed, 200 insertions, 0 deletions
diff --git a/fs/ntfs/lib/decompress_common.c b/fs/ntfs/lib/decompress_common.c
new file mode 100644
index 000000000000..1705face42a3
--- /dev/null
+++ b/fs/ntfs/lib/decompress_common.c
@@ -0,0 +1,200 @@
+// SPDX-License-Identifier: MIT
+/*
+ * decompress_common.c - Code shared by the XPRESS and LZX decompressors
+ *
+ * This is a port of the upstream wimlib "decompress_common.c" which builds
+ * subtable-based Huffman decode tables, as opposed to the older
+ * binary-tree-based format previously used in this library.  The vectorized
+ * (SSE2/AVX2) fill paths are omitted for portability in the kernel.
+ *
+ * Copyright (C) 2022 Eric Biggers
+ */
+
+#include "decompress_common.h"
+
+/* Compute the number of bits with which a subtable must be indexed for a
+ * codeword of length @codeword_len, given that the root table is indexed with
+ * @table_bits bits.
+ */
+static u32 compute_subtable_bits(u32 table_bits,
+				 u32 codeword_len, u16 len_counts[])
+{
+	u32 subtable_bits = codeword_len - table_bits;
+	s32 remainder = (s32)1 << subtable_bits;
+
+	for (;;) {
+		remainder -= len_counts[table_bits + subtable_bits];
+		if (remainder <= 0)
+			break;
+		subtable_bits++;
+		remainder <<= 1;
+	}
+	return subtable_bits;
+}
+
+/* Build the subtables for codewords longer than table_bits. */
+static int build_subtables(u16 decode_table[], u32 num_syms, u32 table_bits,
+			   u16 len_counts[], const u16 sorted_syms[], u32 sym_idx,
+			   u32 decode_table_pos, u32 decode_table_size)
+{
+	u32 subtable_pos = 1U << table_bits;
+	u32 subtable_bits = table_bits;
+	u32 subtable_prefix = (u32)-1;
+	u32 codeword_len = table_bits + 1;
+	u32 codeword = decode_table_pos << 1;
+	u32 prefix;
+	u16 entry;
+	u32 n;
+
+	for (; sym_idx < num_syms; sym_idx++) {
+		while (len_counts[codeword_len] == 0) {
+			codeword_len++;
+			codeword <<= 1;
+		}
+
+		prefix = codeword >> (codeword_len - table_bits);
+
+		if (prefix != subtable_prefix) {
+			subtable_prefix = prefix;
+			subtable_bits = compute_subtable_bits(table_bits, codeword_len,
+							      len_counts);
+			decode_table[subtable_prefix] =
+				MAKE_DECODE_TABLE_ENTRY(subtable_pos, subtable_bits);
+		}
+
+		entry = MAKE_DECODE_TABLE_ENTRY(sorted_syms[sym_idx],
+						codeword_len - table_bits);
+		n = 1U << (subtable_bits - (codeword_len - table_bits));
+
+		/* Defensive bound check: 'lens' is derived from untrusted
+		 * on-disk compressed data, and subtable growth depends on
+		 * its content.  This should never trigger for a correctly
+		 * sized DECODE_TABLE_ENOUGH() value, but turns a wrong value
+		 * into a clean decode failure instead of writing past the
+		 * caller's decode_table[].
+		 */
+		if (unlikely(subtable_pos + n > decode_table_size))
+			return -1;
+
+		do {
+			decode_table[subtable_pos++] = entry;
+		} while (--n);
+
+		len_counts[codeword_len]--;
+		codeword++;
+	}
+
+	return 0;
+}
+
+/*
+ * Given an alphabet of symbols and the length of each symbol's codeword in a
+ * canonical prefix code, build a table for quickly decoding symbols that were
+ * encoded with that code.
+ *
+ * The root table is indexed with 'table_bits' bits.  Codewords not longer than
+ * 'table_bits' are decoded directly from the root table.  Longer codewords are
+ * decoded via subtables: the corresponding root entry is a pointer (the index
+ * of the subtable plus the number of bits with which the subtable is indexed),
+ * and the subtable is indexed with the remaining bits of the codeword.
+ *
+ * Each entry stores both the symbol (high 12 bits) and the codeword length (low
+ * 4 bits), so a single lookup yields the symbol and lets the bitstream be
+ * advanced by the correct number of bits.
+ *
+ * @decode_table:  array in which to build the table (declared with
+ *		   DECODE_TABLE()).  May alias @lens.
+ * @num_syms:      number of symbols in the alphabet.
+ * @table_bits:    log2 of the number of root table entries.
+ * @lens:         array of @num_syms codeword lengths, indexed by symbol.
+ * @max_codeword_len: longest codeword length allowed for this code.
+ * @working_space: temporary array declared with DECODE_TABLE_WORKING_SPACE().
+ * @decode_table_size: number of u16 entries in @decode_table (i.e.
+ *		   ARRAY_SIZE(decode_table) at the call site).  Used only as a
+ *		   defensive bound check against @lens-dependent subtable growth.
+ *
+ * Returns 0 on success, or -1 if the lengths do not form a valid prefix code,
+ * or if building the subtables would overflow @decode_table_size entries.
+ */
+int make_huffman_decode_table(u16 decode_table[], u32 num_syms, u32 table_bits,
+			      const u8 lens[], u32 max_codeword_len,
+			      u16 working_space[], u32 decode_table_size)
+{
+	u16 *const len_counts = &working_space[0];
+	u16 *const offsets = &working_space[1 * (max_codeword_len + 1)];
+	u16 *const sorted_syms = &working_space[2 * (max_codeword_len + 1)];
+	u32 decode_table_pos = 0;
+	u32 sym_idx;
+	u32 codeword_len;
+	s32 remainder = 1;
+	void *entry_ptr = decode_table;
+	u32 len;
+	u32 sym;
+
+	/* Count how many codewords have each length, including 0. */
+	for (len = 0; len <= max_codeword_len; len++)
+		len_counts[len] = 0;
+	for (sym = 0; sym < num_syms; sym++)
+		len_counts[lens[sym]]++;
+
+	/* A codeword of length n should require a proportion of the codespace
+	 * equaling (1/2)^n.  The code is complete iff the codespace is exactly
+	 * filled by the lengths.
+	 */
+	for (len = 1; len <= max_codeword_len; len++) {
+		remainder = (remainder << 1) - len_counts[len];
+		if (unlikely(remainder < 0))
+			return -1;	/* over-subscribed */
+	}
+
+	if (remainder != 0) {
+		/* Incomplete code.  Permitted only if the code is empty. */
+		if (unlikely(remainder != (s32)(1U << max_codeword_len)))
+			return -1;
+
+		/* Empty code: zero the root table so lookups yield symbol 0
+		 * without consuming any bits.
+		 */
+		memset(decode_table, 0, sizeof(decode_table[0]) << table_bits);
+		return 0;
+	}
+
+	/* Sort the symbols primarily by increasing codeword length and
+	 * secondarily by increasing symbol value.
+	 */
+	offsets[0] = 0;
+	for (len = 0; len < max_codeword_len; len++)
+		offsets[len + 1] = offsets[len] + len_counts[len];
+	for (sym = 0; sym < num_syms; sym++)
+		sorted_syms[offsets[lens[sym]]++] = sym;
+
+	/* Fill the root table entries for codewords no longer than table_bits. */
+	sym_idx = offsets[0];
+	codeword_len = 1;
+	for (; codeword_len <= table_bits; codeword_len++) {
+		u32 stores_per_loop = 1U << (table_bits - codeword_len);
+		u32 end_sym_idx = sym_idx + len_counts[codeword_len];
+
+		for (; sym_idx < end_sym_idx; sym_idx++) {
+			u16 v = MAKE_DECODE_TABLE_ENTRY(sorted_syms[sym_idx],
+							codeword_len);
+			u32 n = stores_per_loop;
+			u16 *p = entry_ptr;
+
+			do {
+				*p++ = v;
+			} while (--n);
+			entry_ptr = p;
+		}
+	}
+	decode_table_pos = (u16 *)entry_ptr - decode_table;
+
+	/* If all symbols were processed, no subtables are required. */
+	if (sym_idx == num_syms)
+		return 0;
+
+	/* At least one subtable is required.  Process the remaining symbols. */
+	return build_subtables(decode_table, num_syms, table_bits, len_counts,
+			       sorted_syms, sym_idx, decode_table_pos,
+			       decode_table_size);
+}