av/rtmp/amf/amf.go

/*
NAME
  amf.go

DESCRIPTION
  Action Message Format (AMF) encoding/decoding functions.
  See https://en.wikipedia.org/wiki/Action_Message_Format.

AUTHORS
  Saxon Nelson-Milton <saxon@ausocean.org>
  Dan Kortschak <dan@ausocean.org>
  Jake Lane <jake@ausocean.org>
  Alan Noble <alan@ausocean.org>

LICENSE
  amf.go is Copyright (C) 2017-2019 the Australian Ocean Lab (AusOcean)

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
  for more details.

  You should have received a copy of the GNU General Public License
  along with revid in gpl.txt. If not, see http://www.gnu.org/licenses.

  Derived from librtmp under the GNU Lesser General Public License 2.1
  Copyright (C) 2005-2008 Team XBMC http://www.xbmc.org
  Copyright (C) 2008-2009 Andrej Stepanchuk
  Copyright (C) 2009-2010 Howard Chu
*/

// amf implements Action Message Format (AMF) encoding and decoding.
// See https://en.wikipedia.org/wiki/Action_Message_Format.
package amf

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

// AMF data types.
// NB: we export these sparingly.
const (
	typeNumber = iota
	typeBoolean
	typeString
	TypeObject    // ToDo: consider not exporting this
	typeMovieClip // reserved, not implemented
	TypeNull      // ToDo: consider not exporting this
	typeUndefined
	typeReference
	typeEcmaArray
	TypeObjectEnd // ToDo: consider not exporting this
	typeStrictArray
	typeDate
	typeLongString
	typeUnsupported
	typeRecordset // reserved, not implemented
	typeXmlDoc
	typeTypedObject
	typeAvmplus // reserved, not implemented
	typeInvalid = 0xff
)

// AMF represents an AMF message (object), which is simply a collection of properties.
type Object struct {
	Props []Property // ToDo: consider not exporting this
}

// Property represents an AMF property.
type Property struct {
	name   string
	dtype  uint8
	number float64
	str    string
	obj    Object
}

// AMF errors.
var (
	ErrShortBuffer    = errors.New("amf: short buffer")
	ErrEndOfBuffer    = errors.New("amf: end of buffer")
	ErrInvalidType    = errors.New("amf: invalid type")
	ErrUnimplemented  = errors.New("amf: unimplemented feature")
	ErrDecodingName   = errors.New("amf: name decoding error")
	ErrDecodingString = errors.New("amf: string decoding error")
	ErrUnexpectedEnd  = errors.New("amf: unexpected end")
)

// Basic decoding funcions.
func DecodeInt16(data []byte) uint16 {
	return uint16(binary.BigEndian.Uint16(data))
}

func DecodeInt24(data []byte) uint32 {
	return uint32(data[0])<<16 | uint32(data[1])<<8 | uint32(data[2])
}

func DecodeInt32(data []byte) uint32 {
	return uint32(binary.BigEndian.Uint32(data))
}

func DecodeString(data []byte) string {
	n := DecodeInt16(data)
	return string(data[2 : 2+n])
}

func DecodeLongString(data []byte) string {
	n := DecodeInt32(data)
	return string(data[2 : 2+n])
}

func DecodeNumber(data []byte) float64 {
	return math.Float64frombits(binary.BigEndian.Uint64(data))
}

func DecodeBoolean(data []byte) bool {
	return data[0] != 0
}

// Basic encoding functions.
func EncodeInt24(dst []byte, val int32) ([]byte, error) {
	if len(dst) < 3 {
		return nil, ErrShortBuffer
	}
	dst[0] = byte(val >> 16)
	dst[1] = byte(val >> 8)
	dst[2] = byte(val)
	if len(dst) == 3 {
		return nil, ErrEndOfBuffer
	}
	return dst[3:], nil
}

func EncodeInt32(dst []byte, val int32) ([]byte, error) {
	if len(dst) < 4 {
		return nil, ErrShortBuffer
	}
	binary.BigEndian.PutUint32(dst, uint32(val))
	if len(dst) == 4 {
		return nil, ErrEndOfBuffer
	}
	return dst[4:], nil
}

func EncodeString(dst []byte, val string) ([]byte, error) {
	const typeSize = 1
	if len(val) < 65536 && len(val)+typeSize+binary.Size(int16(0)) > len(dst) {
		return nil, ErrShortBuffer
	}

	if len(val)+typeSize+binary.Size(int32(0)) > len(dst) {
		return nil, ErrShortBuffer
	}

	if len(val) < 65536 {
		dst[0] = typeString
		dst = dst[1:]
		binary.BigEndian.PutUint16(dst[:2], uint16(len(val)))
		dst = dst[2:]
		copy(dst, val)
		if len(dst) == len(val) {
			return nil, ErrEndOfBuffer
		}
		return dst[len(val):], nil
	}

	dst[0] = typeLongString
	dst = dst[1:]
	binary.BigEndian.PutUint32(dst[:4], uint32(len(val)))
	dst = dst[4:]
	copy(dst, val)
	if len(dst) == len(val) {
		return nil, ErrEndOfBuffer
	}
	return dst[len(val):], nil
}

func EncodeNumber(dst []byte, val float64) ([]byte, error) {
	if len(dst) < 9 {
		return nil, ErrShortBuffer
	}
	dst[0] = typeNumber
	dst = dst[1:]
	binary.BigEndian.PutUint64(dst, math.Float64bits(val))
	return dst[8:], nil
}

func EncodeBoolean(dst []byte, val bool) ([]byte, error) {
	if len(dst) < 2 {
		return nil, ErrShortBuffer
	}
	dst[0] = typeBoolean
	if val {
		dst[1] = 1
	}
	if len(dst) == 2 {
		return nil, ErrEndOfBuffer
	}
	return dst[2:], nil

}

func EncodeNamedString(dst []byte, key, val string) ([]byte, error) {
	if 2+len(key) > len(dst) {
		return nil, ErrShortBuffer
	}
	binary.BigEndian.PutUint16(dst[:2], uint16(len(key)))
	dst = dst[2:]
	copy(dst, key)
	if len(key) == len(dst) {
		return nil, ErrEndOfBuffer
	}
	return EncodeString(dst[len(key):], val)
}

func EncodeNamedNumber(dst []byte, key string, val float64) ([]byte, error) {
	if 2+len(key) > len(dst) {
		return nil, ErrShortBuffer
	}
	binary.BigEndian.PutUint16(dst[:2], uint16(len(key)))
	dst = dst[2:]
	copy(dst, key)
	if len(key) == len(dst) {
		return nil, ErrEndOfBuffer
	}
	return EncodeNumber(dst[len(key):], val)
}

func EncodeNamedBoolean(dst []byte, key string, val bool) ([]byte, error) {
	if 2+len(key) > len(dst) {
		return nil, ErrShortBuffer
	}
	binary.BigEndian.PutUint16(dst[:2], uint16(len(key)))
	dst = dst[2:]
	copy(dst, key)
	if len(key) == len(dst) {
		return nil, ErrEndOfBuffer
	}
	return EncodeBoolean(dst[len(key):], val)
}

// Property functions.

func PropSetName(prop *Property, name string) {
	prop.name = name
}

func PropGetNumber(prop *Property) float64 {
	return prop.number
}

func PropGetString(prop *Property) string {
	if prop.dtype == typeString {
		return prop.str
	}
	return ""
}

func PropGetObject(prop *Property, obj *Object) {
	if prop.dtype == TypeObject {
		*obj = prop.obj
	} else {
		*obj = Object{}
	}
}

// PropEncode encodes a property.
func PropEncode(p *Property, dst []byte) ([]byte, error) {
	if p.dtype == typeInvalid {
		return nil, ErrShortBuffer
	}

	if p.dtype != TypeNull && len(p.name)+2+1 >= len(dst) {
		return nil, ErrShortBuffer
	}

	if p.dtype != TypeNull && len(p.name) != 0 {
		binary.BigEndian.PutUint16(dst[:2], uint16(len(p.name)))
		dst = dst[2:]
		copy(dst, p.name)
		dst = dst[len(p.name):]
	}

	var err error
	switch p.dtype {
	case typeNumber:
		dst, err = EncodeNumber(dst, p.number)
	case typeBoolean:
		dst, err = EncodeBoolean(dst, p.number != 0)
	case typeString:
		dst, err = EncodeString(dst, p.str)
	case TypeNull:
		if len(dst) < 2 {
			return nil, ErrShortBuffer
		}
		dst[0] = TypeNull
		dst = dst[1:]
	case TypeObject:
		dst, err = Encode(&p.obj, dst)
	case typeEcmaArray:
		dst, err = EncodeEcmaArray(&p.obj, dst)
	case typeStrictArray:
		dst, err = EncodeArray(&p.obj, dst)
	default:
		dst, err = nil, ErrInvalidType
	}
	return dst, err
}

// PropDecode decodes a property, returning the number of bytes consumed from the supplied buffer.
func PropDecode(prop *Property, data []byte, decodeName bool) (int, error) {
	prop.name = ""

	sz := len(data)
	if len(data) == 0 {
		return 0, ErrEndOfBuffer
	}

	if decodeName {
		if len(data) < 4 {
			return 0, ErrShortBuffer
		}
		n := DecodeInt16(data[:2])
		if int(n) > len(data)-2 {
			return 0, ErrDecodingName
		}

		prop.name = DecodeString(data)
		data = data[2+n:]
	}

	if len(data) == 0 {
		return 0, ErrEndOfBuffer
	}

	prop.dtype = uint8(data[0])
	data = data[1:]

	switch prop.dtype {
	case typeNumber:
		if len(data) < 8 {
			return 0, ErrShortBuffer
		}
		prop.number = DecodeNumber(data[:8])
		data = data[8:]

	case typeBoolean:
		return 0, ErrUnimplemented

	case typeString:
		n := DecodeInt16(data[:2])
		if len(data) < int(n+2) {
			return 0, ErrShortBuffer
		}
		prop.str = DecodeString(data)
		data = data[2+n:]

	case TypeObject:
		n, err := Decode(&prop.obj, data, true)
		if err != nil {
			return 0, err
		}
		data = data[n:]


	case TypeNull, typeUndefined, typeUnsupported:
		prop.dtype = TypeNull

	case typeEcmaArray:
		data = data[4:]
		n, err := Decode(&prop.obj, data, true)
		if err != nil {
			return 0, err
		}
		data = data[n:]

	case TypeObjectEnd:
		return 0, ErrUnexpectedEnd

	default:
		return 0, ErrUnimplemented
	}

	return sz - len(data), nil
}

func PropReset(prop *Property) {
	if prop.dtype == TypeObject || prop.dtype == typeEcmaArray ||
		prop.dtype == typeStrictArray {
		Reset(&prop.obj)
	} else {
		prop.str = ""
	}
	prop.dtype = typeInvalid
}

// Encode serializes an Object into its AMF representation.
func Encode(obj *Object, dst []byte) ([]byte, error) {
	if len(dst) < 5 {
		return nil, ErrShortBuffer
	}

	dst[0] = TypeObject
	dst = dst[1:]

	for i := 0; i < len(obj.Props); i++ {
		var err error
		dst, err = PropEncode(&obj.Props[i], dst)
		if err != nil {
			return nil, err
		}
	}

	if len(dst) < 4 {
		return nil, ErrShortBuffer
	}
	return EncodeInt24(dst, TypeObjectEnd)
}

func EncodeEcmaArray(obj *Object, dst []byte) ([]byte, error) {
	if len(dst) < 5 {
		return nil, ErrShortBuffer
	}

	dst[0] = typeEcmaArray
	dst = dst[1:]
	binary.BigEndian.PutUint32(dst[:4], uint32(len(obj.Props)))
	dst = dst[4:]

	for i := 0; i < len(obj.Props); i++ {
		var err error
		dst, err = PropEncode(&obj.Props[i], dst)
		if err != nil {
			return nil, err
		}
	}

	if len(dst) < 4 {
		return nil, ErrShortBuffer
	}
	return EncodeInt24(dst, TypeObjectEnd)
}

func EncodeArray(obj *Object, dst []byte) ([]byte, error) {
	if len(dst) < 5 {
		return nil, ErrShortBuffer
	}

	dst[0] = typeStrictArray
	dst = dst[1:]
	binary.BigEndian.PutUint32(dst[:4], uint32(len(obj.Props)))
	dst = dst[4:]

	for i := 0; i < len(obj.Props); i++ {
		var err error
		dst, err = PropEncode(&obj.Props[i], dst)
		if err != nil {
			return nil, err
		}
	}

	return dst, nil
}

func Decode(obj *Object, data []byte, decodeName bool) (int, error) {
	sz := len(data)

	obj.Props = obj.Props[:0]
	for len(data) != 0 {
		if len(data) >= 3 && DecodeInt24(data[:3]) == TypeObjectEnd {
			data = data[3:]
			break
		}

		var prop Property
		n, err := PropDecode(&prop, data, decodeName)
		if err != nil {
			return 0, err
		}
		data = data[n:]
		obj.Props = append(obj.Props, prop)
	}

	return sz - len(data), nil
}

func GetProp(obj *Object, name string, idx int32) *Property {
	if idx >= 0 {
		if idx < int32(len(obj.Props)) {
			return &obj.Props[idx]
		}
	} else {
		for i, p := range obj.Props {
			if p.name == name {
				return &obj.Props[i]
			}
		}
	}
	return &Property{}
}

func Reset(obj *Object) {
	for i := range obj.Props {
		PropReset(&obj.Props[i])
	}
	*obj = Object{}
}