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
*/

// Package 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, as defined by the AMF specification.
// NB: we export these sparingly.
const (
	typeNumber      = 0x00
	typeBoolean     = 0x01
	typeString      = 0x02
	TypeObject      = 0x03
	typeMovieClip   = 0x04
	TypeNull        = 0x05
	typeUndefined   = 0x06
	typeReference   = 0x07
	typeEcmaArray   = 0x08
	TypeObjectEnd   = 0x09
	typeStrictArray = 0x0A
	typeDate        = 0x0B
	typeLongString  = 0x0C
	typeUnsupported = 0x0D
	typeRecordset   = 0x0E
	typeXmlDoc      = 0x0F
	typeTypedObject = 0x10
	typeAvmplus     = 0x11
	typeInvalid     = 0xff
)

// AMF represents an AMF message (object), which is simply a collection of properties.
type Object struct {
	Properties []Property
}

// Property represents an AMF property.
type Property struct {
	Name   string
	Type   uint8
	Number float64
	String string
	Object 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")
	ErrPropertyNotFound = errors.New("amf: property not found")
)

// DecodeInt16 decodes a 16-bit integer.
func DecodeInt16(buf []byte) uint16 {
	return uint16(binary.BigEndian.Uint16(buf))
}

// DecodeInt24 decodes a 24-bit integer.
func DecodeInt24(buf []byte) uint32 {
	return uint32(buf[0])<<16 | uint32(buf[1])<<8 | uint32(buf[2])
}

// DecodeInt32 decodes a 32-bit integer.
func DecodeInt32(buf []byte) uint32 {
	return uint32(binary.BigEndian.Uint32(buf))
}

// DecodeString decodes a string that is less than 2^16 bytes long.
func DecodeString(buf []byte) string {
	n := DecodeInt16(buf)
	return string(buf[2 : 2+n])
}

// DecodeLongString decodes a long string.
func DecodeLongString(buf []byte) string {
	n := DecodeInt32(buf)
	return string(buf[2 : 2+n])
}

// DecodeNumber decodes a 64-bit floating-point number.
func DecodeNumber(buf []byte) float64 {
	return math.Float64frombits(binary.BigEndian.Uint64(buf))
}

// DecodeBoolean decodes a boolean.
func DecodeBoolean(buf []byte) bool {
	return buf[0] != 0
}

// EncodeInt24 encodes a 24-bit integer.
func EncodeInt24(buf []byte, val int32) ([]byte, error) {
	if len(buf) < 3 {
		return nil, ErrShortBuffer
	}
	buf[0] = byte(val >> 16)
	buf[1] = byte(val >> 8)
	buf[2] = byte(val)
	if len(buf) == 3 {
		return nil, ErrEndOfBuffer
	}
	return buf[3:], nil
}

// EncodeInt32 encodes a 32-bit integer.
func EncodeInt32(buf []byte, val int32) ([]byte, error) {
	if len(buf) < 4 {
		return nil, ErrShortBuffer
	}
	binary.BigEndian.PutUint32(buf, uint32(val))
	if len(buf) == 4 {
		return nil, ErrEndOfBuffer
	}
	return buf[4:], nil
}

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

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

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

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

// EncodeNumber encodes a 64-bit floating-point number.
func EncodeNumber(buf []byte, val float64) ([]byte, error) {
	if len(buf) < 9 {
		return nil, ErrShortBuffer
	}
	buf[0] = typeNumber
	buf = buf[1:]
	binary.BigEndian.PutUint64(buf, math.Float64bits(val))
	return buf[8:], nil
}

// EncodeBoolean encodes a boolean.
func EncodeBoolean(buf []byte, val bool) ([]byte, error) {
	if len(buf) < 2 {
		return nil, ErrShortBuffer
	}
	buf[0] = typeBoolean
	if val {
		buf[1] = 1
	} else {
		buf[1] = 0
	}
	if len(buf) == 2 {
		return nil, ErrEndOfBuffer
	}
	return buf[2:], nil

}

// EncodeNamedString encodes a named string, where key is the name and val is the string value.
func EncodeNamedString(buf []byte, key, val string) ([]byte, error) {
	if 2+len(key) > len(buf) {
		return nil, ErrShortBuffer
	}
	binary.BigEndian.PutUint16(buf[:2], uint16(len(key)))
	buf = buf[2:]
	copy(buf, key)
	if len(key) == len(buf) {
		return nil, ErrEndOfBuffer
	}
	return EncodeString(buf[len(key):], val)
}

// EncodeNamedNumber encodes a named number, where key is the name and val is the number value.
func EncodeNamedNumber(buf []byte, key string, val float64) ([]byte, error) {
	if 2+len(key) > len(buf) {
		return nil, ErrShortBuffer
	}
	binary.BigEndian.PutUint16(buf[:2], uint16(len(key)))
	buf = buf[2:]
	copy(buf, key)
	if len(key) == len(buf) {
		return nil, ErrEndOfBuffer
	}
	return EncodeNumber(buf[len(key):], val)
}

// EncodeNamedNumber encodes a named boolean, where key is the name and val is the boolean value.
func EncodeNamedBoolean(buf []byte, key string, val bool) ([]byte, error) {
	if 2+len(key) > len(buf) {
		return nil, ErrShortBuffer
	}
	binary.BigEndian.PutUint16(buf[:2], uint16(len(key)))
	buf = buf[2:]
	copy(buf, key)
	if len(key) == len(buf) {
		return nil, ErrEndOfBuffer
	}
	return EncodeBoolean(buf[len(key):], val)
}

// PropEncode encodes a property.
func PropEncode(p *Property, buf []byte) ([]byte, error) {
	if p.Type != TypeNull && len(p.Name)+2+1 >= len(buf) {
		return nil, ErrShortBuffer
	}

	if p.Type != TypeNull && len(p.Name) != 0 {
		binary.BigEndian.PutUint16(buf[:2], uint16(len(p.Name)))
		buf = buf[2:]
		copy(buf, p.Name)
		buf = buf[len(p.Name):]
	}

	switch p.Type {
	case typeNumber:
		return EncodeNumber(buf, p.Number)
	case typeBoolean:
		return EncodeBoolean(buf, p.Number != 0)
	case typeString:
		return EncodeString(buf, p.String)
	case TypeNull:
		if len(buf) < 2 {
			return nil, ErrShortBuffer
		}
		buf[0] = TypeNull
		buf = buf[1:]
	case TypeObject:
		return Encode(&p.Object, buf)
	case typeEcmaArray:
		return EncodeEcmaArray(&p.Object, buf)
	case typeStrictArray:
		return EncodeArray(&p.Object, buf)
	default:
		return nil, ErrInvalidType
	}
	return buf, nil
}

// PropDecode decodes a property, returning the number of bytes consumed from the supplied buffer.
func PropDecode(prop *Property, buf []byte, decodeName bool) (int, error) {
	sz := len(buf)
	if len(buf) == 0 {
		return 0, ErrEndOfBuffer
	}

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

		prop.Name = DecodeString(buf)
		buf = buf[2+n:]
	} else {
		prop.Name = ""
	}

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

	prop.Type = uint8(buf[0])
	buf = buf[1:]

	switch prop.Type {
	case typeNumber:
		if len(buf) < 8 {
			return 0, ErrShortBuffer
		}
		prop.Number = DecodeNumber(buf[:8])
		buf = buf[8:]

	case typeBoolean:
		if len(buf) < 1 {
			return 0, ErrShortBuffer
		}
		prop.Number = float64(uint8(buf[0]))
		buf = buf[1:]

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

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

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

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

	case TypeObjectEnd:
		return 0, ErrUnexpectedEnd

	default:
		return 0, ErrUnimplemented
	}

	return sz - len(buf), nil
}

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

	buf[0] = TypeObject
	buf = buf[1:]

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

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

// EncodeEcmaArray encodes an ECMA array.
func EncodeEcmaArray(obj *Object, buf []byte) ([]byte, error) {
	if len(buf) < 5 {
		return nil, ErrShortBuffer
	}

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

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

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

// EncodeArray encodes an array.
func EncodeArray(obj *Object, buf []byte) ([]byte, error) {
	if len(buf) < 5 {
		return nil, ErrShortBuffer
	}

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

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

	return buf, nil
}

// Decode decodes an object. Property names are only decoded if decodeName is true.
func Decode(obj *Object, buf []byte, decodeName bool) (int, error) {
	sz := len(buf)

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

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

	return sz - len(buf), nil
}

// GetProperty returns a property, either by its index when idx is non-negative, or by its name otherwise.
// If the requested property is not found or the type does not match, an ErrPropertyNotFound error is returned.
func (obj *Object) GetProperty(name string, idx int, typ uint8) (*Property, error) {
	var prop *Property
	if idx >= 0 {
		if idx < len(obj.Properties) {
			prop = &obj.Properties[idx]
		}
	} else {
		for i, p := range obj.Properties {
			if p.Name == name {
				prop = &obj.Properties[i]
				break
			}
		}
	}
	if prop == nil || prop.Type != typ {
		return nil, ErrPropertyNotFound
	}
	return prop, nil
}

// GetNumber is a wrapper for GetProperty that returns a Number property's value, if any.
func (obj *Object) GetNumber(name string, idx int) (float64, error) {
	prop, err := obj.GetProperty(name, idx, typeNumber)
	if err != nil {
		return 0, err
	}
	return prop.Number, nil
}

// GetString is a wrapper for GetProperty that returns a String property's  value, if any.
func (obj *Object) GetString(name string, idx int) (string, error) {
	prop, err := obj.GetProperty(name, idx, typeString)
	if err != nil {
		return "", err
	}
	return prop.String, nil
}

// GetObject is a wrapper for GetProperty that returns an Object property's value, if any.
func (obj *Object) GetObject(name string, idx int) (*Object, error) {
	prop, err := obj.GetProperty(name, idx, TypeObject)
	if err != nil {
		return nil, err
	}
	return &prop.Object, nil
}