494 lines
14 KiB
Go
494 lines
14 KiB
Go
package rdpgfx
|
||
|
||
// ZGFX (RDP8 Bulk Compression) decompressor.
|
||
// Implements the decompression algorithm described in MS-RDPEGFX 2.2.4 / 3.3.8.
|
||
// Based on FreeRDP's reference implementation (libfreerdp/codec/zgfx.c).
|
||
|
||
const zgfxHistorySize = 2500000
|
||
|
||
type zgfxContext struct {
|
||
history []byte
|
||
historyIdx int
|
||
}
|
||
|
||
func newZgfxContext() *zgfxContext {
|
||
return &zgfxContext{
|
||
history: make([]byte, zgfxHistorySize),
|
||
}
|
||
}
|
||
|
||
// Huffman token types
|
||
const (
|
||
tokenLiteral = 0
|
||
tokenMatch = 1
|
||
)
|
||
|
||
type zgfxToken struct {
|
||
prefixLen uint8
|
||
prefixCode uint16
|
||
valueBits uint8
|
||
tokenType uint8
|
||
valueBase uint32
|
||
}
|
||
|
||
// Fixed Huffman table from FreeRDP zgfx.c (Apache-2.0 licensed).
|
||
var zgfxTokenTable = []zgfxToken{
|
||
{1, 0, 8, tokenLiteral, 0}, // 0
|
||
{5, 17, 5, tokenMatch, 0}, // 10001
|
||
{5, 18, 7, tokenMatch, 32}, // 10010
|
||
{5, 19, 9, tokenMatch, 160}, // 10011
|
||
{5, 20, 10, tokenMatch, 672}, // 10100
|
||
{5, 21, 12, tokenMatch, 1696}, // 10101
|
||
{5, 24, 0, tokenLiteral, 0x00}, // 11000
|
||
{5, 25, 0, tokenLiteral, 0x01}, // 11001
|
||
{6, 44, 14, tokenMatch, 5792}, // 101100
|
||
{6, 45, 15, tokenMatch, 22176}, // 101101
|
||
{6, 52, 0, tokenLiteral, 0x02}, // 110100
|
||
{6, 53, 0, tokenLiteral, 0x03}, // 110101
|
||
{6, 54, 0, tokenLiteral, 0xFF}, // 110110
|
||
{7, 92, 18, tokenMatch, 54944}, // 1011100
|
||
{7, 93, 20, tokenMatch, 317088}, // 1011101
|
||
{7, 110, 0, tokenLiteral, 0x04}, // 1101110
|
||
{7, 111, 0, tokenLiteral, 0x05}, // 1101111
|
||
{7, 112, 0, tokenLiteral, 0x06}, // 1110000
|
||
{7, 113, 0, tokenLiteral, 0x07}, // 1110001
|
||
{7, 114, 0, tokenLiteral, 0x08}, // 1110010
|
||
{7, 115, 0, tokenLiteral, 0x09}, // 1110011
|
||
{7, 116, 0, tokenLiteral, 0x0A}, // 1110100
|
||
{7, 117, 0, tokenLiteral, 0x0B}, // 1110101
|
||
{7, 118, 0, tokenLiteral, 0x3A}, // 1110110
|
||
{7, 119, 0, tokenLiteral, 0x3B}, // 1110111
|
||
{7, 120, 0, tokenLiteral, 0x3C}, // 1111000
|
||
{7, 121, 0, tokenLiteral, 0x3D}, // 1111001
|
||
{7, 122, 0, tokenLiteral, 0x3E}, // 1111010
|
||
{7, 123, 0, tokenLiteral, 0x3F}, // 1111011
|
||
{7, 124, 0, tokenLiteral, 0x40}, // 1111100
|
||
{7, 125, 0, tokenLiteral, 0x80}, // 1111101
|
||
{8, 188, 20, tokenMatch, 1365664}, // 10111100
|
||
{8, 189, 21, tokenMatch, 2414240}, // 10111101
|
||
{8, 252, 0, tokenLiteral, 0x0C}, // 11111100
|
||
{8, 253, 0, tokenLiteral, 0x38}, // 11111101
|
||
{8, 254, 0, tokenLiteral, 0x39}, // 11111110
|
||
{8, 255, 0, tokenLiteral, 0x66}, // 11111111
|
||
{9, 380, 22, tokenMatch, 4511392}, // 101111100
|
||
{9, 381, 23, tokenMatch, 8705696}, // 101111101
|
||
{9, 382, 24, tokenMatch, 17094304}, // 101111110
|
||
}
|
||
|
||
// zgfxTokenLut maps the next 9 bits (MSB-first) of the input stream to the
|
||
// matching Huffman token entry. Built once at init(). Replaces the per-bit
|
||
// linear scan over zgfxTokenTable that used to dominate ZGFX hot path
|
||
// profiles — typical RDPGFX traffic decodes thousands of tokens per frame.
|
||
type tokenLutEntry struct {
|
||
prefixLen uint8 // 0 means "no token (truncated input)"
|
||
valueBits uint8
|
||
tokenType uint8
|
||
valueBase uint32
|
||
}
|
||
|
||
var zgfxTokenLut [512]tokenLutEntry
|
||
|
||
func init() {
|
||
for _, t := range zgfxTokenTable {
|
||
// All 9-bit windows whose top prefixLen bits == prefixCode map to t.
|
||
shift := uint(9 - t.prefixLen)
|
||
base := uint32(t.prefixCode) << shift
|
||
span := uint32(1) << shift
|
||
for j := range span {
|
||
idx := base | j
|
||
if zgfxTokenLut[idx].prefixLen == 0 {
|
||
zgfxTokenLut[idx] = tokenLutEntry{
|
||
prefixLen: t.prefixLen,
|
||
valueBits: t.valueBits,
|
||
tokenType: t.tokenType,
|
||
valueBase: t.valueBase,
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// bitReader reads bits MSB-first from a byte slice.
|
||
type bitReader struct {
|
||
data []byte
|
||
bytePos int
|
||
bitPos uint8 // bits remaining in current byte (8..1)
|
||
bitsRemaining uint32 // total decodable bits remaining
|
||
}
|
||
|
||
func newBitReader(data []byte) *bitReader {
|
||
br := &bitReader{data: data}
|
||
if len(data) > 0 {
|
||
br.bitPos = 8
|
||
}
|
||
return br
|
||
}
|
||
|
||
// newBitReaderWithCount creates a reader that tracks total decodable bits.
|
||
// The last byte of RDP8 compressed data encodes the number of padding bits
|
||
// to subtract: bitsAvailable = (len-1)*8 - lastByte.
|
||
func newBitReaderWithCount(data []byte) *bitReader {
|
||
br := &bitReader{data: data}
|
||
if len(data) < 2 {
|
||
return br
|
||
}
|
||
br.bitPos = 8
|
||
paddingBits := uint32(data[len(data)-1])
|
||
totalBits := uint32(len(data)-1) * 8
|
||
if paddingBits > totalBits {
|
||
br.bitsRemaining = 0
|
||
} else {
|
||
br.bitsRemaining = totalBits - paddingBits
|
||
}
|
||
// Exclude the last byte from readable data
|
||
br.data = data[:len(data)-1]
|
||
return br
|
||
}
|
||
|
||
func (br *bitReader) hasBitsRemaining() bool {
|
||
return br.bitsRemaining > 0
|
||
}
|
||
|
||
// peek9 returns up to 9 bits left-aligned as the high bits of a 9-bit value
|
||
// (i.e. bit 8 = next bit out of the stream). Does not advance the reader.
|
||
// avail is the number of valid bits returned; if fewer than 9 bits are
|
||
// available the returned word is zero-padded on the right (LSB).
|
||
func (br *bitReader) peek9() (val uint32, avail uint8) {
|
||
if br.bytePos >= len(br.data) {
|
||
return 0, 0
|
||
}
|
||
// First byte: only the low br.bitPos bits are still unread.
|
||
mask := uint32(1)<<br.bitPos - 1
|
||
bits := uint32(br.data[br.bytePos]) & mask
|
||
avail = br.bitPos
|
||
idx := br.bytePos + 1
|
||
for avail < 9 && idx < len(br.data) {
|
||
bits = (bits << 8) | uint32(br.data[idx])
|
||
avail += 8
|
||
idx++
|
||
}
|
||
// Compute val using the original avail, so the next stream bits land
|
||
// in the high positions of the 9-bit result (zero-padded on the right
|
||
// when fewer than 9 bits are available).
|
||
if avail >= 9 {
|
||
val = (bits >> (uint(avail) - 9)) & 0x1FF
|
||
avail = 9
|
||
} else if avail > 0 {
|
||
val = bits << (9 - avail)
|
||
}
|
||
// Then clamp the reported avail to bitsRemaining. val keeps its high
|
||
// bits valid; lower bits become don't-care padding. decodeToken only
|
||
// accepts a LUT entry whose prefixLen <= the (clamped) avail, so any
|
||
// bits beyond bitsRemaining cannot influence the decoded prefix.
|
||
if uint32(avail) > br.bitsRemaining {
|
||
avail = uint8(br.bitsRemaining)
|
||
}
|
||
return
|
||
}
|
||
|
||
// consumeBits advances the reader by n bits, updating bytePos/bitPos and
|
||
// bitsRemaining. Caller must ensure n <= bitsRemaining (LUT entries
|
||
// already include this guarantee for valid streams).
|
||
func (br *bitReader) consumeBits(n uint8) {
|
||
if uint32(n) > br.bitsRemaining {
|
||
n = uint8(br.bitsRemaining)
|
||
}
|
||
br.bitsRemaining -= uint32(n)
|
||
for n >= br.bitPos {
|
||
n -= br.bitPos
|
||
br.bytePos++
|
||
br.bitPos = 8
|
||
}
|
||
br.bitPos -= n
|
||
if br.bitPos == 0 {
|
||
br.bytePos++
|
||
br.bitPos = 8
|
||
}
|
||
}
|
||
|
||
func (br *bitReader) getBit() uint32 {
|
||
if br.bytePos >= len(br.data) {
|
||
return 0
|
||
}
|
||
br.bitPos--
|
||
bit := uint32((br.data[br.bytePos] >> br.bitPos) & 1)
|
||
if br.bitPos == 0 {
|
||
br.bytePos++
|
||
br.bitPos = 8
|
||
}
|
||
return bit
|
||
}
|
||
|
||
func (br *bitReader) getBits(n uint8) uint32 {
|
||
if n == 0 {
|
||
return 0
|
||
}
|
||
// Fast path: enough bits remain in the current byte.
|
||
if n <= br.bitPos && br.bytePos < len(br.data) {
|
||
br.bitPos -= n
|
||
v := (uint32(br.data[br.bytePos]) >> br.bitPos) & ((1 << n) - 1)
|
||
if br.bitPos == 0 {
|
||
br.bytePos++
|
||
br.bitPos = 8
|
||
}
|
||
return v
|
||
}
|
||
// Slow path: consume bits across byte boundaries byte-at-a-time.
|
||
// This replaces a bit-by-bit loop (up to 24 getBit() calls) with at most
|
||
// ⌈n/8⌉+1 iterations, significantly reducing branch and call overhead.
|
||
var result uint32
|
||
for rem := n; rem > 0; {
|
||
if br.bytePos >= len(br.data) {
|
||
result <<= rem
|
||
break
|
||
}
|
||
take := rem
|
||
if take > br.bitPos {
|
||
take = br.bitPos
|
||
}
|
||
br.bitPos -= take
|
||
rem -= take
|
||
result = (result << take) | ((uint32(br.data[br.bytePos]) >> br.bitPos) & uint32((1<<take)-1))
|
||
if br.bitPos == 0 {
|
||
br.bytePos++
|
||
br.bitPos = 8
|
||
}
|
||
}
|
||
return result
|
||
}
|
||
|
||
// historyWrite copies data into the ring history buffer. Hot path — called
|
||
// for every literal, match expansion, and unencoded raw block. Uses bulk
|
||
// copy() instead of byte-by-byte modulo arithmetic.
|
||
func (z *zgfxContext) historyWrite(data []byte) {
|
||
n := len(data)
|
||
if n == 0 {
|
||
return
|
||
}
|
||
if n >= zgfxHistorySize {
|
||
copy(z.history, data[n-zgfxHistorySize:])
|
||
z.historyIdx = 0
|
||
return
|
||
}
|
||
end := z.historyIdx + n
|
||
if end <= zgfxHistorySize {
|
||
copy(z.history[z.historyIdx:end], data)
|
||
z.historyIdx = end
|
||
if z.historyIdx == zgfxHistorySize {
|
||
z.historyIdx = 0
|
||
}
|
||
} else {
|
||
first := zgfxHistorySize - z.historyIdx
|
||
copy(z.history[z.historyIdx:], data[:first])
|
||
copy(z.history[:n-first], data[first:])
|
||
z.historyIdx = n - first
|
||
}
|
||
}
|
||
|
||
func (z *zgfxContext) outputLiteral(b byte, out *[]byte) {
|
||
z.history[z.historyIdx] = b
|
||
z.historyIdx++
|
||
if z.historyIdx == zgfxHistorySize {
|
||
z.historyIdx = 0
|
||
}
|
||
*out = append(*out, b)
|
||
}
|
||
|
||
// outputMatch copies `count` bytes starting `distance` bytes back in the
|
||
// history into both the output and the history ring. Self-referential
|
||
// matches (count > distance) are handled via byte-wise copy after the first
|
||
// pass, since the pattern grows as it is written.
|
||
func (z *zgfxContext) outputMatch(distance, count int, out *[]byte) {
|
||
if distance <= 0 || count <= 0 {
|
||
return
|
||
}
|
||
o := *out
|
||
base := len(o)
|
||
// Grow output by `count` bytes in a single allocation step.
|
||
if cap(o)-base < count {
|
||
// Standard slice-growth: at least double cap or fit count, whichever larger.
|
||
newCap := max(cap(o)*2, base+count)
|
||
grown := make([]byte, base+count, newCap)
|
||
copy(grown, o)
|
||
o = grown
|
||
} else {
|
||
o = o[:base+count]
|
||
}
|
||
|
||
srcIdx := z.historyIdx - distance
|
||
if srcIdx < 0 {
|
||
srcIdx += zgfxHistorySize
|
||
}
|
||
|
||
if count <= distance {
|
||
// Non-overlapping pattern: 1-2 contiguous copies from history.
|
||
end := srcIdx + count
|
||
if end <= zgfxHistorySize {
|
||
copy(o[base:base+count], z.history[srcIdx:end])
|
||
} else {
|
||
first := zgfxHistorySize - srcIdx
|
||
copy(o[base:base+first], z.history[srcIdx:])
|
||
copy(o[base+first:base+count], z.history[:count-first])
|
||
}
|
||
} else {
|
||
// Self-overlapping: copy the initial `distance` bytes from history,
|
||
// then expand the pattern in-place within the output buffer (the
|
||
// classic LZ77 overlapping copy).
|
||
end := srcIdx + distance
|
||
if end <= zgfxHistorySize {
|
||
copy(o[base:base+distance], z.history[srcIdx:end])
|
||
} else {
|
||
first := zgfxHistorySize - srcIdx
|
||
copy(o[base:base+first], z.history[srcIdx:])
|
||
copy(o[base+first:base+distance], z.history[:distance-first])
|
||
}
|
||
// Overlapping in-place expansion (must be byte-wise; copy() does
|
||
// not guarantee overlapping semantics for src==dst+offset).
|
||
for i := distance; i < count; i++ {
|
||
o[base+i] = o[base+i-distance]
|
||
}
|
||
}
|
||
|
||
*out = o
|
||
z.historyWrite(o[base : base+count])
|
||
}
|
||
|
||
// Decompress decompresses a ZGFX compressed segment payload.
|
||
// The payload must NOT include the 1-byte segment header (flags byte).
|
||
// In RDP8 ZGFX, the last byte of the payload encodes the number of
|
||
// padding bits to subtract from the total bit count.
|
||
//
|
||
// buf is an optional caller-supplied backing buffer (e.g. from a sync.Pool).
|
||
// The returned slice may use buf's backing array or a new one if buf was too
|
||
// small; callers that pool the output must use the RETURNED slice, not buf.
|
||
func (z *zgfxContext) Decompress(data []byte, buf []byte) []byte {
|
||
if len(data) < 2 {
|
||
// Need at least 1 byte of compressed data + 1 byte of padding count
|
||
return nil
|
||
}
|
||
|
||
br := newBitReaderWithCount(data)
|
||
// Use the caller-supplied buffer if it has enough capacity; otherwise fall
|
||
// back to a fresh allocation sized at 3× the compressed input.
|
||
estSize := len(data) * 3
|
||
var out []byte
|
||
if cap(buf) >= estSize {
|
||
out = buf[:0]
|
||
} else {
|
||
out = make([]byte, 0, estSize)
|
||
}
|
||
|
||
for br.hasBitsRemaining() {
|
||
token, ok := z.decodeToken(br)
|
||
if !ok {
|
||
break
|
||
}
|
||
// decodeToken already consumed prefixLen bits.
|
||
|
||
if token.tokenType == tokenLiteral {
|
||
if br.bitsRemaining < uint32(token.valueBits) {
|
||
break
|
||
}
|
||
value := token.valueBase + br.getBits(token.valueBits)
|
||
br.bitsRemaining -= uint32(token.valueBits)
|
||
z.outputLiteral(byte(value), &out)
|
||
} else {
|
||
// Match token
|
||
if br.bitsRemaining < uint32(token.valueBits) {
|
||
break
|
||
}
|
||
distance := int(token.valueBase + br.getBits(token.valueBits))
|
||
br.bitsRemaining -= uint32(token.valueBits)
|
||
|
||
if distance != 0 {
|
||
// Match: copy from history
|
||
count := z.decodeMatchCount(br)
|
||
z.outputMatch(distance, count, &out)
|
||
} else {
|
||
// Unencoded: read raw bytes
|
||
if br.bitsRemaining < 15 {
|
||
break
|
||
}
|
||
rawCount := int(br.getBits(15))
|
||
br.bitsRemaining -= 15
|
||
// Discard remaining bits in current byte to align to byte boundary
|
||
// (equivalent to FreeRDP's cBitsCurrent = 0; BitsCurrent = 0;)
|
||
if br.bitPos < 8 {
|
||
br.bitsRemaining -= uint32(br.bitPos)
|
||
br.bytePos++
|
||
br.bitPos = 8
|
||
}
|
||
if br.bytePos+rawCount > len(br.data) || uint32(rawCount)*8 > br.bitsRemaining {
|
||
break
|
||
}
|
||
rawBytes := br.data[br.bytePos : br.bytePos+rawCount]
|
||
br.bytePos += rawCount
|
||
br.bitsRemaining -= uint32(rawCount) * 8
|
||
z.historyWrite(rawBytes)
|
||
out = append(out, rawBytes...)
|
||
}
|
||
}
|
||
}
|
||
|
||
return out
|
||
}
|
||
|
||
// decodeToken consumes the next Huffman prefix (1..9 bits) from the input
|
||
// stream using the precomputed zgfxTokenLut. Returns the decoded token and
|
||
// true on success; false if the remaining input is too short or malformed.
|
||
//
|
||
// Unlike the old implementation, this version no longer uses a per-bit
|
||
// linear search through the 38-entry token table — the dominant ZGFX cost
|
||
// during real RDPGFX traffic.
|
||
func (z *zgfxContext) decodeToken(br *bitReader) (zgfxToken, bool) {
|
||
val, avail := br.peek9()
|
||
if avail == 0 {
|
||
return zgfxToken{}, false
|
||
}
|
||
e := zgfxTokenLut[val]
|
||
if e.prefixLen == 0 || uint8(avail) < e.prefixLen {
|
||
return zgfxToken{}, false
|
||
}
|
||
br.consumeBits(e.prefixLen)
|
||
return zgfxToken{
|
||
prefixLen: e.prefixLen,
|
||
valueBits: e.valueBits,
|
||
tokenType: e.tokenType,
|
||
valueBase: e.valueBase,
|
||
}, true
|
||
}
|
||
|
||
// decodeMatchCount decodes the match length using FreeRDP's algorithm:
|
||
//
|
||
// 0 → 3
|
||
// 10 + 2 bits → 4 + value (4..7)
|
||
// 110 + 3 bits → 8 + value (8..15)
|
||
// 1110 + 4 bits → 16 + value (16..31)
|
||
// ... and so on (each additional leading 1 doubles the base and adds 1 extra bit)
|
||
func (z *zgfxContext) decodeMatchCount(br *bitReader) int {
|
||
bit := br.getBit()
|
||
br.bitsRemaining--
|
||
if bit == 0 {
|
||
return 3
|
||
}
|
||
|
||
count := 4
|
||
extra := uint8(2)
|
||
|
||
bit = br.getBit()
|
||
br.bitsRemaining--
|
||
for bit == 1 {
|
||
count <<= 1
|
||
extra++
|
||
bit = br.getBit()
|
||
br.bitsRemaining--
|
||
}
|
||
|
||
if br.bitsRemaining < uint32(extra) {
|
||
return count
|
||
}
|
||
count += int(br.getBits(extra))
|
||
br.bitsRemaining -= uint32(extra)
|
||
return count
|
||
}
|