From e3f63a43c964aa7241d9e039d48c2828738c294f Mon Sep 17 00:00:00 2001
From: ShiBen <807629978@qq.com>
Date: Tue, 11 Jan 2022 11:21:18 +0800
Subject: [PATCH] External authentication support send url raw query

Do some dynamic authentication, such as token
---
 internal/core/externalauth.go |  3 +++
 internal/core/hls_muxer.go    |  3 ++-
 internal/core/rtmp_conn.go    | 16 +++++++++-------
 internal/core/rtsp_conn.go    |  6 ++++--
 internal/core/rtsp_session.go |  4 ++--
 rtsp-simple-server.yml        |  1 +
 6 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/internal/core/externalauth.go b/internal/core/externalauth.go
index 751470b1..6d26e88a 100644
--- a/internal/core/externalauth.go
+++ b/internal/core/externalauth.go
@@ -14,6 +14,7 @@ func externalAuth(
 	password string,
 	path string,
 	action string,
+	query string,
 ) error {
 	enc, _ := json.Marshal(struct {
 		IP       string `json:"ip"`
@@ -21,12 +22,14 @@ func externalAuth(
 		Password string `json:"password"`
 		Path     string `json:"path"`
 		Action   string `json:"action"`
+		Query    string `json:"query"`
 	}{
 		IP:       ip,
 		User:     user,
 		Password: password,
 		Path:     path,
 		Action:   action,
+		Query:    query,
 	})
 	res, err := http.Post(ur, "application/json", bytes.NewReader(enc))
 	if err != nil {
diff --git a/internal/core/hls_muxer.go b/internal/core/hls_muxer.go
index 10d43180..99f9366b 100644
--- a/internal/core/hls_muxer.go
+++ b/internal/core/hls_muxer.go
@@ -496,7 +496,8 @@ func (m *hlsMuxer) authenticate(req *http.Request) error {
 			user,
 			pass,
 			m.pathName,
-			"read")
+			"read",
+			req.URL.RawQuery)
 		if err != nil {
 			return pathErrAuthCritical{
 				message: fmt.Sprintf("external authentication failed: %s", err),
diff --git a/internal/core/rtmp_conn.go b/internal/core/rtmp_conn.go
index 2a49a992..1799a427 100644
--- a/internal/core/rtmp_conn.go
+++ b/internal/core/rtmp_conn.go
@@ -29,12 +29,12 @@ const (
 	rtmpConnPauseAfterAuthError = 2 * time.Second
 )
 
-func pathNameAndQuery(inURL *url.URL) (string, url.Values) {
+func pathNameAndQuery(inURL *url.URL) (string, url.Values, string) {
 	// remove leading and trailing slashes inserted by OBS and some other clients
 	tmp := strings.TrimRight(inURL.String(), "/")
 	ur, _ := url.Parse(tmp)
 	pathName := strings.TrimLeft(ur.Path, "/")
-	return pathName, ur.Query()
+	return pathName, ur.Query(), ur.RawQuery
 }
 
 type rtmpConnTrackIDPayloadPair struct {
@@ -217,7 +217,7 @@ func (c *rtmpConn) runInner(ctx context.Context) error {
 }
 
 func (c *rtmpConn) runRead(ctx context.Context) error {
-	pathName, query := pathNameAndQuery(c.conn.URL())
+	pathName, query, rawQuery := pathNameAndQuery(c.conn.URL())
 
 	res := c.pathManager.onReaderSetupPlay(pathReaderSetupPlayReq{
 		author:   c,
@@ -226,7 +226,7 @@ func (c *rtmpConn) runRead(ctx context.Context) error {
 			pathIPs []interface{},
 			pathUser conf.Credential,
 			pathPass conf.Credential) error {
-			return c.authenticate(pathName, pathIPs, pathUser, pathPass, "read", query)
+			return c.authenticate(pathName, pathIPs, pathUser, pathPass, "read", query, rawQuery)
 		},
 	})
 
@@ -462,7 +462,7 @@ func (c *rtmpConn) runPublish(ctx context.Context) error {
 		tracks = append(tracks, audioTrack)
 	}
 
-	pathName, query := pathNameAndQuery(c.conn.URL())
+	pathName, query, rawQuery := pathNameAndQuery(c.conn.URL())
 
 	res := c.pathManager.onPublisherAnnounce(pathPublisherAnnounceReq{
 		author:   c,
@@ -471,7 +471,7 @@ func (c *rtmpConn) runPublish(ctx context.Context) error {
 			pathIPs []interface{},
 			pathUser conf.Credential,
 			pathPass conf.Credential) error {
-			return c.authenticate(pathName, pathIPs, pathUser, pathPass, "publish", query)
+			return c.authenticate(pathName, pathIPs, pathUser, pathPass, "publish", query, rawQuery)
 		},
 	})
 
@@ -599,6 +599,7 @@ func (c *rtmpConn) authenticate(
 	pathPass conf.Credential,
 	action string,
 	query url.Values,
+	rawQuery string,
 ) error {
 	if c.externalAuthenticationURL != "" {
 		err := externalAuth(
@@ -607,7 +608,8 @@ func (c *rtmpConn) authenticate(
 			query.Get("user"),
 			query.Get("pass"),
 			pathName,
-			action)
+			action,
+			rawQuery)
 		if err != nil {
 			return pathErrAuthCritical{
 				message: fmt.Sprintf("external authentication failed: %s", err),
diff --git a/internal/core/rtsp_conn.go b/internal/core/rtsp_conn.go
index 5c48cd3a..59f899a0 100644
--- a/internal/core/rtsp_conn.go
+++ b/internal/core/rtsp_conn.go
@@ -108,6 +108,7 @@ func (c *rtspConn) authenticate(
 	pathPass conf.Credential,
 	action string,
 	req *base.Request,
+	query string,
 ) error {
 	if c.externalAuthenticationURL != "" {
 		username := ""
@@ -126,7 +127,8 @@ func (c *rtspConn) authenticate(
 			username,
 			password,
 			pathName,
-			action)
+			action,
+			query)
 		if err != nil {
 			c.authFailures++
 
@@ -247,7 +249,7 @@ func (c *rtspConn) onDescribe(ctx *gortsplib.ServerHandlerOnDescribeCtx,
 			pathIPs []interface{},
 			pathUser conf.Credential,
 			pathPass conf.Credential) error {
-			return c.authenticate(ctx.Path, pathIPs, pathUser, pathPass, "read", ctx.Req)
+			return c.authenticate(ctx.Path, pathIPs, pathUser, pathPass, "read", ctx.Req, ctx.Query)
 		},
 	})
 
diff --git a/internal/core/rtsp_session.go b/internal/core/rtsp_session.go
index 5f956d83..c3d7dffb 100644
--- a/internal/core/rtsp_session.go
+++ b/internal/core/rtsp_session.go
@@ -161,7 +161,7 @@ func (s *rtspSession) onAnnounce(c *rtspConn, ctx *gortsplib.ServerHandlerOnAnno
 			pathIPs []interface{},
 			pathUser conf.Credential,
 			pathPass conf.Credential) error {
-			return c.authenticate(ctx.Path, pathIPs, pathUser, pathPass, "publish", ctx.Req)
+			return c.authenticate(ctx.Path, pathIPs, pathUser, pathPass, "publish", ctx.Req, ctx.Query)
 		},
 	})
 
@@ -220,7 +220,7 @@ func (s *rtspSession) onSetup(c *rtspConn, ctx *gortsplib.ServerHandlerOnSetupCt
 				pathIPs []interface{},
 				pathUser conf.Credential,
 				pathPass conf.Credential) error {
-				return c.authenticate(ctx.Path, pathIPs, pathUser, pathPass, "read", ctx.Req)
+				return c.authenticate(ctx.Path, pathIPs, pathUser, pathPass, "read", ctx.Req, ctx.Query)
 			},
 		})
 
diff --git a/rtsp-simple-server.yml b/rtsp-simple-server.yml
index b98f95db..45d75794 100644
--- a/rtsp-simple-server.yml
+++ b/rtsp-simple-server.yml
@@ -26,6 +26,7 @@ readBufferCount: 512
 #   "password": "password",
 #   "path": "path",
 #   "action": "read|publish"
+#   "query": "url's raw query"
 # }
 # If the response code is 20x, authentication is accepted, otherwise
 # it is discarded.