From 807bcc25a4c89a458ecf5fd4fca08f42011bd710 Mon Sep 17 00:00:00 2001 From: Iwasaki Yudai Date: Thu, 24 Aug 2017 14:40:28 +0900 Subject: [PATCH] Refine API of webtty package --- backend/localcommand/local_command.go | 4 -- server/handlers.go | 24 +++-------- server/slave.go | 2 +- server/ws_wrapper.go | 33 ++++++++++++++++ webtty/errors.go | 5 ++- webtty/master.go | 57 +++------------------------ webtty/message_types.go | 2 + webtty/option.go | 15 +++++-- webtty/slave.go | 6 ++- webtty/webtty.go | 31 +++++++-------- 10 files changed, 82 insertions(+), 97 deletions(-) create mode 100644 server/ws_wrapper.go diff --git a/backend/localcommand/local_command.go b/backend/localcommand/local_command.go index 449e620..4beca86 100644 --- a/backend/localcommand/local_command.go +++ b/backend/localcommand/local_command.go @@ -123,10 +123,6 @@ func (lcmd *LocalCommand) ResizeTerminal(width int, height int) error { } } -func (lcmd *LocalCommand) GetTerminalSize() (int, int, error) { - return pty.Getsize(lcmd.pty) -} - func (lcmd *LocalCommand) closeTimeoutC() <-chan time.Time { if lcmd.closeTimeout >= 0 { return time.After(lcmd.closeTimeout) diff --git a/server/handlers.go b/server/handlers.go index 674ecab..f1a2f9b 100644 --- a/server/handlers.go +++ b/server/handlers.go @@ -148,29 +148,17 @@ func (server *Server) processWSConn(ctx context.Context, conn *websocket.Conn) e if server.options.EnableReconnect { opts = append(opts, webtty.WithReconnect(server.options.ReconnectTime)) } - if server.options.Width > 0 || server.options.Height > 0 { - width, height, err := slave.GetTerminalSize() - if err != nil { - return errors.Wrapf(err, "failed to get default terminal size") - } - if server.options.Width > 0 { - width = server.options.Width - } - if server.options.Height > 0 { - height = server.options.Height - } - err = slave.ResizeTerminal(width, height) - if err != nil { - return errors.Wrapf(err, "failed to resize terminal") - } - - opts = append(opts, webtty.WithFixedSize(server.options.Width, server.options.Height)) + if server.options.Width > 0 { + opts = append(opts, webtty.WithFixedColumns(server.options.Width)) + } + if server.options.Height > 0 { + opts = append(opts, webtty.WithFixedRows(server.options.Height)) } if server.options.Preferences != nil { opts = append(opts, webtty.WithMasterPreferences(server.options.Preferences)) } - tty, err := webtty.New(conn, slave, opts...) + tty, err := webtty.New(&wsWrapper{conn}, slave, opts...) if err != nil { return errors.Wrapf(err, "failed to create webtty") } diff --git a/server/slave.go b/server/slave.go index d259abf..77d0973 100644 --- a/server/slave.go +++ b/server/slave.go @@ -8,7 +8,7 @@ import ( type Slave interface { webtty.Slave - GetTerminalSize() (width int, height int, err error) + Close() error } type Factory interface { diff --git a/server/ws_wrapper.go b/server/ws_wrapper.go new file mode 100644 index 0000000..819c393 --- /dev/null +++ b/server/ws_wrapper.go @@ -0,0 +1,33 @@ +package server + +import ( + "github.com/gorilla/websocket" +) + +type wsWrapper struct { + *websocket.Conn +} + +func (wsw *wsWrapper) Write(p []byte) (n int, err error) { + writer, err := wsw.Conn.NextWriter(websocket.TextMessage) + if err != nil { + return 0, err + } + defer writer.Close() + return writer.Write(p) +} + +func (wsw *wsWrapper) Read(p []byte) (n int, err error) { + for { + msgType, reader, err := wsw.Conn.NextReader() + if err != nil { + return 0, err + } + + if msgType != websocket.TextMessage { + continue + } + + return reader.Read(p) + } +} diff --git a/webtty/errors.go b/webtty/errors.go index ea2fcee..0ea6473 100644 --- a/webtty/errors.go +++ b/webtty/errors.go @@ -5,6 +5,9 @@ import ( ) var ( - ErrSlaveClosed = errors.New("slave closed") + // ErrSlaveClosed indicates the function has exited by the slave + ErrSlaveClosed = errors.New("slave closed") + + // ErrSlaveClosed is returned when the slave connection is closed. ErrMasterClosed = errors.New("master closed") ) diff --git a/webtty/master.go b/webtty/master.go index 3312ebf..7f30f47 100644 --- a/webtty/master.go +++ b/webtty/master.go @@ -1,55 +1,8 @@ -/* -Copyright (c) 2013 The Gorilla WebSocket Authors. All rights reserved. - - Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - -Redistributions of source code must retain the above copyright notice, this -list of conditions and the following disclaimer. - - Redistributions in binary form must reproduce the above copyright notice, -this list of conditions and the following disclaimer in the documentation -and/or other materials provided with the distribution. - - THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND -ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED -WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -*/ - package webtty -// Master represents a PTY master, usually it's a websocket connection. -type Master interface { - WriteMessage(messageType int, data []byte) error - ReadMessage() (messageType int, p []byte, err error) -} - -// The message types are defined in RFC 6455, section 11.8. -const ( - // TextMessage denotes a text data message. The text message payload is - // interpreted as UTF-8 encoded text data. - WSTextMessage = 1 - - // BinaryMessage denotes a binary data message. - WSBinaryMessage = 2 - - // CloseMessage denotes a close control message. The optional message - // payload contains a numeric code and text. Use the FormatCloseMessage - // function to format a close message payload. - WSCloseMessage = 8 - - // PingMessage denotes a ping control message. The optional message payload - // is UTF-8 encoded text. - WSPingMessage = 9 - - // PongMessage denotes a ping control message. The optional message payload - // is UTF-8 encoded text. - WSPongMessage = 10 +import ( + "io" ) + +// Master represents a PTY master, usually it's a websocket connection. +type Master io.ReadWriter diff --git a/webtty/message_types.go b/webtty/message_types.go index 8b1721e..c437a42 100644 --- a/webtty/message_types.go +++ b/webtty/message_types.go @@ -1,5 +1,7 @@ package webtty +// Protocols defines the name of this protocol, +// which is supposed to be used to the subprotocol of Websockt streams. var Protocols = []string{"webtty"} const ( diff --git a/webtty/option.go b/webtty/option.go index 5dc407a..1618e89 100644 --- a/webtty/option.go +++ b/webtty/option.go @@ -17,11 +17,18 @@ func WithPermitWrite() Option { } } -// WithFixedSize sets a fixed size to TTY master. -func WithFixedSize(width int, height int) Option { +// WithFixedColumns sets a fixed width to TTY master. +func WithFixedColumns(columns int) Option { return func(wt *WebTTY) error { - wt.width = width - wt.height = height + wt.columns = columns + return nil + } +} + +// WithFixedRows sets a fixed height to TTY master. +func WithFixedRows(rows int) Option { + return func(wt *WebTTY) error { + wt.rows = rows return nil } } diff --git a/webtty/slave.go b/webtty/slave.go index a054aad..b9893ba 100644 --- a/webtty/slave.go +++ b/webtty/slave.go @@ -6,8 +6,12 @@ import ( // Slave represents a PTY slave, typically it's a local command. type Slave interface { - io.ReadWriteCloser + io.ReadWriter + // WindowTitleVariables returns any values that can be used to fill out + // the title of a terminal. WindowTitleVariables() map[string]interface{} + + // ResizeTerminal sets a new size of the terminal. ResizeTerminal(columns int, rows int) error } diff --git a/webtty/webtty.go b/webtty/webtty.go index c4df7d3..ed4f8b9 100644 --- a/webtty/webtty.go +++ b/webtty/webtty.go @@ -9,7 +9,7 @@ import ( "github.com/pkg/errors" ) -// WebTTY bridges sets of a PTY slave and its PTY master. +// WebTTY bridges a PTY slave and its PTY master. // To support text-based streams and side channel commands such as // terminal resizing, WebTTY uses an original protocol. type WebTTY struct { @@ -20,9 +20,9 @@ type WebTTY struct { windowTitle []byte permitWrite bool - width int - height int - reconnect int // in milliseconds + columns int + rows int + reconnect int // in seconds masterPrefs []byte bufferSize int @@ -39,8 +39,8 @@ func New(masterConn Master, slave Slave, options ...Option) (*WebTTY, error) { slave: slave, permitWrite: false, - width: 0, - height: 0, + columns: 0, + rows: 0, bufferSize: 1024, } @@ -52,11 +52,12 @@ func New(masterConn Master, slave Slave, options ...Option) (*WebTTY, error) { return wt, nil } -// Run starts the WebTTY. +// Run starts the main process of the WebTTY. // This method blocks until the context is canceled. // Note that the master and slave are left intact even // after the context is canceled. Closing them is caller's // responsibility. +// If the connection to one end gets closed, returns ErrSlaveClosed or ErrMasterClosed. func (wt *WebTTY) Run(ctx context.Context) error { err := wt.sendInitializeMessage() if err != nil { @@ -84,16 +85,14 @@ func (wt *WebTTY) Run(ctx context.Context) error { go func() { errs <- func() error { + buffer := make([]byte, wt.bufferSize) for { - typ, data, err := wt.masterConn.ReadMessage() + n, err := wt.masterConn.Read(buffer) if err != nil { return ErrMasterClosed } - if typ != WSTextMessage { - continue - } - err = wt.handleMasterReadEvent(data) + err = wt.handleMasterReadEvent(buffer[:n]) if err != nil { return err } @@ -148,7 +147,7 @@ func (wt *WebTTY) masterWrite(data []byte) error { wt.writeMutex.Lock() defer wt.writeMutex.Unlock() - err := wt.masterConn.WriteMessage(WSTextMessage, data) + _, err := wt.masterConn.Write(data) if err != nil { return errors.Wrapf(err, "failed to write to master") } @@ -183,7 +182,7 @@ func (wt *WebTTY) handleMasterReadEvent(data []byte) error { } case ResizeTerminal: - if wt.width != 0 && wt.height != 0 { + if wt.columns != 0 && wt.rows != 0 { break } @@ -196,12 +195,12 @@ func (wt *WebTTY) handleMasterReadEvent(data []byte) error { if err != nil { return errors.Wrapf(err, "received malformed data for terminal resize") } - rows := wt.height + rows := wt.rows if rows == 0 { rows = int(args.Rows) } - columns := wt.width + columns := wt.columns if columns == 0 { columns = int(args.Columns) }