improve examples (#778)

examples/server-record-format-h264-to-disk/main.go

@@ -0,0 +1,167 @@
+package main
+
+import (
+	"fmt"
+	"log"
+	"sync"
+
+	"github.com/pion/rtp"
+
+	"github.com/bluenviron/gortsplib/v4"
+	"github.com/bluenviron/gortsplib/v4/pkg/base"
+	"github.com/bluenviron/gortsplib/v4/pkg/description"
+	"github.com/bluenviron/gortsplib/v4/pkg/format"
+	"github.com/bluenviron/gortsplib/v4/pkg/format/rtph264"
+)
+
+// This example shows how to
+// 1. create a RTSP server which accepts plain connections.
+// 2. allow a single client to publish a stream, containing a H264 format.
+// 3. save the content of the H264 media in a file in MPEG-TS format.
+
+type serverHandler struct {
+	server      *gortsplib.Server
+	mutex       sync.Mutex
+	publisher   *gortsplib.ServerSession
+	media       *description.Media
+	format      *format.H264
+	rtpDec      *rtph264.Decoder
+	mpegtsMuxer *mpegtsMuxer
+}
+
+// called when a connection is opened.
+func (sh *serverHandler) OnConnOpen(ctx *gortsplib.ServerHandlerOnConnOpenCtx) {
+	log.Printf("conn opened")
+}
+
+// called when a connection is closed.
+func (sh *serverHandler) OnConnClose(ctx *gortsplib.ServerHandlerOnConnCloseCtx) {
+	log.Printf("conn closed (%v)", ctx.Error)
+}
+
+// called when a session is opened.
+func (sh *serverHandler) OnSessionOpen(ctx *gortsplib.ServerHandlerOnSessionOpenCtx) {
+	log.Printf("session opened")
+}
+
+// called when a session is closed.
+func (sh *serverHandler) OnSessionClose(ctx *gortsplib.ServerHandlerOnSessionCloseCtx) {
+	log.Printf("session closed")
+
+	sh.mutex.Lock()
+	defer sh.mutex.Unlock()
+
+	sh.publisher = nil
+	sh.mpegtsMuxer.close()
+}
+
+// called when receiving an ANNOUNCE request.
+func (sh *serverHandler) OnAnnounce(ctx *gortsplib.ServerHandlerOnAnnounceCtx) (*base.Response, error) {
+	log.Printf("ANNOUNCE request")
+
+	sh.mutex.Lock()
+	defer sh.mutex.Unlock()
+
+	if sh.publisher != nil {
+		sh.publisher.Close()
+		sh.mpegtsMuxer.close()
+	}
+
+	// find the H264 media and format
+	var forma *format.H264
+	medi := ctx.Description.FindFormat(&forma)
+	if medi == nil {
+		return &base.Response{
+			StatusCode: base.StatusBadRequest,
+		}, fmt.Errorf("H264 media not found")
+	}
+
+	// setup RTP -> H264 decoder
+	rtpDec, err := forma.CreateDecoder()
+	if err != nil {
+		panic(err)
+	}
+
+	// setup H264 -> MPEGTS muxer
+	mpegtsMuxer := &mpegtsMuxer{
+		fileName: "mystream.ts",
+		sps:      forma.SPS,
+		pps:      forma.PPS,
+	}
+	err = mpegtsMuxer.initialize()
+	if err != nil {
+		return &base.Response{
+			StatusCode: base.StatusBadRequest,
+		}, err
+	}
+
+	sh.publisher = ctx.Session
+	sh.media = medi
+	sh.format = forma
+	sh.rtpDec = rtpDec
+	sh.mpegtsMuxer = mpegtsMuxer
+
+	return &base.Response{
+		StatusCode: base.StatusOK,
+	}, nil
+}
+
+// called when receiving a SETUP request.
+func (sh *serverHandler) OnSetup(ctx *gortsplib.ServerHandlerOnSetupCtx) (*base.Response, *gortsplib.ServerStream, error) {
+	// prevent readers from using the server.
+	if ctx.Session.State() == gortsplib.ServerSessionStateInitial {
+		return &base.Response{
+			StatusCode: base.StatusNotImplemented,
+		}, nil, nil
+	}
+
+	log.Printf("SETUP request")
+
+	return &base.Response{
+		StatusCode: base.StatusOK,
+	}, nil, nil
+}
+
+// called when receiving a RECORD request.
+func (sh *serverHandler) OnRecord(ctx *gortsplib.ServerHandlerOnRecordCtx) (*base.Response, error) {
+	log.Printf("RECORD request")
+
+	// called when receiving a RTP packet
+	ctx.Session.OnPacketRTP(sh.media, sh.format, func(pkt *rtp.Packet) {
+		// decode timestamp
+		pts, ok := ctx.Session.PacketPTS2(sh.media, pkt)
+		if !ok {
+			return
+		}
+
+		au, err := sh.rtpDec.Decode(pkt)
+		if err != nil {
+			return
+		}
+
+		// encode H264 access unit into MPEG-TS
+		sh.mpegtsMuxer.writeH264(au, pts)
+	})
+
+	return &base.Response{
+		StatusCode: base.StatusOK,
+	}, nil
+}
+
+func main() {
+	// configure the server
+	h := &serverHandler{}
+	h.server = &gortsplib.Server{
+		Handler:           h,
+		RTSPAddress:       ":8554",
+		UDPRTPAddress:     ":8000",
+		UDPRTCPAddress:    ":8001",
+		MulticastIPRange:  "224.1.0.0/16",
+		MulticastRTPPort:  8002,
+		MulticastRTCPPort: 8003,
+	}
+
+	// start server and wait until a fatal error
+	log.Printf("server is ready on %s", h.server.RTSPAddress)
+	panic(h.server.StartAndWait())
+}

examples/server-record-format-h264-to-disk/mpegts_muxer.go

@@ -0,0 +1,109 @@
+package main
+
+import (
+	"bufio"
+	"os"
+
+	"github.com/bluenviron/mediacommon/v2/pkg/codecs/h264"
+	"github.com/bluenviron/mediacommon/v2/pkg/formats/mpegts"
+)
+
+// mpegtsMuxer allows to save a H264 stream into a MPEG-TS file.
+type mpegtsMuxer struct {
+	fileName string
+	sps      []byte
+	pps      []byte
+
+	f            *os.File
+	b            *bufio.Writer
+	w            *mpegts.Writer
+	track        *mpegts.Track
+	dtsExtractor *h264.DTSExtractor
+}
+
+// initialize initializes a mpegtsMuxer.
+func (e *mpegtsMuxer) initialize() error {
+	var err error
+	e.f, err = os.Create(e.fileName)
+	if err != nil {
+		return err
+	}
+	e.b = bufio.NewWriter(e.f)
+
+	e.track = &mpegts.Track{
+		Codec: &mpegts.CodecH264{},
+	}
+
+	e.w = &mpegts.Writer{W: e.b, Tracks: []*mpegts.Track{e.track}}
+	err = e.w.Initialize()
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// close closes all the mpegtsMuxer resources.
+func (e *mpegtsMuxer) close() {
+	e.b.Flush()
+	e.f.Close()
+}
+
+// writeH264 writes a H264 access unit into MPEG-TS.
+func (e *mpegtsMuxer) writeH264(au [][]byte, pts int64) error {
+	var filteredAU [][]byte
+
+	nonIDRPresent := false
+	idrPresent := false
+
+	for _, nalu := range au {
+		typ := h264.NALUType(nalu[0] & 0x1F)
+		switch typ {
+		case h264.NALUTypeSPS:
+			e.sps = nalu
+			continue
+
+		case h264.NALUTypePPS:
+			e.pps = nalu
+			continue
+
+		case h264.NALUTypeAccessUnitDelimiter:
+			continue
+
+		case h264.NALUTypeIDR:
+			idrPresent = true
+
+		case h264.NALUTypeNonIDR:
+			nonIDRPresent = true
+		}
+
+		filteredAU = append(filteredAU, nalu)
+	}
+
+	au = filteredAU
+
+	if au == nil || (!nonIDRPresent && !idrPresent) {
+		return nil
+	}
+
+	// add SPS and PPS before every access unit that contains an IDR
+	if idrPresent {
+		au = append([][]byte{e.sps, e.pps}, au...)
+	}
+
+	if e.dtsExtractor == nil {
+		// skip samples silently until we find one with a IDR
+		if !idrPresent {
+			return nil
+		}
+		e.dtsExtractor = h264.NewDTSExtractor()
+	}
+
+	dts, err := e.dtsExtractor.Extract(au, pts)
+	if err != nil {
+		return err
+	}
+
+	// encode into MPEG-TS
+	return e.w.WriteH264(e.track, pts, dts, au)
+}