Skip to content

haoqu.ma / gotty

Go to a project
Commit 25627da8 by Iwasaki Yudai
Options

Restructure handler function

1 parent 7e2befa1
Inline Side-by-side
Showing with 133 additions and 106 deletions
app/app.go
...@@ -7,6 +7,7 @@ import ( ...@@ -7,6 +7,7 @@ import (
"log" "log"
"math/big" "math/big"
"net/http" "net/http"
"os"
"os/exec" "os/exec"
"strconv" "strconv"
"strings" "strings"
...@@ -21,6 +22,8 @@ import ( ...@@ -21,6 +22,8 @@ import (
type App struct { type App struct {
options Options options Options
upgrader *websocket.Upgrader
} }
type Options struct { type Options struct {
...@@ -35,6 +38,12 @@ type Options struct { ...@@ -35,6 +38,12 @@ type Options struct {
func New(options Options) *App { func New(options Options) *App {
return &App{ return &App{
options: options, options: options,
upgrader: &websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
Subprotocols: []string{"gotty"},
},
} }
} }
...@@ -77,7 +86,7 @@ func (app *App) Run() error { ...@@ -77,7 +86,7 @@ func (app *App) Run() error {
fs := http.StripPrefix(path, http.FileServer(&assetfs.AssetFS{Asset: Asset, AssetDir: AssetDir, Prefix: "bindata"})) fs := http.StripPrefix(path, http.FileServer(&assetfs.AssetFS{Asset: Asset, AssetDir: AssetDir, Prefix: "bindata"}))
http.Handle(path, fs) http.Handle(path, fs)
http.HandleFunc(path+"ws", app.generateHandler()) http.HandleFunc(path+"ws", app.handler)
endpoint := app.options.Address + ":" + app.options.Port endpoint := app.options.Address + ":" + app.options.Port
log.Printf("Server is running at %s, command: %s", endpoint+path, strings.Join(app.options.Command, " ")) log.Printf("Server is running at %s, command: %s", endpoint+path, strings.Join(app.options.Command, " "))
...@@ -94,120 +103,142 @@ func (app *App) Run() error { ...@@ -94,120 +103,142 @@ func (app *App) Run() error {
return nil return nil
} }
func (app *App) generateHandler() func(w http.ResponseWriter, r *http.Request) { func (app *App) handler(w http.ResponseWriter, r *http.Request) {
return func(w http.ResponseWriter, r *http.Request) { log.Printf("New client connected: %s", r.RemoteAddr)
log.Printf("New client connected: %s", r.RemoteAddr)
upgrader := websocket.Upgrader{ if r.Method != "GET" {
ReadBufferSize: 1024, http.Error(w, "Method not allowed", 405)
WriteBufferSize: 1024, return
Subprotocols: []string{"gotty"}, }
}
if r.Method != "GET" { conn, err := app.upgrader.Upgrade(w, r, nil)
http.Error(w, "Method not allowed", 405) if err != nil {
return log.Print("Failed to upgrade connection")
} return
}
cmd := exec.Command(app.options.Command[0], app.options.Command[1:]...)
ptyIo, err := pty.Start(cmd)
if err != nil {
log.Print("Failed to execute command")
return
}
log.Printf("Command is running for client %s with PID %d", r.RemoteAddr, cmd.Process.Pid)
context := &clientContext{
request: r,
connection: conn,
command: cmd,
pty: ptyIo,
}
app.goHandleConnection(context)
}
func (app *App) goHandleConnection(context *clientContext) {
exit := make(chan bool, 2)
go func() {
defer func() { exit <- true }()
conn, err := upgrader.Upgrade(w, r, nil) app.processSend(context)
}()
go func() {
defer func() { exit <- true }()
app.processReceive(context)
}()
go func() {
<-exit
context.command.Wait()
context.connection.Close()
log.Printf("Connection closed: %s", context.request.RemoteAddr)
}()
}
func (app *App) processSend(context *clientContext) {
buf := make([]byte, 1024)
utf8f := utf8reader.New(context.pty)
for {
size, err := utf8f.Read(buf)
if err != nil { if err != nil {
log.Print("Failed to upgrade connection") log.Printf("Command exited for: %s", context.request.RemoteAddr)
return return
} }
cmd := exec.Command(app.options.Command[0], app.options.Command[1:]...) writer, err := context.connection.NextWriter(websocket.TextMessage)
fio, err := pty.Start(cmd)
log.Printf("Command is running for client %s with PID %d", r.RemoteAddr, cmd.Process.Pid)
if err != nil { if err != nil {
log.Print("Failed to execute command")
return return
} }
exit := make(chan bool, 2) writer.Write(buf[:size])
writer.Close()
go func() { }
defer func() { exit <- true }() }
buf := make([]byte, 1024) func (app *App) processReceive(context *clientContext) {
utf8f := utf8reader.New(fio) for {
_, data, err := context.connection.ReadMessage()
if err != nil {
return
}
for { switch data[0] {
size, err := utf8f.Read(buf) case Input:
if err != nil { if !app.options.PermitWrite {
log.Printf("Command exited for: %s", r.RemoteAddr) break
return }
}
writer, err := conn.NextWriter(websocket.TextMessage) _, err := context.pty.Write(data[1:])
if err != nil { if err != nil {
return return
} }
writer.Write(buf[:size]) case ResizeTerminal:
writer.Close() var args argResizeTerminal
err = json.Unmarshal(data[1:], &args)
if err != nil {
log.Print("Malformed remote command")
return
} }
}()
window := struct {
go func() { row uint16
defer func() { exit <- true }() col uint16
x uint16
for { y uint16
_, data, err := conn.ReadMessage() }{
if err != nil { uint16(args.Rows),
return uint16(args.Columns),
} 0,
0,
switch data[0] {
case Input:
if !app.options.PermitWrite {
break
}
_, err := fio.Write(data[1:])
if err != nil {
return
}
case ResizeTerminal:
var args argResizeTerminal
err = json.Unmarshal(data[1:], &args)
if err != nil {
log.Print("Malformed remote command")
return
}
window := struct {
row uint16
col uint16
x uint16
y uint16
}{
uint16(args.Rows),
uint16(args.Columns),
0,
0,
}
syscall.Syscall(
syscall.SYS_IOCTL,
fio.Fd(),
syscall.TIOCSWINSZ,
uintptr(unsafe.Pointer(&window)),
)
default:
log.Print("Unknown message type")
return
}
} }
}() syscall.Syscall(
syscall.SYS_IOCTL,
go func() { context.pty.Fd(),
<-exit syscall.TIOCSWINSZ,
cmd.Wait() uintptr(unsafe.Pointer(&window)),
conn.Close() )
log.Printf("Connection closed: %s", r.RemoteAddr)
}() default:
log.Print("Unknown message type")
return
}
}
}
func generateRandomString(length int) string {
const base = 36
size := big.NewInt(base)
n := make([]byte, length)
for i, _ := range n {
c, _ := rand.Int(rand.Reader, size)
n[i] = strconv.FormatInt(c.Int64(), base)[0]
} }
return string(n)
} }
const ( const (
...@@ -220,13 +251,9 @@ type argResizeTerminal struct { ...@@ -220,13 +251,9 @@ type argResizeTerminal struct {
Rows float64 Rows float64
} }
func generateRandomString(length int) string { type clientContext struct {
const base = 36 request *http.Request
size := big.NewInt(base) connection *websocket.Conn
n := make([]byte, length) command *exec.Cmd
for i, _ := range n { pty *os.File
c, _ := rand.Int(rand.Reader, size)
n[i] = strconv.FormatInt(c.Int64(), base)[0]
}
return string(n)
} }
Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!