stream-server/rtmp/chunk.go

338 lines
10 KiB
Go

package rtmp
import (
"net"
"time"
"encoding/binary"
"errors"
)
// data intake object, only needed since most of the work is
// just keeping track of the previous values for higher chunk formats to reuse
type ChunkStream struct {
timestamp uint32
last_msg_strm_id uint32
last_msg_len uint32
last_msg_type uint8
timedelta uint32
}
// actual message object, metadata from the chunk header
// data either spans the chunk or split across multiple with format 3
type Message struct {
data []byte
curr_bytes_read uint32
timestamp uint32
msg_type uint8
msg_len uint32
}
// just here because I got tired of constantly assigning byte buffers and didn't want to
// abstract it, plus helped figuring out how to use pointers and memory allocation on
// restricted memory
type ChunkBuffers struct {
time []byte
msg_len []byte
msg_typeid []byte
msg_streamid []byte
fmt_csid_byte []byte
csid_true []byte
}
// reads the initial variable size header that defines the chunk's format and chunkstream id
// 5.3.1.1
func read_basic_header(conn net.Conn, chunk_bufs_ptr *ChunkBuffers) (format uint8, csid uint32, err error) {
if _, err = conn.Read(chunk_bufs_ptr.fmt_csid_byte); err != nil {
return
}
format = uint8(chunk_bufs_ptr.fmt_csid_byte[0] >> 6) // get first 2 bits for format 0-3
switch chunk_bufs_ptr.fmt_csid_byte[0] & 0x3f { // last 6 bits 0-63
case 0: // csid 0 is invalid, means true csid is the next byte, + 64 for the 6 bits prior
if _, err = conn.Read(chunk_bufs_ptr.csid_true[1:]); err != nil {
return
}
csid = uint32(chunk_bufs_ptr.csid_true[1]) + 64
case 1: // csid 1 is invalid, means true csid is in the next 2 bytes, reverse order (little endian) and the 64, reconstruct
if _, err = conn.Read(chunk_bufs_ptr.csid_true); err != nil {
return
}
csid = uint32(binary.LittleEndian.Uint16(chunk_bufs_ptr.csid_true)) + 64
default: // by default true csid was read
csid = uint32(chunk_bufs_ptr.fmt_csid_byte[0] & 0x3f)
}
return
}
func read_time(conn net.Conn, chunk_bufs_ptr *ChunkBuffers) (time uint32, extended_time bool, err error) {
chunk_bufs_ptr.time[0] = 0
if _, err = conn.Read(chunk_bufs_ptr.time[1:]); err != nil {
return
}
time = binary.BigEndian.Uint32(chunk_bufs_ptr.time)
if time ^ 0xffffff == 0 {
extended_time = true
}
return
}
func read_msg_len(conn net.Conn, chunk_bufs_ptr *ChunkBuffers) (msg_len uint32, err error) {
if _, err = conn.Read(chunk_bufs_ptr.msg_len[1:]); err != nil {
return
}
msg_len = binary.BigEndian.Uint32(chunk_bufs_ptr.msg_len)
return
}
func read_msg_typeid(conn net.Conn, chunk_bufs_ptr *ChunkBuffers) (msg_type uint8, err error) {
if _, err = conn.Read(chunk_bufs_ptr.msg_typeid); err != nil {
return
}
msg_type = chunk_bufs_ptr.msg_typeid[0]
return
}
func read_msg_streamid(conn net.Conn, chunk_bufs_ptr *ChunkBuffers) (msg_streamid uint32, err error) {
if _, err = conn.Read(chunk_bufs_ptr.msg_streamid); err != nil {
return
}
msg_streamid = binary.LittleEndian.Uint32(chunk_bufs_ptr.msg_streamid)
return
}
func read_time_extd(conn net.Conn, chunk_bufs_ptr *ChunkBuffers) (time uint32, err error) {
if _, err = conn.Read(chunk_bufs_ptr.time); err != nil {
return
}
time = binary.BigEndian.Uint32(chunk_bufs_ptr.time)
return
}
// msg header format 0, all data, typically for the start of a new chunkstream or reset, or all data changes
// 5.3.1.2.1
func read_message_header_0(conn net.Conn, chnk_stream_ptr *ChunkStream, msg_ptr *Message, chunk_bufs_ptr *ChunkBuffers) (error) {
var extended_time bool
var err error
// first bit of data is time in the first 3 bytes, if max value, defined as extended time
// true time value is at end of header (technically after but its still effectively in the header)
chnk_stream_ptr.timestamp, extended_time, err = read_time(conn, chunk_bufs_ptr)
if err != nil {
return err
}
msg_ptr.msg_len, err = read_msg_len(conn, chunk_bufs_ptr)
if err != nil {
return err
}
msg_ptr.data = make([]byte, msg_ptr.msg_len)
chnk_stream_ptr.last_msg_len = msg_ptr.msg_len // assign data for different formats
msg_ptr.msg_type, err = read_msg_typeid(conn, chunk_bufs_ptr)
if err != nil {
return err
}
chnk_stream_ptr.last_msg_type = msg_ptr.msg_type
chnk_stream_ptr.last_msg_strm_id, err = read_msg_streamid(conn, chunk_bufs_ptr)
if err != nil {
return err
}
if extended_time {
chnk_stream_ptr.timestamp, err = read_time_extd(conn, chunk_bufs_ptr)
if err != nil {
return err
}
}
msg_ptr.timestamp = chnk_stream_ptr.timestamp
return nil
}
// format 1, typically for the 2nd message in a previously init chunkstream with most other data changed
// no longer time but timedelta. timestamp for the message is computed from the stored timestamp when the stream was opened
// otherwise the same except the msg_stream_id is reused (unchanged)
// 5.3.1.2.2
func read_message_header_1(conn net.Conn, chnk_stream_ptr *ChunkStream, msg_ptr *Message, chunk_bufs_ptr *ChunkBuffers) (error) {
var extended_time bool
var err error
chnk_stream_ptr.timedelta, extended_time, err = read_time(conn, chunk_bufs_ptr)
if err != nil {
return err
}
msg_ptr.msg_len, err = read_msg_len(conn, chunk_bufs_ptr)
if err != nil {
return err
}
msg_ptr.data = make([]byte, msg_ptr.msg_len)
chnk_stream_ptr.last_msg_len = msg_ptr.msg_len
msg_ptr.msg_type, err = read_msg_typeid(conn, chunk_bufs_ptr)
if err != nil {
return err
}
chnk_stream_ptr.last_msg_type = msg_ptr.msg_type
if extended_time {
chnk_stream_ptr.timedelta, err = read_time_extd(conn, chunk_bufs_ptr)
if err != nil {
return err
}
}
chnk_stream_ptr.timestamp += chnk_stream_ptr.timedelta
msg_ptr.timestamp = chnk_stream_ptr.timestamp
return nil
}
// format 2, typically for when format 1 can be used but even the message len and type are identical
// 5.3.1.2.3
func read_message_header_2(conn net.Conn, chnk_stream_ptr *ChunkStream, msg_ptr *Message, chunk_bufs_ptr *ChunkBuffers) (error) {
var extended_time bool
var err error
if chnk_stream_ptr == nil {
panic(err)
}
chnk_stream_ptr.timedelta, extended_time, err = read_time(conn, chunk_bufs_ptr)
if err != nil {
return err
}
msg_ptr.msg_len = chnk_stream_ptr.last_msg_len
msg_ptr.data = make([]byte, msg_ptr.msg_len)
msg_ptr.msg_type = chnk_stream_ptr.last_msg_type
if extended_time {
chnk_stream_ptr.timedelta, err = read_time_extd(conn, chunk_bufs_ptr)
if err != nil {
return err
}
}
chnk_stream_ptr.timestamp += chnk_stream_ptr.timedelta
msg_ptr.timestamp = chnk_stream_ptr.timestamp
return nil
}
// read the actual chunkdata given the space for the message that has been constructed using the chunk headers and
// configured peer chunk size
func read_chunk_data(conn net.Conn, msg_ptr *Message, chnk_size uint32) (uint32, error) {
bytes_left := msg_ptr.msg_len - msg_ptr.curr_bytes_read
var buffer_end uint32
if bytes_left < chnk_size {
buffer_end = bytes_left + msg_ptr.curr_bytes_read
} else {
buffer_end = chnk_size + msg_ptr.curr_bytes_read
}
n, err := conn.Read(msg_ptr.data[msg_ptr.curr_bytes_read:buffer_end])
if err != nil {
return 0, err
}
msg_ptr.curr_bytes_read += uint32(n)
return uint32(n), nil
}
// wrapper function to handle all formats and chunkstreams returns a pointer to a Message (can be nil if not completed)
// n_read was there to handle ack win, but doesn't really seem to matter at all
func ReadChunk(conn net.Conn, open_chnkstrms map[uint32]*ChunkStream, open_msgs map[uint32]*Message, chnk_size uint32, chunk_bufs_ptr *ChunkBuffers, n_read *uint32) (*Message, error){
conn.SetDeadline(time.Now().Add(10 * time.Second))
var chnkstream_ptr *ChunkStream
var msg_ptr *Message
format, csid, err := read_basic_header(conn, chunk_bufs_ptr)
if err != nil {
return nil, err
}
if csid < 64 {
*n_read += 1
} else if csid < 320 {
*n_read += 2
} else {
*n_read += 3
}
switch format {
case 0:
// format 0, assume new chunkstream which will override and reset any previous chunkstream regardless
chnkstream_ptr = new(ChunkStream)
open_chnkstrms[csid] = chnkstream_ptr // assign the newly made chunkstream to the map of those currently open
msg_ptr = new(Message) // cannot be a continued message since the stream just started
if err := read_message_header_0(conn, chnkstream_ptr, msg_ptr, chunk_bufs_ptr); err != nil {
return nil, err
}
*n_read += 11
if msg_ptr.timestamp > 0xffffff {
*n_read += 4
}
case 1:
// format 1, chunkstream with this csid MUST have been opened before, but message is just starting since
// message metadata is being declared
chnkstream_ptr = open_chnkstrms[csid]
if chnkstream_ptr == nil {
return nil, errors.New("Chunkstream was not opened, bad obs")
}
msg_ptr = new(Message)
if err := read_message_header_1(conn, chnkstream_ptr, msg_ptr, chunk_bufs_ptr); err != nil {
return nil, err
}
*n_read += 7
if msg_ptr.timestamp > 0xffffff {
*n_read += 4
}
case 2:
// format 2, same as 1
chnkstream_ptr = open_chnkstrms[csid]
if chnkstream_ptr == nil {
return nil, errors.New("Chunkstream was not opened, bad obs")
}
msg_ptr = new(Message)
if err := read_message_header_2(conn, chnkstream_ptr, msg_ptr, chunk_bufs_ptr); err != nil {
return nil, err
}
*n_read += 3
if msg_ptr.timestamp > 0xffffff {
*n_read += 4
}
case 3:
// format 3 has no header, only happens when either a previous message over the csid was too large for the chunksize
// or a new message with the same exact metadata including timedelta is being sent.
chnkstream_ptr = open_chnkstrms[csid]
if chnkstream_ptr == nil {
return nil, errors.New("Chunkstream was not opened, bad obs")
}
msg_prev, ok := open_msgs[chnkstream_ptr.last_msg_strm_id]
if !ok {
msg_ptr = new(Message)
msg_ptr.msg_type = chnkstream_ptr.last_msg_type
msg_ptr.msg_len = chnkstream_ptr.last_msg_len
msg_ptr.data = make([]byte, msg_ptr.msg_len)
chnkstream_ptr.timestamp += chnkstream_ptr.timedelta
msg_ptr.timestamp = chnkstream_ptr.timestamp
} else {
msg_ptr = msg_prev
}
}
n, err := read_chunk_data(conn, msg_ptr, chnk_size)
if err != nil {
return nil, err
}
*n_read += n
conn.SetDeadline(time.Time{})
if msg_ptr.curr_bytes_read < msg_ptr.msg_len { // if the full message was not read, save it as an open message for the next chunk
open_msgs[chnkstream_ptr.last_msg_strm_id] = msg_ptr
return nil, nil
} else { // else a full message was read, remove it from the list of open messages if it was there at all and return it
delete(open_msgs, chnkstream_ptr.last_msg_strm_id)
return msg_ptr, nil
}
}