[*]修复会导致程序崩溃的bug

wwj-saber · Nov 27, 2021 · ad6cf27 · ad6cf27
1 parent 03c9edd
commit ad6cf27
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/lib/grdp/protocol/pdu/pdu.go b/lib/grdp/protocol/pdu/pdu.go
@@ -3,6 +3,8 @@ package pdu
 import (
 	"bytes"
 	"encoding/hex"
+	"errors"
+	"fmt"
 
 	"kscan/lib/grdp/core"
 	"kscan/lib/grdp/emission"
@@ -138,6 +140,14 @@ func (c *Client) connect(data *gcc.ClientCoreData, userId uint16, channelId uint
 }
 
 func (c *Client) recvDemandActivePDU(s []byte) {
+	defer func() {
+		if e := recover(); e != nil {
+			err := errors.New(fmt.Sprint("recv demand active pdu error: ", e))
+			glog.Debug(err, e)
+			return
+		}
+	}()
+
 	glog.Debug("PDU recvDemandActivePDU", hex.EncodeToString(s))
 	r := bytes.NewReader(s)
 	pdu, err := readPDU(r)

diff --git a/lib/grdp/protocol/sec/sec.go b/lib/grdp/protocol/sec/sec.go
@@ -7,6 +7,7 @@ import (
 	"crypto/sha1"
 	"encoding/hex"
 	"errors"
+	"fmt"
 	"io"
 	"math/big"
 	"unicode/utf16"
@@ -357,6 +358,14 @@ func (s *SEC) readEncryptedPayload(data []byte, checkSum bool) []byte {
 
 }
 func (s *SEC) writeEncryptedPayload(data []byte, checkSum bool) []byte {
+	defer func() {
+		if e := recover(); e != nil {
+			err := errors.New(fmt.Sprint("write encrypted payload error: ", e))
+			glog.Debug(err, e)
+			return
+		}
+	}()
+
 	if s.nbEncryptedPacket == 4096 {
 
 	}
@@ -624,6 +633,15 @@ func (e *ClientSecurityExchangePDU) serialize() []byte {
 	return buff.Bytes()
 }
 func (c *Client) sendClientRandom() {
+	defer func() {
+		if e := recover(); e != nil {
+			err := errors.New(fmt.Sprint("send client random error: ", e))
+			glog.Debug(err, e)
+			_ = c.Close()
+			return
+		}
+	}()
+
 	glog.Info("send Client Random")
 
 	clientRandom := core.Random(32)