refactor WS authentication.

internal/session/auth.go

@@ -3,35 +3,54 @@ package session
 import (
 	"fmt"
 	"net/http"
+	"strings"
 
+	"demodesk/neko/internal/types"
 	"demodesk/neko/internal/utils"
 )
 
-// TODO: Refactor
-func (manager *SessionManagerCtx) Authenticate(r *http.Request) (string, string, bool, error) {
-	ip := r.RemoteAddr
+const (
+	token_name = "password"
+)
 
-	//if ws.conf.Proxy {
-	//	ip = utils.ReadUserIP(r)
-	//}
+func (manager *SessionManagerCtx) Authenticate(r *http.Request) (types.Session, error) {
+	token := getToken(r)
+	if token == "" {
+		return nil, fmt.Errorf("no password provided")
+	}
+
+	isAdmin := (token == manager.config.AdminPassword)
+	isUser := (token == manager.config.Password)
+
+	if !isAdmin && !isUser {
+		return nil, fmt.Errorf("invalid password")
+	}
 
 	id, err := utils.NewUID(32)
 	if err != nil {
-		return "", ip, false, err
+		return nil, err
 	}
 
-	passwords, ok := r.URL.Query()["password"]
-	if !ok || len(passwords[0]) < 1 {
-		return "", ip, false, fmt.Errorf("no password provided")
-	}
-
-	if passwords[0] == manager.config.AdminPassword {
-		return id, ip, true, nil
-	}
-
-	if passwords[0] == manager.config.Password {
-		return id, ip, false, nil
-	}
-
-	return "", ip, false, fmt.Errorf("invalid password: %s", passwords[0])
+	return manager.New(id, isAdmin), nil
+}
+
+func getToken(r *http.Request) string {
+	// Get token from query
+	if token := r.URL.Query().Get(token_name); token != "" {
+		return token
+	}
+
+	// Get token from authorization header
+	bearer := r.Header.Get("Authorization")
+	if len(bearer) > 7 && strings.ToUpper(bearer[0:6]) == "BEARER" {
+		return bearer[7:]
+	}
+
+	// Get token from cookie
+	cookie, err := r.Cookie(token_name)
+	if err == nil {
+		return cookie.Value
+	}
+
+	return ""
 }

internal/session/manager.go

@@ -30,23 +30,16 @@ type SessionManagerCtx struct {
 	emmiter events.EventEmmiter
 }
 
-func (manager *SessionManagerCtx) New(id string, admin bool, socket types.WebSocket) types.Session {
+func (manager *SessionManagerCtx) New(id string, admin bool) types.Session {
 	session := &SessionCtx{
 		id:        id,
 		admin:     admin,
 		manager:   manager,
-		socket:    socket,
 		logger:    manager.logger.With().Str("id", id).Logger(),
 		connected: false,
 	}
 
 	manager.members[id] = session
-	manager.emmiter.Emit("created", session)
-
-	if !manager.capture.Streaming() && len(manager.members) > 0 {
-		manager.capture.StartStream()
-	}
-
 	return session
 }
 
@@ -66,10 +59,6 @@ func (manager *SessionManagerCtx) Destroy(id string) error {
 		delete(manager.members, id)
 		err := session.destroy()
 
-		if !manager.capture.Streaming() && len(manager.members) <= 0 {
-			manager.capture.StopStream()
-		}
-
 		manager.emmiter.Emit("destroy", id)
 		return err
 	}
@@ -164,12 +153,22 @@ func (manager *SessionManagerCtx) OnHostCleared(listener func(session types.Sess
 
 func (manager *SessionManagerCtx) OnDestroy(listener func(id string)) {
 	manager.emmiter.On("destroy", func(payload ...interface{}) {
+		// Stop streaming, if everyone left
+		if manager.capture.Streaming() && len(manager.members) == 0 {
+			manager.capture.StopStream()
+		}
+
 		listener(payload[0].(string))
 	})
 }
 
 func (manager *SessionManagerCtx) OnCreated(listener func(session types.Session)) {
 	manager.emmiter.On("created", func(payload ...interface{}) {
+		// Start streaming, when first joins
+		if !manager.capture.Streaming() {
+			manager.capture.StartStream()
+		}
+	
 		listener(payload[0].(*SessionCtx))
 	})
 }

internal/session/session.go

@@ -62,6 +62,7 @@ func (session *SessionCtx) SetName(name string) {
 
 func (session *SessionCtx) SetSocket(socket types.WebSocket) {
 	session.socket = socket
+	session.manager.emmiter.Emit("created", session)
 }
 
 func (session *SessionCtx) SetPeer(peer types.Peer) {