nl-im-service/internal/mediaserver/rtmp_amf0.go

/**
 * package mediaserver
 *
 * AMF0 (Action Message Format 0) 编解码实现
 * 用于 RTMP 命令消息的序列化和反序列化
 */
package mediaserver

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

// AMF0 数据类型标记
const (
	AMF0_NUMBER        = 0x00 // 8 bytes double
	AMF0_BOOLEAN       = 0x01 // 1 byte
	AMF0_STRING        = 0x02 // 2 bytes length + data
	AMF0_OBJECT        = 0x03 // key-value pairs
	AMF0_MOVIECLIP     = 0x04 // reserved
	AMF0_NULL          = 0x05 // no data
	AMF0_UNDEFINED     = 0x06 // no data
	AMF0_REFERENCE     = 0x07 // 2 bytes
	AMF0_ECMA_ARRAY    = 0x08 // associative array
	AMF0_OBJECT_END    = 0x09 // object end marker
	AMF0_STRICT_ARRAY  = 0x0A // strict array
	AMF0_DATE          = 0x0B // 8 bytes double + 2 bytes timezone
	AMF0_LONG_STRING   = 0x0C // 4 bytes length + data
	AMF0_UNSUPPORTED   = 0x0D
	AMF0_RECORDSET     = 0x0E // reserved
	AMF0_XML_DOCUMENT  = 0x0F
	AMF0_TYPED_OBJECT  = 0x10
	AMF0_AVMPLUS       = 0x11 // switch to AMF3
)

// AMF0Object 表示一个 AMF0 对象
type AMF0Object map[string]interface{}

// AMF0Encoder AMF0 编码器
type AMF0Encoder struct {
	buf *bytes.Buffer
}

// AMF0Decoder AMF0 解码器
type AMF0Decoder struct {
	data   []byte
	offset int
}

// NewAMF0Encoder 创建 AMF0 编码器
func NewAMF0Encoder() *AMF0Encoder {
	return &AMF0Encoder{
		buf: new(bytes.Buffer),
	}
}

// NewAMF0Decoder 创建 AMF0 解码器
func NewAMF0Decoder(data []byte) *AMF0Decoder {
	return &AMF0Decoder{
		data:   data,
		offset: 0,
	}
}

// Bytes 获取编码后的字节
func (e *AMF0Encoder) Bytes() []byte {
	return e.buf.Bytes()
}

// Reset 重置编码器
func (e *AMF0Encoder) Reset() {
	e.buf.Reset()
}

// EncodeNumber 编码数字
func (e *AMF0Encoder) EncodeNumber(val float64) {
	e.buf.WriteByte(AMF0_NUMBER)
	bits := math.Float64bits(val)
	binary.Write(e.buf, binary.BigEndian, bits)
}

// EncodeBoolean 编码布尔值
func (e *AMF0Encoder) EncodeBoolean(val bool) {
	e.buf.WriteByte(AMF0_BOOLEAN)
	if val {
		e.buf.WriteByte(1)
	} else {
		e.buf.WriteByte(0)
	}
}

// EncodeString 编码字符串
func (e *AMF0Encoder) EncodeString(val string) {
	data := []byte(val)
	if len(data) > 0xFFFF {
		// Long string
		e.buf.WriteByte(AMF0_LONG_STRING)
		binary.Write(e.buf, binary.BigEndian, uint32(len(data)))
	} else {
		e.buf.WriteByte(AMF0_STRING)
		binary.Write(e.buf, binary.BigEndian, uint16(len(data)))
	}
	e.buf.Write(data)
}

// EncodeNull 编码 null
func (e *AMF0Encoder) EncodeNull() {
	e.buf.WriteByte(AMF0_NULL)
}

// EncodeObject 编码对象
func (e *AMF0Encoder) EncodeObject(obj AMF0Object) {
	e.buf.WriteByte(AMF0_OBJECT)
	for key, val := range obj {
		// 写入属性名（不带类型标记）
		binary.Write(e.buf, binary.BigEndian, uint16(len(key)))
		e.buf.WriteString(key)
		// 写入值
		e.EncodeValue(val)
	}
	// 写入对象结束标记
	e.buf.Write([]byte{0, 0, AMF0_OBJECT_END})
}

// EncodeEcmaArray 编码 ECMA 数组
func (e *AMF0Encoder) EncodeEcmaArray(obj AMF0Object) {
	e.buf.WriteByte(AMF0_ECMA_ARRAY)
	binary.Write(e.buf, binary.BigEndian, uint32(len(obj)))
	for key, val := range obj {
		binary.Write(e.buf, binary.BigEndian, uint16(len(key)))
		e.buf.WriteString(key)
		e.EncodeValue(val)
	}
	e.buf.Write([]byte{0, 0, AMF0_OBJECT_END})
}

// EncodeValue 编码任意值
func (e *AMF0Encoder) EncodeValue(val interface{}) {
	switch v := val.(type) {
	case float64:
		e.EncodeNumber(v)
	case float32:
		e.EncodeNumber(float64(v))
	case int:
		e.EncodeNumber(float64(v))
	case int64:
		e.EncodeNumber(float64(v))
	case int32:
		e.EncodeNumber(float64(v))
	case uint32:
		e.EncodeNumber(float64(v))
	case bool:
		e.EncodeBoolean(v)
	case string:
		e.EncodeString(v)
	case nil:
		e.EncodeNull()
	case AMF0Object:
		e.EncodeObject(v)
	case map[string]interface{}:
		e.EncodeObject(AMF0Object(v))
	default:
		e.EncodeNull()
	}
}

// Remaining 返回剩余未解析的字节数
func (d *AMF0Decoder) Remaining() int {
	return len(d.data) - d.offset
}

// DecodeAll 解码所有值
func (d *AMF0Decoder) DecodeAll() ([]interface{}, error) {
	var values []interface{}
	for d.Remaining() > 0 {
		val, err := d.DecodeValue()
		if err != nil {
			break
		}
		values = append(values, val)
	}
	return values, nil
}

// DecodeValue 解码一个值
func (d *AMF0Decoder) DecodeValue() (interface{}, error) {
	if d.Remaining() < 1 {
		return nil, fmt.Errorf("数据不足")
	}

	marker := d.data[d.offset]
	d.offset++

	switch marker {
	case AMF0_NUMBER:
		return d.decodeNumber()
	case AMF0_BOOLEAN:
		return d.decodeBoolean()
	case AMF0_STRING:
		return d.decodeString()
	case AMF0_OBJECT:
		return d.decodeObject()
	case AMF0_NULL, AMF0_UNDEFINED:
		return nil, nil
	case AMF0_ECMA_ARRAY:
		return d.decodeEcmaArray()
	case AMF0_STRICT_ARRAY:
		return d.decodeStrictArray()
	case AMF0_LONG_STRING:
		return d.decodeLongString()
	case AMF0_DATE:
		return d.decodeDate()
	default:
		return nil, fmt.Errorf("未知的 AMF0 类型: 0x%02X", marker)
	}
}

// decodeNumber 解码数字
func (d *AMF0Decoder) decodeNumber() (float64, error) {
	if d.Remaining() < 8 {
		return 0, fmt.Errorf("数据不足以解码 number")
	}
	bits := binary.BigEndian.Uint64(d.data[d.offset : d.offset+8])
	d.offset += 8
	return math.Float64frombits(bits), nil
}

// decodeBoolean 解码布尔值
func (d *AMF0Decoder) decodeBoolean() (bool, error) {
	if d.Remaining() < 1 {
		return false, fmt.Errorf("数据不足以解码 boolean")
	}
	val := d.data[d.offset] != 0
	d.offset++
	return val, nil
}

// decodeString 解码字符串
func (d *AMF0Decoder) decodeString() (string, error) {
	if d.Remaining() < 2 {
		return "", fmt.Errorf("数据不足以解码 string length")
	}
	length := int(binary.BigEndian.Uint16(d.data[d.offset : d.offset+2]))
	d.offset += 2

	if d.Remaining() < length {
		return "", fmt.Errorf("数据不足以解码 string data")
	}
	str := string(d.data[d.offset : d.offset+length])
	d.offset += length
	return str, nil
}

// decodeLongString 解码长字符串
func (d *AMF0Decoder) decodeLongString() (string, error) {
	if d.Remaining() < 4 {
		return "", fmt.Errorf("数据不足以解码 long string length")
	}
	length := int(binary.BigEndian.Uint32(d.data[d.offset : d.offset+4]))
	d.offset += 4

	if d.Remaining() < length {
		return "", fmt.Errorf("数据不足以解码 long string data")
	}
	str := string(d.data[d.offset : d.offset+length])
	d.offset += length
	return str, nil
}

// decodeObject 解码对象
func (d *AMF0Decoder) decodeObject() (AMF0Object, error) {
	obj := make(AMF0Object)

	for {
		// 读取属性名长度
		if d.Remaining() < 2 {
			return nil, fmt.Errorf("数据不足以解码对象属性名长度")
		}
		nameLen := int(binary.BigEndian.Uint16(d.data[d.offset : d.offset+2]))
		d.offset += 2

		// 检查是否到达对象结束
		if nameLen == 0 {
			if d.Remaining() < 1 || d.data[d.offset] != AMF0_OBJECT_END {
				return nil, fmt.Errorf("对象结束标记缺失")
			}
			d.offset++
			break
		}

		// 读取属性名
		if d.Remaining() < nameLen {
			return nil, fmt.Errorf("数据不足以解码对象属性名")
		}
		name := string(d.data[d.offset : d.offset+nameLen])
		d.offset += nameLen

		// 读取属性值
		val, err := d.DecodeValue()
		if err != nil {
			return nil, fmt.Errorf("解码对象属性值失败: %w", err)
		}
		obj[name] = val
	}

	return obj, nil
}

// decodeEcmaArray 解码 ECMA 数组
func (d *AMF0Decoder) decodeEcmaArray() (AMF0Object, error) {
	if d.Remaining() < 4 {
		return nil, fmt.Errorf("数据不足以解码 ECMA 数组长度")
	}
	// 读取数组长度（但实际不使用，因为以 object end 结束）
	d.offset += 4

	return d.decodeObject()
}

// decodeStrictArray 解码严格数组
func (d *AMF0Decoder) decodeStrictArray() ([]interface{}, error) {
	if d.Remaining() < 4 {
		return nil, fmt.Errorf("数据不足以解码严格数组长度")
	}
	length := int(binary.BigEndian.Uint32(d.data[d.offset : d.offset+4]))
	d.offset += 4

	arr := make([]interface{}, length)
	for i := 0; i < length; i++ {
		val, err := d.DecodeValue()
		if err != nil {
			return nil, fmt.Errorf("解码数组元素失败: %w", err)
		}
		arr[i] = val
	}
	return arr, nil
}

// decodeDate 解码日期
func (d *AMF0Decoder) decodeDate() (float64, error) {
	if d.Remaining() < 10 {
		return 0, fmt.Errorf("数据不足以解码 date")
	}
	bits := binary.BigEndian.Uint64(d.data[d.offset : d.offset+8])
	d.offset += 8
	// 跳过时区信息 (2 bytes)
	d.offset += 2
	return math.Float64frombits(bits), nil
}

// EncodeAMF0 编码多个值为 AMF0 格式
func EncodeAMF0(values ...interface{}) []byte {
	encoder := NewAMF0Encoder()
	for _, val := range values {
		encoder.EncodeValue(val)
	}
	return encoder.Bytes()
}

// DecodeAMF0 从 AMF0 格式解码多个值
func DecodeAMF0(data []byte) ([]interface{}, error) {
	decoder := NewAMF0Decoder(data)
	return decoder.DecodeAll()
}

// EncodeConnectResult 编码 connect 响应
func EncodeConnectResult(transactionID float64) []byte {
	encoder := NewAMF0Encoder()

	// _result
	encoder.EncodeString("_result")
	encoder.EncodeNumber(transactionID)

	// Properties
	encoder.EncodeObject(AMF0Object{
		"fmsVer":       "FMS/3,5,7,7009",
		"capabilities": float64(31),
		"mode":         float64(1),
	})

	// Information
	encoder.EncodeObject(AMF0Object{
		"level":          "status",
		"code":           "NetConnection.Connect.Success",
		"description":    "Connection succeeded.",
		"objectEncoding": float64(0),
	})

	return encoder.Bytes()
}

// EncodeCreateStreamResult 编码 createStream 响应
func EncodeCreateStreamResult(transactionID float64, streamID float64) []byte {
	encoder := NewAMF0Encoder()
	encoder.EncodeString("_result")
	encoder.EncodeNumber(transactionID)
	encoder.EncodeNull()
	encoder.EncodeNumber(streamID)
	return encoder.Bytes()
}

// EncodeOnStatus 编码 onStatus 消息
func EncodeOnStatus(code, level, description string) []byte {
	encoder := NewAMF0Encoder()
	encoder.EncodeString("onStatus")
	encoder.EncodeNumber(0)
	encoder.EncodeNull()
	encoder.EncodeObject(AMF0Object{
		"level":       level,
		"code":        code,
		"description": description,
	})
	return encoder.Bytes()
}

// EncodeOnBWDone 编码 onBWDone 消息
func EncodeOnBWDone() []byte {
	encoder := NewAMF0Encoder()
	encoder.EncodeString("onBWDone")
	encoder.EncodeNumber(0)
	encoder.EncodeNull()
	return encoder.Bytes()
}