gotty/server/handlers.go

240 lines
5.9 KiB
Go
Raw Permalink Normal View History

2017-02-25 22:37:07 +00:00
package server
import (
"bytes"
"context"
"encoding/json"
"fmt"
"log"
"net/http"
"net/url"
"sync/atomic"
"github.com/gorilla/websocket"
"github.com/pkg/errors"
"github.com/yudai/gotty/webtty"
)
2017-08-13 04:40:00 +00:00
func (server *Server) generateHandleWS(ctx context.Context, cancel context.CancelFunc, counter *counter) http.HandlerFunc {
once := new(int64)
2017-08-13 04:40:00 +00:00
go func() {
select {
case <-counter.timer().C:
cancel()
case <-ctx.Done():
}
}()
2017-02-25 22:37:07 +00:00
return func(w http.ResponseWriter, r *http.Request) {
if server.options.Once {
success := atomic.CompareAndSwapInt64(once, 0, 1)
if !success {
2017-02-25 22:37:07 +00:00
http.Error(w, "Server is shutting down", http.StatusServiceUnavailable)
return
}
}
2017-08-13 04:40:00 +00:00
num := counter.add(1)
2017-02-25 22:37:07 +00:00
closeReason := "unknown reason"
defer func() {
2017-08-13 04:40:00 +00:00
num := counter.done()
2017-02-25 22:37:07 +00:00
log.Printf(
"Connection closed by %s: %s, connections: %d/%d",
closeReason, r.RemoteAddr, num, server.options.MaxConnection,
2017-02-25 22:37:07 +00:00
)
if server.options.Once {
cancel()
}
}()
if int64(server.options.MaxConnection) != 0 {
2017-08-13 04:40:00 +00:00
if num > server.options.MaxConnection {
2017-02-25 22:37:07 +00:00
closeReason = "exceeding max number of connections"
return
}
}
log.Printf("New client connected: %s, connections: %d/%d", r.RemoteAddr, num, server.options.MaxConnection)
2017-02-25 22:37:07 +00:00
if r.Method != "GET" {
http.Error(w, "Method not allowed", 405)
return
}
conn, err := server.upgrader.Upgrade(w, r, nil)
if err != nil {
2017-08-21 06:36:23 +00:00
closeReason = err.Error()
2017-02-25 22:37:07 +00:00
return
}
defer conn.Close()
2019-01-29 14:44:20 +00:00
err = server.processWSConn(ctx, conn, r.Header)
2017-02-25 22:37:07 +00:00
switch err {
case ctx.Err():
closeReason = "cancelation"
case webtty.ErrSlaveClosed:
closeReason = server.factory.Name()
case webtty.ErrMasterClosed:
closeReason = "client"
default:
closeReason = fmt.Sprintf("an error: %s", err)
}
}
}
2019-01-29 14:44:20 +00:00
func (server *Server) processWSConn(ctx context.Context, conn *websocket.Conn, header http.Header) error {
2017-02-25 22:37:07 +00:00
typ, initLine, err := conn.ReadMessage()
if err != nil {
return errors.Wrapf(err, "failed to authenticate websocket connection")
}
if typ != websocket.TextMessage {
return errors.New("failed to authenticate websocket connection: invalid message type")
}
var init InitMessage
err = json.Unmarshal(initLine, &init)
if err != nil {
return errors.Wrapf(err, "failed to authenticate websocket connection")
}
if init.AuthToken != server.options.Credential {
return errors.New("failed to authenticate websocket connection")
}
queryPath := "?"
if server.options.PermitArguments && init.Arguments != "" {
queryPath = init.Arguments
}
2019-01-29 14:44:20 +00:00
headerArguments := header.Get("Arguments")
if server.options.PermitArguments && headerArguments != "" {
queryPath = queryPath + "&" + headerArguments
}
2017-02-25 22:37:07 +00:00
query, err := url.Parse(queryPath)
if err != nil {
return errors.Wrapf(err, "failed to parse arguments")
}
params := query.Query()
var slave Slave
2019-01-29 14:44:20 +00:00
2017-02-25 22:37:07 +00:00
slave, err = server.factory.New(params)
if err != nil {
return errors.Wrapf(err, "failed to create backend")
}
defer slave.Close()
titleVars := server.titleVariables(
[]string{"server", "master", "slave"},
map[string]map[string]interface{}{
"server": server.options.TitleVariables,
"master": map[string]interface{}{
"remote_addr": conn.RemoteAddr(),
},
"slave": slave.WindowTitleVariables(),
},
)
titleBuf := new(bytes.Buffer)
err = server.titleTemplate.Execute(titleBuf, titleVars)
if err != nil {
return errors.Wrapf(err, "failed to fill window title template")
}
opts := []webtty.Option{
webtty.WithWindowTitle(titleBuf.Bytes()),
}
if server.options.PermitWrite {
opts = append(opts, webtty.WithPermitWrite())
}
if server.options.EnableReconnect {
opts = append(opts, webtty.WithReconnect(server.options.ReconnectTime))
}
2017-08-24 05:40:28 +00:00
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))
2017-02-25 22:37:07 +00:00
}
if server.options.Preferences != nil {
opts = append(opts, webtty.WithMasterPreferences(server.options.Preferences))
}
2017-08-24 05:40:28 +00:00
tty, err := webtty.New(&wsWrapper{conn}, slave, opts...)
2017-02-25 22:37:07 +00:00
if err != nil {
return errors.Wrapf(err, "failed to create webtty")
}
err = tty.Run(ctx)
return err
}
func (server *Server) handleIndex(w http.ResponseWriter, r *http.Request) {
titleVars := server.titleVariables(
[]string{"server", "master"},
map[string]map[string]interface{}{
"server": server.options.TitleVariables,
"master": map[string]interface{}{
"remote_addr": r.RemoteAddr,
},
},
)
titleBuf := new(bytes.Buffer)
err := server.titleTemplate.Execute(titleBuf, titleVars)
if err != nil {
http.Error(w, "Internal Server Error", 500)
return
}
indexVars := map[string]interface{}{
"title": titleBuf.String(),
}
indexBuf := new(bytes.Buffer)
err = server.indexTemplate.Execute(indexBuf, indexVars)
if err != nil {
http.Error(w, "Internal Server Error", 500)
return
}
w.Write(indexBuf.Bytes())
}
func (server *Server) handleAuthToken(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/javascript")
// @TODO hashing?
w.Write([]byte("var gotty_auth_token = '" + server.options.Credential + "';"))
}
func (server *Server) handleConfig(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/javascript")
w.Write([]byte("var gotty_term = '" + server.options.Term + "';"))
}
2017-02-25 22:37:07 +00:00
// titleVariables merges maps in a specified order.
// varUnits are name-keyed maps, whose names will be iterated using order.
func (server *Server) titleVariables(order []string, varUnits map[string]map[string]interface{}) map[string]interface{} {
titleVars := map[string]interface{}{}
for _, name := range order {
vars, ok := varUnits[name]
if !ok {
panic("title variable name error")
}
for key, val := range vars {
titleVars[key] = val
}
}
// safe net for conflicted keys
for _, name := range order {
titleVars[name] = varUnits[name]
}
return titleVars
}