Refine API of webtty package

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