test RTP timestamps

Overlord-Client/cmd/agent/webrtcpub/h264_track.go

@@ -0,0 +1,85 @@
+//go:build overlord_webrtc
+
+package webrtcpub
+
+import (
+	"crypto/rand"
+	"encoding/binary"
+	"errors"
+	"sync"
+	"time"
+
+	"github.com/pion/rtp"
+	"github.com/pion/rtp/codecs"
+	"github.com/pion/webrtc/v4"
+)
+
+const (
+	h264ClockRate = 90_000
+	h264RTPMTU    = 1200
+)
+
+type rtpPacketSink interface {
+	WriteRTP(*rtp.Packet) error
+}
+
+type h264TrackWriter struct {
+	mu          sync.Mutex
+	t           rtpPacketSink
+	packetizer  rtp.Packetizer
+	timestamp   uint32
+	lastCapture time.Time
+}
+
+func newH264TrackWriter(track *webrtc.TrackLocalStaticRTP) *h264TrackWriter {
+	return newH264TrackWriterWithSink(track, randomRTPVideoTimestamp())
+}
+
+func newH264TrackWriterWithSink(sink rtpPacketSink, initialTimestamp uint32) *h264TrackWriter {
+	return &h264TrackWriter{
+		t: sink,
+		packetizer: rtp.NewPacketizer(
+			h264RTPMTU, 0, initialTimestamp, &codecs.H264Payloader{},
+			rtp.NewRandomSequencer(), h264ClockRate,
+		),
+		timestamp: initialTimestamp,
+	}
+}
+
+func randomRTPVideoTimestamp() uint32 {
+	var value [4]byte
+	if _, err := rand.Read(value[:]); err == nil {
+		return binary.BigEndian.Uint32(value[:])
+	}
+	return uint32(time.Now().UnixNano())
+}
+
+func (w *h264TrackWriter) WriteH264(nalu []byte, capturedAt time.Time) error {
+	if len(nalu) == 0 {
+		return nil
+	}
+	if capturedAt.IsZero() {
+		capturedAt = time.Now()
+	}
+
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if !w.lastCapture.IsZero() && capturedAt.After(w.lastCapture) {
+		ticks := uint64(capturedAt.Sub(w.lastCapture).Seconds() * h264ClockRate)
+		if ticks == 0 {
+			ticks = 1
+		}
+		w.timestamp += uint32(ticks)
+	}
+	w.lastCapture = capturedAt
+
+	packets := w.packetizer.Packetize(nalu, 0)
+	var writeErrs []error
+	for _, packet := range packets {
+		packet.Timestamp = w.timestamp
+		if err := w.t.WriteRTP(packet); err != nil {
+			writeErrs = append(writeErrs, err)
+		}
+	}
+	return errors.Join(writeErrs...)
+}