From e2ba39922cf7e958432217fc122951ced1a1769c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Tue, 16 Nov 2021 23:48:40 +0100
Subject: [PATCH] added control protection.

---
 docs/changelog.md                         |  1 +
 docs/getting-started/configuration.md     |  3 ++
 server/internal/types/config/websocket.go |  9 +++++
 server/internal/websocket/websocket.go    | 41 +++++++++++++++++++++++
 4 files changed, 54 insertions(+)

diff --git a/docs/changelog.md b/docs/changelog.md
index d53ed05b..dce92124 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -5,6 +5,7 @@
 ### New Features
 - Lock controls for users, globally.
 - Ability to set locks from config `NEKO_LOCKS=control login`.
+- Added control protection - users can gain control only if at least one admin is in the room `NEKO_CONTROL_PROTECTION=true`.
 
 ### Misc
 - ARM-based images not bound to Raspberry Pi only.
diff --git a/docs/getting-started/configuration.md b/docs/getting-started/configuration.md
index d81fac1b..270648c3 100644
--- a/docs/getting-started/configuration.md
+++ b/docs/getting-started/configuration.md
@@ -68,6 +68,9 @@ NEKO_ICESERVERS:
 NEKO_LOCKS:
   - Resources, that will be locked when starting, separated by whitespace.
   - Currently supported: control login
+NEKO_CONTROL_PROTECTION:
+  - Control protection means, users can gain control only if at least one admin is in the room.
+  - e.g. false
 ```
 
 ## Agruments
diff --git a/server/internal/types/config/websocket.go b/server/internal/types/config/websocket.go
index b21ef5ca..56450a63 100644
--- a/server/internal/types/config/websocket.go
+++ b/server/internal/types/config/websocket.go
@@ -10,6 +10,8 @@ type WebSocket struct {
 	AdminPassword string
 	Proxy         bool
 	Locks         []string
+
+	ControlProtection bool
 }
 
 func (WebSocket) Init(cmd *cobra.Command) error {
@@ -33,6 +35,11 @@ func (WebSocket) Init(cmd *cobra.Command) error {
 		return err
 	}
 
+	cmd.PersistentFlags().Bool("control_protection", true, "control protection means, users can gain control only if at least one admin is in the room")
+	if err := viper.BindPFlag("control_protection", cmd.PersistentFlags().Lookup("control_protection")); err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -41,4 +48,6 @@ func (s *WebSocket) Set() {
 	s.AdminPassword = viper.GetString("password_admin")
 	s.Proxy = viper.GetBool("proxy")
 	s.Locks = viper.GetStringSlice("locks")
+
+	s.ControlProtection = viper.GetBool("control_protection")
 }
diff --git a/server/internal/websocket/websocket.go b/server/internal/websocket/websocket.go
index d60de436..c7643e37 100644
--- a/server/internal/websocket/websocket.go
+++ b/server/internal/websocket/websocket.go
@@ -18,6 +18,8 @@ import (
 	"m1k1o/neko/internal/utils"
 )
 
+const CONTROL_PROTECTION_SESSION = "control_protection"
+
 func New(sessions types.SessionManager, remote types.RemoteManager, broadcast types.BroadcastManager, webrtc types.WebRTCManager, conf *config.WebSocket) *WebSocketHandler {
 	logger := log.With().Str("module", "websocket").Logger()
 
@@ -30,6 +32,12 @@ func New(sessions types.SessionManager, remote types.RemoteManager, broadcast ty
 		logger.Info().Msgf("locked resources: %+v", conf.Locks)
 	}
 
+	// if control protection is enabled
+	if conf.ControlProtection {
+		locks["control"] = CONTROL_PROTECTION_SESSION
+		logger.Info().Msgf("control locked on behalf of control protection")
+	}
+
 	return &WebSocketHandler{
 		logger:   logger,
 		shutdown: make(chan interface{}),
@@ -84,6 +92,23 @@ func (ws *WebSocketHandler) Start() {
 		} else {
 			ws.logger.Debug().Str("id", id).Msg("session connected")
 		}
+
+		// if control protection is enabled and at least one admin
+		// and if room was locked on behalf control protection, unlock
+		sess, ok := ws.handler.locked["control"]
+		if ok && ws.conf.ControlProtection && sess == CONTROL_PROTECTION_SESSION && len(ws.sessions.Admins()) > 0 {
+			delete(ws.handler.locked, "control")
+			ws.logger.Info().Msgf("control unlocked on behalf of control protection")
+
+			if err := ws.sessions.Broadcast(
+				message.AdminLock{
+					Event:    event.ADMIN_UNLOCK,
+					ID:       id,
+					Resource: "control",
+				}, nil); err != nil {
+				ws.logger.Warn().Err(err).Msgf("broadcasting event %s has failed", event.ADMIN_UNLOCK)
+			}
+		}
 	})
 
 	ws.sessions.OnDestroy(func(id string, session types.Session) {
@@ -92,6 +117,22 @@ func (ws *WebSocketHandler) Start() {
 		} else {
 			ws.logger.Debug().Str("id", id).Msg("session destroyed")
 		}
+
+		// if control protection is enabled and no admin
+		if ws.conf.ControlProtection && len(ws.sessions.Admins()) == 0 {
+			ws.handler.locked["control"] = CONTROL_PROTECTION_SESSION
+			ws.logger.Info().Msgf("control locked and released on behalf of control protection")
+			ws.handler.adminRelease(id, session)
+
+			if err := ws.sessions.Broadcast(
+				message.AdminLock{
+					Event:    event.ADMIN_LOCK,
+					ID:       id,
+					Resource: "control",
+				}, nil); err != nil {
+				ws.logger.Warn().Err(err).Msgf("broadcasting event %s has failed", event.ADMIN_LOCK)
+			}
+		}
 	})
 
 	ws.wg.Add(1)