monibuca/plugin/gb28181/pkg/transceiver.go

package gb28181

import (
	"errors"
	"fmt"
	"net"
	"os"
	"strings"

	"github.com/pion/rtp"
	"m7s.live/v5"
	"m7s.live/v5/pkg"
	"m7s.live/v5/pkg/task"
	"m7s.live/v5/pkg/util"
	rtp2 "m7s.live/v5/plugin/rtp/pkg"
)

const (
	StartCodePS        = 0x000001ba
	StartCodeSYS       = 0x000001bb
	StartCodeMAP       = 0x000001bc
	StartCodeVideo     = 0x000001e0
	StartCodeAudio     = 0x000001c0
	PrivateStreamCode  = 0x000001bd
	MEPGProgramEndCode = 0x000001b9
)

type PSPublisher struct {
	*m7s.Publisher
	*util.BufReader
	Receiver Receiver
}

var ErrRTPReceiveLost = errors.New("rtp receive lost")

type Receiver struct {
	task.Task
	rtp.Packet
	FeedChan   chan []byte
	psm        util.Memory
	dump       *os.File
	dumpLen    []byte
	psVideo    PSVideo
	psAudio    PSAudio
	RTPReader  *rtp2.TCP
	ListenAddr string
	listener   net.Listener
	StreamMode string // 数据流传输模式（UDP:udp传输/TCP-ACTIVE：tcp主动模式/TCP-PASSIVE：tcp被动模式）
}

func NewPSPublisher(puber *m7s.Publisher) *PSPublisher {
	ret := &PSPublisher{
		Publisher: puber,
	}
	ret.Receiver.FeedChan = make(chan []byte, 10)
	ret.BufReader = util.NewBufReaderChan(ret.Receiver.FeedChan)
	ret.Receiver.psVideo.SetAllocator(ret.Allocator)
	ret.Receiver.psAudio.SetAllocator(ret.Allocator)
	return ret
}

func (p *PSPublisher) ReadPayload() (payload util.Memory, err error) {
	payloadlen, err := p.ReadBE(2)
	if err != nil {
		return
	}
	return p.ReadBytes(payloadlen)
}

func (p *PSPublisher) Demux() {
	var payload util.Memory
	defer p.Info("demux exit")
	for {
		code, err := p.ReadBE32(4)
		if err != nil {
			return
		}
		p.Trace("demux", "code", code)
		switch code {
		case StartCodePS:
			var psl byte
			if err = p.Skip(9); err != nil {
				return
			}
			psl, err = p.ReadByte()
			if err != nil {
				return
			}
			psl &= 0x07
			if err = p.Skip(int(psl)); err != nil {
				return
			}
		case StartCodeVideo:
			payload, err = p.ReadPayload()
			var annexB *pkg.AnnexB
			annexB, err = p.Receiver.psVideo.parsePESPacket(payload)
			if annexB != nil {
				err = p.WriteVideo(annexB)
			}
		case StartCodeAudio:
			payload, err = p.ReadPayload()
			var audioFrame pkg.IAVFrame
			audioFrame, err = p.Receiver.psAudio.parsePESPacket(payload)
			if audioFrame != nil {
				err = p.WriteAudio(audioFrame)
			}
		case StartCodeMAP:
			p.decProgramStreamMap()
		case StartCodeSYS, PrivateStreamCode:
			p.ReadPayload()
		default:
			p.ReadPayload()
		}
	}
}

func (dec *PSPublisher) decProgramStreamMap() (err error) {
	dec.Receiver.psm, err = dec.ReadPayload()
	if err != nil {
		return err
	}
	var programStreamInfoLen, programStreamMapLen, elementaryStreamInfoLength uint32
	var streamType, elementaryStreamID byte
	reader := dec.Receiver.psm.NewReader()
	reader.Skip(2)
	programStreamInfoLen, err = reader.ReadBE(2)
	reader.Skip(int(programStreamInfoLen))
	programStreamMapLen, err = reader.ReadBE(2)
	for programStreamMapLen > 0 {
		streamType, err = reader.ReadByte()
		elementaryStreamID, err = reader.ReadByte()
		if elementaryStreamID >= 0xe0 && elementaryStreamID <= 0xef {
			dec.Receiver.psVideo.streamType = streamType
		} else if elementaryStreamID >= 0xc0 && elementaryStreamID <= 0xdf {
			dec.Receiver.psAudio.streamType = streamType
		}
		elementaryStreamInfoLength, err = reader.ReadBE(2)
		reader.Skip(int(elementaryStreamInfoLength))
		programStreamMapLen -= 4 + elementaryStreamInfoLength
	}
	return nil
}

func (p *Receiver) ReadRTP(rtp util.Buffer) (err error) {
	lastSeq := p.SequenceNumber
	if err = p.Unmarshal(rtp); err != nil {
		p.Error("unmarshal error", "err", err)
		return
	}
	if lastSeq == 0 || p.SequenceNumber == lastSeq+1 {
		if p.Enabled(p, task.TraceLevel) {
			p.Trace("rtp", "len", rtp.Len(), "seq", p.SequenceNumber, "payloadType", p.PayloadType, "ssrc", p.SSRC)
		}
		copyData := make([]byte, len(p.Payload))
		copy(copyData, p.Payload)
		p.FeedChan <- copyData
		return
	}
	return ErrRTPReceiveLost
}

func (p *Receiver) Start() (err error) {
	if strings.ToUpper(p.StreamMode) == "TCP-ACTIVE" {
		// TCP主动模式不需要监听，直接返回
		p.Info("TCP-ACTIVE mode, no need to listen")
		return nil
	}
	// TCP被动模式
	p.listener, err = net.Listen("tcp4", p.ListenAddr)
	if err != nil {
		p.Error("start listen", "err", err)
		return errors.New("start listen,err" + err.Error())
	}
	p.Info("start listen", "addr", p.ListenAddr)
	return
}

func (p *Receiver) Dispose() {
	if p.listener != nil {
		p.listener.Close()
	}
	if p.RTPReader != nil {
		p.RTPReader.Close()
	}
	if p.FeedChan != nil {
		close(p.FeedChan)
	}
}

func (p *Receiver) Go() error {
	if strings.ToUpper(p.StreamMode) == "TCP-ACTIVE" {
		// TCP主动模式，主动连接设备
		addr := p.ListenAddr
		if !strings.Contains(addr, ":") {
			addr = ":" + addr
		}
		if strings.HasPrefix(addr, ":") {
			p.Error("invalid address, missing IP", "addr", addr)
			return fmt.Errorf("invalid address %s, missing IP", addr)
		}
		p.Info("TCP-ACTIVE mode, connecting to device", "addr", addr)
		conn, err := net.Dial("tcp", addr)
		if err != nil {
			p.Error("connect to device failed", "err", err)
			return err
		}
		p.RTPReader = (*rtp2.TCP)(conn.(*net.TCPConn))
		p.Info("connected to device", "addr", conn.RemoteAddr())
		return p.RTPReader.Read(p.ReadRTP)
	}
	// TCP被动模式
	p.Info("start accept")
	conn, err := p.listener.Accept()
	if err != nil {
		p.Error("accept", "err", err)
		return err
	}
	p.RTPReader = (*rtp2.TCP)(conn.(*net.TCPConn))
	p.Info("accept", "addr", conn.RemoteAddr())
	return p.RTPReader.Read(p.ReadRTP)
}