diff --git a/app/app.go b/app/app.go index 998c9db..ed2e250 100644 --- a/app/app.go +++ b/app/app.go @@ -14,6 +14,7 @@ import ( "os/exec" "strconv" "strings" + "sync" "text/template" "github.com/braintree/manners" @@ -243,6 +244,7 @@ func (app *App) handleWS(w http.ResponseWriter, r *http.Request) { connection: conn, command: cmd, pty: ptyIo, + writeMutex: &sync.Mutex{}, } context.goHandleClient() diff --git a/app/client_context.go b/app/client_context.go index 63ad67e..ca38d23 100644 --- a/app/client_context.go +++ b/app/client_context.go @@ -8,6 +8,7 @@ import ( "os" "os/exec" "strings" + "sync" "syscall" "unsafe" @@ -20,6 +21,7 @@ type clientContext struct { connection *websocket.Conn command *exec.Cmd pty *os.File + writeMutex *sync.Mutex } const ( @@ -96,19 +98,24 @@ func (context *clientContext) processSend() { for { size, err := context.pty.Read(buf) - safeMessage := base64.StdEncoding.EncodeToString([]byte(buf[:size])) if err != nil { log.Printf("Command exited for: %s", context.request.RemoteAddr) return } - - err = context.connection.WriteMessage(websocket.TextMessage, append([]byte{Output}, []byte(safeMessage)...)) - if err != nil { + safeMessage := base64.StdEncoding.EncodeToString([]byte(buf[:size])) + if err = context.write(append([]byte{Output}, []byte(safeMessage)...)); err != nil { + log.Printf(err.Error()) return } } } +func (context *clientContext) write(data []byte) error { + context.writeMutex.Lock() + defer context.writeMutex.Unlock() + return context.connection.WriteMessage(websocket.TextMessage, data) +} + func (context *clientContext) sendInitialize() error { hostname, _ := os.Hostname() titleVars := ContextVars{ @@ -118,34 +125,34 @@ func (context *clientContext) sendInitialize() error { RemoteAddr: context.request.RemoteAddr, } + context.writeMutex.Lock() writer, err := context.connection.NextWriter(websocket.TextMessage) if err != nil { + context.writeMutex.Unlock() return err } writer.Write([]byte{SetWindowTitle}) if err = context.app.titleTemplate.Execute(writer, titleVars); err != nil { + context.writeMutex.Unlock() return err } writer.Close() + context.writeMutex.Unlock() htermPrefs := make(map[string]interface{}) for key, value := range context.app.options.Preferences { htermPrefs[strings.Replace(key, "_", "-", -1)] = value } prefs, _ := json.Marshal(htermPrefs) - context.connection.WriteMessage( - websocket.TextMessage, - append([]byte{SetPreferences}, prefs...), - ) - + if err := context.write(append([]byte{SetPreferences}, prefs...)); err != nil { + return err + } if context.app.options.EnableReconnect { reconnect, _ := json.Marshal(context.app.options.ReconnectTime) - context.connection.WriteMessage( - websocket.TextMessage, - append([]byte{SetReconnect}, reconnect...), - ) + if err := context.write(append([]byte{SetReconnect}, reconnect...)); err != nil { + return err + } } - return nil } @@ -153,6 +160,11 @@ func (context *clientContext) processReceive() { for { _, data, err := context.connection.ReadMessage() if err != nil { + log.Print(err.Error()) + return + } + if len(data) == 0 { + log.Print("An error has occured") return } @@ -168,7 +180,10 @@ func (context *clientContext) processReceive() { } case Ping: - context.connection.WriteMessage(websocket.TextMessage, []byte{Pong}) + if err := context.write([]byte{Pong}); err != nil { + log.Print(err.Error()) + return + } case ResizeTerminal: var args argResizeTerminal err = json.Unmarshal(data[1:], &args)