改进版---增加自定义参数的k8s-web-terminal

感觉比之前的完善，

支持自定义端口，自定义kubeconfig文件路径，自定义权证api地址。

package main

import (
	"encoding/json"
	"flag"
	"fmt"

	"log"
	"net/http"
	"os"
	"path/filepath"
	"sync"

	"github.com/gorilla/websocket"
	"k8s.io/api/core/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/remotecommand"
)

var (
	clientset *kubernetes.Clientset
	kconfig   *rest.Config
	apiUrl    *string
)

// message from web socket client
type xtermMessage struct {
	MsgType string `json:"type"`  // 类型:resize客户端调整终端, input客户端输入
	Input   string `json:"input"` // msgtype=input情况下使用
	Rows    uint16 `json:"rows"`  // msgtype=resize情况下使用
	Cols    uint16 `json:"cols"`  // msgtype=resize情况下使用
}

type WSStreamHandler struct {
	conn        *websocket.Conn
	resizeEvent chan remotecommand.TerminalSize
	rbuf        []byte
	cond        *sync.Cond

	sync.Mutex
}

// Run start a loop to fetch from ws client and store the data in byte buffer
func (h *WSStreamHandler) Run(userToken string) {

	accessUrl := fmt.Sprintf("%s?token=%s", *apiUrl, userToken)
	respToken, err := http.Get(accessUrl)
	if respToken != nil {
		defer respToken.Body.Close()
	}

	if err != nil {
		log.Println(err)
		return
	}

	if respToken.StatusCode != 200 {
		h.Write([]byte("亲，你无权进入此容器。bye~bye~"))
		h.conn.Close()
		return
	}

	for {
		_, p, err := h.conn.ReadMessage()
		if err != nil {
			log.Println("ws ReadMessage err: ", err)
			// 新增如果前端关闭，后端任然继续读取数据将会报错
			// panic: repeated read on failed websocket connection。
			h.conn.Close()
			return
		}
		xmsg := xtermMessage{}
		if err := json.Unmarshal(p, &xmsg); err != nil {
			log.Println("json.Unmarshal err: ", err)
		}

		switch xmsg.MsgType {
		case "input":
			{
				h.Lock()
				// log.Printf("reading input: %s", string(xmsg.Input))
				h.rbuf = append(h.rbuf, xmsg.Input...)
				h.cond.Signal()
				h.Unlock()
			}
		case "resize":
			{
				ev := remotecommand.TerminalSize{
					Width:  xmsg.Cols,
					Height: xmsg.Rows}
				h.resizeEvent <- ev
			}
		default:
			log.Println("other xmsg.MsgType: not input or resize.")
		}
	}
}

func (h *WSStreamHandler) Read(b []byte) (size int, err error) {
	h.Lock()
	for len(h.rbuf) == 0 {
		h.cond.Wait()
	}
	size = copy(b, h.rbuf)
	h.rbuf = h.rbuf[size:]
	h.Unlock()
	return
}

func (h *WSStreamHandler) Write(b []byte) (size int, err error) {
	size = len(b)
	err = h.conn.WriteMessage(websocket.TextMessage, b)

	return
}

func (h *WSStreamHandler) Next() (size *remotecommand.TerminalSize) {
	ret := <-h.resizeEvent
	size = &ret

	return
}

func wsHandler(resp http.ResponseWriter, req *http.Request) {
	var (
		conn          *websocket.Conn
		sshReq        *rest.Request
		podName       string
		podNs         string
		containerName string
		executor      remotecommand.Executor
		handler       *WSStreamHandler
		err           error
	)

	// 解析GET参数
	if err = req.ParseForm(); err != nil {
		return
	}

	var userToken string = req.Form.Get("userToken")
	podNs = req.Form.Get("namespace")
	podName = req.Form.Get("pod")
	containerName = req.Form.Get("container")

	// 得到websocket长连接
	upgrader := websocket.Upgrader{}
	if conn, err = upgrader.Upgrade(resp, req, nil); err != nil {
		log.Fatalf("error creating ws conn: %v", err)
	}

	sshReq = clientset.CoreV1().RESTClient().Post().
		Resource("pods").
		Name(podName).
		Namespace(podNs).
		SubResource("exec").
		VersionedParams(&v1.PodExecOptions{
			Container: containerName,
			Command:   []string{"sh"},
			Stdin:     true,
			Stdout:    true,
			Stderr:    true,
			TTY:       true,
		}, scheme.ParameterCodec)
	// 创建到容器的连接
	if executor, err = remotecommand.NewSPDYExecutor(kconfig, "POST", sshReq.URL()); err != nil {
		log.Fatalf("error creating spdy executor: %v", err)
	}

	log.Println("connectingi to pod...")
	handler = &WSStreamHandler{
		conn:        conn,
		resizeEvent: make(chan remotecommand.TerminalSize)}
	handler.cond = sync.NewCond(handler)

	// run loop to fetch data from ws client
	go handler.Run(userToken)

	err = executor.Stream(remotecommand.StreamOptions{
		Stdin:             handler,
		Stdout:            handler,
		Stderr:            handler,
		TerminalSizeQueue: handler,
		Tty:               true,
	})
	if err != nil {
		log.Println("error executor: ", err)
		handler.Write([]byte("亲，你选择的容器不存在！bye~bye~"))
		handler.conn.Close()
		return
	}

	return
}

func main() {

	staticDir := "./public"
	fs := http.FileServer(http.Dir(staticDir))
	http.Handle("/", fs)

	var err error
	cfgpath := flag.String("kubecfg", filepath.Join(homeDir(), ".kube", "config"), "kubeconfig绝对路径")
	httpPort := flag.String("port", "8000", "监听端口号")
	apiUrl = flag.String("api", "http://127.0.0.1/api.html", "权限api地址")

	flag.Parse()

	log.Println("加载kubeconfig的路径:", *cfgpath)
	log.Println("端口号:", *httpPort)
	log.Println("权限认证api:", *apiUrl)

	if kconfig, err = clientcmd.BuildConfigFromFlags("", *cfgpath); err != nil {
		log.Fatalf("error creating k8s config: %v", err)
	}

	if clientset, err = kubernetes.NewForConfig(kconfig); err != nil {
		log.Fatalf("error creating clientset: %v", err)
	}

	http.HandleFunc("/terminal", wsHandler)

	log.Println("server running...")
	http.ListenAndServe(":"+*httpPort, nil)
	log.Println("server stopped...")
}

func homeDir() string {
	if h := os.Getenv("HOME"); h != "" {
		return h
	}

	return os.Getenv("USERPROFILE") // windows
}