From 161aaeec8decedcdf2dcd51106538ca297ae3823 Mon Sep 17 00:00:00 2001
From: Gabe Kangas <gabek@real-ity.com>
Date: Mon, 1 Jun 2020 17:15:11 -0700
Subject: [PATCH] Reject invalid streaming key

---
 config.go                                   |  7 +--
 config/{config.yaml => config-example.yaml} |  3 +-
 handler.go                                  |  6 +++
 main.go                                     | 41 ------------------
 rtmp.go                                     | 48 +++++++++++++++++++++
 5 files changed, 60 insertions(+), 45 deletions(-)
 rename config/{config.yaml => config-example.yaml} (86%)
 create mode 100644 rtmp.go

diff --git a/config.go b/config.go
index bb1f22f38..56d61af81 100644
--- a/config.go
+++ b/config.go
@@ -20,8 +20,9 @@ type Config struct {
 }
 
 type VideoSettings struct {
-	ResolutionWidth      int `yaml:"resolutionWidth"`
-	ChunkLengthInSeconds int `yaml:"chunkLengthInSeconds"`
+	ResolutionWidth      int    `yaml:"resolutionWidth"`
+	ChunkLengthInSeconds int    `yaml:"chunkLengthInSeconds"`
+	StreamingKey         string `yaml:"streamingKey"`
 }
 
 // MaxNumberOnDisk must be at least as large as MaxNumberInPlaylist
@@ -39,7 +40,7 @@ func getConfig() Config {
 	filePath := "config/config.yaml"
 
 	if !fileExists(filePath) {
-		log.Fatal("ERROR: valid config/config.yaml is required")
+		log.Fatal("ERROR: valid config/config.yaml is required.  Copy config/config-example.yaml to config/config.yaml and edit.")
 	}
 
 	yamlFile, err := ioutil.ReadFile(filePath)
diff --git a/config/config.yaml b/config/config-example.yaml
similarity index 86%
rename from config/config.yaml
rename to config/config-example.yaml
index 78aa3567c..dee8a0c73 100644
--- a/config/config.yaml
+++ b/config/config-example.yaml
@@ -6,11 +6,12 @@ webServerPort: 8080
 videoSettings:
   resolutionWidth: 900
   chunkLengthInSeconds: 4
+  streamingKey: abc123
 
 files:
   maxNumberInPlaylist: 30
   maxNumberOnDisk: 60
 
 ipfs:
-  enabled: true
+  enabled: false
   gateway: https://ipfs.io
\ No newline at end of file
diff --git a/handler.go b/handler.go
index d05e2ab16..25fdcaa9e 100644
--- a/handler.go
+++ b/handler.go
@@ -41,6 +41,12 @@ func (h *Handler) OnCreateStream(timestamp uint32, cmd *rtmpmsg.NetConnectionCre
 func (h *Handler) OnPublish(timestamp uint32, cmd *rtmpmsg.NetStreamPublish) error {
 	// log.Printf("OnPublish: %#v", cmd)
 
+	log.Println("Incoming stream connected.")
+
+	if cmd.PublishingName != configuration.VideoSettings.StreamingKey {
+		return errors.New("Invalid streaming key!  Rejecting incoming stream.")
+	}
+
 	// (example) Reject a connection when PublishingName is empty
 	if cmd.PublishingName == "" {
 		return errors.New("PublishingName is empty")
diff --git a/main.go b/main.go
index 6b1877682..add7ad56b 100644
--- a/main.go
+++ b/main.go
@@ -1,15 +1,11 @@
 package main
 
 import (
-	"io"
-	"net"
 	"net/http"
 	"strconv"
 
 	icore "github.com/ipfs/interface-go-ipfs-core"
-	"github.com/sirupsen/logrus"
 	log "github.com/sirupsen/logrus"
-	"github.com/yutopp/go-rtmp"
 )
 
 var ipfs icore.CoreAPI
@@ -56,40 +52,3 @@ func startChatServer() {
 
 	log.Fatal(http.ListenAndServe(":"+strconv.Itoa(configuration.WebServerPort), nil))
 }
-
-func startRTMPService() {
-	port := 1935
-	log.Printf("RTMP server is listening for incoming stream on port %d.\n", port)
-
-	tcpAddr, err := net.ResolveTCPAddr("tcp", ":"+strconv.Itoa(port))
-	if err != nil {
-		log.Panicf("Failed: %+v", err)
-	}
-
-	listener, err := net.ListenTCP("tcp", tcpAddr)
-	if err != nil {
-		log.Panicf("Failed: %+v", err)
-	}
-
-	srv := rtmp.NewServer(&rtmp.ServerConfig{
-		OnConnect: func(conn net.Conn) (io.ReadWriteCloser, *rtmp.ConnConfig) {
-			l := log.StandardLogger()
-			l.SetLevel(logrus.WarnLevel)
-
-			h := &Handler{}
-
-			return conn, &rtmp.ConnConfig{
-				Handler: h,
-
-				ControlState: rtmp.StreamControlStateConfig{
-					DefaultBandwidthWindowSize: 6 * 1024 * 1024 / 8,
-				},
-
-				Logger: l,
-			}
-		},
-	})
-	if err := srv.Serve(listener); err != nil {
-		log.Panicf("Failed: %+v", err)
-	}
-}
diff --git a/rtmp.go b/rtmp.go
new file mode 100644
index 000000000..98ffbaf87
--- /dev/null
+++ b/rtmp.go
@@ -0,0 +1,48 @@
+package main
+
+import (
+	"io"
+	"net"
+	"strconv"
+
+	"github.com/sirupsen/logrus"
+	log "github.com/sirupsen/logrus"
+	"github.com/yutopp/go-rtmp"
+)
+
+func startRTMPService() {
+	port := 1935
+	log.Printf("RTMP server is listening for incoming stream on port %d.\n", port)
+
+	tcpAddr, err := net.ResolveTCPAddr("tcp", ":"+strconv.Itoa(port))
+	if err != nil {
+		log.Panicf("Failed: %+v", err)
+	}
+
+	listener, err := net.ListenTCP("tcp", tcpAddr)
+	if err != nil {
+		log.Panicf("Failed: %+v", err)
+	}
+
+	srv := rtmp.NewServer(&rtmp.ServerConfig{
+		OnConnect: func(conn net.Conn) (io.ReadWriteCloser, *rtmp.ConnConfig) {
+			l := log.StandardLogger()
+			l.SetLevel(logrus.WarnLevel)
+
+			h := &Handler{}
+
+			return conn, &rtmp.ConnConfig{
+				Handler: h,
+
+				ControlState: rtmp.StreamControlStateConfig{
+					DefaultBandwidthWindowSize: 6 * 1024 * 1024 / 8,
+				},
+
+				Logger: l,
+			}
+		},
+	})
+	if err := srv.Serve(listener); err != nil {
+		log.Panicf("Failed: %+v", err)
+	}
+}