stream-server/rtmp/chunk2.go

package rtmp

import (
	"net"
	"encoding/binary"
	"time"
)

// counterpart to the read_basic_header from chunk.go, basically the reverse of that
func write_basic_header(chunk_buffer *[]byte, format uint8, csid uint32) {
	if csid < 64 {
		(*chunk_buffer)[0] = (format << 6) + uint8(csid)
	} else if csid < 320 {
		(*chunk_buffer)[0] = format << 6
		(*chunk_buffer)[1] = uint8(csid - 64)
	} else {
		(*chunk_buffer)[0] = (format << 6) + 1
		binary.LittleEndian.PutUint16((*chunk_buffer)[1:3], uint16(csid - 64))
	}
}

// only consider writing format 0 for new, or format 3 for cut off. As the server only has to send
// protocol messages, no real point optimizing for a demo
func write_message_header_0(chunk_buffer *[]byte, msg_strmid uint32, msg_ptr *Message, basic_header_size uint32) {
	msg_len_inter := make([]byte, 4) // uint32 to uint24 hack
	binary.BigEndian.PutUint32(msg_len_inter, msg_ptr.msg_len)
	copy((*chunk_buffer)[basic_header_size + 3:basic_header_size + 6],  msg_len_inter[1:]) // skip timestamp since it will always be 0 for protocol and control msgs
	(*chunk_buffer)[basic_header_size + 6] = msg_ptr.msg_type
	binary.LittleEndian.PutUint32((*chunk_buffer)[basic_header_size + 7:basic_header_size + 11], msg_strmid)
}

// write the actual data
func write_chunk_data(chunk_buffer *[]byte, data []byte, bytes_buffered *uint32, chunk_data_size uint32, header_size uint32) {
	copy((*chunk_buffer)[header_size:header_size+chunk_data_size], data[*bytes_buffered:*bytes_buffered + chunk_data_size])
	*bytes_buffered += chunk_data_size
}

// chunk writing wrapper
func WriteChunk(conn net.Conn, csid uint32, msg_strmid uint32, msg_ptr *Message, chunk_size uint32) (error) {
	conn.SetDeadline(time.Now().Add(10 * time.Second))

	var i uint32 = 0 // track number of bytes written
	for i < msg_ptr.msg_len { // until message is fully sent
		var chunk_data_size uint32
		// check if remaining message needs to be chunked or not
		if msg_ptr.msg_len - i > chunk_size {
			chunk_data_size = chunk_size
		} else {
			chunk_data_size = msg_ptr.msg_len - i
		}

		basic_header_size := uint32(1)
		if csid >= 64 {
			basic_header_size += 1
		}
		if csid >= 320 {
			basic_header_size += 1
		}

		var chunk_buffer []byte
		var header_size uint32
		if i == 0 { // if chunkstream is to be opened/reset (format 0 = message header)
			chunk_buffer = make([]byte, basic_header_size + 11 + chunk_data_size)
			write_basic_header(&chunk_buffer, 0, csid)
			write_message_header_0(&chunk_buffer, msg_strmid, msg_ptr, basic_header_size)
			header_size = basic_header_size + 11
		} else {
			chunk_buffer = make([]byte, basic_header_size + chunk_data_size)
			write_basic_header(&chunk_buffer, 3, csid)
			header_size = basic_header_size
		}
		write_chunk_data(&chunk_buffer, msg_ptr.data, &i, chunk_data_size, header_size)
		if _, err := conn.Write(chunk_buffer); err != nil { // expunge buffer over network connection
			return err
		}
	}
	conn.SetDeadline(time.Time{})

	return nil
}