support encrypted streams with SRTP and MIKEY (#520) (#809)

server_stream_media.go

@@ -1,6 +1,8 @@
 package gortsplib
 
 import (
+	"crypto/rand"
+	"fmt"
 	"sync/atomic"
 
 	"github.com/bluenviron/gortsplib/v4/pkg/description"
@@ -12,6 +14,7 @@ type serverStreamMedia struct {
 	media   *description.Media
 	trackID int
 
+	srtpOutCtx      *wrappedSRTPContext
 	formats         map[uint8]*serverStreamFormat
 	multicastWriter *serverMulticastWriter
 	bytesSent       *uint64
@@ -40,6 +43,30 @@ func (sm *serverStreamMedia) initialize() error {
 		sm.formats[forma.PayloadType()] = sf
 	}
 
+	if sm.st.Server.TLSConfig != nil {
+		srtpOutKey := make([]byte, srtpKeyLength)
+		_, err := rand.Read(srtpOutKey)
+		if err != nil {
+			return err
+		}
+
+		ssrcs := make([]uint32, len(sm.formats))
+		n := 0
+		for _, cf := range sm.formats {
+			ssrcs[n] = cf.localSSRC
+			n++
+		}
+
+		sm.srtpOutCtx = &wrappedSRTPContext{
+			key:   srtpOutKey,
+			ssrcs: ssrcs,
+		}
+		err = sm.srtpOutCtx.initialize()
+		if err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
@@ -54,35 +81,75 @@ func (sm *serverStreamMedia) close() {
 }
 
 func (sm *serverStreamMedia) writePacketRTCP(pkt rtcp.Packet) error {
-	byts, err := pkt.Marshal()
+	plain, err := pkt.Marshal()
 	if err != nil {
 		return err
 	}
 
-	le := len(byts)
+	maxPlainPacketSize := sm.st.Server.MaxPacketSize
+	if sm.srtpOutCtx != nil {
+		maxPlainPacketSize -= srtcpOverhead
+	}
+
+	if len(plain) > maxPlainPacketSize {
+		return fmt.Errorf("packet is too big")
+	}
+
+	var encr []byte
+	if sm.srtpOutCtx != nil {
+		encr = make([]byte, sm.st.Server.MaxPacketSize)
+		encr, err = sm.srtpOutCtx.encryptRTCP(encr, plain, nil)
+		if err != nil {
+			return err
+		}
+	}
+
+	encrLen := uint64(len(encr))
+	plainLen := uint64(len(plain))
 
 	// send unicast
 	for r := range sm.st.activeUnicastReaders {
-		if _, ok := r.setuppedMedias[sm.media]; ok {
-			err := r.writePacketRTCPEncoded(sm.media, byts)
-			if err != nil {
-				r.onStreamWriteError(err)
-				continue
+		if sm, ok := r.setuppedMedias[sm.media]; ok {
+			if r.setuppedSecure {
+				err := sm.writePacketRTCPEncoded(encr)
+				if err != nil {
+					r.onStreamWriteError(err)
+					continue
+				}
+
+				atomic.AddUint64(sm.bytesSent, encrLen)
+			} else {
+				err := sm.writePacketRTCPEncoded(plain)
+				if err != nil {
+					r.onStreamWriteError(err)
+					continue
+				}
+
+				atomic.AddUint64(sm.bytesSent, plainLen)
 			}
 
-			atomic.AddUint64(sm.bytesSent, uint64(le))
 			atomic.AddUint64(sm.rtcpPacketsSent, 1)
 		}
 	}
 
 	// send multicast
 	if sm.multicastWriter != nil {
-		err := sm.multicastWriter.writePacketRTCP(byts)
-		if err != nil {
-			return err
+		if sm.srtpOutCtx != nil {
+			err := sm.multicastWriter.writePacketRTCP(encr)
+			if err != nil {
+				return err
+			}
+
+			atomic.AddUint64(sm.bytesSent, encrLen)
+		} else {
+			err := sm.multicastWriter.writePacketRTCP(plain)
+			if err != nil {
+				return err
+			}
+
+			atomic.AddUint64(sm.bytesSent, plainLen)
 		}
 
-		atomic.AddUint64(sm.bytesSent, uint64(le))
 		atomic.AddUint64(sm.rtcpPacketsSent, 1)
 	}