From 00201af40cbd77d30d658bec4f2e04935dc34075 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:00:24 +0100
Subject: [PATCH] implement config locks.

---
 docs/changelog.md                         |  1 +
 docs/getting-started/configuration.md     |  3 +++
 server/internal/types/config/websocket.go |  7 +++++++
 server/internal/websocket/websocket.go    | 11 ++++++++++-
 4 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/docs/changelog.md b/docs/changelog.md
index 2c662569..d53ed05b 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -4,6 +4,7 @@
 
 ### New Features
 - Lock controls for users, globally.
+- Ability to set locks from config `NEKO_LOCKS=control login`.
 
 ### 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 fab42fd1..d81fac1b 100644
--- a/docs/getting-started/configuration.md
+++ b/docs/getting-started/configuration.md
@@ -65,6 +65,9 @@ NEKO_ICESERVERS:
   - Describes multiple STUN and TURN server that can be used by the ICEAgent to establish a connection with a peer
   - e.g. '[{"urls": ["turn:turn.example.com:19302", "stun:stun.example.com:19302"], "username": "name", "credential": "password"}, {"urls": ["stun:stun.example2.com:19302"]}]'
   - [More information](https://developer.mozilla.org/en-US/docs/Web/API/RTCIceServer)
+NEKO_LOCKS:
+  - Resources, that will be locked when starting, separated by whitespace.
+  - Currently supported: control login
 ```
 
 ## Agruments
diff --git a/server/internal/types/config/websocket.go b/server/internal/types/config/websocket.go
index 099f34bd..b21ef5ca 100644
--- a/server/internal/types/config/websocket.go
+++ b/server/internal/types/config/websocket.go
@@ -9,6 +9,7 @@ type WebSocket struct {
 	Password      string
 	AdminPassword string
 	Proxy         bool
+	Locks         []string
 }
 
 func (WebSocket) Init(cmd *cobra.Command) error {
@@ -27,6 +28,11 @@ func (WebSocket) Init(cmd *cobra.Command) error {
 		return err
 	}
 
+	cmd.PersistentFlags().StringSlice("locks", []string{}, "resources, that will be locked when starting (control, login)")
+	if err := viper.BindPFlag("locks", cmd.PersistentFlags().Lookup("locks")); err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -34,4 +40,5 @@ func (s *WebSocket) Set() {
 	s.Password = viper.GetString("password")
 	s.AdminPassword = viper.GetString("password_admin")
 	s.Proxy = viper.GetBool("proxy")
+	s.Locks = viper.GetStringSlice("locks")
 }
diff --git a/server/internal/websocket/websocket.go b/server/internal/websocket/websocket.go
index 3a8c10a3..d60de436 100644
--- a/server/internal/websocket/websocket.go
+++ b/server/internal/websocket/websocket.go
@@ -21,6 +21,15 @@ import (
 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()
 
+	locks := make(map[string]string)
+	for _, lock := range conf.Locks {
+		locks[lock] = "" // empty session ID
+	}
+
+	if len(conf.Locks) > 0 {
+		logger.Info().Msgf("locked resources: %+v", conf.Locks)
+	}
+
 	return &WebSocketHandler{
 		logger:   logger,
 		shutdown: make(chan interface{}),
@@ -39,7 +48,7 @@ func New(sessions types.SessionManager, remote types.RemoteManager, broadcast ty
 			sessions:  sessions,
 			webrtc:    webrtc,
 			banned:    make(map[string]bool),
-			locked:    make(map[string]string),
+			locked:    locks,
 		},
 		conns: 0,
 	}