From ad5f6b8db5ed5b4040041f1fb8270ee84e4b70c0 Mon Sep 17 00:00:00 2001
From: videogame hacker <half-kh-hacker@hackery.site>
Date: Thu, 23 Feb 2023 16:35:15 +0000
Subject: [PATCH] Add streams database and password auth

---
 wish-server/.env          |  1 +
 wish-server/.gitignore    |  1 +
 wish-server/db_schema.sql |  4 +++
 wish-server/go.mod        |  1 +
 wish-server/go.sum        |  2 ++
 wish-server/main.go       | 56 ++++++++++++++++++++++++++++-----------
 6 files changed, 49 insertions(+), 16 deletions(-)
 create mode 100644 wish-server/.gitignore
 create mode 100644 wish-server/db_schema.sql

diff --git a/wish-server/.env b/wish-server/.env
index 0d8751c..aa4c0f3 100644
--- a/wish-server/.env
+++ b/wish-server/.env
@@ -1 +1,2 @@
 HTTP_ADDRESS=127.0.0.1:3001
+DATABASE_URL=./streams.db
diff --git a/wish-server/.gitignore b/wish-server/.gitignore
new file mode 100644
index 0000000..4523c5d
--- /dev/null
+++ b/wish-server/.gitignore
@@ -0,0 +1 @@
+/streams.db
diff --git a/wish-server/db_schema.sql b/wish-server/db_schema.sql
new file mode 100644
index 0000000..b9456bc
--- /dev/null
+++ b/wish-server/db_schema.sql
@@ -0,0 +1,4 @@
+CREATE TABLE streams (
+  stream TEXT NOT NULL,
+  password TEXT NOT NULL
+) STRICT;
diff --git a/wish-server/go.mod b/wish-server/go.mod
index 16394df..4e25f6e 100644
--- a/wish-server/go.mod
+++ b/wish-server/go.mod
@@ -9,6 +9,7 @@ require (
 
 require (
 	github.com/google/uuid v1.3.0 // indirect
+	github.com/mattn/go-sqlite3 v1.14.16 // indirect
 	github.com/pion/datachannel v1.5.5 // indirect
 	github.com/pion/dtls/v2 v2.2.4 // indirect
 	github.com/pion/logging v0.2.2 // indirect
diff --git a/wish-server/go.sum b/wish-server/go.sum
index 5b40257..c05bcae 100644
--- a/wish-server/go.sum
+++ b/wish-server/go.sum
@@ -24,6 +24,8 @@ github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwA
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
+github.com/mattn/go-sqlite3 v1.14.16 h1:yOQRA0RpS5PFz/oikGwBEqvAWhWg5ufRz4ETLjwpU1Y=
+github.com/mattn/go-sqlite3 v1.14.16/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg=
 github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
 github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU=
 github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
diff --git a/wish-server/main.go b/wish-server/main.go
index 8eb046a..aee2ec0 100644
--- a/wish-server/main.go
+++ b/wish-server/main.go
@@ -10,9 +10,13 @@ import (
 	"strings"
 	"sync"
 
+	"database/sql"
+
 	"github.com/joho/godotenv"
 	"github.com/pion/interceptor"
 	"github.com/pion/webrtc/v3"
+
+	_ "github.com/mattn/go-sqlite3"
 )
 
 type WebRTCStream struct {
@@ -23,6 +27,7 @@ var (
 	runningStreams     map[string]WebRTCStream
 	runningStreamsLock sync.Mutex
 	api                *webrtc.API
+	db                 *sql.DB
 )
 
 func main() {
@@ -30,6 +35,32 @@ func main() {
 		log.Fatal(err)
 	}
 
+	db = setupDatabase()
+	defer db.Close()
+
+	api = setupWebRTC()
+	runningStreams = map[string]WebRTCStream{}
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/api/wish-server/whip", withCors(HandleWHIP))
+	mux.HandleFunc("/api/wish-server/whep", withCors(HandleWHEP))
+
+	log.Fatal((&http.Server{
+		Handler: mux,
+		Addr:    os.Getenv("HTTP_ADDRESS"),
+	}).ListenAndServe())
+}
+
+func setupDatabase() *sql.DB {
+	db, err := sql.Open("sqlite3", os.Getenv("DATABASE_URL"))
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return db
+}
+
+func setupWebRTC() *webrtc.API {
 	mediaEngine := &webrtc.MediaEngine{}
 	if err := mediaEngine.RegisterDefaultCodecs(); err != nil {
 		log.Fatal(err)
@@ -42,21 +73,11 @@ func main() {
 
 	settingEngine := webrtc.SettingEngine{}
 
-	api = webrtc.NewAPI(
+	return webrtc.NewAPI(
 		webrtc.WithMediaEngine(mediaEngine),
 		webrtc.WithInterceptorRegistry(interceptorRegistry),
 		webrtc.WithSettingEngine(settingEngine),
 	)
-
-	runningStreams = map[string]WebRTCStream{}
-	mux := http.NewServeMux()
-	mux.HandleFunc("/api/wish-server/whip", withCors(HandleWHIP))
-	mux.HandleFunc("/api/wish-server/whep", withCors(HandleWHEP))
-
-	log.Fatal((&http.Server{
-		Handler: mux,
-		Addr:    os.Getenv("HTTP_ADDRESS"),
-	}).ListenAndServe())
 }
 
 func withCors(next func(w http.ResponseWriter, r *http.Request)) http.HandlerFunc {
@@ -107,14 +128,17 @@ func getTracksForStream(streamName string) (
 }
 
 func HandleWHIP(res http.ResponseWriter, req *http.Request) {
-	// TODO: SQLite3 to store stream keys (split auth by colon)
-
-	streamName := req.Header.Get("Authorization")
-	if streamName == "" {
-		logHTTPError(res, "Stream name was not set", http.StatusBadRequest)
+	authorization := req.Header.Get("Authorization")
+	if authorization == "" {
+		logHTTPError(res, "Authorization was not set", http.StatusBadRequest)
 		return
 	}
 
+	streamName, streamPassword, _ := strings.Cut(authorization, ":")
+	if db.QueryRow("SELECT * FROM streams WHERE stream = ? AND password = ?", streamName, streamPassword).Err() != nil {
+		logHTTPError(res, "Invalid stream authorization", http.StatusUnauthorized)
+	}
+
 	offer, err := io.ReadAll(req.Body)
 	if err != nil {
 		logHTTPError(res, err.Error(), http.StatusInternalServerError)