diff --git a/connector.go b/connector.go index d972e5e..c906116 100644 --- a/connector.go +++ b/connector.go @@ -54,7 +54,7 @@ func CreateZapLogger(config LoggingConfig) *zap.SugaredLogger { // ExposedProcess contains Process along with connected clients and cached output. type ExposedProcess struct { *Process - Clients *xsync.MapOf[string, chan interface{}] + Clients *xsync.MapOf[chan interface{}, string] Console string ConsoleLock sync.RWMutex } @@ -175,7 +175,7 @@ func InitializeConnector(config *Config) *Connector { func (connector *Connector) AddProcess(proc *Process) { process := &ExposedProcess{ Process: proc, - Clients: xsync.NewMapOf[string, chan interface{}](), + Clients: xsync.NewMapOf[chan interface{}, string](), Console: "", } connector.Processes.Store(process.Name, process) @@ -198,7 +198,7 @@ func (connector *Connector) AddProcess(proc *Process) { process.Console = strings.Join(truncate[len(truncate)-2500:], "\n") } process.Console = process.Console + "\n" + m - process.Clients.Range(func(key string, connection chan interface{}) bool { + process.Clients.Range(func(connection chan interface{}, token string) bool { connection <- m return true }) diff --git a/endpoints_misc.go b/endpoints_misc.go index e746758..b36e2f6 100644 --- a/endpoints_misc.go +++ b/endpoints_misc.go @@ -331,22 +331,21 @@ func consoleEndpoint(connector *Connector, w http.ResponseWriter, r *http.Reques process.ConsoleLock.RLock() defer process.ConsoleLock.RUnlock() writeChannel <- process.Console - process.Clients.Store(token, writeChannel) + process.Clients.Store(writeChannel, token) })() // Read messages from the user and execute them. for { - client, _ := process.Clients.Load(token) - // Another client has connected with the same token. Terminate existing connection. - if client != writeChannel { + _, ok := process.Clients.Load(writeChannel) // If gone, stop reading messages from client. + if !ok { break } // Read messages from the user. _, message, err := c.ReadMessage() if err != nil { - process.Clients.Delete(token) + process.Clients.Delete(writeChannel) break // The WebSocket connection has terminated. } else if _, ok := connector.Authenticator.GetUsers().Load(user); !ok && r.RemoteAddr != "@" { - process.Clients.Delete(token) + process.Clients.Delete(writeChannel) c.Close() break } diff --git a/process.go b/process.go index 8206bcc..d9b9b21 100644 --- a/process.go +++ b/process.go @@ -149,8 +149,8 @@ func (process *Process) MonitorProcess(connector *Connector) error { "stopped/crashed, and has now been removed.") if process, loaded := connector.Processes.LoadAndDelete(process.Name); loaded { <-time.After(5 * time.Second) - process.Clients.Range(func(key string, ws chan interface{}) bool { - ws <- nil + process.Clients.Range(func(connection chan interface{}, token string) bool { + connection <- nil return true }) process.Clients.Clear()