123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549 |
- package msgp
-
- import (
- "fmt"
- "math"
- )
-
- const (
- // Complex64Extension is the extension number used for complex64
- Complex64Extension = 3
-
- // Complex128Extension is the extension number used for complex128
- Complex128Extension = 4
-
- // TimeExtension is the extension number used for time.Time
- TimeExtension = 5
- )
-
- // our extensions live here
- var extensionReg = make(map[int8]func() Extension)
-
- // RegisterExtension registers extensions so that they
- // can be initialized and returned by methods that
- // decode `interface{}` values. This should only
- // be called during initialization. f() should return
- // a newly-initialized zero value of the extension. Keep in
- // mind that extensions 3, 4, and 5 are reserved for
- // complex64, complex128, and time.Time, respectively,
- // and that MessagePack reserves extension types from -127 to -1.
- //
- // For example, if you wanted to register a user-defined struct:
- //
- // msgp.RegisterExtension(10, func() msgp.Extension { &MyExtension{} })
- //
- // RegisterExtension will panic if you call it multiple times
- // with the same 'typ' argument, or if you use a reserved
- // type (3, 4, or 5).
- func RegisterExtension(typ int8, f func() Extension) {
- switch typ {
- case Complex64Extension, Complex128Extension, TimeExtension:
- panic(fmt.Sprint("msgp: forbidden extension type:", typ))
- }
- if _, ok := extensionReg[typ]; ok {
- panic(fmt.Sprint("msgp: RegisterExtension() called with typ", typ, "more than once"))
- }
- extensionReg[typ] = f
- }
-
- // ExtensionTypeError is an error type returned
- // when there is a mis-match between an extension type
- // and the type encoded on the wire
- type ExtensionTypeError struct {
- Got int8
- Want int8
- }
-
- // Error implements the error interface
- func (e ExtensionTypeError) Error() string {
- return fmt.Sprintf("msgp: error decoding extension: wanted type %d; got type %d", e.Want, e.Got)
- }
-
- // Resumable returns 'true' for ExtensionTypeErrors
- func (e ExtensionTypeError) Resumable() bool { return true }
-
- func errExt(got int8, wanted int8) error {
- return ExtensionTypeError{Got: got, Want: wanted}
- }
-
- // Extension is the interface fulfilled
- // by types that want to define their
- // own binary encoding.
- type Extension interface {
- // ExtensionType should return
- // a int8 that identifies the concrete
- // type of the extension. (Types <0 are
- // officially reserved by the MessagePack
- // specifications.)
- ExtensionType() int8
-
- // Len should return the length
- // of the data to be encoded
- Len() int
-
- // MarshalBinaryTo should copy
- // the data into the supplied slice,
- // assuming that the slice has length Len()
- MarshalBinaryTo([]byte) error
-
- UnmarshalBinary([]byte) error
- }
-
- // RawExtension implements the Extension interface
- type RawExtension struct {
- Data []byte
- Type int8
- }
-
- // ExtensionType implements Extension.ExtensionType, and returns r.Type
- func (r *RawExtension) ExtensionType() int8 { return r.Type }
-
- // Len implements Extension.Len, and returns len(r.Data)
- func (r *RawExtension) Len() int { return len(r.Data) }
-
- // MarshalBinaryTo implements Extension.MarshalBinaryTo,
- // and returns a copy of r.Data
- func (r *RawExtension) MarshalBinaryTo(d []byte) error {
- copy(d, r.Data)
- return nil
- }
-
- // UnmarshalBinary implements Extension.UnmarshalBinary,
- // and sets r.Data to the contents of the provided slice
- func (r *RawExtension) UnmarshalBinary(b []byte) error {
- if cap(r.Data) >= len(b) {
- r.Data = r.Data[0:len(b)]
- } else {
- r.Data = make([]byte, len(b))
- }
- copy(r.Data, b)
- return nil
- }
-
- // WriteExtension writes an extension type to the writer
- func (mw *Writer) WriteExtension(e Extension) error {
- l := e.Len()
- var err error
- switch l {
- case 0:
- o, err := mw.require(3)
- if err != nil {
- return err
- }
- mw.buf[o] = mext8
- mw.buf[o+1] = 0
- mw.buf[o+2] = byte(e.ExtensionType())
- case 1:
- o, err := mw.require(2)
- if err != nil {
- return err
- }
- mw.buf[o] = mfixext1
- mw.buf[o+1] = byte(e.ExtensionType())
- case 2:
- o, err := mw.require(2)
- if err != nil {
- return err
- }
- mw.buf[o] = mfixext2
- mw.buf[o+1] = byte(e.ExtensionType())
- case 4:
- o, err := mw.require(2)
- if err != nil {
- return err
- }
- mw.buf[o] = mfixext4
- mw.buf[o+1] = byte(e.ExtensionType())
- case 8:
- o, err := mw.require(2)
- if err != nil {
- return err
- }
- mw.buf[o] = mfixext8
- mw.buf[o+1] = byte(e.ExtensionType())
- case 16:
- o, err := mw.require(2)
- if err != nil {
- return err
- }
- mw.buf[o] = mfixext16
- mw.buf[o+1] = byte(e.ExtensionType())
- default:
- switch {
- case l < math.MaxUint8:
- o, err := mw.require(3)
- if err != nil {
- return err
- }
- mw.buf[o] = mext8
- mw.buf[o+1] = byte(uint8(l))
- mw.buf[o+2] = byte(e.ExtensionType())
- case l < math.MaxUint16:
- o, err := mw.require(4)
- if err != nil {
- return err
- }
- mw.buf[o] = mext16
- big.PutUint16(mw.buf[o+1:], uint16(l))
- mw.buf[o+3] = byte(e.ExtensionType())
- default:
- o, err := mw.require(6)
- if err != nil {
- return err
- }
- mw.buf[o] = mext32
- big.PutUint32(mw.buf[o+1:], uint32(l))
- mw.buf[o+5] = byte(e.ExtensionType())
- }
- }
- // we can only write directly to the
- // buffer if we're sure that it
- // fits the object
- if l <= mw.bufsize() {
- o, err := mw.require(l)
- if err != nil {
- return err
- }
- return e.MarshalBinaryTo(mw.buf[o:])
- }
- // here we create a new buffer
- // just large enough for the body
- // and save it as the write buffer
- err = mw.flush()
- if err != nil {
- return err
- }
- buf := make([]byte, l)
- err = e.MarshalBinaryTo(buf)
- if err != nil {
- return err
- }
- mw.buf = buf
- mw.wloc = l
- return nil
- }
-
- // peek at the extension type, assuming the next
- // kind to be read is Extension
- func (m *Reader) peekExtensionType() (int8, error) {
- p, err := m.R.Peek(2)
- if err != nil {
- return 0, err
- }
- spec := sizes[p[0]]
- if spec.typ != ExtensionType {
- return 0, badPrefix(ExtensionType, p[0])
- }
- if spec.extra == constsize {
- return int8(p[1]), nil
- }
- size := spec.size
- p, err = m.R.Peek(int(size))
- if err != nil {
- return 0, err
- }
- return int8(p[size-1]), nil
- }
-
- // peekExtension peeks at the extension encoding type
- // (must guarantee at least 1 byte in 'b')
- func peekExtension(b []byte) (int8, error) {
- spec := sizes[b[0]]
- size := spec.size
- if spec.typ != ExtensionType {
- return 0, badPrefix(ExtensionType, b[0])
- }
- if len(b) < int(size) {
- return 0, ErrShortBytes
- }
- // for fixed extensions,
- // the type information is in
- // the second byte
- if spec.extra == constsize {
- return int8(b[1]), nil
- }
- // otherwise, it's in the last
- // part of the prefix
- return int8(b[size-1]), nil
- }
-
- // ReadExtension reads the next object from the reader
- // as an extension. ReadExtension will fail if the next
- // object in the stream is not an extension, or if
- // e.Type() is not the same as the wire type.
- func (m *Reader) ReadExtension(e Extension) (err error) {
- var p []byte
- p, err = m.R.Peek(2)
- if err != nil {
- return
- }
- lead := p[0]
- var read int
- var off int
- switch lead {
- case mfixext1:
- if int8(p[1]) != e.ExtensionType() {
- err = errExt(int8(p[1]), e.ExtensionType())
- return
- }
- p, err = m.R.Peek(3)
- if err != nil {
- return
- }
- err = e.UnmarshalBinary(p[2:])
- if err == nil {
- _, err = m.R.Skip(3)
- }
- return
-
- case mfixext2:
- if int8(p[1]) != e.ExtensionType() {
- err = errExt(int8(p[1]), e.ExtensionType())
- return
- }
- p, err = m.R.Peek(4)
- if err != nil {
- return
- }
- err = e.UnmarshalBinary(p[2:])
- if err == nil {
- _, err = m.R.Skip(4)
- }
- return
-
- case mfixext4:
- if int8(p[1]) != e.ExtensionType() {
- err = errExt(int8(p[1]), e.ExtensionType())
- return
- }
- p, err = m.R.Peek(6)
- if err != nil {
- return
- }
- err = e.UnmarshalBinary(p[2:])
- if err == nil {
- _, err = m.R.Skip(6)
- }
- return
-
- case mfixext8:
- if int8(p[1]) != e.ExtensionType() {
- err = errExt(int8(p[1]), e.ExtensionType())
- return
- }
- p, err = m.R.Peek(10)
- if err != nil {
- return
- }
- err = e.UnmarshalBinary(p[2:])
- if err == nil {
- _, err = m.R.Skip(10)
- }
- return
-
- case mfixext16:
- if int8(p[1]) != e.ExtensionType() {
- err = errExt(int8(p[1]), e.ExtensionType())
- return
- }
- p, err = m.R.Peek(18)
- if err != nil {
- return
- }
- err = e.UnmarshalBinary(p[2:])
- if err == nil {
- _, err = m.R.Skip(18)
- }
- return
-
- case mext8:
- p, err = m.R.Peek(3)
- if err != nil {
- return
- }
- if int8(p[2]) != e.ExtensionType() {
- err = errExt(int8(p[2]), e.ExtensionType())
- return
- }
- read = int(uint8(p[1]))
- off = 3
-
- case mext16:
- p, err = m.R.Peek(4)
- if err != nil {
- return
- }
- if int8(p[3]) != e.ExtensionType() {
- err = errExt(int8(p[3]), e.ExtensionType())
- return
- }
- read = int(big.Uint16(p[1:]))
- off = 4
-
- case mext32:
- p, err = m.R.Peek(6)
- if err != nil {
- return
- }
- if int8(p[5]) != e.ExtensionType() {
- err = errExt(int8(p[5]), e.ExtensionType())
- return
- }
- read = int(big.Uint32(p[1:]))
- off = 6
-
- default:
- err = badPrefix(ExtensionType, lead)
- return
- }
-
- p, err = m.R.Peek(read + off)
- if err != nil {
- return
- }
- err = e.UnmarshalBinary(p[off:])
- if err == nil {
- _, err = m.R.Skip(read + off)
- }
- return
- }
-
- // AppendExtension appends a MessagePack extension to the provided slice
- func AppendExtension(b []byte, e Extension) ([]byte, error) {
- l := e.Len()
- var o []byte
- var n int
- switch l {
- case 0:
- o, n = ensure(b, 3)
- o[n] = mext8
- o[n+1] = 0
- o[n+2] = byte(e.ExtensionType())
- return o[:n+3], nil
- case 1:
- o, n = ensure(b, 3)
- o[n] = mfixext1
- o[n+1] = byte(e.ExtensionType())
- n += 2
- case 2:
- o, n = ensure(b, 4)
- o[n] = mfixext2
- o[n+1] = byte(e.ExtensionType())
- n += 2
- case 4:
- o, n = ensure(b, 6)
- o[n] = mfixext4
- o[n+1] = byte(e.ExtensionType())
- n += 2
- case 8:
- o, n = ensure(b, 10)
- o[n] = mfixext8
- o[n+1] = byte(e.ExtensionType())
- n += 2
- case 16:
- o, n = ensure(b, 18)
- o[n] = mfixext16
- o[n+1] = byte(e.ExtensionType())
- n += 2
- default:
- switch {
- case l < math.MaxUint8:
- o, n = ensure(b, l+3)
- o[n] = mext8
- o[n+1] = byte(uint8(l))
- o[n+2] = byte(e.ExtensionType())
- n += 3
- case l < math.MaxUint16:
- o, n = ensure(b, l+4)
- o[n] = mext16
- big.PutUint16(o[n+1:], uint16(l))
- o[n+3] = byte(e.ExtensionType())
- n += 4
- default:
- o, n = ensure(b, l+6)
- o[n] = mext32
- big.PutUint32(o[n+1:], uint32(l))
- o[n+5] = byte(e.ExtensionType())
- n += 6
- }
- }
- return o, e.MarshalBinaryTo(o[n:])
- }
-
- // ReadExtensionBytes reads an extension from 'b' into 'e'
- // and returns any remaining bytes.
- // Possible errors:
- // - ErrShortBytes ('b' not long enough)
- // - ExtensionTypeError{} (wire type not the same as e.Type())
- // - TypeError{} (next object not an extension)
- // - InvalidPrefixError
- // - An umarshal error returned from e.UnmarshalBinary
- func ReadExtensionBytes(b []byte, e Extension) ([]byte, error) {
- l := len(b)
- if l < 3 {
- return b, ErrShortBytes
- }
- lead := b[0]
- var (
- sz int // size of 'data'
- off int // offset of 'data'
- typ int8
- )
- switch lead {
- case mfixext1:
- typ = int8(b[1])
- sz = 1
- off = 2
- case mfixext2:
- typ = int8(b[1])
- sz = 2
- off = 2
- case mfixext4:
- typ = int8(b[1])
- sz = 4
- off = 2
- case mfixext8:
- typ = int8(b[1])
- sz = 8
- off = 2
- case mfixext16:
- typ = int8(b[1])
- sz = 16
- off = 2
- case mext8:
- sz = int(uint8(b[1]))
- typ = int8(b[2])
- off = 3
- if sz == 0 {
- return b[3:], e.UnmarshalBinary(b[3:3])
- }
- case mext16:
- if l < 4 {
- return b, ErrShortBytes
- }
- sz = int(big.Uint16(b[1:]))
- typ = int8(b[3])
- off = 4
- case mext32:
- if l < 6 {
- return b, ErrShortBytes
- }
- sz = int(big.Uint32(b[1:]))
- typ = int8(b[5])
- off = 6
- default:
- return b, badPrefix(ExtensionType, lead)
- }
-
- if typ != e.ExtensionType() {
- return b, errExt(typ, e.ExtensionType())
- }
-
- // the data of the extension starts
- // at 'off' and is 'sz' bytes long
- if len(b[off:]) < sz {
- return b, ErrShortBytes
- }
- tot := off + sz
- return b[tot:], e.UnmarshalBinary(b[off:tot])
- }
|