fix: incorrect disk info on unix-like os

XZB-1248 · Sep 17, 2022 · 9d9a992 · 9d9a992
1 parent 98b1beb
commit 9d9a992
Show file tree

Hide file tree

Showing 15 changed files with 369 additions and 302 deletions.
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -11,7 +11,7 @@ jobs:
 
     strategy:
       matrix:
-        go-version: [ 1.17 ]
+        go-version: [ 1.18.1 ]
 
     steps:
       - uses: actions/checkout@v3
@@ -57,7 +57,7 @@ jobs:
     strategy:
       matrix:
         node-version: [ 16.x ]
-        go-version: [ 1.17 ]
+        go-version: [ 1.18.1 ]
 
     steps:
     - uses: actions/checkout@v3

diff --git a/client/common/common.go b/client/common/common.go
@@ -14,55 +14,76 @@ import (
 
 type Conn struct {
 	*ws.Conn
-	Secret []byte
+	secret    []byte
+	secretHex string
 }
 
 var WSConn *Conn
-var WSLock = sync.Mutex{}
-var HTTP = req.C().SetUserAgent(`SPARK COMMIT: ` + config.COMMIT)
+var Mutex = &sync.Mutex{}
+var HTTP = CreateClient()
 
 const MaxMessageSize = 32768 + 1024
 
-func SendData(data []byte, wsConn *Conn) error {
-	WSLock.Lock()
-	defer WSLock.Unlock()
+func CreateConn(wsConn *ws.Conn, secret []byte) *Conn {
+	return &Conn{
+		Conn:      wsConn,
+		secret:    secret,
+		secretHex: hex.EncodeToString(secret),
+	}
+}
+
+func CreateClient() *req.Client {
+	return req.C().SetUserAgent(`SPARK COMMIT: ` + config.COMMIT)
+}
+
+func (wsConn *Conn) SendData(data []byte) error {
+	Mutex.Lock()
+	defer Mutex.Unlock()
 	if WSConn == nil {
 		return errors.New(`${i18n|wsClosed}`)
 	}
-	wsConn.SetWriteDeadline(time.Now().Add(5 * time.Second))
+	wsConn.SetWriteDeadline(Now.Add(5 * time.Second))
 	defer wsConn.SetWriteDeadline(time.Time{})
 	return wsConn.WriteMessage(ws.BinaryMessage, data)
 }
 
-func SendPack(pack interface{}, wsConn *Conn) error {
-	WSLock.Lock()
-	defer WSLock.Unlock()
+func (wsConn *Conn) SendPack(pack interface{}) error {
+	Mutex.Lock()
+	defer Mutex.Unlock()
 	data, err := utils.JSON.Marshal(pack)
 	if err != nil {
 		return err
 	}
-	data, err = utils.Encrypt(data, wsConn.Secret)
+	data, err = utils.Encrypt(data, wsConn.secret)
 	if err != nil {
 		return err
 	}
 	if len(data) > MaxMessageSize {
 		_, err = HTTP.R().
 			SetBody(data).
-			SetHeader(`Secret`, hex.EncodeToString(wsConn.Secret)).
+			SetHeader(`Secret`, wsConn.secretHex).
 			Send(`POST`, config.GetBaseURL(false)+`/ws`)
 		return err
 	}
 	if WSConn == nil {
 		return errors.New(`${i18n|wsClosed}`)
 	}
-	wsConn.SetWriteDeadline(time.Now().Add(5 * time.Second))
+	wsConn.SetWriteDeadline(Now.Add(5 * time.Second))
 	defer wsConn.SetWriteDeadline(time.Time{})
 	return wsConn.WriteMessage(ws.BinaryMessage, data)
 }
 
-func SendCb(pack, prev modules.Packet, wsConn *Conn) error {
+func (wsConn *Conn) SendCallback(pack, prev modules.Packet) error {
 	if len(prev.Event) > 0 {
 		pack.Event = prev.Event
 	}
-	return SendPack(pack, wsConn)
+	return wsConn.SendPack(pack)
+}
+
+func (wsConn *Conn) GetSecret() []byte {
+	return wsConn.secret
+}
+
+func (wsConn *Conn) GetSecretHex() string {
+	return wsConn.secretHex
 }
diff --git a/client/common/time.go b/client/common/time.go
@@ -0,0 +1,16 @@
+package common
+
+import "time"
+
+var Now time.Time = time.Now()
+var Unix int64 = Now.Unix()
+
+// To prevent call time.Now().Unix() too often.
+func init() {
+	go func() {
+		for now := range time.NewTicker(time.Second).C {
+			Now = now
+			Unix = now.Unix()
+		}
+	}()
+}
diff --git a/client/core/core.go b/client/core/core.go
@@ -26,45 +26,18 @@ var stop bool
 var (
 	errNoSecretHeader = errors.New(`can not find secret header`)
 )
-var handlers = map[string]func(pack modules.Packet, wsConn *common.Conn){
-	`ping`:           ping,
-	`offline`:        offline,
-	`lock`:           lock,
-	`logoff`:         logoff,
-	`hibernate`:      hibernate,
-	`suspend`:        suspend,
-	`restart`:        restart,
-	`shutdown`:       shutdown,
-	`screenshot`:     screenshot,
-	`initTerminal`:   initTerminal,
-	`inputTerminal`:  inputTerminal,
-	`resizeTerminal`: resizeTerminal,
-	`pingTerminal`:   pingTerminal,
-	`killTerminal`:   killTerminal,
-	`listFiles`:      listFiles,
-	`fetchFile`:      fetchFile,
-	`removeFiles`:    removeFiles,
-	`uploadFiles`:    uploadFiles,
-	`uploadTextFile`: uploadTextFile,
-	`listProcesses`:  listProcesses,
-	`killProcess`:    killProcess,
-	`initDesktop`:    initDesktop,
-	`pingDesktop`:    pingDesktop,
-	`killDesktop`:    killDesktop,
-	`getDesktop`:     getDesktop,
-}
 
 func Start() {
 	for !stop {
 		var err error
 		if common.WSConn != nil {
-			common.WSLock.Lock()
+			common.Mutex.Lock()
 			common.WSConn.Close()
-			common.WSLock.Unlock()
+			common.Mutex.Unlock()
 		}
-		common.WSLock.Lock()
+		common.Mutex.Lock()
 		common.WSConn, err = connectWS()
-		common.WSLock.Unlock()
+		common.Mutex.Unlock()
 		if err != nil && !stop {
 			golog.Error(`Connection error: `, err)
 			<-time.After(3 * time.Second)
@@ -105,7 +78,7 @@ func connectWS() (*common.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &common.Conn{Conn: wsConn, Secret: secret}, nil
+	return common.CreateConn(wsConn, secret), nil
 }
 
 func reportWS(wsConn *common.Conn) error {
@@ -114,18 +87,18 @@ func reportWS(wsConn *common.Conn) error {
 		return err
 	}
 	pack := modules.CommonPack{Act: `report`, Data: *device}
-	err = common.SendPack(pack, wsConn)
+	err = wsConn.SendPack(pack)
 	common.WSConn.SetWriteDeadline(time.Time{})
 	if err != nil {
 		return err
 	}
-	common.WSConn.SetReadDeadline(time.Now().Add(5 * time.Second))
+	common.WSConn.SetReadDeadline(common.Now.Add(5 * time.Second))
 	_, data, err := common.WSConn.ReadMessage()
 	common.WSConn.SetReadDeadline(time.Time{})
 	if err != nil {
 		return err
 	}
-	data, err = utils.Decrypt(data, common.WSConn.Secret)
+	data, err = utils.Decrypt(data, common.WSConn.GetSecret())
 	if err != nil {
 		return err
 	}
@@ -148,7 +121,7 @@ func checkUpdate(wsConn *common.Conn) error {
 		SetQueryParam(`os`, runtime.GOOS).
 		SetQueryParam(`arch`, runtime.GOARCH).
 		SetQueryParam(`commit`, config.COMMIT).
-		SetHeader(`Secret`, hex.EncodeToString(wsConn.Secret)).
+		SetHeader(`Secret`, wsConn.GetSecretHex()).
 		Send(`POST`, config.GetBaseURL(false)+`/api/client/update`)
 	if err != nil {
 		return err
@@ -189,7 +162,7 @@ func handleWS(wsConn *common.Conn) error {
 			golog.Error(err)
 			return nil
 		}
-		data, err = utils.Decrypt(data, wsConn.Secret)
+		data, err = utils.Decrypt(data, wsConn.GetSecret())
 		if err != nil {
 			golog.Error(err)
 			errCount++
@@ -220,7 +193,7 @@ func handleWS(wsConn *common.Conn) error {
 
 func handleAct(pack modules.Packet, wsConn *common.Conn) {
 	if act, ok := handlers[pack.Act]; !ok {
-		common.SendCb(modules.Packet{Code: 1, Msg: `${i18n|actionNotImplemented}`}, pack, wsConn)
+		wsConn.SendCallback(modules.Packet{Code: 1, Msg: `${i18n|actionNotImplemented}`}, pack)
 	} else {
 		defer func() {
 			if r := recover(); r != nil {

diff --git a/client/core/device.go b/client/core/device.go
@@ -147,17 +147,20 @@ func GetRAMInfo() (modules.IO, error) {
 }
 
 func GetDiskInfo() (modules.IO, error) {
+	devices := map[string]struct{}{}
 	result := modules.IO{}
-	disk.IOCounters()
-	disks, err := disk.Partitions(true)
+	disks, err := disk.Partitions(false)
 	if err != nil {
 		return result, nil
 	}
 	for i := 0; i < len(disks); i++ {
-		stat, err := disk.Usage(disks[i].Mountpoint)
-		if err == nil {
-			result.Total += stat.Total
-			result.Used += stat.Used
+		if _, ok := devices[disks[i].Device]; !ok {
+			devices[disks[i].Device] = struct{}{}
+			stat, err := disk.Usage(disks[i].Mountpoint)
+			if err == nil {
+				result.Total += stat.Total
+				result.Used += stat.Used
+			}
 		}
 	}
 	result.Usage = float64(result.Used) / float64(result.Total) * 100