Files
lkm/rtc/rtc_stream.go
2025-07-27 15:24:00 +08:00

113 lines
3.1 KiB
Go

package rtc
import (
"github.com/lkmio/avformat"
"github.com/lkmio/avformat/collections"
"github.com/lkmio/avformat/utils"
"github.com/lkmio/lkm/stream"
"github.com/pion/interceptor"
"github.com/pion/webrtc/v4"
"net"
)
var (
webrtcApi *webrtc.API
SupportedCodecs = map[utils.AVCodecID]interface{}{
utils.AVCodecIdH264: webrtc.MimeTypeH264,
utils.AVCodecIdH265: webrtc.MimeTypeH265,
utils.AVCodecIdAV1: webrtc.MimeTypeAV1,
utils.AVCodecIdVP8: webrtc.MimeTypeVP8,
utils.AVCodecIdVP9: webrtc.MimeTypeVP9,
utils.AVCodecIdOPUS: webrtc.MimeTypeOpus,
utils.AVCodecIdPCMALAW: webrtc.MimeTypePCMA,
utils.AVCodecIdPCMMULAW: webrtc.MimeTypePCMU,
utils.AVCodecIdADPCMG722: webrtc.MimeTypeG722,
}
)
type transStream struct {
stream.BaseTransStream
segments []stream.TransStreamSegment
}
func (t *transStream) Input(packet *avformat.AVPacket, trackIndex int) ([]*collections.ReferenceCounter[[]byte], int64, bool, error) {
t.ClearOutStreamBuffer()
var data *collections.ReferenceCounter[[]byte]
if utils.AVMediaTypeAudio == packet.MediaType {
data = collections.NewReferenceCounter(packet.Data)
} else if utils.AVMediaTypeVideo == packet.MediaType {
annexBData := avformat.AVCCPacket2AnnexB(t.FindTrackWithStreamIndex(packet.Index).Stream, packet)
data = collections.NewReferenceCounter(annexBData)
}
duration := int64(uint32(packet.GetDuration(1000)))
key := utils.AVMediaTypeVideo == packet.MediaType && packet.Key
if t.HasVideo() && stream.AppConfig.GOPCache {
// 遇到视频关键帧, 丢弃前一帧缓存
if key {
t.segments = t.segments[:0]
}
t.segments = append(t.segments, stream.TransStreamSegment{
Data: []*collections.ReferenceCounter[[]byte]{data},
TS: duration,
Key: key,
Index: trackIndex,
})
}
t.AppendOutStreamBuffer(data)
return t.OutBuffer[:t.OutBufferSize], duration, key, nil
}
func (t *transStream) ReadKeyFrameBuffer() ([]stream.TransStreamSegment, error) {
return t.segments, nil
}
func (t *transStream) WriteHeader() error {
return nil
}
func InitConfig() {
setting := webrtc.SettingEngine{}
var ips []string
ips = append(ips, stream.AppConfig.PublicIP)
udpListener, err := net.ListenUDP("udp", &net.UDPAddr{
IP: net.ParseIP(stream.AppConfig.ListenIP),
Port: stream.AppConfig.WebRtc.Port,
})
if err != nil {
panic(err)
}
// 设置公网ip和监听端口
setting.SetICEUDPMux(webrtc.NewICEUDPMux(nil, udpListener))
setting.SetNAT1To1IPs(ips, webrtc.ICECandidateTypeHost)
// 注册音视频编码器
m := &webrtc.MediaEngine{}
if err := m.RegisterDefaultCodecs(); err != nil {
panic(err)
}
i := &interceptor.Registry{}
if err := webrtc.RegisterDefaultInterceptors(m, i); err != nil {
panic(err)
}
webrtcApi = webrtc.NewAPI(webrtc.WithMediaEngine(m), webrtc.WithInterceptorRegistry(i), webrtc.WithSettingEngine(setting))
}
func NewTransStream() stream.TransStream {
t := &transStream{}
return t
}
func TransStreamFactory(_ stream.Source, _ stream.TransStreamProtocol, _ []*stream.Track, _ stream.Sink) (stream.TransStream, error) {
return NewTransStream(), nil
}