mirror of
https://github.com/AlexxIT/go2rtc.git
synced 2025-10-25 17:20:27 +08:00
Total rework RTMP client
This commit is contained in:
Alexey Khit
408
pkg/rtmp/rtmp.go
Normal file
408
pkg/rtmp/rtmp.go
Normal file
@@ -0,0 +1,408 @@
|
||||
package rtmp
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
|
||||
"github.com/AlexxIT/go2rtc/pkg/flv/amf"
|
||||
)
|
||||
|
||||
const (
|
||||
MsgSetPacketSize = 1
|
||||
MsgServerBandwidth = 5
|
||||
MsgClientBandwidth = 6
|
||||
MsgCommand = 20
|
||||
|
||||
//MsgAck = 3
|
||||
//MsgControl = 4
|
||||
//MsgAudioPacket = 8
|
||||
//MsgVideoPacket = 9
|
||||
//MsgDataExt = 15
|
||||
//MsgCommandExt = 17
|
||||
//MsgData = 18
|
||||
)
|
||||
|
||||
var ErrResponse = errors.New("rtmp: wrong response")
|
||||
|
||||
// rtmp - implements flv.Transport
|
||||
type rtmp struct {
|
||||
url string
|
||||
app string
|
||||
stream string
|
||||
pktSize uint32
|
||||
|
||||
headers map[uint32]*header
|
||||
|
||||
conn net.Conn
|
||||
rd io.Reader
|
||||
}
|
||||
|
||||
type header struct {
|
||||
msgTime uint32
|
||||
msgSize uint32
|
||||
msgType byte
|
||||
}
|
||||
|
||||
func (c *rtmp) ReadTag() (byte, uint32, []byte, error) {
|
||||
hdrType, sid, err := c.readHeader()
|
||||
if err != nil {
|
||||
return 0, 0, nil, err
|
||||
}
|
||||
|
||||
// storing header information for support header type 3
|
||||
hdr, ok := c.headers[sid]
|
||||
if !ok {
|
||||
hdr = &header{}
|
||||
c.headers[sid] = hdr
|
||||
}
|
||||
|
||||
var b []byte
|
||||
|
||||
// https://en.wikipedia.org/wiki/Real-Time_Messaging_Protocol#Packet_structure
|
||||
switch hdrType {
|
||||
case 0: // 12 byte header (full header)
|
||||
if b, err = c.readSize(11); err != nil {
|
||||
return 0, 0, nil, err
|
||||
}
|
||||
_ = b[7]
|
||||
hdr.msgTime = Uint24(b) // timestamp
|
||||
hdr.msgSize = Uint24(b[3:]) // msgdatalen
|
||||
hdr.msgType = b[6] // msgtypeid
|
||||
_ = binary.BigEndian.Uint32(b[7:]) // msgsid
|
||||
|
||||
case 1: // 8 bytes - like type b00, not including message ID (4 last bytes)
|
||||
if b, err = c.readSize(7); err != nil {
|
||||
return 0, 0, nil, err
|
||||
}
|
||||
_ = b[6]
|
||||
hdr.msgTime = Uint24(b) // timestamp
|
||||
hdr.msgSize = Uint24(b[3:]) // msgdatalen
|
||||
hdr.msgType = b[6] // msgtypeid
|
||||
|
||||
case 2: // 4 bytes - Basic Header and timestamp (3 bytes) are included
|
||||
if b, err = c.readSize(3); err != nil {
|
||||
return 0, 0, nil, err
|
||||
}
|
||||
hdr.msgTime = Uint24(b) // timestamp
|
||||
|
||||
case 3: // 1 byte - only the Basic Header is included
|
||||
// use here hdr from previous msg with same session ID (sid)
|
||||
}
|
||||
|
||||
timeMS := hdr.msgTime
|
||||
if timeMS == 0xFFFFFF {
|
||||
if b, err = c.readSize(4); err != nil {
|
||||
return 0, 0, nil, err
|
||||
}
|
||||
timeMS = binary.BigEndian.Uint32(b)
|
||||
}
|
||||
|
||||
//log.Printf("[rtmp] hdrType=%d sid=%d msdTime=%d msgSize=%d msgType=%d", hdrType, sid, hdr.msgTime, hdr.msgSize, hdr.msgType)
|
||||
|
||||
// 1. Response zero size
|
||||
if hdr.msgSize == 0 {
|
||||
return hdr.msgType, timeMS, nil, nil
|
||||
}
|
||||
|
||||
b = make([]byte, hdr.msgSize)
|
||||
|
||||
// 2. Response small packet
|
||||
if c.pktSize == 0 || hdr.msgSize < c.pktSize {
|
||||
if _, err = io.ReadFull(c.rd, b); err != nil {
|
||||
return 0, 0, nil, err
|
||||
}
|
||||
return hdr.msgType, timeMS, b, nil
|
||||
}
|
||||
|
||||
// 3. Response big packet
|
||||
var i0 uint32
|
||||
for i1 := c.pktSize; i1 < hdr.msgSize; i1 += c.pktSize {
|
||||
if _, err = io.ReadFull(c.rd, b[i0:i1]); err != nil {
|
||||
return 0, 0, nil, err
|
||||
}
|
||||
|
||||
if _, _, err = c.readHeader(); err != nil {
|
||||
return 0, 0, nil, err
|
||||
}
|
||||
|
||||
if hdr.msgTime == 0xFFFFFF {
|
||||
if _, err = c.readSize(4); err != nil {
|
||||
return 0, 0, nil, err
|
||||
}
|
||||
}
|
||||
|
||||
i0 = i1
|
||||
}
|
||||
|
||||
if _, err = io.ReadFull(c.rd, b[i0:]); err != nil {
|
||||
return 0, 0, nil, err
|
||||
}
|
||||
|
||||
return hdr.msgType, timeMS, b, nil
|
||||
}
|
||||
|
||||
func (c *rtmp) Close() error {
|
||||
return c.conn.Close()
|
||||
}
|
||||
|
||||
func (c *rtmp) init() error {
|
||||
if err := c.handshake(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := c.sendConfig(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
c.headers = map[uint32]*header{}
|
||||
|
||||
if err := c.sendConnect(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := c.sendPlay(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *rtmp) handshake() error {
|
||||
// simple handshake without real random and check response
|
||||
const randomSize = 4 + 4 + 1528
|
||||
|
||||
b := make([]byte, 1+randomSize)
|
||||
b[0] = 0x03
|
||||
if _, err := c.conn.Write(b); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if _, err := io.ReadFull(c.rd, b); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if b[0] != 3 {
|
||||
return errors.New("rtmp: wrong handshake")
|
||||
}
|
||||
|
||||
if _, err := c.conn.Write(b[1:]); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if _, err := io.ReadFull(c.rd, b[1:]); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *rtmp) sendConfig() error {
|
||||
b := make([]byte, 5)
|
||||
binary.BigEndian.PutUint32(b, 65536)
|
||||
if err := c.sendRequest(MsgSetPacketSize, 0, b[:4]); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
binary.BigEndian.PutUint32(b, 2500000)
|
||||
if err := c.sendRequest(MsgServerBandwidth, 0, b[:4]); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
binary.BigEndian.PutUint32(b, 10000000) // ack size
|
||||
b[4] = 2 // limit type
|
||||
if err := c.sendRequest(MsgClientBandwidth, 0, b); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *rtmp) sendConnect() error {
|
||||
msg := amf.AMF{}
|
||||
msg.WriteString("connect")
|
||||
msg.WriteNumber(1)
|
||||
msg.WriteObject(map[string]any{
|
||||
"app": c.app,
|
||||
"flashVer": "MAC 32,0,0,0",
|
||||
"tcUrl": c.url,
|
||||
"fpad": false, // ?
|
||||
"capabilities": 15, // ?
|
||||
"audioCodecs": 4071, // ?
|
||||
"videoCodecs": 252, // ?
|
||||
"videoFunction": 1, // ?
|
||||
})
|
||||
|
||||
if err := c.sendRequest(MsgCommand, 0, msg.Bytes()); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
s, err := c.waitCode()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if s != "NetConnection.Connect.Success" {
|
||||
return errors.New("rtmp: wrong code: " + s)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *rtmp) sendPlay() error {
|
||||
msg := amf.NewWriter()
|
||||
msg.WriteString("createStream")
|
||||
msg.WriteNumber(2)
|
||||
msg.WriteNull()
|
||||
|
||||
if err := c.sendRequest(MsgCommand, 0, msg.Bytes()); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
args, err := c.waitResponse()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if len(args) < 4 {
|
||||
return ErrResponse
|
||||
}
|
||||
|
||||
sid, _ := args[3].(float64)
|
||||
|
||||
msg = amf.NewWriter()
|
||||
msg.WriteString("play")
|
||||
msg.WriteNumber(3)
|
||||
msg.WriteNull()
|
||||
msg.WriteString(c.stream)
|
||||
|
||||
if err = c.sendRequest(MsgCommand, uint32(sid), msg.Bytes()); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
s, err := c.waitCode()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
switch s {
|
||||
case "NetStream.Play.Start", "NetStream.Play.Reset":
|
||||
return nil
|
||||
}
|
||||
|
||||
return errors.New("rtmp: wrong code: " + s)
|
||||
}
|
||||
|
||||
func (c *rtmp) sendRequest(msgType byte, streamID uint32, payload []byte) error {
|
||||
n := len(payload)
|
||||
b := make([]byte, 12+n)
|
||||
_ = b[12]
|
||||
|
||||
switch msgType {
|
||||
case MsgSetPacketSize, MsgServerBandwidth, MsgClientBandwidth:
|
||||
b[0] = 0x02 // chunk ID
|
||||
case MsgCommand:
|
||||
if streamID == 0 {
|
||||
b[0] = 0x03 // chunk ID
|
||||
} else {
|
||||
b[0] = 0x08 // chunk ID
|
||||
}
|
||||
}
|
||||
|
||||
//PutUint24(b[1:], 0) // timestamp
|
||||
PutUint24(b[4:], uint32(n)) // payload size
|
||||
b[7] = msgType // message type
|
||||
binary.BigEndian.PutUint32(b[8:], streamID) // message stream ID
|
||||
copy(b[12:], payload)
|
||||
|
||||
if _, err := c.conn.Write(b); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *rtmp) readHeader() (byte, uint32, error) {
|
||||
b, err := c.readSize(1)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
hdrType := b[0] >> 6
|
||||
sid := uint32(b[0] & 0b111111)
|
||||
|
||||
switch sid {
|
||||
case 0:
|
||||
if b, err = c.readSize(1); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
sid = 64 + uint32(b[0])
|
||||
case 1:
|
||||
if b, err = c.readSize(2); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
sid = 64 + uint32(binary.BigEndian.Uint16(b))
|
||||
}
|
||||
|
||||
return hdrType, sid, nil
|
||||
}
|
||||
|
||||
func (c *rtmp) readSize(n uint32) ([]byte, error) {
|
||||
b := make([]byte, n)
|
||||
if _, err := io.ReadAtLeast(c.rd, b, int(n)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
func (c *rtmp) waitResponse() ([]any, error) {
|
||||
for {
|
||||
msgType, _, b, err := c.ReadTag()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
switch msgType {
|
||||
case MsgSetPacketSize:
|
||||
c.pktSize = binary.BigEndian.Uint32(b)
|
||||
|
||||
case MsgCommand:
|
||||
return amf.NewReader(b).ReadItems()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *rtmp) waitCode() (string, error) {
|
||||
args, err := c.waitResponse()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if len(args) < 4 {
|
||||
return "", ErrResponse
|
||||
}
|
||||
|
||||
m, _ := args[3].(map[string]any)
|
||||
if m == nil {
|
||||
return "", ErrResponse
|
||||
}
|
||||
|
||||
v, _ := m["code"]
|
||||
if v == nil {
|
||||
return "", ErrResponse
|
||||
}
|
||||
|
||||
s, _ := v.(string)
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func PutUint24(b []byte, v uint32) {
|
||||
_ = b[2]
|
||||
b[0] = byte(v >> 16)
|
||||
b[1] = byte(v >> 8)
|
||||
b[2] = byte(v)
|
||||
}
|
||||
|
||||
func Uint24(b []byte) uint32 {
|
||||
_ = b[2]
|
||||
return uint32(b[0])<<16 | uint32(b[1])<<8 | uint32(b[2])
|
||||
}
|
||||
Reference in New Issue
Block a user