From 3defbd60db0edf92787d57d753477647258ec38a Mon Sep 17 00:00:00 2001
From: Alexey Khit <alexey.khit@gmail.com>
Date: Sun, 23 Jul 2023 17:09:09 +0300
Subject: [PATCH] Add deadline handler for SRTP server

---
 pkg/core/helpers.go | 12 +++++++++---
 pkg/srtp/session.go | 26 ++++++++++++++++----------
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/pkg/core/helpers.go b/pkg/core/helpers.go
index b3871a6d..9a9bf11c 100644
--- a/pkg/core/helpers.go
+++ b/pkg/core/helpers.go
@@ -1,12 +1,18 @@
 package core
 
 import (
-	cryptorand "crypto/rand"
-	"github.com/rs/zerolog/log"
+	"crypto/rand"
 	"runtime"
 	"strconv"
 	"strings"
 	"time"
+
+	"github.com/rs/zerolog/log"
+)
+
+const (
+	ConnDialTimeout = time.Second * 3
+	ConnDeadline    = time.Second * 3
 )
 
 // Now90000 - timestamp for Video (clock rate = 90000 samples per second)
@@ -19,7 +25,7 @@ const symbols = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ-
 // RandString base10 - numbers, base16 - hex, base36 - digits+letters, base64 - URL safe symbols
 func RandString(size, base byte) string {
 	b := make([]byte, size)
-	if _, err := cryptorand.Read(b); err != nil {
+	if _, err := rand.Read(b); err != nil {
 		panic(err)
 	}
 	for i := byte(0); i < size; i++ {
diff --git a/pkg/srtp/session.go b/pkg/srtp/session.go
index 853ada05..82988c0a 100644
--- a/pkg/srtp/session.go
+++ b/pkg/srtp/session.go
@@ -1,11 +1,12 @@
 package srtp
 
 import (
+	"time"
+
 	"github.com/AlexxIT/go2rtc/pkg/core"
 	"github.com/pion/rtcp"
 	"github.com/pion/rtp"
 	"github.com/pion/srtp/v2"
-	"time"
 )
 
 type Session struct {
@@ -19,6 +20,8 @@ type Session struct {
 	Track *core.Receiver
 	Recv  uint32
 
+	Deadline *time.Timer
+
 	lastSequence  uint32
 	lastTimestamp uint32
 	//lastPacket    *rtp.Packet
@@ -28,17 +31,16 @@ type Session struct {
 	totalLost uint32
 }
 
-func (s *Session) SetKeys(
-	localKey, localSalt, remoteKey, remoteSalt []byte,
-) (err error) {
-	if s.localCtx, err = srtp.CreateContext(
-		localKey, localSalt, GuessProfile(localKey),
-	); err != nil {
+func (s *Session) LastTime() time.Time {
+	return s.lastTime
+}
+
+func (s *Session) SetKeys(localKey, localSalt, remoteKey, remoteSalt []byte) (err error) {
+	s.localCtx, err = srtp.CreateContext(localKey, localSalt, GuessProfile(localKey))
+	if err != nil {
 		return
 	}
-	s.remoteCtx, err = srtp.CreateContext(
-		remoteKey, remoteSalt, GuessProfile(remoteKey),
-	)
+	s.remoteCtx, err = srtp.CreateContext(remoteKey, remoteSalt, GuessProfile(remoteKey))
 	return
 }
 
@@ -56,6 +58,10 @@ func (s *Session) HandleRTP(data []byte) (err error) {
 		return
 	}
 
+	if s.Deadline != nil {
+		s.Deadline.Reset(core.ConnDeadline)
+	}
+
 	now := time.Now()
 
 	// https://www.ietf.org/rfc/rfc3550.txt