add command line option --protocols to set what protocols are available

bluenviron · Jan 15, 2020 · f607fa8 · f607fa8
1 parent e78f3f2
commit f607fa8
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 9 deletions.
diff --git a/README.md b/README.md
@@ -72,12 +72,13 @@ rtsp-simple-server v0.0.0
 RTSP server.
 
 Flags:
-  --help            Show context-sensitive help (also try --help-long and --help-man).
-  --version         print rtsp-simple-server version
-  --rtsp-port=8554  port of the RTSP TCP listener
-  --rtp-port=8000   port of the RTP UDP listener
-  --rtcp-port=8001  port of the RTCP UDP listener
-  --publish-key=""  optional authentication key required to publish
+  --help                 Show context-sensitive help (also try --help-long and --help-man).
+  --version              print rtsp-simple-server version
+  --protocols="udp,tcp"  supported protocols
+  --rtsp-port=8554       port of the RTSP TCP listener
+  --rtp-port=8000        port of the RTP UDP listener
+  --rtcp-port=8001       port of the RTCP UDP listener
+  --publish-key=""       optional authentication key required to publish
 ```
 
 

diff --git a/client.go b/client.go
@@ -386,6 +386,18 @@ func (c *client) handleRequest(req *rtsp.Request) bool {
 				}
 				return false
 			}() {
+				if _, ok := c.p.protocols[_STREAM_PROTOCOL_UDP]; !ok {
+					c.log("ERR: udp streaming is disabled")
+					c.rconn.WriteResponse(&rtsp.Response{
+						StatusCode: 461,
+						Status:     "Unsupported Transport",
+						Headers: map[string]string{
+							"CSeq": cseq,
+						},
+					})
+					return false
+				}
+
 				rtpPort, rtcpPort := th.getClientPorts()
 				if rtpPort == 0 || rtcpPort == 0 {
 					c.writeResError(req, fmt.Errorf("transport header does not have valid client ports (%s)", transportStr))
@@ -447,6 +459,18 @@ func (c *client) handleRequest(req *rtsp.Request) bool {
 
 				// play via TCP
 			} else if _, ok := th["RTP/AVP/TCP"]; ok {
+				if _, ok := c.p.protocols[_STREAM_PROTOCOL_TCP]; !ok {
+					c.log("ERR: tcp streaming is disabled")
+					c.rconn.WriteResponse(&rtsp.Response{
+						StatusCode: 461,
+						Status:     "Unsupported Transport",
+						Headers: map[string]string{
+							"CSeq": cseq,
+						},
+					})
+					return false
+				}
+
 				if c.path != "" && path != c.path {
 					c.writeResError(req, fmt.Errorf("path has changed"))
 					return false
@@ -529,7 +553,19 @@ func (c *client) handleRequest(req *rtsp.Request) bool {
 					return true
 				}
 				return false
-			}(); ok {
+			}() {
+				if _, ok := c.p.protocols[_STREAM_PROTOCOL_UDP]; !ok {
+					c.log("ERR: udp streaming is disabled")
+					c.rconn.WriteResponse(&rtsp.Response{
+						StatusCode: 461,
+						Status:     "Unsupported Transport",
+						Headers: map[string]string{
+							"CSeq": cseq,
+						},
+					})
+					return false
+				}
+
 				rtpPort, rtcpPort := th.getClientPorts()
 				if rtpPort == 0 || rtcpPort == 0 {
 					c.writeResError(req, fmt.Errorf("transport header does not have valid client ports (%s)", transportStr))
@@ -580,6 +616,18 @@ func (c *client) handleRequest(req *rtsp.Request) bool {
 
 				// record via TCP
 			} else if _, ok := th["RTP/AVP/TCP"]; ok {
+				if _, ok := c.p.protocols[_STREAM_PROTOCOL_TCP]; !ok {
+					c.log("ERR: tcp streaming is disabled")
+					c.rconn.WriteResponse(&rtsp.Response{
+						StatusCode: 461,
+						Status:     "Unsupported Transport",
+						Headers: map[string]string{
+							"CSeq": cseq,
+						},
+					})
+					return false
+				}
+
 				var interleaved string
 				err = func() error {
 					c.p.mutex.Lock()

diff --git a/main.go b/main.go
@@ -6,6 +6,7 @@ import (
 	"net"
 	"os"
 	"regexp"
+	"strings"
 	"sync"
 
 	"gopkg.in/alecthomas/kingpin.v2"
@@ -40,6 +41,7 @@ func (s streamProtocol) String() string {
 }
 
 type program struct {
+	protocols  map[streamProtocol]struct{}
 	rtspPort   int
 	rtpPort    int
 	rtcpPort   int
@@ -52,7 +54,24 @@ type program struct {
 	publishers map[string]*client
 }
 
-func newProgram(rtspPort int, rtpPort int, rtcpPort int, publishKey string) (*program, error) {
+func newProgram(protocolsStr string, rtspPort int, rtpPort int, rtcpPort int, publishKey string) (*program, error) {
+	protocols := make(map[streamProtocol]struct{})
+	for _, proto := range strings.Split(protocolsStr, ",") {
+		switch proto {
+		case "udp":
+			protocols[_STREAM_PROTOCOL_UDP] = struct{}{}
+
+		case "tcp":
+			protocols[_STREAM_PROTOCOL_TCP] = struct{}{}
+
+		default:
+			return nil, fmt.Errorf("unsupported protocol: %s", proto)
+		}
+	}
+	if len(protocols) == 0 {
+		return nil, fmt.Errorf("no protocols supplied")
+	}
+
 	if publishKey != "" {
 		if !regexp.MustCompile("^[a-zA-Z0-9]+$").MatchString(publishKey) {
 			return nil, fmt.Errorf("publish key must be alphanumeric")
@@ -62,6 +81,7 @@ func newProgram(rtspPort int, rtpPort int, rtcpPort int, publishKey string) (*pr
 	log.Printf("rtsp-simple-server %s", Version)
 
 	p := &program{
+		protocols:  protocols,
 		rtspPort:   rtspPort,
 		rtpPort:    rtpPort,
 		rtcpPort:   rtcpPort,
@@ -128,6 +148,7 @@ func main() {
 
 	version := kingpin.Flag("version", "print rtsp-simple-server version").Bool()
 
+	protocols := kingpin.Flag("protocols", "supported protocols").Default("udp,tcp").String()
 	rtspPort := kingpin.Flag("rtsp-port", "port of the RTSP TCP listener").Default("8554").Int()
 	rtpPort := kingpin.Flag("rtp-port", "port of the RTP UDP listener").Default("8000").Int()
 	rtcpPort := kingpin.Flag("rtcp-port", "port of the RTCP UDP listener").Default("8001").Int()
@@ -140,7 +161,7 @@ func main() {
 		os.Exit(0)
 	}
 
-	p, err := newProgram(*rtspPort, *rtpPort, *rtcpPort, *publishKey)
+	p, err := newProgram(*protocols, *rtspPort, *rtpPort, *rtcpPort, *publishKey)
 	if err != nil {
 		log.Fatal("ERR: ", err)
 	}