Implement websocket client connection authorization

Signed-off-by: billfort <[email protected]>

.gitignore

@@ -11,13 +11,15 @@ cscope*
 *.suo
 .DS_Store
 *~
+*.bak
 
 /build/
 /certs/
 /ChainDB/
 /ChainDB.config
 /Log/
 /vendor/
+/test/
 
 /nknd
 /nknc

api/common/errcode/errcode.go

@@ -9,7 +9,7 @@ const (
 	ILLEGAL_DATAFORMAT       ErrCode = 41003
 	INVALID_METHOD           ErrCode = 42001
 	INVALID_PARAMS           ErrCode = 42002
-	INVALID_TOKEN            ErrCode = 42003
+	INVALID_SIGNATURE        ErrCode = 42003
 	INVALID_JSON             ErrCode = 42004
 	INVALID_TRANSACTION      ErrCode = 43001
 	INVALID_ASSET            ErrCode = 43002
@@ -58,7 +58,7 @@ var ErrMessage = map[ErrCode]string{
 	ILLEGAL_DATAFORMAT:      "ILLEGAL DATAFORMAT",
 	INVALID_METHOD:          "INVALID METHOD",
 	INVALID_PARAMS:          "INVALID PARAMS",
-	INVALID_TOKEN:           "VERIFY TOKEN ERROR",
+	INVALID_SIGNATURE:       "VERIFY SIGNATURE ERROR",
 	INVALID_JSON:            "INVALID JSON",
 	INVALID_TRANSACTION:     "INVALID TRANSACTION",
 	INVALID_ASSET:           "INVALID ASSET",

api/websocket/server/server.go

@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"compress/zlib"
 	"context"
+	"crypto/rand"
+	"crypto/sha256"
 	"crypto/tls"
 	"encoding/hex"
 	"encoding/json"
@@ -192,6 +194,50 @@ func (ws *WsServer) registryMethod() {
 			return api.RespPacking(api.NodeInfo(wsAddr, rpcAddr, pubkey, id), errcode.WRONG_NODE)
 		}
 
+		// client auth
+		signature, okSig := cmd["Signature"]
+		clientSalt, okSalt := cmd["ClientSalt"]
+
+		if okSig && okSalt { // if client send ClientSalt and Signature, then check it
+			strSignature, typeOk := signature.(string) // interface type assertion
+			if !typeOk {
+				return api.RespPacking(err.Error(), errcode.INVALID_PARAMS)
+			}
+			byteSignature, err := hex.DecodeString(strSignature)
+			if err != nil {
+				return api.RespPacking(err.Error(), errcode.ILLEGAL_DATAFORMAT)
+			}
+
+			strClientSalt, typeOk := clientSalt.(string) // interface type assertion
+			if !typeOk {
+				return api.RespPacking(err.Error(), errcode.INVALID_PARAMS)
+			}
+			byteClientSalt, err := hex.DecodeString(strClientSalt)
+			if err != nil {
+				return api.RespPacking(err.Error(), errcode.ILLEGAL_DATAFORMAT)
+			}
+
+			sess := cmd["session"].(*session.Session)
+			challenge := sess.Challenge[:]
+			challenge = append(challenge, byteClientSalt...)
+			hash := sha256.Sum256(challenge)
+
+			err = crypto.Verify(pubKey, hash[:], byteSignature)
+			if err != nil { // fail verify challenge signature
+				go func() {
+					log.Warning("Client signature is not right, close its conneciton now")
+					time.Sleep(3 * time.Second)       // sleep several second, let response reach client
+					ws.SessionList.CloseSession(sess) // close this session
+				}()
+
+				return api.RespPacking(nil, errcode.INVALID_SIGNATURE)
+			} else {
+				log.Infof("client auth pass")
+			}
+		} else {
+			log.Infof("client doesn't send signature, it should be old version sdk")
+		}
+
 		newSessionID := hex.EncodeToString(clientID)
 		session, err := ws.SessionList.ChangeSessionToClient(cmd["Userid"].(string), newSessionID)
 		if err != nil {
@@ -287,6 +333,13 @@ func (ws *WsServer) websocketHandler(w http.ResponseWriter, r *http.Request) {
 		return nil
 	})
 
+	// client auth
+	err = ws.sendClientAuthChallenge(sess)
+	if err != nil {
+		log.Error("send client auth challenge: ", err)
+		return
+	}
+
 	done := make(chan struct{})
 	defer close(done)
 	go func() {
@@ -422,6 +475,7 @@ func (ws *WsServer) OnDataHandle(curSession *session.Session, messageType int, b
 	}
 	req["Userid"] = curSession.GetSessionId()
 	req["IsTls"] = r.TLS != nil
+	req["session"] = curSession
 	ret := action.handler(ws, req, r.Context())
 	resp := api.ResponsePack(ret["error"].(errcode.ErrCode))
 	resp["Action"] = actionName
@@ -452,14 +506,15 @@ func (ws *WsServer) deleteTxHashs(sSessionId string) {
 	}
 }
 
-func (ws *WsServer) respondToSession(session *session.Session, resp map[string]interface{}) {
+func (ws *WsServer) respondToSession(session *session.Session, resp map[string]interface{}) error {
 	resp["Desc"] = errcode.ErrMessage[resp["Error"].(errcode.ErrCode)]
 	data, err := json.Marshal(resp)
 	if err != nil {
 		log.Error("Websocket response:", err)
-		return
+		return err
 	}
-	session.SendText(data)
+	err = session.SendText(data)
+	return err
 }
 
 func (ws *WsServer) respondToId(sSessionId string, resp map[string]interface{}) {
@@ -570,3 +625,17 @@ func (ws *WsServer) sendInboundRelayMessageToClient(v interface{}) {
 		log.Error("Decode relay message failed")
 	}
 }
+
+// client auth, generate challenge
+func (ws *WsServer) sendClientAuthChallenge(sess *session.Session) error {
+	resp := api.ResponsePack(errcode.SUCCESS)
+	resp["Action"] = "authChallenge"
+
+	challenge := make([]byte, 32)
+	rand.Reader.Read(challenge)
+	resp["Challenge"] = hex.EncodeToString(challenge)
+	sess.Challenge = challenge // save this challenge for verifying later.
+
+	err := ws.respondToSession(sess, resp)
+	return err
+}

api/websocket/session/session.go

@@ -24,6 +24,8 @@ type Session struct {
 
 	wsLock sync.Mutex
 	ws     *websocket.Conn
+
+	Challenge []byte // client auth, authorization challenge
 }
 
 func (s *Session) GetSessionId() string {