Refine API of webtty package

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)

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")
 	}

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 {

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)
+	}
+}

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")
 )

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

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 (

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
 	}
 }

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
 }

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)
 		}