Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
11 kB · 421 lines
Go
at fable51/moq
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422// Package main contains an example.package api
import ( "context" "crypto/tls" "fmt" "io" "net" "strings" "time"
"github.com/bluenviron/gortmplib" "github.com/bluenviron/gortmplib/pkg/h264conf" "github.com/bluenviron/gortmplib/pkg/message" "github.com/bluenviron/gortsplib/v5/pkg/format" "golang.org/x/sync/errgroup" "stream.place/streamplace/pkg/config" "stream.place/streamplace/pkg/log" "stream.place/streamplace/pkg/media")
// This example shows how to:// 1. create a RTMP server// 2. accept a stream from a reader.// 3. broadcast the stream to readers.
var RTMPTimeout = 10 * time.Second
const RTMPPrefix = "/live/"
func (a *StreamplaceAPI) HandleRTMPPublisher(ctx context.Context, sc *gortmplib.ServerConn) error { err := sc.RW.(net.Conn).SetReadDeadline(time.Now().Add(RTMPTimeout)) if err != nil { return err }
if !strings.HasPrefix(sc.URL.Path, RTMPPrefix) { return fmt.Errorf("RTMP publisher is not allowed to publish to %s (must start with %s)", sc.URL.String(), RTMPPrefix) } streamKey := strings.TrimPrefix(sc.URL.Path, RTMPPrefix) mediaSigner, err := a.MakeMediaSigner(ctx, streamKey) if err != nil { return fmt.Errorf("failed to make media signer: %w", err) }
streamer := mediaSigner.Streamer() ctx = log.WithLogValues(ctx, "streamer", streamer) session := &media.RTMPSession{ EventChan: make(chan any, 1024), MediaSigner: mediaSigner, } a.rtmpSessionsLock.Lock() a.rtmpSessions[streamer] = session a.rtmpSessionsLock.Unlock()
defer func() { a.rtmpSessionsLock.Lock() delete(a.rtmpSessions, streamer) a.rtmpSessionsLock.Unlock() close(session.EventChan) }()
g, ctx := errgroup.WithContext(ctx) r, err := initializeRTMPReader(ctx, sc, session) if err != nil { return err }
stopRead := context.AfterFunc(ctx, func() { _ = sc.RW.(net.Conn).Close() }) defer stopRead() g.Go(func() error { for { if ctx.Err() != nil { return ctx.Err() } err = sc.RW.(net.Conn).SetReadDeadline(time.Now().Add(RTMPTimeout)) if err != nil { return err } err = r.Read() if err != nil { return err } } })
g.Go(func() error { return a.MediaManager.RTMPIngest(ctx, fmt.Sprintf("rtmp://%s/live/%s", a.rtmpInternalPlaybackAddr, streamer), mediaSigner) })
return g.Wait()}
// initializeRTMPReader is shared by live publishers and the unpublished shadow.func initializeRTMPReader(ctx context.Context, sc *gortmplib.ServerConn, session *media.RTMPSession) (*gortmplib.Reader, error) { r := &gortmplib.Reader{Conn: sc} if err := r.Initialize(); err != nil { return nil, err } send := func(event any) { select { case session.EventChan <- event: case <-ctx.Done(): } } for _, track := range r.Tracks() { switch track := track.(type) { case *format.H264: session.VideoTrack = track r.OnDataH264(track, func(pts, dts time.Duration, au [][]byte) { send(&media.RTMPH264Data{AU: au, PTS: pts, DTS: dts}) }) case *format.MPEG4Audio: session.AudioTrack = track r.OnDataMPEG4Audio(track, func(pts time.Duration, au []byte) { send(&media.RTMPAACData{AU: au, PTS: pts}) }) default: return nil, fmt.Errorf("unsupported track type: %T", track) } } return r, nil}
func (a *StreamplaceAPI) HandleRTMPPlayback(ctx context.Context, sc *gortmplib.ServerConn) error { if !strings.HasPrefix(sc.URL.Path, RTMPPrefix) { return fmt.Errorf("RTMP publisher is not allowed to publish to %s (must start with %s)", sc.URL.String(), RTMPPrefix) } streamer := strings.TrimPrefix(sc.URL.Path, RTMPPrefix) a.rtmpSessionsLock.Lock() session, ok := a.rtmpSessions[streamer] a.rtmpSessionsLock.Unlock() if !ok { return fmt.Errorf("RTMP session not found for streamer %s", streamer) }
return relayRTMPSession(ctx, sc, session)}
// relayRTMPSession writes a publisher's session out to conn, the internal// playback connection the ingest pipeline reads with rtmp2src, until the// session closes or ctx ends.func relayRTMPSession(ctx context.Context, conn gortmplib.Conn, session *media.RTMPSession) error { tracks := []format.Format{} if session.VideoTrack != nil { tracks = append(tracks, session.VideoTrack) } if session.AudioTrack != nil { tracks = append(tracks, session.AudioTrack) } w := &gortmplib.Writer{ Conn: conn, Tracks: tracks, } err := w.Initialize() if err != nil { return err } for { select { case <-ctx.Done(): return ctx.Err() case event := <-session.EventChan: if event == nil { return fmt.Errorf("RTMP session closed: %w", io.EOF) } switch event := event.(type) { case *media.RTMPH264Data: // gortmplib's reader hands on an AVC sequence header as an AU // of just SPS and PPS, including the stream's first one, which // it replays after reading ahead for the tracks. Written back as // a frame, it takes the timestamp of the keyframe after it, and // h264parse drops that keyframe's PTS, which fails the // segmenter's muxer. Write it as the sequence header it was. if msg, ok := h264ConfigMessage(event.AU, event.DTS); ok { if err := conn.Write(msg); err != nil { return fmt.Errorf("error writing H264 config: %w", err) } continue } err := w.WriteH264(session.VideoTrack, event.PTS, event.DTS, event.AU) if err != nil { return fmt.Errorf("error writing H264: %w", err) } case *media.RTMPAACData: err := w.WriteMPEG4Audio(session.AudioTrack, event.PTS, event.AU) if err != nil { return fmt.Errorf("error writing MPEG4Audio: %w", err) } default: return fmt.Errorf("unsupported event type: %T", event) } } }}
func (a *StreamplaceAPI) HandleRTMPPublishConn(ctx context.Context, conn net.Conn) error { err := conn.SetReadDeadline(time.Now().Add(RTMPTimeout)) if err != nil { return err }
sc := &gortmplib.ServerConn{ RW: conn, } err = sc.Initialize() if err != nil { return err }
err = sc.Accept() if err != nil { return err }
if sc.Publish { return a.HandleRTMPPublisher(ctx, sc) } return fmt.Errorf("RTMP playback is not allowed")}
func (a *StreamplaceAPI) HandleRTMPPlaybackConn(ctx context.Context, conn net.Conn) error { err := conn.SetReadDeadline(time.Now().Add(RTMPTimeout)) if err != nil { return err }
sc := &gortmplib.ServerConn{ RW: conn, } err = sc.Initialize() if err != nil { return err }
err = sc.Accept() if err != nil { return err }
if !sc.Publish { return a.HandleRTMPPlayback(ctx, sc) } return fmt.Errorf("RTMP playback is not allowed")}
func (a *StreamplaceAPI) ServeRTMP(ctx context.Context) error { ln, err := net.Listen("tcp", a.CLI.RTMPAddr) if err != nil { return fmt.Errorf("failed to listen: %w", err) } defer ln.Close()
go func() { <-ctx.Done() ln.Close() }()
log.Log(ctx, "rtmp server starting", "addr", a.CLI.RTMPAddr)
g, ctx := errgroup.WithContext(ctx) g.Go(func() error { return a.ServeRTMPInternalPlayback(ctx) }) g.Go(func() error { for { if ctx.Err() != nil { return ctx.Err() } conn, err := ln.Accept() if err != nil { return fmt.Errorf("error accepting RTMP connection: %w", err) } go func() { err := a.HandleRTMPPublishConn(ctx, conn) if err != nil { log.Error(ctx, "error handling RTMP publish connection", "error", err) } }() } })
return g.Wait()}
// Serve RTMP internal playback server for gstreamer to pull fromfunc (a *StreamplaceAPI) ServeRTMPInternalPlayback(ctx context.Context) error { ln, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { return fmt.Errorf("failed to listen: %w", err) } addr := ln.Addr().String() defer ln.Close()
_, port, err := net.SplitHostPort(addr) if err != nil { return fmt.Errorf("failed to split host and port: %w", err) }
go func() { <-ctx.Done() ln.Close() }()
a.rtmpInternalPlaybackAddr = fmt.Sprintf("127.0.0.1:%s", port)
log.Log(ctx, "rtmp internal playback server starting", "addr", a.rtmpInternalPlaybackAddr)
// Accept loop in a goroutine so we can select on context.Done for { if ctx.Err() != nil { return ctx.Err() } conn, err := ln.Accept() if err != nil { return fmt.Errorf("error accepting RTMP connection: %w", err) }
go func() { err := a.HandleRTMPPlaybackConn(ctx, conn) if err != nil { log.Error(ctx, "error handling RTMP internal playback connection", "error", err) } }() }}
func (a *StreamplaceAPI) ServeRTMPS(ctx context.Context, cli *config.CLI) error { var tlsConfig *tls.Config if a.ACME != nil { tlsConfig = a.ACME.TLSConfig() } else { cert, err := tls.LoadX509KeyPair(cli.TLSCertPath, cli.TLSKeyPath) if err != nil { return fmt.Errorf("failed to load TLS certificate: %w", err) } tlsConfig = &tls.Config{ Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS12, } }
ln, err := tls.Listen("tcp", cli.RTMPSAddr, tlsConfig) if err != nil { return fmt.Errorf("failed to create RTMPS listener: %w", err) }
log.Log(ctx, "rtmps server starting", "addr", cli.RTMPAddr)
go func() { <-ctx.Done() ln.Close() }()
g, ctx := errgroup.WithContext(ctx) g.Go(func() error { return a.ServeRTMPInternalPlayback(ctx) }) g.Go(func() error { for { if ctx.Err() != nil { return ctx.Err() } conn, err := ln.Accept() if err != nil { return fmt.Errorf("error accepting RTMP connection: %w", err) } go func() { err := a.HandleRTMPPublishConn(ctx, conn) if err != nil { log.Error(ctx, "error handling RTMP publish connection", "error", err) } }() } })
return g.Wait()}
// h264ConfigMessage returns the FLV sequence header for an H.264 access unit// that holds nothing but parameter sets (one SPS and one PPS), or false for// any other AU.func h264ConfigMessage(au [][]byte, dts time.Duration) (*message.Video, bool) { var conf h264conf.Conf for _, nalu := range au { if len(nalu) == 0 { return nil, false } switch nalu[0] & 0x1f { case 7: // SPS if conf.SPS != nil { return nil, false } conf.SPS = nalu case 8: // PPS if conf.PPS != nil { return nil, false } conf.PPS = nalu default: return nil, false } } if conf.SPS == nil || conf.PPS == nil { return nil, false } buf, err := conf.Marshal() if err != nil { return nil, false } return &message.Video{ ChunkStreamID: message.VideoChunkStreamID, MessageStreamID: 0x1000000, Codec: message.CodecH264, IsKeyFrame: true, Type: message.VideoTypeConfig, Payload: buf, DTS: dts, }, true}