123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206 |
- 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()
- }
- }
|