From abe3f5c89f13b1a633549b0c83f9fb62cbf1b5df Mon Sep 17 00:00:00 2001
From: Iwasaki Yudai <yudai@arielworks.com>
Date: Sun, 16 Aug 2015 19:00:09 -0700
Subject: [PATCH] Don't send invalid UTF-8 string

When a rune stands over the buffer boundary, invalid UTF-8 bytes were
sent to the client. Now the logic checks if the tail of the buffer can
be decoded as a valid rune. When it's an invalid byte sequence, the logic
keeps it in the buffer and send it in the next loop with the following byte sequence.
---
 app/app.go | 41 +++++++++++++++++++++++++++++++++++------
 1 file changed, 35 insertions(+), 6 deletions(-)

diff --git a/app/app.go b/app/app.go
index f0a224c..4d8be5b 100644
--- a/app/app.go
+++ b/app/app.go
@@ -7,6 +7,7 @@ import (
 	"os/exec"
 	"strings"
 	"syscall"
+	"unicode/utf8"
 	"unsafe"
 
 	"github.com/elazarl/go-bindata-assetfs"
@@ -53,8 +54,8 @@ func (app *App) generateHandler() func(w http.ResponseWriter, r *http.Request) {
 		log.Printf("New client connected: %s", r.RemoteAddr)
 
 		upgrader := websocket.Upgrader{
-			ReadBufferSize:  0,
-			WriteBufferSize: 0,
+			ReadBufferSize:  1024,
+			WriteBufferSize: 1024,
 			Subprotocols:    []string{"gotty"},
 		}
 
@@ -81,19 +82,47 @@ func (app *App) generateHandler() func(w http.ResponseWriter, r *http.Request) {
 		go func() {
 			defer func() { exit <- true }()
 
-			buf := make([]byte, 512)
+			buf := make([]byte, 1024)
+			leftOver := 0
 			for {
-				len, err := fio.Read(buf)
+				size, err := fio.Read(buf[leftOver:])
+				size += leftOver
+
 				if err != nil {
 					log.Printf("command exited for: %s", r.RemoteAddr)
 					return
 				}
-				writer, err := conn.NextWriter(1)
+
+				writer, err := conn.NextWriter(websocket.TextMessage)
 				if err != nil {
 					return
 				}
-				writer.Write(buf[0:len])
+
+				// UTF-8 Boundary check
+				for leftOver = 0; leftOver < utf8.UTFMax; leftOver++ {
+					re, _ := utf8.DecodeLastRune(
+						buf[:size-leftOver],
+					)
+
+					if re != utf8.RuneError {
+						break
+					}
+					// Invalid UTF rune
+				}
+
+				if leftOver == utf8.UTFMax-1 {
+					re, _ := utf8.DecodeLastRune(buf[:size-leftOver])
+					if re == utf8.RuneError {
+						log.Fatal("UTF8 Boundary error.")
+					}
+				}
+
+				writer.Write(buf[:size-leftOver])
 				writer.Close()
+
+				for i := 0; i < leftOver; i++ {
+					buf[i] = buf[size-leftOver+i]
+				}
 			}
 		}()