From 1b518b94fddc01baad33f7a00babd80d5074c021 Mon Sep 17 00:00:00 2001
From: Alexey Khit <alexey.khit@gmail.com>
Date: Sun, 27 Nov 2022 10:08:37 +0300
Subject: [PATCH] Add support auth for RTSP server

---
 cmd/rtsp/rtsp.go | 14 ++++++++++----
 pkg/rtsp/conn.go | 16 ++++++++++++++++
 pkg/tcp/auth.go  | 18 ++++++++++++++++++
 3 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/cmd/rtsp/rtsp.go b/cmd/rtsp/rtsp.go
index 2da8b570..1c7d600e 100644
--- a/cmd/rtsp/rtsp.go
+++ b/cmd/rtsp/rtsp.go
@@ -14,7 +14,9 @@ import (
 func Init() {
 	var conf struct {
 		Mod struct {
-			Listen string `yaml:"listen"`
+			Listen   string `yaml:"listen"`
+			Username string `yaml:"username"`
+			Password string `yaml:"password"`
 		} `yaml:"rtsp"`
 	}
 
@@ -52,7 +54,12 @@ func Init() {
 			if err != nil {
 				return
 			}
-			go tcpHandler(conn)
+
+			c := rtsp.NewServer(conn)
+			if conf.Mod.Username != "" {
+				c.Auth(conf.Mod.Username, conf.Mod.Password)
+			}
+			go tcpHandler(c)
 		}
 	}()
 }
@@ -121,13 +128,12 @@ func rtspHandler(url string) (streamer.Producer, error) {
 	return conn, nil
 }
 
-func tcpHandler(c net.Conn) {
+func tcpHandler(conn *rtsp.Conn) {
 	var name string
 	var closer func()
 
 	trace := log.Trace().Enabled()
 
-	conn := rtsp.NewServer(c)
 	conn.Listen(func(msg interface{}) {
 		if trace {
 			switch msg := msg.(type) {
diff --git a/pkg/rtsp/conn.go b/pkg/rtsp/conn.go
index 02dc597f..aebc0280 100644
--- a/pkg/rtsp/conn.go
+++ b/pkg/rtsp/conn.go
@@ -99,6 +99,11 @@ func NewServer(conn net.Conn) *Conn {
 	return c
 }
 
+func (c *Conn) Auth(username, password string) {
+	info := url.UserPassword(username, password)
+	c.auth = tcp.NewAuth(info)
+}
+
 func (c *Conn) parseURI() (err error) {
 	c.URL, err = url.Parse(c.uri)
 	if err != nil {
@@ -490,6 +495,17 @@ func (c *Conn) Accept() error {
 
 		c.Fire(req)
 
+		if !c.auth.Validate(req) {
+			res := &tcp.Response{
+				Status: "401 Unauthorized",
+				Header: map[string][]string{"Www-Authenticate": {`Basic realm="go2rtc"`}},
+			}
+			if err = c.Response(res); err != nil {
+				return err
+			}
+			continue
+		}
+
 		// Receiver: OPTIONS > DESCRIBE > SETUP... > PLAY > TEARDOWN
 		// Sender: OPTIONS > ANNOUNCE > SETUP... > RECORD > TEARDOWN
 		switch req.Method {
diff --git a/pkg/tcp/auth.go b/pkg/tcp/auth.go
index 27cb2e94..0aa2c45b 100644
--- a/pkg/tcp/auth.go
+++ b/pkg/tcp/auth.go
@@ -80,6 +80,24 @@ func (a *Auth) Write(req *Request) {
 	}
 }
 
+func (a *Auth) Validate(req *Request) bool {
+	if a == nil {
+		return true
+	}
+
+	header := req.Header.Get("Authorization")
+	if header == "" {
+		return false
+	}
+
+	if a.Method == AuthUnknown {
+		a.Method = AuthBasic
+		a.header = "Basic " + B64(a.user, a.pass)
+	}
+
+	return header == a.header
+}
+
 func Between(s, sub1, sub2 string) string {
 	i := strings.Index(s, sub1)
 	if i < 0 {