open-standards/wai/live-conformance/_rfc8878_zstd.py

"""A Zstandard decoder written from RFC 8878, in pure Python, for the prior-carriage verifier.

It decodes one frame against a raw-content dictionary (RFC 8878 §5: the dictionary's bytes are
history before the frame's first byte, and the repeat offsets start at 1, 4 and 8), checks the
frame's content size and its XXH64 checksum when they are stated, and refuses anything it does
not read with `ZstdError`. It shares no code with wai-rs, which decodes with the Zstandard
library. It is written for the corpus's small frames, not for speed.
"""


class ZstdError(Exception):
    pass


MAGIC = 0xFD2FB528
MAX_BLOCK = 128 * 1024


def frame_window(frame):
    """(Window_Size, header length, content size or None, checksum flag) of a frame header."""
    if len(frame) < 6 or int.from_bytes(frame[:4], "little") != MAGIC:
        raise ZstdError("no frame magic")
    fhd = frame[4]
    if fhd & 0x08:
        raise ZstdError("reserved bit set")
    single = bool(fhd & 0x20)
    did = (0, 1, 2, 4)[fhd & 3]
    fcs = {0: 1 if single else 0, 1: 2, 2: 4, 3: 8}[fhd >> 6]
    at = 5
    window = None
    if not single:
        wd = frame[at]
        at += 1
        base = 1 << (10 + (wd >> 3))
        window = base + (base // 8) * (wd & 7)
    at += did
    if len(frame) < at + fcs:
        raise ZstdError("header truncated")
    size = int.from_bytes(frame[at:at + fcs], "little") if fcs else None
    if fcs == 2:
        size += 256
    at += fcs
    if window is None:
        window = size
    return window, at, size, bool(fhd & 0x04)


# --- bit readers --------------------------------------------------------------------------------

class Forward:
    """Bits read from the least significant bit of each byte, first byte first (§4.1.1)."""

    def __init__(self, data):
        self.v = int.from_bytes(data, "little")
        self.n = len(data) * 8
        self.pos = 0

    def read(self, k):
        if self.pos + k > self.n:
            raise ZstdError("table description past its end")
        x = (self.v >> self.pos) & ((1 << k) - 1)
        self.pos += k
        return x

    def peek(self, k):
        return (self.v >> self.pos) & ((1 << k) - 1)


class Backward:
    """A bitstream read from its end (§4.1): the highest set bit of the last byte ends it."""

    def __init__(self, data):
        if not data or data[-1] == 0:
            raise ZstdError("bitstream without its end mark")
        self.v = int.from_bytes(data, "little")
        self.left = self.v.bit_length() - 1

    def read(self, k):
        if k == 0:
            return 0
        self.left -= k
        if self.left >= 0:
            return (self.v >> self.left) & ((1 << k) - 1)
        return (self.v << -self.left) & ((1 << k) - 1)

    def peek(self, k):
        if self.left >= k:
            return (self.v >> (self.left - k)) & ((1 << k) - 1)
        return (self.v << (k - self.left)) & ((1 << k) - 1)

    def overflowed(self):
        return self.left < 0


# --- FSE (§4.1) ---------------------------------------------------------------------------------

def read_distribution(br, max_symbol, max_log):
    """An FSE table description: (accuracy log, normalised counts)."""
    log = br.read(4) + 5
    if log > max_log:
        raise ZstdError("accuracy log too large")
    remaining = (1 << log) + 1
    threshold = 1 << log
    bits = log + 1
    counts = []
    previous_zero = False
    while remaining > 1 and len(counts) <= max_symbol:
        if previous_zero:
            while True:
                flag = br.read(2)
                counts.extend([0] * flag)
                if flag != 3:
                    break
            if len(counts) > max_symbol:
                break
        top = (2 * threshold - 1) - remaining
        low = br.peek(bits - 1)
        if low < top:
            value = low
            br.read(bits - 1)
        else:
            value = br.peek(bits)
            if value >= threshold:
                value -= top
            br.read(bits)
        count = value - 1
        remaining -= -count if count < 0 else count
        counts.append(count)
        previous_zero = count == 0
        while remaining < threshold:
            bits -= 1
            threshold >>= 1
    if remaining != 1 or len(counts) > max_symbol + 1:
        raise ZstdError("distribution does not sum")
    return log, counts


def build_table(log, counts):
    """The decoding table: per state, (symbol, bits, baseline)."""
    size = 1 << log
    symbol = [0] * size
    high = size - 1
    for s, c in enumerate(counts):
        if c == -1:
            symbol[high] = s
            high -= 1
    step = (size >> 1) + (size >> 3) + 3
    pos = 0
    for s, c in enumerate(counts):
        for _ in range(max(c, 0)):
            symbol[pos] = s
            pos = (pos + step) & (size - 1)
            while pos > high:
                pos = (pos + step) & (size - 1)
    if pos != 0:
        raise ZstdError("distribution does not spread")
    nxt = [1 if c == -1 else c for c in counts]
    table = []
    for u in range(size):
        s = symbol[u]
        x = nxt[s]
        nxt[s] += 1
        nb = log - (x.bit_length() - 1)
        table.append((s, nb, (x << nb) - size))
    return table


# --- Huffman literals (§4.2) --------------------------------------------------------------------

def huffman_weights(data):
    """(table, max bits, bytes read) from a Huffman tree description."""
    if not data:
        raise ZstdError("no tree description")
    head = data[0]
    if head < 128:
        stream = data[1:1 + head]
        if len(stream) != head:
            raise ZstdError("tree description truncated")
        fb = Forward(stream)
        log, counts = read_distribution(fb, 255, 6)
        table = build_table(log, counts)
        used = (fb.pos + 7) // 8
        bb = Backward(stream[used:])
        s1, s2 = bb.read(log), bb.read(log)
        weights = []
        # Two states, alternating, until a state update reads past the stream's start; the
        # other state then gives the last weight (§4.2.1.2).
        while True:
            sym, nb, base = table[s1]
            weights.append(sym)
            s1 = base + bb.read(nb)
            if bb.overflowed():
                weights.append(table[s2][0])
                break
            sym, nb, base = table[s2]
            weights.append(sym)
            s2 = base + bb.read(nb)
            if bb.overflowed():
                weights.append(table[s1][0])
                break
            if len(weights) > 255:
                raise ZstdError("too many weights")
        used = 1 + head
    else:
        n = head - 127
        used = 1 + (n + 1) // 2
        if len(data) < used:
            raise ZstdError("tree description truncated")
        weights = []
        for i in range(n):
            byte = data[1 + i // 2]
            weights.append(byte >> 4 if i % 2 == 0 else byte & 15)
    if any(w > 11 for w in weights):
        raise ZstdError("weight too large")
    total = sum(1 << (w - 1) for w in weights if w)
    if total == 0:
        raise ZstdError("no weights")
    maxbits = total.bit_length()
    rest = (1 << maxbits) - total
    if rest & (rest - 1):
        raise ZstdError("weights do not complete a tree")
    weights.append(rest.bit_length())
    if maxbits > 11:
        raise ZstdError("code too long")
    table = [None] * (1 << maxbits)
    pos = 0
    for w in range(1, 12):
        for s, ws in enumerate(weights):
            if ws == w:
                nb = maxbits + 1 - w
                for _ in range(1 << (w - 1)):
                    table[pos] = (s, nb)
                    pos += 1
    if pos != len(table):
        raise ZstdError("weights do not fill the table")
    return (table, maxbits), used


def huffman_stream(data, tree, count):
    table, maxbits = tree
    br = Backward(data)
    out = bytearray()
    for _ in range(count):
        s, nb = table[br.peek(maxbits)]
        br.read(nb)
        out.append(s)
    if br.left != 0:
        raise ZstdError("Huffman stream not consumed exactly")
    return bytes(out)


def literals(block, state):
    """(literal bytes, bytes read) of a Literals_Section (§3.1.1.3.1)."""
    b0 = block[0]
    kind, fmt = b0 & 3, (b0 >> 2) & 3
    if kind in (0, 1):
        if fmt in (0, 2):
            hl, size = 1, b0 >> 3
        elif fmt == 1:
            hl, size = 2, int.from_bytes(block[:2], "little") >> 4
        else:
            hl, size = 3, int.from_bytes(block[:3], "little") >> 4
        if kind == 0:
            if len(block) < hl + size:
                raise ZstdError("raw literals truncated")
            return block[hl:hl + size], hl + size
        if len(block) < hl + 1:
            raise ZstdError("RLE literals truncated")
        return bytes([block[hl]]) * size, hl + 1
    lhc = int.from_bytes(block[:5].ljust(5, b"\0"), "little")
    if fmt in (0, 1):
        hl, size, csize, streams = 3, (lhc >> 4) & 0x3FF, (lhc >> 14) & 0x3FF, 1 if fmt == 0 else 4
    elif fmt == 2:
        hl, size, csize, streams = 4, (lhc >> 4) & 0x3FFF, (lhc >> 18) & 0x3FFF, 4
    else:
        hl, size, csize, streams = 5, (lhc >> 4) & 0x3FFFF, (lhc >> 22) & 0x3FFFF, 4
    body = block[hl:hl + csize]
    if len(body) != csize:
        raise ZstdError("compressed literals truncated")
    at = 0
    if kind == 2:
        state["huffman"], at = huffman_weights(body)
    elif state.get("huffman") is None:
        raise ZstdError("treeless literals with no previous tree")
    tree = state["huffman"]
    if streams == 1:
        out = huffman_stream(body[at:], tree, size)
    else:
        if len(body) < at + 6:
            raise ZstdError("jump table truncated")
        s1, s2, s3 = (int.from_bytes(body[at + 2 * i:at + 2 * i + 2], "little") for i in range(3))
        at += 6
        seg = (size + 3) // 4
        parts = []
        for i, n in enumerate((s1, s2, s3, len(body) - at - s1 - s2 - s3)):
            if n < 0:
                raise ZstdError("jump table past the literals")
            parts.append(huffman_stream(body[at:at + n], tree, seg if i < 3 else size - 3 * seg))
            at += n
        out = b"".join(parts)
    return out, hl + csize


# --- sequences (§3.1.1.3.2) ---------------------------------------------------------------------

LL_BASE = list(range(16)) + [16, 18, 20, 22, 24, 28, 32, 40, 48, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384,
                             32768, 65536]
LL_BITS = [0] * 16 + [1, 1, 1, 1, 2, 2, 3, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
ML_BASE = [n + 3 for n in range(32)] + [35, 37, 39, 41, 43, 47, 51, 59, 67, 83, 99, 131, 259, 515, 1027, 2051,
                                        4099, 8195, 16387, 32771, 65539]
ML_BITS = [0] * 32 + [1, 1, 1, 1, 2, 2, 3, 3, 4, 4, 5, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
LL_DEFAULT = (6, [4, 3, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 2, 1, 1, 1, 1, 1,
                  -1, -1, -1, -1])
ML_DEFAULT = (6, [1, 4, 3, 2, 2, 2, 2, 2, 2] + [1] * 37 + [-1] * 7)
OF_DEFAULT = (5, [1, 1, 1, 1, 1, 1, 2, 2, 2] + [1] * 15 + [-1] * 5)
KINDS = (("ll", 35, 9, LL_DEFAULT), ("of", 31, 8, OF_DEFAULT), ("ml", 52, 9, ML_DEFAULT))


def sequences(block, state):
    """The sequences of a Sequences_Section: [(literals length, offset value, match length)]."""
    if not block:
        raise ZstdError("sequences section missing")
    b0 = block[0]
    if b0 == 0:
        if len(block) != 1:
            raise ZstdError("bytes after an empty sequences section")
        return []
    if b0 < 128:
        n, at = b0, 1
    elif b0 < 255:
        n, at = ((b0 - 128) << 8) + block[1], 2
    else:
        n, at = block[1] + (block[2] << 8) + 0x7F00, 3
    modes = block[at]
    at += 1
    if modes & 3:
        raise ZstdError("reserved bits in the modes")
    tables = {}
    for (name, maxsym, maxlog, default), shift in zip(KINDS, (6, 4, 2)):
        mode = (modes >> shift) & 3
        if mode == 0:
            tables[name] = build_table(*default)
        elif mode == 1:
            if block[at] > maxsym:
                raise ZstdError("RLE symbol out of range")
            tables[name] = [(block[at], 0, 0)]
            at += 1
        elif mode == 2:
            fb = Forward(block[at:])
            log, counts = read_distribution(fb, maxsym, maxlog)
            tables[name] = build_table(log, counts)
            at += (fb.pos + 7) // 8
        else:
            if state.get(name) is None:
                raise ZstdError("repeat mode with no previous table")
            tables[name] = state[name]
        state[name] = tables[name]
    br = Backward(block[at:])

    def log_of(t):
        return len(t).bit_length() - 1

    ll, of, ml = (br.read(log_of(tables[k])) for k in ("ll", "of", "ml"))
    out = []
    for i in range(n):
        ofc = tables["of"][of][0]
        mlc = tables["ml"][ml][0]
        llc = tables["ll"][ll][0]
        if ofc > 31 or mlc > 52 or llc > 35:
            raise ZstdError("code out of range")
        offset = (1 << ofc) + br.read(ofc)
        match = ML_BASE[mlc] + br.read(ML_BITS[mlc])
        lit = LL_BASE[llc] + br.read(LL_BITS[llc])
        out.append((lit, offset, match))
        if i + 1 < n:
            for k in ("ll", "ml", "of"):
                s, nb, base = tables[k][{"ll": ll, "ml": ml, "of": of}[k]]
                v = base + br.read(nb)
                if k == "ll":
                    ll = v
                elif k == "ml":
                    ml = v
                else:
                    of = v
    if br.left != 0:
        raise ZstdError("sequence bitstream not consumed exactly")
    return out


def execute(seqs, lits, history, reps):
    """Appends the block's output to `history` (dictionary and frame so far)."""
    at = 0
    for lit, value, match in seqs:
        if at + lit > len(lits):
            raise ZstdError("literals overrun")
        history += lits[at:at + lit]
        at += lit
        if value > 3:
            offset = value - 3
            reps[:] = [offset, reps[0], reps[1]]
        else:
            idx = value + (1 if lit == 0 else 0)
            if idx == 1:
                offset = reps[0]
            elif idx == 2:
                offset = reps[1]
                reps[:] = [offset, reps[0], reps[2]]
            elif idx == 3:
                offset = reps[2]
                reps[:] = [offset, reps[0], reps[1]]
            else:
                offset = reps[0] - 1
                reps[:] = [offset, reps[0], reps[1]]
        if offset <= 0 or offset > len(history):
            raise ZstdError("offset past the history")
        start = len(history) - offset
        for j in range(match):
            history.append(history[start + j])
    history += lits[at:]


# --- XXH64 --------------------------------------------------------------------------------------

M64 = (1 << 64) - 1
P1, P2, P3, P4, P5 = (0x9E3779B185EBCA87, 0xC2B2AE3D27D4EB4F, 0x165667B19E3779F9, 0x85EBCA77C2B2AE63,
                      0x27D4EB2F165667C5)


def _rotl(x, r):
    return ((x << r) | (x >> (64 - r))) & M64


def _round(acc, lane):
    return (_rotl((acc + lane * P2) & M64, 31) * P1) & M64


def xxh64(data, seed=0):
    n = len(data)
    i = 0
    if n >= 32:
        v = [(seed + P1 + P2) & M64, (seed + P2) & M64, seed, (seed - P1) & M64]
        while i + 32 <= n:
            for k in range(4):
                v[k] = _round(v[k], int.from_bytes(data[i + 8 * k:i + 8 * k + 8], "little"))
            i += 32
        h = (_rotl(v[0], 1) + _rotl(v[1], 7) + _rotl(v[2], 12) + _rotl(v[3], 18)) & M64
        for k in range(4):
            h = ((h ^ _round(0, v[k])) * P1 + P4) & M64
    else:
        h = (seed + P5) & M64
    h = (h + n) & M64
    while i + 8 <= n:
        h ^= _round(0, int.from_bytes(data[i:i + 8], "little"))
        h = (_rotl(h, 27) * P1 + P4) & M64
        i += 8
    if i + 4 <= n:
        h ^= (int.from_bytes(data[i:i + 4], "little") * P1) & M64
        h = (_rotl(h, 23) * P2 + P3) & M64
        i += 4
    while i < n:
        h ^= (data[i] * P5) & M64
        h = (_rotl(h, 11) * P1) & M64
        i += 1
    h ^= h >> 33
    h = (h * P2) & M64
    h ^= h >> 29
    h = (h * P3) & M64
    h ^= h >> 32
    return h


# --- a frame ------------------------------------------------------------------------------------

def decode(frame, dictionary=b"", limit=None):
    """The content of one whole frame, decoded against `dictionary` as raw content; nothing may
    follow the frame. With `limit`, a frame whose content would pass it raises `TooLarge`."""
    window, at, size, checksum = frame_window(frame)
    history = bytearray(dictionary)
    base = len(history)
    reps = [1, 4, 8]
    state = {}
    while True:
        if len(frame) < at + 3:
            raise ZstdError("block header truncated")
        h = int.from_bytes(frame[at:at + 3], "little")
        at += 3
        last, kind, bsize = h & 1, (h >> 1) & 3, h >> 3
        if bsize > min(window, MAX_BLOCK) and kind != 1:
            raise ZstdError("block too large")
        if kind == 0:
            if len(frame) < at + bsize:
                raise ZstdError("raw block truncated")
            history += frame[at:at + bsize]
            at += bsize
        elif kind == 1:
            if len(frame) < at + 1:
                raise ZstdError("RLE block truncated")
            history += bytes([frame[at]]) * bsize
            at += 1
        elif kind == 2:
            block = frame[at:at + bsize]
            if len(block) != bsize:
                raise ZstdError("compressed block truncated")
            lits, used = literals(block, state)
            execute(sequences(block[used:], state), lits, history, reps)
            at += bsize
        else:
            raise ZstdError("reserved block type")
        if limit is not None and len(history) - base > limit:
            raise TooLarge()
        if last:
            break
    content = bytes(history[base:])
    if size is not None and len(content) != size:
        raise ZstdError("content size not as stated")
    if checksum:
        if frame[at:at + 4] != (xxh64(content) & 0xFFFFFFFF).to_bytes(4, "little"):
            raise ZstdError("checksum mismatch")
        at += 4
    if at != len(frame):
        raise ZstdError("bytes after the frame")
    return content


class TooLarge(Exception):
    pass