914 lines
23 KiB
Go
914 lines
23 KiB
Go
package capnp
|
|
|
|
import (
|
|
"bufio"
|
|
"encoding/binary"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"sync"
|
|
|
|
"zombiezen.com/go/capnproto2/internal/packed"
|
|
)
|
|
|
|
// Security limits. Matches C++ implementation.
|
|
const (
|
|
defaultTraverseLimit = 64 << 20 // 64 MiB
|
|
defaultDepthLimit = 64
|
|
|
|
maxStreamSegments = 512
|
|
|
|
defaultDecodeLimit = 64 << 20 // 64 MiB
|
|
)
|
|
|
|
const maxDepth = ^uint(0)
|
|
|
|
// A Message is a tree of Cap'n Proto objects, split into one or more
|
|
// segments of contiguous memory. The only required field is Arena.
|
|
// A Message is safe to read from multiple goroutines.
|
|
type Message struct {
|
|
// rlimit must be first so that it is 64-bit aligned.
|
|
// See sync/atomic docs.
|
|
rlimit ReadLimiter
|
|
rlimitInit sync.Once
|
|
|
|
Arena Arena
|
|
|
|
// CapTable is the indexed list of the clients referenced in the
|
|
// message. Capability pointers inside the message will use this table
|
|
// to map pointers to Clients. The table is usually populated by the
|
|
// RPC system.
|
|
//
|
|
// See https://capnproto.org/encoding.html#capabilities-interfaces for
|
|
// more details on the capability table.
|
|
CapTable []Client
|
|
|
|
// TraverseLimit limits how many total bytes of data are allowed to be
|
|
// traversed while reading. Traversal is counted when a Struct or
|
|
// List is obtained. This means that calling a getter for the same
|
|
// sub-struct multiple times will cause it to be double-counted. Once
|
|
// the traversal limit is reached, pointer accessors will report
|
|
// errors. See https://capnproto.org/encoding.html#amplification-attack
|
|
// for more details on this security measure.
|
|
//
|
|
// If not set, this defaults to 64 MiB.
|
|
TraverseLimit uint64
|
|
|
|
// DepthLimit limits how deeply-nested a message structure can be.
|
|
// If not set, this defaults to 64.
|
|
DepthLimit uint
|
|
|
|
// mu protects the following fields:
|
|
mu sync.Mutex
|
|
segs map[SegmentID]*Segment
|
|
firstSeg Segment // Preallocated first segment. msg is non-nil once initialized.
|
|
}
|
|
|
|
// NewMessage creates a message with a new root and returns the first
|
|
// segment. It is an error to call NewMessage on an arena with data in it.
|
|
func NewMessage(arena Arena) (msg *Message, first *Segment, err error) {
|
|
msg = &Message{Arena: arena}
|
|
switch arena.NumSegments() {
|
|
case 0:
|
|
first, err = msg.allocSegment(wordSize)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
case 1:
|
|
first, err = msg.Segment(0)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
if len(first.data) > 0 {
|
|
return nil, nil, errHasData
|
|
}
|
|
default:
|
|
return nil, nil, errHasData
|
|
}
|
|
if first.ID() != 0 {
|
|
return nil, nil, errors.New("capnp: arena allocated first segment with non-zero ID")
|
|
}
|
|
seg, _, err := alloc(first, wordSize) // allocate root
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
if seg != first {
|
|
return nil, nil, errors.New("capnp: arena didn't allocate first word in first segment")
|
|
}
|
|
return msg, first, nil
|
|
}
|
|
|
|
// Reset resets a message to use a different arena, allowing a single
|
|
// Message to be reused for reading multiple messages. This invalidates
|
|
// any existing pointers in the Message, so use with caution.
|
|
func (m *Message) Reset(arena Arena) {
|
|
m.mu.Lock()
|
|
m.Arena = arena
|
|
m.CapTable = nil
|
|
m.segs = nil
|
|
m.firstSeg = Segment{}
|
|
m.mu.Unlock()
|
|
if m.TraverseLimit == 0 {
|
|
m.ReadLimiter().Reset(defaultTraverseLimit)
|
|
} else {
|
|
m.ReadLimiter().Reset(m.TraverseLimit)
|
|
}
|
|
}
|
|
|
|
// Root returns the pointer to the message's root object.
|
|
//
|
|
// Deprecated: Use RootPtr.
|
|
func (m *Message) Root() (Pointer, error) {
|
|
p, err := m.RootPtr()
|
|
return p.toPointer(), err
|
|
}
|
|
|
|
// RootPtr returns the pointer to the message's root object.
|
|
func (m *Message) RootPtr() (Ptr, error) {
|
|
s, err := m.Segment(0)
|
|
if err != nil {
|
|
return Ptr{}, err
|
|
}
|
|
return s.root().PtrAt(0)
|
|
}
|
|
|
|
// SetRoot sets the message's root object to p.
|
|
//
|
|
// Deprecated: Use SetRootPtr.
|
|
func (m *Message) SetRoot(p Pointer) error {
|
|
return m.SetRootPtr(toPtr(p))
|
|
}
|
|
|
|
// SetRootPtr sets the message's root object to p.
|
|
func (m *Message) SetRootPtr(p Ptr) error {
|
|
s, err := m.Segment(0)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return s.root().SetPtr(0, p)
|
|
}
|
|
|
|
// AddCap appends a capability to the message's capability table and
|
|
// returns its ID.
|
|
func (m *Message) AddCap(c Client) CapabilityID {
|
|
n := CapabilityID(len(m.CapTable))
|
|
m.CapTable = append(m.CapTable, c)
|
|
return n
|
|
}
|
|
|
|
// ReadLimiter returns the message's read limiter. Useful if you want
|
|
// to reset the traversal limit while reading.
|
|
func (m *Message) ReadLimiter() *ReadLimiter {
|
|
m.rlimitInit.Do(func() {
|
|
if m.TraverseLimit == 0 {
|
|
m.rlimit.limit = defaultTraverseLimit
|
|
} else {
|
|
m.rlimit.limit = m.TraverseLimit
|
|
}
|
|
})
|
|
return &m.rlimit
|
|
}
|
|
|
|
func (m *Message) depthLimit() uint {
|
|
if m.DepthLimit != 0 {
|
|
return m.DepthLimit
|
|
}
|
|
return defaultDepthLimit
|
|
}
|
|
|
|
// NumSegments returns the number of segments in the message.
|
|
func (m *Message) NumSegments() int64 {
|
|
return int64(m.Arena.NumSegments())
|
|
}
|
|
|
|
// Segment returns the segment with the given ID.
|
|
func (m *Message) Segment(id SegmentID) (*Segment, error) {
|
|
if isInt32Bit && id > maxInt32 {
|
|
return nil, errSegment32Bit
|
|
}
|
|
if int64(id) >= m.Arena.NumSegments() {
|
|
return nil, errSegmentOutOfBounds
|
|
}
|
|
m.mu.Lock()
|
|
if seg := m.segment(id); seg != nil {
|
|
m.mu.Unlock()
|
|
return seg, nil
|
|
}
|
|
data, err := m.Arena.Data(id)
|
|
if err != nil {
|
|
m.mu.Unlock()
|
|
return nil, err
|
|
}
|
|
seg := m.setSegment(id, data)
|
|
m.mu.Unlock()
|
|
return seg, nil
|
|
}
|
|
|
|
// segment returns the segment with the given ID.
|
|
// The caller must be holding m.mu.
|
|
func (m *Message) segment(id SegmentID) *Segment {
|
|
if m.segs == nil {
|
|
if id == 0 && m.firstSeg.msg != nil {
|
|
return &m.firstSeg
|
|
}
|
|
return nil
|
|
}
|
|
return m.segs[id]
|
|
}
|
|
|
|
// setSegment creates or updates the Segment with the given ID.
|
|
// The caller must be holding m.mu.
|
|
func (m *Message) setSegment(id SegmentID, data []byte) *Segment {
|
|
if m.segs == nil {
|
|
if id == 0 {
|
|
m.firstSeg = Segment{
|
|
id: id,
|
|
msg: m,
|
|
data: data,
|
|
}
|
|
return &m.firstSeg
|
|
}
|
|
m.segs = make(map[SegmentID]*Segment)
|
|
if m.firstSeg.msg != nil {
|
|
m.segs[0] = &m.firstSeg
|
|
}
|
|
} else if seg := m.segs[id]; seg != nil {
|
|
seg.data = data
|
|
return seg
|
|
}
|
|
seg := &Segment{
|
|
id: id,
|
|
msg: m,
|
|
data: data,
|
|
}
|
|
m.segs[id] = seg
|
|
return seg
|
|
}
|
|
|
|
// allocSegment creates or resizes an existing segment such that
|
|
// cap(seg.Data) - len(seg.Data) >= sz.
|
|
func (m *Message) allocSegment(sz Size) (*Segment, error) {
|
|
m.mu.Lock()
|
|
if m.segs == nil && m.firstSeg.msg != nil {
|
|
m.segs = make(map[SegmentID]*Segment)
|
|
m.segs[0] = &m.firstSeg
|
|
}
|
|
id, data, err := m.Arena.Allocate(sz, m.segs)
|
|
if err != nil {
|
|
m.mu.Unlock()
|
|
return nil, err
|
|
}
|
|
if isInt32Bit && id > maxInt32 {
|
|
m.mu.Unlock()
|
|
return nil, errSegment32Bit
|
|
}
|
|
seg := m.setSegment(id, data)
|
|
m.mu.Unlock()
|
|
return seg, nil
|
|
}
|
|
|
|
// alloc allocates sz zero-filled bytes. It prefers using s, but may
|
|
// use a different segment in the same message if there's not sufficient
|
|
// capacity.
|
|
func alloc(s *Segment, sz Size) (*Segment, Address, error) {
|
|
sz = sz.padToWord()
|
|
if sz > maxSize-wordSize {
|
|
return nil, 0, errOverflow
|
|
}
|
|
|
|
if !hasCapacity(s.data, sz) {
|
|
var err error
|
|
s, err = s.msg.allocSegment(sz)
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
}
|
|
|
|
addr := Address(len(s.data))
|
|
end, ok := addr.addSize(sz)
|
|
if !ok {
|
|
return nil, 0, errOverflow
|
|
}
|
|
space := s.data[len(s.data):end]
|
|
s.data = s.data[:end]
|
|
for i := range space {
|
|
space[i] = 0
|
|
}
|
|
return s, addr, nil
|
|
}
|
|
|
|
// An Arena loads and allocates segments for a Message.
|
|
type Arena interface {
|
|
// NumSegments returns the number of segments in the arena.
|
|
// This must not be larger than 1<<32.
|
|
NumSegments() int64
|
|
|
|
// Data loads the data for the segment with the given ID. IDs are in
|
|
// the range [0, NumSegments()).
|
|
// must be tightly packed in the range [0, NumSegments()).
|
|
Data(id SegmentID) ([]byte, error)
|
|
|
|
// Allocate selects a segment to place a new object in, creating a
|
|
// segment or growing the capacity of a previously loaded segment if
|
|
// necessary. If Allocate does not return an error, then the
|
|
// difference of the capacity and the length of the returned slice
|
|
// must be at least minsz. segs is a map of segment slices returned
|
|
// by the Data method keyed by ID (although the length of these slices
|
|
// may have changed by previous allocations). Allocate must not
|
|
// modify segs.
|
|
//
|
|
// If Allocate creates a new segment, the ID must be one larger than
|
|
// the last segment's ID or zero if it is the first segment.
|
|
//
|
|
// If Allocate returns an previously loaded segment's ID, then the
|
|
// arena is responsible for preserving the existing data in the
|
|
// returned byte slice.
|
|
Allocate(minsz Size, segs map[SegmentID]*Segment) (SegmentID, []byte, error)
|
|
}
|
|
|
|
type singleSegmentArena []byte
|
|
|
|
// SingleSegment returns a new arena with an expanding single-segment
|
|
// buffer. b can be used to populate the segment for reading or to
|
|
// reserve memory of a specific size. A SingleSegment arena does not
|
|
// return errors unless you attempt to access another segment.
|
|
func SingleSegment(b []byte) Arena {
|
|
ssa := new(singleSegmentArena)
|
|
*ssa = b
|
|
return ssa
|
|
}
|
|
|
|
func (ssa *singleSegmentArena) NumSegments() int64 {
|
|
return 1
|
|
}
|
|
|
|
func (ssa *singleSegmentArena) Data(id SegmentID) ([]byte, error) {
|
|
if id != 0 {
|
|
return nil, errSegmentOutOfBounds
|
|
}
|
|
return *ssa, nil
|
|
}
|
|
|
|
func (ssa *singleSegmentArena) Allocate(sz Size, segs map[SegmentID]*Segment) (SegmentID, []byte, error) {
|
|
data := []byte(*ssa)
|
|
if segs[0] != nil {
|
|
data = segs[0].data
|
|
}
|
|
if len(data)%int(wordSize) != 0 {
|
|
return 0, nil, errors.New("capnp: segment size is not a multiple of word size")
|
|
}
|
|
if hasCapacity(data, sz) {
|
|
return 0, data, nil
|
|
}
|
|
inc, err := nextAlloc(int64(len(data)), int64(maxSegmentSize()), sz)
|
|
if err != nil {
|
|
return 0, nil, fmt.Errorf("capnp: alloc %d bytes: %v", sz, err)
|
|
}
|
|
buf := make([]byte, len(data), cap(data)+inc)
|
|
copy(buf, data)
|
|
*ssa = buf
|
|
return 0, *ssa, nil
|
|
}
|
|
|
|
type roSingleSegment []byte
|
|
|
|
func (ss roSingleSegment) NumSegments() int64 {
|
|
return 1
|
|
}
|
|
|
|
func (ss roSingleSegment) Data(id SegmentID) ([]byte, error) {
|
|
if id != 0 {
|
|
return nil, errSegmentOutOfBounds
|
|
}
|
|
return ss, nil
|
|
}
|
|
|
|
func (ss roSingleSegment) Allocate(sz Size, segs map[SegmentID]*Segment) (SegmentID, []byte, error) {
|
|
return 0, nil, errors.New("capnp: segment is read-only")
|
|
}
|
|
|
|
type multiSegmentArena [][]byte
|
|
|
|
// MultiSegment returns a new arena that allocates new segments when
|
|
// they are full. b can be used to populate the buffer for reading or
|
|
// to reserve memory of a specific size.
|
|
func MultiSegment(b [][]byte) Arena {
|
|
msa := new(multiSegmentArena)
|
|
*msa = b
|
|
return msa
|
|
}
|
|
|
|
// demuxArena slices b into a multi-segment arena.
|
|
func demuxArena(hdr streamHeader, data []byte) (Arena, error) {
|
|
segs := make([][]byte, int(hdr.maxSegment())+1)
|
|
for i := range segs {
|
|
sz, err := hdr.segmentSize(uint32(i))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
segs[i], data = data[:sz:sz], data[sz:]
|
|
}
|
|
return MultiSegment(segs), nil
|
|
}
|
|
|
|
func (msa *multiSegmentArena) NumSegments() int64 {
|
|
return int64(len(*msa))
|
|
}
|
|
|
|
func (msa *multiSegmentArena) Data(id SegmentID) ([]byte, error) {
|
|
if int64(id) >= int64(len(*msa)) {
|
|
return nil, errSegmentOutOfBounds
|
|
}
|
|
return (*msa)[id], nil
|
|
}
|
|
|
|
func (msa *multiSegmentArena) Allocate(sz Size, segs map[SegmentID]*Segment) (SegmentID, []byte, error) {
|
|
var total int64
|
|
for i, data := range *msa {
|
|
id := SegmentID(i)
|
|
if s := segs[id]; s != nil {
|
|
data = s.data
|
|
}
|
|
if hasCapacity(data, sz) {
|
|
return id, data, nil
|
|
}
|
|
total += int64(cap(data))
|
|
if total < 0 {
|
|
// Overflow.
|
|
return 0, nil, fmt.Errorf("capnp: alloc %d bytes: message too large", sz)
|
|
}
|
|
}
|
|
n, err := nextAlloc(total, 1<<63-1, sz)
|
|
if err != nil {
|
|
return 0, nil, fmt.Errorf("capnp: alloc %d bytes: %v", sz, err)
|
|
}
|
|
buf := make([]byte, 0, n)
|
|
id := SegmentID(len(*msa))
|
|
*msa = append(*msa, buf)
|
|
return id, buf, nil
|
|
}
|
|
|
|
// nextAlloc computes how much more space to allocate given the number
|
|
// of bytes allocated in the entire message and the requested number of
|
|
// bytes. It will always return a multiple of wordSize. max must be a
|
|
// multiple of wordSize. The sum of curr and the returned size will
|
|
// always be less than max.
|
|
func nextAlloc(curr, max int64, req Size) (int, error) {
|
|
if req == 0 {
|
|
return 0, nil
|
|
}
|
|
maxinc := int64(1<<32 - 8) // largest word-aligned Size
|
|
if isInt32Bit {
|
|
maxinc = 1<<31 - 8 // largest word-aligned int
|
|
}
|
|
if int64(req) > maxinc {
|
|
return 0, errors.New("allocation too large")
|
|
}
|
|
req = req.padToWord()
|
|
want := curr + int64(req)
|
|
if want <= curr || want > max {
|
|
return 0, errors.New("allocation overflows message size")
|
|
}
|
|
new := curr
|
|
double := new + new
|
|
switch {
|
|
case want < 1024:
|
|
next := (1024 - curr + 7) &^ 7
|
|
if next < curr {
|
|
return int((curr + 7) &^ 7), nil
|
|
}
|
|
return int(next), nil
|
|
case want > double:
|
|
return int(req), nil
|
|
default:
|
|
for 0 < new && new < want {
|
|
new += new / 4
|
|
}
|
|
if new <= 0 {
|
|
return int(req), nil
|
|
}
|
|
delta := new - curr
|
|
if delta > maxinc {
|
|
return int(maxinc), nil
|
|
}
|
|
return int((delta + 7) &^ 7), nil
|
|
}
|
|
}
|
|
|
|
// A Decoder represents a framer that deserializes a particular Cap'n
|
|
// Proto input stream.
|
|
type Decoder struct {
|
|
r io.Reader
|
|
|
|
segbuf [msgHeaderSize]byte
|
|
hdrbuf []byte
|
|
|
|
reuse bool
|
|
buf []byte
|
|
msg Message
|
|
arena roSingleSegment
|
|
|
|
// Maximum number of bytes that can be read per call to Decode.
|
|
// If not set, a reasonable default is used.
|
|
MaxMessageSize uint64
|
|
}
|
|
|
|
// NewDecoder creates a new Cap'n Proto framer that reads from r.
|
|
func NewDecoder(r io.Reader) *Decoder {
|
|
return &Decoder{r: r}
|
|
}
|
|
|
|
// NewPackedDecoder creates a new Cap'n Proto framer that reads from a
|
|
// packed stream r.
|
|
func NewPackedDecoder(r io.Reader) *Decoder {
|
|
return NewDecoder(packed.NewReader(bufio.NewReader(r)))
|
|
}
|
|
|
|
// Decode reads a message from the decoder stream.
|
|
func (d *Decoder) Decode() (*Message, error) {
|
|
maxSize := d.MaxMessageSize
|
|
if maxSize == 0 {
|
|
maxSize = defaultDecodeLimit
|
|
}
|
|
if _, err := io.ReadFull(d.r, d.segbuf[:]); err != nil {
|
|
return nil, err
|
|
}
|
|
maxSeg := binary.LittleEndian.Uint32(d.segbuf[:])
|
|
if maxSeg > maxStreamSegments {
|
|
return nil, errTooManySegments
|
|
}
|
|
hdrSize := streamHeaderSize(maxSeg)
|
|
if hdrSize > maxSize || hdrSize > (1<<31-1) {
|
|
return nil, errDecodeLimit
|
|
}
|
|
d.hdrbuf = resizeSlice(d.hdrbuf, int(hdrSize))
|
|
copy(d.hdrbuf, d.segbuf[:])
|
|
if _, err := io.ReadFull(d.r, d.hdrbuf[msgHeaderSize:]); err != nil {
|
|
return nil, err
|
|
}
|
|
hdr, _, err := parseStreamHeader(d.hdrbuf)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
total, err := hdr.totalSize()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// TODO(someday): if total size is greater than can fit in one buffer,
|
|
// attempt to allocate buffer per segment.
|
|
if total > maxSize-hdrSize || total > (1<<31-1) {
|
|
return nil, errDecodeLimit
|
|
}
|
|
if !d.reuse {
|
|
buf := make([]byte, int(total))
|
|
if _, err := io.ReadFull(d.r, buf); err != nil {
|
|
return nil, err
|
|
}
|
|
arena, err := demuxArena(hdr, buf)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &Message{Arena: arena}, nil
|
|
}
|
|
d.buf = resizeSlice(d.buf, int(total))
|
|
if _, err := io.ReadFull(d.r, d.buf); err != nil {
|
|
return nil, err
|
|
}
|
|
var arena Arena
|
|
if hdr.maxSegment() == 0 {
|
|
d.arena = d.buf[:len(d.buf):len(d.buf)]
|
|
arena = &d.arena
|
|
} else {
|
|
var err error
|
|
arena, err = demuxArena(hdr, d.buf)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
d.msg.Reset(arena)
|
|
return &d.msg, nil
|
|
}
|
|
|
|
func resizeSlice(b []byte, size int) []byte {
|
|
if cap(b) < size {
|
|
return make([]byte, size)
|
|
}
|
|
return b[:size]
|
|
}
|
|
|
|
// ReuseBuffer causes the decoder to reuse its buffer on subsequent decodes.
|
|
// The decoder may return messages that cannot handle allocations.
|
|
func (d *Decoder) ReuseBuffer() {
|
|
d.reuse = true
|
|
}
|
|
|
|
// Unmarshal reads an unpacked serialized stream into a message. No
|
|
// copying is performed, so the objects in the returned message read
|
|
// directly from data.
|
|
func Unmarshal(data []byte) (*Message, error) {
|
|
if len(data) == 0 {
|
|
return nil, io.EOF
|
|
}
|
|
hdr, data, err := parseStreamHeader(data)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if tot, err := hdr.totalSize(); err != nil {
|
|
return nil, err
|
|
} else if tot > uint64(len(data)) {
|
|
return nil, io.ErrUnexpectedEOF
|
|
}
|
|
arena, err := demuxArena(hdr, data)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &Message{Arena: arena}, nil
|
|
}
|
|
|
|
// UnmarshalPacked reads a packed serialized stream into a message.
|
|
func UnmarshalPacked(data []byte) (*Message, error) {
|
|
if len(data) == 0 {
|
|
return nil, io.EOF
|
|
}
|
|
data, err := packed.Unpack(nil, data)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return Unmarshal(data)
|
|
}
|
|
|
|
// MustUnmarshalRoot reads an unpacked serialized stream and returns
|
|
// its root pointer. If there is any error, it panics.
|
|
//
|
|
// Deprecated: Use MustUnmarshalRootPtr.
|
|
func MustUnmarshalRoot(data []byte) Pointer {
|
|
msg, err := Unmarshal(data)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
p, err := msg.Root()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return p
|
|
}
|
|
|
|
// MustUnmarshalRootPtr reads an unpacked serialized stream and returns
|
|
// its root pointer. If there is any error, it panics.
|
|
func MustUnmarshalRootPtr(data []byte) Ptr {
|
|
msg, err := Unmarshal(data)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
p, err := msg.RootPtr()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return p
|
|
}
|
|
|
|
// An Encoder represents a framer for serializing a particular Cap'n
|
|
// Proto stream.
|
|
type Encoder struct {
|
|
w io.Writer
|
|
hdrbuf []byte
|
|
bufs [][]byte
|
|
|
|
packed bool
|
|
packbuf []byte
|
|
}
|
|
|
|
// NewEncoder creates a new Cap'n Proto framer that writes to w.
|
|
func NewEncoder(w io.Writer) *Encoder {
|
|
return &Encoder{w: w}
|
|
}
|
|
|
|
// NewPackedEncoder creates a new Cap'n Proto framer that writes to a
|
|
// packed stream w.
|
|
func NewPackedEncoder(w io.Writer) *Encoder {
|
|
return &Encoder{w: w, packed: true}
|
|
}
|
|
|
|
// Encode writes a message to the encoder stream.
|
|
func (e *Encoder) Encode(m *Message) error {
|
|
nsegs := m.NumSegments()
|
|
if nsegs == 0 {
|
|
return errMessageEmpty
|
|
}
|
|
e.bufs = append(e.bufs[:0], nil) // first element is placeholder for header
|
|
maxSeg := uint32(nsegs - 1)
|
|
hdrSize := streamHeaderSize(maxSeg)
|
|
if uint64(cap(e.hdrbuf)) < hdrSize {
|
|
e.hdrbuf = make([]byte, 0, hdrSize)
|
|
}
|
|
e.hdrbuf = appendUint32(e.hdrbuf[:0], maxSeg)
|
|
for i := int64(0); i < nsegs; i++ {
|
|
s, err := m.Segment(SegmentID(i))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
n := len(s.data)
|
|
if int64(n) > int64(maxSize) {
|
|
return errSegmentTooLarge
|
|
}
|
|
e.hdrbuf = appendUint32(e.hdrbuf, uint32(Size(n)/wordSize))
|
|
e.bufs = append(e.bufs, s.data)
|
|
}
|
|
if len(e.hdrbuf)%int(wordSize) != 0 {
|
|
e.hdrbuf = appendUint32(e.hdrbuf, 0)
|
|
}
|
|
e.bufs[0] = e.hdrbuf
|
|
if e.packed {
|
|
return e.writePacked(e.bufs)
|
|
}
|
|
return e.write(e.bufs)
|
|
}
|
|
|
|
func (e *Encoder) writePacked(bufs [][]byte) error {
|
|
for _, b := range bufs {
|
|
e.packbuf = packed.Pack(e.packbuf[:0], b)
|
|
if _, err := e.w.Write(e.packbuf); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *Message) segmentSizes() ([]Size, error) {
|
|
nsegs := m.NumSegments()
|
|
sizes := make([]Size, nsegs)
|
|
for i := int64(0); i < nsegs; i++ {
|
|
s, err := m.Segment(SegmentID(i))
|
|
if err != nil {
|
|
return sizes[:i], err
|
|
}
|
|
n := len(s.data)
|
|
if int64(n) > int64(maxSize) {
|
|
return sizes[:i], errSegmentTooLarge
|
|
}
|
|
sizes[i] = Size(n)
|
|
}
|
|
return sizes, nil
|
|
}
|
|
|
|
// Marshal concatenates the segments in the message into a single byte
|
|
// slice including framing.
|
|
func (m *Message) Marshal() ([]byte, error) {
|
|
// Compute buffer size.
|
|
// TODO(light): error out if too many segments
|
|
nsegs := m.NumSegments()
|
|
if nsegs == 0 {
|
|
return nil, errMessageEmpty
|
|
}
|
|
maxSeg := uint32(nsegs - 1)
|
|
hdrSize := streamHeaderSize(maxSeg)
|
|
sizes, err := m.segmentSizes()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// TODO(light): error out if too large
|
|
total := uint64(hdrSize) + totalSize(sizes)
|
|
|
|
// Fill in buffer.
|
|
buf := make([]byte, hdrSize, total)
|
|
// TODO: remove marshalStreamHeader and inline.
|
|
marshalStreamHeader(buf, sizes)
|
|
for i := int64(0); i < nsegs; i++ {
|
|
s, err := m.Segment(SegmentID(i))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
buf = append(buf, s.data...)
|
|
}
|
|
return buf, nil
|
|
}
|
|
|
|
// MarshalPacked marshals the message in packed form.
|
|
func (m *Message) MarshalPacked() ([]byte, error) {
|
|
data, err := m.Marshal()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
buf := make([]byte, 0, len(data))
|
|
buf = packed.Pack(buf, data)
|
|
return buf, nil
|
|
}
|
|
|
|
// Stream header sizes.
|
|
const (
|
|
msgHeaderSize = 4
|
|
segHeaderSize = 4
|
|
)
|
|
|
|
// streamHeaderSize returns the size of the header, given the
|
|
// first 32-bit number.
|
|
func streamHeaderSize(n uint32) uint64 {
|
|
return (msgHeaderSize + segHeaderSize*(uint64(n)+1) + 7) &^ 7
|
|
}
|
|
|
|
// marshalStreamHeader marshals the sizes into the byte slice, which
|
|
// must be of size streamHeaderSize(len(sizes) - 1).
|
|
//
|
|
// TODO: remove marshalStreamHeader and inline.
|
|
func marshalStreamHeader(b []byte, sizes []Size) {
|
|
binary.LittleEndian.PutUint32(b, uint32(len(sizes)-1))
|
|
for i, sz := range sizes {
|
|
loc := msgHeaderSize + i*segHeaderSize
|
|
binary.LittleEndian.PutUint32(b[loc:], uint32(sz/Size(wordSize)))
|
|
}
|
|
}
|
|
|
|
// appendUint32 appends a uint32 to a byte slice and returns the
|
|
// new slice.
|
|
func appendUint32(b []byte, v uint32) []byte {
|
|
b = append(b, 0, 0, 0, 0)
|
|
binary.LittleEndian.PutUint32(b[len(b)-4:], v)
|
|
return b
|
|
}
|
|
|
|
type streamHeader struct {
|
|
b []byte
|
|
}
|
|
|
|
// parseStreamHeader parses the header of the stream framing format.
|
|
func parseStreamHeader(data []byte) (h streamHeader, tail []byte, err error) {
|
|
if uint64(len(data)) < streamHeaderSize(0) {
|
|
return streamHeader{}, nil, io.ErrUnexpectedEOF
|
|
}
|
|
maxSeg := binary.LittleEndian.Uint32(data)
|
|
// TODO(light): check int
|
|
hdrSize := streamHeaderSize(maxSeg)
|
|
if uint64(len(data)) < hdrSize {
|
|
return streamHeader{}, nil, io.ErrUnexpectedEOF
|
|
}
|
|
return streamHeader{b: data}, data[hdrSize:], nil
|
|
}
|
|
|
|
func (h streamHeader) maxSegment() uint32 {
|
|
return binary.LittleEndian.Uint32(h.b)
|
|
}
|
|
|
|
func (h streamHeader) segmentSize(i uint32) (Size, error) {
|
|
s := binary.LittleEndian.Uint32(h.b[msgHeaderSize+i*segHeaderSize:])
|
|
sz, ok := wordSize.times(int32(s))
|
|
if !ok {
|
|
return 0, errSegmentTooLarge
|
|
}
|
|
return sz, nil
|
|
}
|
|
|
|
func (h streamHeader) totalSize() (uint64, error) {
|
|
var sum uint64
|
|
for i := uint64(0); i <= uint64(h.maxSegment()); i++ {
|
|
x, err := h.segmentSize(uint32(i))
|
|
if err != nil {
|
|
return sum, err
|
|
}
|
|
sum += uint64(x)
|
|
}
|
|
return sum, nil
|
|
}
|
|
|
|
func hasCapacity(b []byte, sz Size) bool {
|
|
return sz <= Size(cap(b)-len(b))
|
|
}
|
|
|
|
func totalSize(s []Size) uint64 {
|
|
var sum uint64
|
|
for _, sz := range s {
|
|
sum += uint64(sz)
|
|
}
|
|
return sum
|
|
}
|
|
|
|
const (
|
|
maxInt32 = 0x7fffffff
|
|
maxInt = int(^uint(0) >> 1)
|
|
|
|
isInt32Bit = maxInt == maxInt32
|
|
)
|
|
|
|
// maxSegmentSize returns the maximum permitted size of a single segment
|
|
// on this platform.
|
|
//
|
|
// This is effectively a compile-time constant, but can't be represented
|
|
// as a constant because it requires a conditional. It is trivially
|
|
// inlinable and optimizable, so should act like one.
|
|
func maxSegmentSize() Size {
|
|
if isInt32Bit {
|
|
return Size(maxInt32 - 7)
|
|
} else {
|
|
return maxSize - 7
|
|
}
|
|
}
|
|
|
|
var (
|
|
errSegmentOutOfBounds = errors.New("capnp: segment ID out of bounds")
|
|
errSegment32Bit = errors.New("capnp: segment ID larger than 31 bits")
|
|
errMessageEmpty = errors.New("capnp: marshalling an empty message")
|
|
errHasData = errors.New("capnp: NewMessage called on arena with data")
|
|
errSegmentTooLarge = errors.New("capnp: segment too large")
|
|
errTooManySegments = errors.New("capnp: too many segments to decode")
|
|
errDecodeLimit = errors.New("capnp: message too large")
|
|
)
|