go-hhc-cli/hhc_telws/telws_handler.go

package hhc_telws

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"git.swzry.com/zry/gorillaws2netconn"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/pascaldekloe/jwt"
	"github.com/tjfoc/gmsm/sm2"
	"io"
	"net/http"
	"time"
)

const JWT_KEY_LENGTH = 256

type CliBackendHandler struct {
	jwtKey             []byte
	errLog             func(err error)
	authHandlers       map[string]TelwsAuthHandler
	authMethodList     []string
	gingrp             *gin.RouterGroup
	jwtttl             time.Duration
	sm2key             *sm2.PrivateKey
	sm2pubkey          string
	upgrader           *websocket.Upgrader
	sessHandlerFactory TelwsSessionHandlerFactory
}

func NewCliBackendHandler(gingrp *gin.RouterGroup, tsfactory TelwsSessionHandlerFactory) *CliBackendHandler {
	cbh := &CliBackendHandler{
		gingrp:             gingrp,
		authMethodList:     make([]string, 0),
		authHandlers:       make(map[string]TelwsAuthHandler),
		sessHandlerFactory: tsfactory,
	}
	cbh.upgrader = &websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
	cbh.gingrp.GET("/", cbh.handleGetTelwsInfo)
	cbh.gingrp.POST("/login.satori", cbh.handleLogin)
	cbh.gingrp.GET("/ws.satori", cbh.handleWebSocket)
	return cbh
}

func (cbh *CliBackendHandler) InitSecure(jwtttl time.Duration) error {
	cbh.jwtttl = jwtttl
	jbuf := make([]byte, JWT_KEY_LENGTH)
	_, err := io.ReadFull(rand.Reader, jbuf)
	if err != nil {
		return err
	}
	cbh.jwtKey = jbuf
	pk, err := sm2.GenerateKey()
	if err != nil {
		return err
	}
	cbh.sm2key = pk
	xenc := hex.EncodeToString(pk.X.Bytes())
	yenc := hex.EncodeToString(pk.Y.Bytes())
	cbh.sm2pubkey = SM2PubKeyCombine(xenc, yenc)
	return nil
}

func (cbh *CliBackendHandler) AddAuthHandler(ah TelwsAuthHandler) {
	ahname := ah.Register(cbh.jwtKey, cbh.jwtttl)
	cbh.authHandlers[ahname] = ah
	cbh.authMethodList = append(cbh.authMethodList, ahname)
}

func (cbh *CliBackendHandler) SetErrorLogFunction(fe func(err error)) {
	cbh.errLog = fe
}

func (cbh *CliBackendHandler) handleGetTelwsInfo(ctx *gin.Context) {
	ctx.JSON(200, gin.H{
		"protocol":     "telws",
		"protocol_ver": "1.0",
		"auth_methods": cbh.authMethodList,
		"encrypt": map[string]string{
			"method":     "sm2",
			"pubkey":     cbh.sm2pubkey,
			"cipherMode": "C1C3C2",
			"header":     "04",
		},
	})
}

type postJson_Login struct {
	AuthData string `json:"authData"`
}

func (cbh *CliBackendHandler) handleLogin(ctx *gin.Context) {
	var pjl postJson_Login
	err := ctx.BindJSON(&pjl)
	if err != nil {
		ctx.JSON(200, gin.H{
			"suc":   false,
			"etype": "input_json_error",
		})
		return
	}
	data := pjl.AuthData
	if data == "" {
		ctx.JSON(200, gin.H{
			"suc":   false,
			"etype": "no_cipher_text",
		})
		return
	}
	hdec, err := hex.DecodeString(data)
	if err != nil {
		ctx.JSON(200, gin.H{
			"suc":   false,
			"etype": "hex_decode",
		})
		return
	}
	pt, err := sm2.Decrypt(cbh.sm2key, hdec)
	if err != nil {
		ctx.JSON(200, gin.H{
			"suc":   false,
			"etype": "sm2_decrypt",
		})
		return
	}
	var jdata map[string]string
	err = json.Unmarshal(pt, &jdata)
	if err != nil {
		ctx.JSON(200, gin.H{
			"suc":   false,
			"etype": "decrypted_json_error",
		})
		return
	}
	met, ok := jdata["method"]
	if !ok {
		ctx.JSON(200, gin.H{
			"suc":   false,
			"etype": "no_auth_method_specified",
		})
		return
	}
	ah, ok := cbh.authHandlers[met]
	if !ok {
		ctx.JSON(200, gin.H{
			"suc":   false,
			"etype": "unsupported_auth_method",
		})
		return
	}
	suc, jwtstr, emsg := ah.Login(jdata)
	if !suc {
		ctx.JSON(200, gin.H{
			"suc":            false,
			"etype":          "auth_failed",
			"auth_error_msg": emsg,
		})
		return
	}
	ctx.JSON(200, gin.H{
		"suc": true,
		"jwt": jwtstr,
	})
}

func (cbh *CliBackendHandler) handleWebSocket(ctx *gin.Context) {
	jwtstr := ctx.Request.FormValue("token")
	ws, err := cbh.upgrader.Upgrade(ctx.Writer, ctx.Request, nil)
	if err != nil {
		fmt.Println("WS Error:", err.Error())
		return
	}
	defer ws.Close()
	nconn := &gorillaws2netconn.NetConn4Gorilla{WS: ws, WriteMessageType: websocket.BinaryMessage}
	vclm, err := jwt.HMACCheck([]byte(jwtstr), cbh.jwtKey)
	if err != nil {
		if cbh.errLog != nil {
			cbh.errLog(err)
		}
		_, _ = fmt.Fprintf(nconn, "Websocket auth failed: %s\r\n", err.Error())
		ws.Close()
		return
	}
	tscc := WrapNewTelwsSessionContext()
	tscc.SetJwtClaims(vclm)
	tscc.SetConn(nconn)
	sh := cbh.sessHandlerFactory.GetSessionHandler()
	err = sh.HandleConnection(tscc.GetUserContext())
	if err != nil {
		if cbh.errLog != nil {
			cbh.errLog(err)
		}
		_, _ = fmt.Fprintf(nconn, "Websocket connection end by error: %s\r\n", err.Error())
		ws.Close()
		return
	} else {
		_, _ = fmt.Fprintf(nconn, "Websocket connection closed.\r\n")
		ws.Close()
	}
}