From 039d69ba4216f837c9150a8ae0dc3ef8ea2482cc Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Tue, 19 Jan 2021 14:28:46 +0100
Subject: [PATCH] rewrite extraction of control attribute

---
 pkg/base/url.go    | 64 +++++++++++++++++++++++++++++-----------------
 serverconf_test.go | 23 ++++++++++++-----
 serverconn.go      | 18 ++-----------
 3 files changed, 59 insertions(+), 46 deletions(-)

diff --git a/pkg/base/url.go b/pkg/base/url.go
index 83e74450..f7aa5f3f 100644
--- a/pkg/base/url.go
+++ b/pkg/base/url.go
@@ -3,8 +3,46 @@ package base
 import (
 	"fmt"
 	"net/url"
+	"strconv"
 )
 
+func stringsReverseIndex(s, substr string) int {
+	for i := len(s) - 1 - len(substr); i >= 0; i-- {
+		if s[i:i+len(substr)] == substr {
+			return i
+		}
+	}
+	return -1
+}
+
+// PathSplitControlAttribute splits a path and query from a control attribute.
+func PathSplitControlAttribute(pathAndQuery string) (int, string, bool) {
+	i := stringsReverseIndex(pathAndQuery, "/trackID=")
+
+	// URL doesn't contain trackID - we assume it's track 0
+	if i < 0 {
+		return 0, pathAndQuery, true
+	}
+
+	tmp, err := strconv.ParseInt(pathAndQuery[i+len("/trackID="):], 10, 64)
+	if err != nil || tmp < 0 {
+		return 0, "", false
+	}
+	trackID := int(tmp)
+
+	return trackID, pathAndQuery[:i], true
+}
+
+// PathSplitQuery splits a path from a query.
+func PathSplitQuery(pathAndQuery string) (string, string) {
+	i := stringsReverseIndex(pathAndQuery, "?")
+	if i >= 0 {
+		return pathAndQuery[:i], pathAndQuery[i:]
+	}
+
+	return pathAndQuery, ""
+}
+
 // URL is a RTSP URL.
 // This is basically an HTTP URL with some additional functions to handle
 // control attributes.
@@ -67,19 +105,12 @@ func (u *URL) CloneWithoutCredentials() *URL {
 
 // RTSPPath returns the path of a RTSP URL.
 func (u *URL) RTSPPath() (string, bool) {
-	var path string
-	if u.RawPath != "" {
-		path = u.RawPath
-	} else {
-		path = u.Path
-	}
-
-	// remove leading slash
-	if len(path) == 0 || path[0] != '/' {
+	pathAndQuery, ok := u.RTSPPathAndQuery()
+	if !ok {
 		return "", false
 	}
-	path = path[1:]
 
+	path, _ := PathSplitQuery(pathAndQuery)
 	return path, true
 }
 
@@ -104,19 +135,6 @@ func (u *URL) RTSPPathAndQuery() (string, bool) {
 	return pathAndQuery, true
 }
 
-// SetRTSPPathAndQuery sets the path and query of a RTSP URL.
-func (u *URL) SetRTSPPathAndQuery(pathAndQuery string) {
-	nu1 := &URL{
-		Scheme: u.Scheme,
-		Opaque: u.Opaque,
-		User:   u.User,
-		Host:   u.Host,
-		Path:   "/",
-	}
-	nu2, _ := ParseURL(nu1.String() + pathAndQuery)
-	*u = *nu2
-}
-
 // AddControlAttribute adds a control attribute to a RTSP url.
 func (u *URL) AddControlAttribute(controlPath string) {
 	if controlPath[0] != '?' {
diff --git a/serverconf_test.go b/serverconf_test.go
index 557425cc..6e119327 100644
--- a/serverconf_test.go
+++ b/serverconf_test.go
@@ -170,15 +170,24 @@ func (ts *testServ) handleConn(conn *ServerConn) {
 	}
 
 	onSetup := func(req *base.Request, th *headers.Transport, trackID int) (*base.Response, error) {
-		reqPath, ok := req.URL.RTSPPath()
-		if !ok {
-			return &base.Response{
-				StatusCode: base.StatusBadRequest,
-			}, fmt.Errorf("invalid path (%s)", req.URL)
-		}
-
 		switch conn.State() {
 		case ServerConnStateInitial, ServerConnStatePrePlay:
+			pathAndQuery, ok := req.URL.RTSPPathAndQuery()
+			if !ok {
+				return &base.Response{
+					StatusCode: base.StatusBadRequest,
+				}, fmt.Errorf("invalid path (%s)", req.URL)
+			}
+
+			_, pathAndQuery, ok = base.PathSplitControlAttribute(pathAndQuery)
+			if !ok {
+				return &base.Response{
+					StatusCode: base.StatusBadRequest,
+				}, fmt.Errorf("invalid path (%s)", req.URL)
+			}
+
+			reqPath, _ := base.PathSplitQuery(pathAndQuery)
+
 			if reqPath != "teststream" {
 				return &base.Response{
 					StatusCode: base.StatusBadRequest,
diff --git a/serverconn.go b/serverconn.go
index a1817503..ca26919c 100644
--- a/serverconn.go
+++ b/serverconn.go
@@ -6,7 +6,6 @@ import (
 	"errors"
 	"fmt"
 	"net"
-	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -500,23 +499,10 @@ func (sc *ServerConn) handleRequest(req *base.Request) (*base.Response, error) {
 
 			trackID, err := func() (int, error) {
 				if th.Mode == nil || *th.Mode == headers.TransportModePlay {
-					i := strings.Index(pathAndQuery, "/trackID=")
-
-					// URL doesn't contain trackID - we assume it's track 0
-					if i < 0 {
-						return 0, nil
-					}
-
-					tmp, err := strconv.ParseInt(pathAndQuery[i+len("/trackID="):], 10, 64)
-					if err != nil || tmp < 0 {
+					trackID, _, ok := base.PathSplitControlAttribute(pathAndQuery)
+					if !ok {
 						return 0, fmt.Errorf("invalid track (%s)", pathAndQuery)
 					}
-					trackID := int(tmp)
-
-					// remove track ID from path
-					nu := req.URL.Clone()
-					nu.SetRTSPPathAndQuery(pathAndQuery[:i])
-					req.URL = nu
 
 					return trackID, nil
 				}