lychee-slicer/asar/pickle/pickle.go

package pickle

import (
	"encoding/binary"
	"errors"
	"fmt"
	"math"
)

const (
	SIZE_UINT32        = 4
	PAYLOAD_UNIT       = 64
	CAPACITY_READ_ONLY = 1 << 53
)

func alignInt(i, alignment int) int {
	return i + ((alignment - (i % alignment)) % alignment)
}

type PickleIterator struct {
	payload       []byte
	payloadOffset int
	readIndex     int
	endIndex      int
}

func NewPickleIterator(pickle *Pickle) *PickleIterator {
	return &PickleIterator{
		payload:       pickle.header,
		payloadOffset: pickle.headerSize,
		readIndex:     0,
		endIndex:      pickle.getPayloadSize(),
	}
}

func (pi *PickleIterator) ReadBool() bool {
	return pi.ReadInt() != 0
}

func (pi *PickleIterator) ReadInt() int32 {
	buf := pi.readBytes(binary.Size(int32(0)))
	return int32(binary.LittleEndian.Uint32(buf))
}

func (pi *PickleIterator) ReadUInt32() uint32 {
	buf := pi.readBytes(binary.Size(uint32(0)))
	return binary.LittleEndian.Uint32(buf)
}

func (pi *PickleIterator) ReadInt64() int64 {
	buf := pi.readBytes(binary.Size(int64(0)))
	return int64(binary.LittleEndian.Uint64(buf))
}

func (pi *PickleIterator) ReadUInt64() uint64 {
	buf := pi.readBytes(binary.Size(uint64(0)))
	return binary.LittleEndian.Uint64(buf)
}

func (pi *PickleIterator) ReadFloat() float32 {
	return math.Float32frombits(pi.ReadUInt32())
}

func (pi *PickleIterator) readDouble() float64 {
	return math.Float64frombits(pi.ReadUInt64())
}

func (pi *PickleIterator) ReadString() string {
	length := pi.ReadInt()
	return string(pi.readBytes(int(length)))
}

func (pi *PickleIterator) readBytes(length int) []byte {
	readPayloadOffset := pi.getReadPayloadOffsetAndAdvance(length)
	return pi.payload[readPayloadOffset : readPayloadOffset+length]
}

func (pi *PickleIterator) getReadPayloadOffsetAndAdvance(length int) int {
	if length > pi.endIndex-pi.readIndex {
		pi.readIndex = pi.endIndex
		panic(errors.New("failed to read data with length of " + fmt.Sprint(length)))
	}
	readPayloadOffset := pi.payloadOffset + pi.readIndex
	pi.advance(length)
	return readPayloadOffset
}

func (pi *PickleIterator) advance(size int) {
	alignedSize := alignInt(size, SIZE_UINT32)
	if pi.endIndex-pi.readIndex < alignedSize {
		pi.readIndex = pi.endIndex
	} else {
		pi.readIndex += alignedSize
	}
}

type Pickle struct {
	header              []byte
	headerSize          int
	capacityAfterHeader int
	writeOffset         int
}

func newPickle(buffer []byte) *Pickle {
	p := &Pickle{}
	if buffer != nil {
		p.header = buffer
		p.headerSize = len(buffer) - p.getPayloadSize()
		p.capacityAfterHeader = CAPACITY_READ_ONLY
		p.writeOffset = 0
		if p.headerSize > len(buffer) {
			p.headerSize = 0
		}
		if p.headerSize != alignInt(p.headerSize, SIZE_UINT32) {
			p.headerSize = 0
		}
		if p.headerSize == 0 {
			p.header = make([]byte, 0)
		}
	} else {
		p.header = make([]byte, 0)
		p.headerSize = SIZE_UINT32
		p.capacityAfterHeader = 0
		p.writeOffset = 0
		p.resize(PAYLOAD_UNIT)
		p.setPayloadSize(0)
	}
	return p
}

func CreateEmpty() *Pickle {
	return newPickle(nil)
}

func CreateFromBuffer(buffer []byte) *Pickle {
	return newPickle(buffer)
}

func (p *Pickle) GetHeader() []byte {
	return p.header
}

func (p *Pickle) GetHeaderSize() int {
	return p.headerSize
}

func (p *Pickle) CreateIterator() *PickleIterator {
	return NewPickleIterator(p)
}

func (p *Pickle) ToBuffer() []byte {
	return p.header[:p.headerSize+p.getPayloadSize()]
}

func (p *Pickle) WriteBool(value bool) {
	if value {
		p.WriteInt(1)
	}
	p.WriteInt(0)
}

func (p *Pickle) WriteInt(value int32) {
	p.writeBytes(value, binary.LittleEndian)
}

func (p *Pickle) WriteUInt32(value uint32) {
	p.writeBytes(value, binary.LittleEndian)
}

func (p *Pickle) WriteInt64(value int64) {
	p.writeBytes(value, binary.LittleEndian)
}

func (p *Pickle) WriteUInt64(value uint64) {
	p.writeBytes(value, binary.LittleEndian)
}

func (p *Pickle) WriteFloat(value float32) {
	p.writeBytes(math.Float32bits(value), binary.LittleEndian)
}

func (p *Pickle) WriteDouble(value float64) {
	p.writeBytes(math.Float64bits(value), binary.LittleEndian)
}

func (p *Pickle) WriteString(value string) {
	length := len(value)
	p.WriteInt(int32(length))
	p.writeBytes([]byte(value), binary.LittleEndian)
}

func (p *Pickle) setPayloadSize(payloadSize int) {
	binary.LittleEndian.PutUint32(p.header[:4], uint32(payloadSize))
}

func (p *Pickle) getPayloadSize() int {
	return int(binary.LittleEndian.Uint32(p.header[:4]))
}

func (p *Pickle) writeBytes(data any, byteOrder binary.ByteOrder) {
	length := binary.Size(data)
	if length == -1 {
		panic(errors.New("unsupported data type"))
	}
	dataLength := alignInt(length, SIZE_UINT32)
	newSize := p.writeOffset + dataLength
	if newSize > p.capacityAfterHeader {
		p.resize(max(p.capacityAfterHeader*2, newSize))
	}
	binary.Encode(p.header[p.headerSize+p.writeOffset:], byteOrder, data)

	endOffset := p.headerSize + p.writeOffset + length
	for i := endOffset; i < endOffset+dataLength-length; i++ {
		p.header[i] = 0
	}

	p.setPayloadSize(newSize)
	p.writeOffset = newSize
}

func (p *Pickle) resize(newCapacity int) {
	newCapacity = alignInt(newCapacity, PAYLOAD_UNIT)
	p.header = append(p.header, make([]byte, newCapacity)...)
	p.capacityAfterHeader = newCapacity
}