通过websocket 实现与容器的交互

package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"

	app "./console"
)

func main() {
	var address *string = flag.String("Address", "", "server listen address")
	var port *string = flag.String("Port", "10001", "server listen port")
	var sessionKey *string = flag.String("SessionKey", "_auth_user_id", "user serssion key")
	flag.Parse()
	options := app.DefaultOptions
	options.Address = *address
	options.Port = *port
	options.SessionKey = *sessionKey

	app, err := app.New(nil, &options)
	registerSignals(app)
	err = app.Run()
	if err != nil {
		exit(err, 4)
	}
}

func exit(err error, code int) {
	if err != nil {
		fmt.Println(err)
	}
	os.Exit(code)
}

func registerSignals(app *app.App) {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(
		sigChan,
		syscall.SIGINT,
		syscall.SIGTERM,
	)
	go func() {
		for {
			s := <-sigChan
			switch s {
			case syscall.SIGINT, syscall.SIGTERM:
				if app.Exit() {
					fmt.Println("Send ^C to force exit.")
				} else {
					os.Exit(5)
				}
			}
		}
	}()
}

package app

import (
	"crypto/md5"
	"encoding/hex"
	"encoding/json"
	"log"
	"net/http"
	"os/exec"
	"sync"
	"text/template"

	"github.com/braintree/manners"
	"github.com/gorilla/websocket"
	"github.com/kr/pty"
	"github.com/yudai/umutex"
)

type App struct {
	command []string
	options *Options

	upgrader *websocket.Upgrader

	titleTemplate *template.Template

	onceMutex *umutex.UnblockingMutex
}

type Options struct {
	Address         string                 `hcl:"address"`
	Port            string                 `hcl:"port"`
	PermitWrite     bool                   `hcl:"permit_write"`
	IndexFile       string                 `hcl:"index_file"`
	TitleFormat     string                 `hcl:"title_format"`
	EnableReconnect bool                   `hcl:"enable_reconnect"`
	ReconnectTime   int                    `hcl:"reconnect_time"`
	PermitArguments bool                   `hcl:"permit_arguments"`
	CloseSignal     int                    `hcl:"close_signal"`
	Preferences     HtermPrefernces        `hcl:"preferences"`
	RawPreferences  map[string]interface{} `hcl:"preferences"`
	SessionKey      string                 `hcl:"session_key"`
}

var Version = "0.0.1"

var DefaultOptions = Options{
	Address:         "",
	Port:            "10001",
	PermitWrite:     true,
	IndexFile:       "",
	TitleFormat:     "DTTY Command",
	EnableReconnect: true,
	ReconnectTime:   10,
	CloseSignal:     1, // syscall.SIGHUP
	Preferences:     HtermPrefernces{},
	SessionKey:      "_auth_user_id",
}

type InitMessage struct {
	T_id string `json:t_id`
	S_id string `json:s_id`
	C_id string `json:c_id`
	Md5  string `json:md5`
}

func checkSameOrigin(r *http.Request) bool {
	return true
}

func New(command []string, options *Options) (*App, error) {
	titleTemplate, _ := template.New("title").Parse(options.TitleFormat)
	return &App{
		command: command,
		options: options,
		upgrader: &websocket.Upgrader{
			ReadBufferSize:  1024,
			WriteBufferSize: 1024,
			CheckOrigin:     checkSameOrigin,
		},
		titleTemplate: titleTemplate,
		onceMutex:     umutex.New(),
	}, nil
}

func (app *App) Run() error {
	wsHandler := http.HandlerFunc(app.handleWS)
	siteHandler := wrapHeaders(http.Handler(http.NewServeMux()))
	wsMux := http.NewServeMux()
	wsMux.Handle("/", siteHandler)
	wsMux.Handle("/ws", wsHandler)
	siteHandler = (http.Handler(wsMux))

	siteHandler = wrapLogger(siteHandler)

	s := manners.NewWithServer(&http.Server{
		Addr:    app.options.Address + ":" + app.options.Port,
		Handler: siteHandler,
	})
	log.Printf("Start server on port " + app.options.Port)
	log.Fatal(s.ListenAndServe())
	log.Printf("Exiting...")

	return nil
}

func (app *App) handleWS(w http.ResponseWriter, r *http.Request) {
	log.Printf("New client connected: %s", r.RemoteAddr)

	if r.Method != "GET" {
		http.Error(w, "Method not allowed", 405)
		return
	}
	conn, err := app.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("Failed to upgrade connection: " + err.Error())
		return
	}

	//defer conn.Close()

	_, stream, err := conn.ReadMessage()
	if err != nil {
		log.Print("Failed to authenticate websocket connection " + err.Error())
		conn.Close()
		return
	}

	message := string(stream)
	log.Print("message=", message)

	var init InitMessage

	err = json.Unmarshal(stream, &init)

	//todo auth
	if init.C_id == "" {
		log.Print("Parameter is error:" + init.C_id)
		conn.WriteMessage(websocket.TextMessage, []byte("Parameter is error !"))
		conn.Close()
		return
	}
	key := init.T_id + "_" + init.S_id + "_" + init.C_id + "_yunpingtai"
	md5 := md5Func(key)
	log.Print(key + ":" + md5)
	if md5 != init.Md5 {
		log.Print("Auth is not allowed !")
		conn.WriteMessage(websocket.TextMessage, []byte("Auth is not allowed!"))
		conn.Close()
		return
	}

	cmd := exec.Command("docker", "exec", "-ti", init.C_id, "/bin/bash")
	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{
		app:        app,
		request:    r,
		connection: conn,
		command:    cmd,
		pty:        ptyIo,
		writeMutex: &sync.Mutex{},
	}

	context.goHandleClient()
}

func (app *App) Exit() (firstCall bool) {
	manners.Close()
	return true
}

func wrapLogger(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		rw := &responseWrapper{w, 200}
		handler.ServeHTTP(rw, r)
		log.Printf("%s %d %s %s", r.RemoteAddr, rw.status, r.Method, r.URL.Path)
	})
}

func wrapHeaders(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Server", "GoTTY/"+Version)
		handler.ServeHTTP(w, r)
	})
}

func md5Func(str string) string {
	h := md5.New()
	h.Write([]byte(str))
	cipherStr := h.Sum(nil)
	return hex.EncodeToString(cipherStr)
}