466 lines
18 KiB
Plaintext
466 lines
18 KiB
Plaintext
from stdint import *
|
||
import zlib.zchecksum as zchecksum
|
||
import zlib.zdef as zdef
|
||
import zlib.zhuff as zhuff
|
||
import string
|
||
import memhub
|
||
import t, c
|
||
|
||
|
||
@t.Object
|
||
class zinflate_stream:
|
||
pool: rawptr[memhub.MemManager]
|
||
reader: zdef.zbit_reader
|
||
|
||
window: BYTEPTR
|
||
window_pos: t.CInt
|
||
window_size: t.CInt
|
||
|
||
wbits: t.CInt
|
||
is_finished: t.CInt
|
||
is_initialized: t.CInt
|
||
|
||
header_parsed: t.CInt
|
||
is_gzip: t.CInt
|
||
|
||
adler: t.CUInt32T
|
||
crc: t.CUInt32T
|
||
|
||
output: BYTEPTR
|
||
output_len: t.CSizeT
|
||
output_cap: t.CSizeT
|
||
|
||
input_data: t.CConst | BYTEPTR
|
||
input_len: t.CSizeT
|
||
input_pos: t.CSizeT
|
||
|
||
max_length: t.CSizeT
|
||
|
||
def __init__(self):
|
||
pass
|
||
|
||
def output_byte(self, byte: BYTE):
|
||
if self.max_length > 0 and self.output_len >= self.max_length: return
|
||
|
||
if self.output_len >= self.output_cap:
|
||
new_cap: t.CSizeT = self.output_cap * 2
|
||
if new_cap < 256: new_cap = 256
|
||
new_buf: BYTEPTR = BYTEPTR(zdef.zdef_alloc(self.pool, new_cap))
|
||
if not new_buf: return
|
||
string.memcpy(new_buf, self.output, self.output_len)
|
||
zdef.zdef_free(self.pool, self.output)
|
||
self.output = new_buf
|
||
self.output_cap = new_cap
|
||
self.output[self.output_len] = byte
|
||
self.output_len += 1
|
||
self.window[self.window_pos & (zdef.ZDEFLATE_WINDOW_SIZE - 1)] = byte
|
||
self.window_pos += 1
|
||
|
||
|
||
def parse_header(self) -> t.CInt:
|
||
if self.header_parsed: return 0
|
||
if self.wbits < 0:
|
||
self.header_parsed = 1
|
||
self.is_gzip = 0
|
||
return 0
|
||
if self.input_pos + 2 > self.input_len: return -3
|
||
b0: BYTE = self.input_data[self.input_pos]
|
||
b1: BYTE = self.input_data[self.input_pos + 1]
|
||
if b0 == 0x1F and b1 == 0x8B:
|
||
self.is_gzip = 1
|
||
if self.input_pos + 10 > self.input_len: return -3
|
||
self.input_pos += 10
|
||
# FLG 位于 gzip 头第 4 字节（偏移 3），即 self.input_pos - 7
|
||
flg: BYTE = self.input_data[self.input_pos - 7]
|
||
if flg & 0x04:
|
||
if self.input_pos + 2 > self.input_len: return -3
|
||
xlen: UINT = self.input_data[self.input_pos] | (self.input_data[self.input_pos + 1] << 8)
|
||
self.input_pos += 2 + xlen
|
||
if flg & 0x08:
|
||
while self.input_pos < self.input_len and self.input_data[self.input_pos] != 0: self.input_pos += 1
|
||
self.input_pos += 1
|
||
if flg & 0x10:
|
||
while self.input_pos < self.input_len and self.input_data[self.input_pos] != 0: self.input_pos += 1
|
||
self.input_pos += 1
|
||
if flg & 0x02:
|
||
self.input_pos += 2
|
||
if self.input_pos > self.input_len: return -3
|
||
else:
|
||
self.is_gzip = 0
|
||
if (b0 * 256 + b1) % 31 != 0: return -3
|
||
cm: t.CInt = b0 & 0x0F
|
||
if cm != 8: return -3
|
||
self.input_pos += 2
|
||
self.header_parsed = 1
|
||
return 0
|
||
|
||
|
||
def read_bits_from_input(self, nbits: t.CInt, out: UINTPTR) -> t.CInt:
|
||
val: UINT = 0
|
||
for i in range(nbits):
|
||
if self.input_pos >= self.input_len: return -3
|
||
if self.input_data[self.input_pos] & (UINT(1) << (self.reader.bit_pos)):
|
||
val |= (UINT(1) << i)
|
||
self.reader.bit_pos += 1
|
||
if self.reader.bit_pos == 8:
|
||
self.reader.bit_pos = 0
|
||
self.input_pos += 1
|
||
c.Set(c.Deref(out), val)
|
||
return 0
|
||
|
||
|
||
def read_huffman(self, dt: rawptr[zhuff.zhuff_decode_tree]) -> t.CInt:
|
||
node: int = dt.root
|
||
while dt.nodes[node].symbol == -1:
|
||
bit: UINT
|
||
if self.read_bits_from_input(1, c.Addr(bit)) != 0: return -1
|
||
dir: int = t.CInt(bit)
|
||
if dt.nodes[node].children[dir] == -1: return -3
|
||
node = dt.nodes[node].children[dir]
|
||
return dt.nodes[node].symbol
|
||
|
||
def inflate_block_stored(self) -> t.CInt:
|
||
if self.reader.bit_pos != 0:
|
||
self.reader.bit_pos = 0
|
||
self.input_pos += 1
|
||
if self.input_pos + 4 > self.input_len: return -3
|
||
length: UINT = self.input_data[self.input_pos] | (self.input_data[self.input_pos + 1] << 8)
|
||
nlen: UINT = self.input_data[self.input_pos + 2] | (self.input_data[self.input_pos + 3] << 8)
|
||
self.input_pos += 4
|
||
if (length ^ nlen) != 0xFFFF: return -3
|
||
if self.input_pos + length > self.input_len: return -3
|
||
i: t.CUnsignedInt
|
||
for i in range(length):
|
||
self.output_byte(self.input_data[self.input_pos])
|
||
self.input_pos += 1
|
||
return 0
|
||
|
||
|
||
def inflate_block_fixed(self) -> t.CInt:
|
||
lit_tree = zhuff.zhuff_tree()
|
||
dist_tree = zhuff.zhuff_tree()
|
||
lit_tree.build_fixed_lit_tree()
|
||
dist_tree.build_fixed_dist_tree()
|
||
|
||
lit_dt: rawptr[zhuff.zhuff_decode_tree] = zdef.zdef_alloc(self.pool, zhuff.zhuff_decode_tree.__sizeof__())
|
||
dist_dt: rawptr[zhuff.zhuff_decode_tree] = zdef.zdef_alloc(self.pool, zhuff.zhuff_decode_tree.__sizeof__())
|
||
if not lit_dt or not dist_dt:
|
||
if lit_dt: zdef.zdef_free(self.pool, lit_dt)
|
||
if dist_dt: zdef.zdef_free(self.pool, dist_dt)
|
||
return -4
|
||
string.memset(lit_dt, 0, zhuff.zhuff_decode_tree.__sizeof__())
|
||
string.memset(dist_dt, 0, zhuff.zhuff_decode_tree.__sizeof__())
|
||
lit_dt.build_decode_tree(c.Addr(lit_tree))
|
||
dist_dt.build_decode_tree(c.Addr(dist_tree))
|
||
|
||
result: t.CInt = 0
|
||
while True:
|
||
if self.max_length > 0 and self.output_len >= self.max_length: break
|
||
|
||
sym: t.CInt = self.read_huffman(lit_dt)
|
||
if sym < 0:
|
||
result = -3
|
||
break
|
||
if sym == 256: break
|
||
if sym < 256:
|
||
self.output_byte(BYTE(sym))
|
||
else:
|
||
len_idx: t.CInt = sym - 257
|
||
if len_idx < 0 or len_idx >= 29:
|
||
result = -3
|
||
break
|
||
length: t.CInt = zdef.zdeflate_len_base[len_idx]
|
||
extra: t.CInt = zdef.zdeflate_len_extra_bits[len_idx]
|
||
if extra > 0:
|
||
extra_val: UINT
|
||
if self.read_bits_from_input(extra, c.Addr(extra_val)) != 0:
|
||
result = -3
|
||
break
|
||
length += extra_val
|
||
|
||
dist_sym: t.CInt = self.read_huffman(dist_dt)
|
||
if dist_sym < 0 or dist_sym >= 30:
|
||
result = -3
|
||
break
|
||
dist: t.CInt = zdef.zdeflate_dist_base[dist_sym]
|
||
extra = zdef.zdeflate_dist_extra_bits[dist_sym]
|
||
if extra > 0:
|
||
extra_val: UINT
|
||
if self.read_bits_from_input(extra, c.Addr(extra_val)) != 0:
|
||
result = -3
|
||
break
|
||
dist += extra_val
|
||
|
||
for i in range(length):
|
||
if self.max_length > 0 and self.output_len >= self.max_length: break
|
||
src_pos: t.CInt = (self.window_pos - dist) & (zdef.ZDEFLATE_WINDOW_SIZE - 1)
|
||
byte: BYTE = self.window[src_pos]
|
||
self.output_byte(byte)
|
||
|
||
zdef.zdef_free(self.pool, lit_dt)
|
||
zdef.zdef_free(self.pool, dist_dt)
|
||
return result
|
||
|
||
|
||
def inflate_block_dynamic(self) -> t.CInt:
|
||
hlit_val: UINT
|
||
hdist_val: UINT
|
||
hclen_val: UINT
|
||
if self.read_bits_from_input(5, c.Addr(hlit_val)) != 0: return -3
|
||
if self.read_bits_from_input(5, c.Addr(hdist_val)) != 0: return -3
|
||
if self.read_bits_from_input(4, c.Addr(hclen_val)) != 0: return -3
|
||
|
||
hlit: t.CInt = t.CInt(hlit_val) + 257
|
||
hdist: t.CInt = t.CInt(hdist_val) + 1
|
||
hclen: t.CInt = t.CInt(hclen_val) + 4
|
||
|
||
cl_lengths: t.CArray[t.CInt, 19]
|
||
string.memset(cl_lengths, 0, cl_lengths.__sizeof__())
|
||
for i in range(hclen):
|
||
len_val: UINT
|
||
if self.read_bits_from_input(3, c.Addr(len_val)) != 0: return -3
|
||
cl_lengths[zdef.zdeflate_codelen_order[i]] = t.CInt(len_val)
|
||
|
||
cl_tree = zhuff.zhuff_tree()
|
||
cl_tree.build_tree_from_lengths(cl_lengths, 19, 7)
|
||
|
||
cl_dt = zhuff.zhuff_decode_tree()
|
||
cl_dt.build_decode_tree(c.Addr(cl_tree))
|
||
|
||
total: t.CInt = hlit + hdist
|
||
all_lengths: t.CInt | t.CPtr = zdef.zdef_alloc(self.pool, total * int.__sizeof__())
|
||
if not all_lengths: return -4
|
||
|
||
idx: t.CInt = 0
|
||
while idx < total:
|
||
sym: t.CInt = self.read_huffman(c.Addr(cl_dt))
|
||
if sym < 0 or sym > 18:
|
||
zdef.zdef_free(self.pool, all_lengths)
|
||
return -3
|
||
if sym < 16:
|
||
all_lengths[idx] = sym
|
||
idx += 1
|
||
elif sym == 16:
|
||
rep: UINT
|
||
if self.read_bits_from_input(2, c.Addr(rep)) != 0:
|
||
zdef.zdef_free(self.pool, all_lengths)
|
||
return -3
|
||
rep += 3
|
||
if idx == 0 or idx + rep > total:
|
||
zdef.zdef_free(self.pool, all_lengths)
|
||
return -3
|
||
prev: t.CInt = all_lengths[idx - 1]
|
||
i: t.CUnsignedInt
|
||
for i in range(rep):
|
||
all_lengths[idx] = prev
|
||
idx += 1
|
||
elif sym == 17:
|
||
rep: t.CUnsignedInt
|
||
if self.read_bits_from_input(3, c.Addr(rep)) != 0:
|
||
zdef.zdef_free(self.pool, all_lengths)
|
||
return -3
|
||
rep += 3
|
||
if idx + rep > total:
|
||
zdef.zdef_free(self.pool, all_lengths)
|
||
return -3
|
||
i: t.CUnsignedInt
|
||
for i in range(rep):
|
||
all_lengths[idx] = 0
|
||
idx += 1
|
||
else:
|
||
rep: t.CUnsignedInt
|
||
if self.read_bits_from_input(7, c.Addr(rep)) != 0:
|
||
zdef.zdef_free(self.pool, all_lengths)
|
||
return -3
|
||
rep += 11
|
||
if idx + rep > total:
|
||
zdef.zdef_free(self.pool, all_lengths)
|
||
return -3
|
||
i: t.CUnsignedInt
|
||
for i in range(rep):
|
||
all_lengths[idx] = 0
|
||
idx += 1
|
||
|
||
lit_tree = zhuff.zhuff_tree()
|
||
dist_tree = zhuff.zhuff_tree()
|
||
|
||
lit_lengths: t.CArray[t.CInt, 288]
|
||
string.memset(lit_lengths, 0, lit_lengths.__sizeof__())
|
||
for i in range(hlit):
|
||
lit_lengths[i] = all_lengths[i]
|
||
lit_tree.build_tree_from_lengths(lit_lengths, hlit, zdef.ZDEFLATE_MAX_BITS)
|
||
|
||
dist_lengths: t.CArray[t.CInt, 32]
|
||
string.memset(dist_lengths, 0, dist_lengths.__sizeof__())
|
||
for i in range(hdist):
|
||
dist_lengths[i] = all_lengths[hlit + i]
|
||
dist_tree.build_tree_from_lengths(dist_lengths, hdist, zdef.ZDEFLATE_MAX_BITS)
|
||
|
||
zdef.zdef_free(self.pool, all_lengths)
|
||
|
||
lit_dt = zhuff.zhuff_decode_tree()
|
||
dist_dt = zhuff.zhuff_decode_tree()
|
||
lit_dt.build_decode_tree(c.Addr(lit_tree))
|
||
dist_dt.build_decode_tree(c.Addr(dist_tree))
|
||
|
||
while True:
|
||
if self.max_length > 0 and self.output_len >= self.max_length: return 0
|
||
sym: t.CInt = self.read_huffman(c.Addr(lit_dt))
|
||
if sym < 0: return -3
|
||
if sym == 256: return 0
|
||
if sym < 256:
|
||
self.output_byte(BYTE(sym))
|
||
else:
|
||
len_idx: t.CInt = sym - 257
|
||
if len_idx < 0 or len_idx >= 29: return -3
|
||
length: t.CInt = zdef.zdeflate_len_base[len_idx]
|
||
extra: t.CInt = zdef.zdeflate_len_extra_bits[len_idx]
|
||
if extra > 0:
|
||
extra_val: UINT
|
||
if self.read_bits_from_input(extra, c.Addr(extra_val)) != 0: return -3
|
||
length += extra_val
|
||
dist_sym: t.CInt = self.read_huffman(c.Addr(dist_dt))
|
||
if dist_sym < 0 or dist_sym >= 30: return -3
|
||
dist: t.CInt = zdef.zdeflate_dist_base[dist_sym]
|
||
extra = zdef.zdeflate_dist_extra_bits[dist_sym]
|
||
if extra > 0:
|
||
extra_val: UINT
|
||
if self.read_bits_from_input(extra, c.Addr(extra_val)) != 0: return -3
|
||
dist += extra_val
|
||
for i in range(length):
|
||
if self.max_length > 0 and self.output_len >= self.max_length: return 0
|
||
src_pos: t.CInt = (self.window_pos - dist) & (zdef.ZDEFLATE_WINDOW_SIZE - 1)
|
||
byte: BYTE = self.window[src_pos]
|
||
self.output_byte(byte)
|
||
|
||
def inflate_stream(self) -> t.CInt:
|
||
err: t.CInt = self.parse_header()
|
||
if err != 0: return err
|
||
while not self.is_finished:
|
||
bfinal_val: UINT
|
||
btype_val: UINT
|
||
if self.read_bits_from_input(1, c.Addr(bfinal_val)) != 0: return -3
|
||
if self.read_bits_from_input(2, c.Addr(btype_val)) != 0: return -3
|
||
bfinal: t.CInt = t.CInt(bfinal_val)
|
||
btype: t.CInt = t.CInt(btype_val)
|
||
if btype == 0:
|
||
err = self.inflate_block_stored()
|
||
elif btype == 1:
|
||
err = self.inflate_block_fixed()
|
||
elif btype == 2:
|
||
err = self.inflate_block_dynamic()
|
||
else:
|
||
return -3
|
||
if err != 0: return err
|
||
if bfinal: self.is_finished = 1
|
||
if self.max_length > 0 and self.output_len >= self.max_length: return 0
|
||
if self.is_gzip:
|
||
if self.reader.bit_pos != 0:
|
||
self.reader.bit_pos = 0
|
||
self.input_pos += 1
|
||
if self.input_pos + 8 <= self.input_len:
|
||
self.crc = ((t.CUInt32T(self.input_data[self.input_pos + 0]) << 0) |
|
||
(t.CUInt32T(self.input_data[self.input_pos + 1]) << 8) |
|
||
(t.CUInt32T(self.input_data[self.input_pos + 2]) << 16) |
|
||
(t.CUInt32T(self.input_data[self.input_pos + 3]) << 24))
|
||
self.input_pos += 8
|
||
elif self.wbits >= 0:
|
||
if self.reader.bit_pos != 0:
|
||
self.reader.bit_pos = 0
|
||
self.input_pos += 1
|
||
if self.input_pos + 4 <= self.input_len:
|
||
self.adler = ((t.CUInt32T(self.input_data[self.input_pos + 0]) << 24) |
|
||
(t.CUInt32T(self.input_data[self.input_pos + 1]) << 16) |
|
||
(t.CUInt32T(self.input_data[self.input_pos + 2]) << 8) |
|
||
(t.CUInt32T(self.input_data[self.input_pos + 3]) << 0))
|
||
self.input_pos += 4
|
||
return 0
|
||
|
||
# out 是"指针的指针"（i8**）。不要写成 `t.CUInt8T | t.CPtr[t.CPtr]`：
|
||
# 联合里嵌 `t.CPtr[...]` 下标时，定义侧按 PtrCount 累加得 i8**，而 stub/declare
|
||
# 侧（StubGen/VpiReader）只把整段等于 `t.CPtr`/`CPtr` 的段认作 CPtr 标记，
|
||
# 于是坍缩成首段 `t.CUInt8T` = i8 → 跨模块 ABI 不一致（调用方只传 1 字节，
|
||
# 被调用方按 8 字节指针写入 → ASan heap-buffer-overflow）。
|
||
# 改用"指针别名 | t.CPtr"：两侧都按"追加 1 层指针"解析，同为 i8**。
|
||
# 同形先例：includes/w32/win32sync.vp:65 `Context: VOIDPTR | t.CPtr`、
|
||
# includes/posix.vp:199 `argv: CHARPTR | t.CPtr`。
|
||
def decompress(self, data: UINT8PTR, length: t.CSizeT,
|
||
max_length: t.CSizeT, out: BYTEPTR | t.CPtr, out_len: t.CSizeT | t.CPtr) -> t.CInt:
|
||
if not self.is_initialized: return -2
|
||
self.input_data = data
|
||
self.input_len = length
|
||
self.input_pos = 0
|
||
self.reader.bit_pos = 0
|
||
self.output_len = 0
|
||
self.max_length = max_length
|
||
err: t.CInt = self.inflate_stream()
|
||
if err != 0: return err
|
||
c.Set(c.Deref(out_len), self.output_len)
|
||
if self.output_len == 0:
|
||
c.Set(c.Deref(out), UINT8PTR(zdef.zdef_alloc(self.pool, 1)))
|
||
else:
|
||
dst: UINT8PTR = UINT8PTR(zdef.zdef_alloc(self.pool, self.output_len))
|
||
c.Set(c.Deref(out), dst)
|
||
if dst:
|
||
string.memcpy(dst, self.output, self.output_len)
|
||
return 0
|
||
|
||
def set_dictionary(self, dict: UINT8PTR, length: t.CSizeT) -> t.CInt:
|
||
if not dict: return -2
|
||
use: t.CSizeT = length
|
||
if use > zdef.ZDEFLATE_WINDOW_SIZE:
|
||
dict += length - zdef.ZDEFLATE_WINDOW_SIZE
|
||
use = zdef.ZDEFLATE_WINDOW_SIZE
|
||
string.memcpy(self.window, dict, use)
|
||
self.window_pos = t.CInt(use)
|
||
return 0
|
||
|
||
def copy(self) -> zinflate_stream | t.CPtr:
|
||
z: zinflate_stream | t.CPtr = zdef.zdef_alloc(self.pool, zinflate_stream.__sizeof__())
|
||
if not z: return None
|
||
string.memcpy(z, self, zinflate_stream.__sizeof__())
|
||
z.window = BYTEPTR(zdef.zdef_alloc(self.pool, zdef.ZDEFLATE_WINDOW_SIZE))
|
||
z.output = BYTEPTR(zdef.zdef_alloc(self.pool, self.output_cap))
|
||
if not z.window or not z.output:
|
||
zinflate_stream.destroy(z)
|
||
return None
|
||
string.memcpy(z.window, self.window, zdef.ZDEFLATE_WINDOW_SIZE)
|
||
string.memcpy(z.output, self.output, self.output_cap)
|
||
return z
|
||
|
||
def destroy(self):
|
||
if not self: return
|
||
if self.window: zdef.zdef_free(self.pool, self.window)
|
||
if self.output: zdef.zdef_free(self.pool, self.output)
|
||
zdef.zdef_free(self.pool, self)
|
||
|
||
|
||
def zinflate_one_shot(pool: rawptr[memhub.MemManager], data: UINT8PTR, length: t.CSizeT,
|
||
wbits: t.CInt, bufsize: t.CSizeT, out_len: t.CSizeT | t.CPtr) -> t.CUInt8T | t.CPtr:
|
||
s: zinflate_stream | t.CPtr = zinflate_create(pool, wbits)
|
||
if not s: return None
|
||
out: t.CUInt8T | t.CPtr = None
|
||
err: t.CInt = s.decompress(data, length, 0, c.Addr(out), out_len)
|
||
s.destroy()
|
||
if err != 0: return None
|
||
return out
|
||
|
||
|
||
def zinflate_create(pool: rawptr[memhub.MemManager], wbits: t.CInt) -> zinflate_stream | t.CPtr:
|
||
s: zinflate_stream | t.CPtr = zdef.zdef_alloc(pool, zinflate_stream.__sizeof__())
|
||
if not s: return None
|
||
string.memset(s, 0, zinflate_stream.__sizeof__())
|
||
s.pool = pool
|
||
s.wbits = wbits
|
||
s.is_initialized = 1
|
||
s.adler = 1
|
||
s.window = BYTEPTR(zdef.zdef_alloc(pool, zdef.ZDEFLATE_WINDOW_SIZE))
|
||
s.output = BYTEPTR(zdef.zdef_alloc(pool, 256))
|
||
s.output_cap = 256
|
||
s.output_len = 0
|
||
s.max_length = 0
|
||
if not s.window or not s.output:
|
||
s.destroy()
|
||
return None
|
||
return s
|