NET-1778: scale test changes (#920)

* fix UpdateHost call issue

* add unzip and change encrypt for peerUpdate message

* add old decrypt func back for compatibility

* change broker keepalive

* fix mq server restart connection issue

Author: yabinma

auth/auth.go

@@ -37,6 +37,10 @@ func isTokenExpired(tokenString string) bool {
 	return true
 }
 
+func CleanJwtToken() {
+	jwtToken = ""
+}
+
 // Authenticate authenticates with netmaker api to permit subsequent interactions with the api
 func Authenticate(server *config.Server, host *config.Config) (string, error) {
 	if jwtToken != "" && !isTokenExpired(jwtToken) {

functions/daemon.go

@@ -1,9 +1,14 @@
 package functions
 
 import (
+	"bytes"
+	"compress/gzip"
 	"context"
+	"crypto/aes"
+	"crypto/cipher"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"os"
 	"os/signal"
@@ -341,7 +346,7 @@ func setupMQTT(server *config.Server) error {
 	opts.SetAutoReconnect(true)
 	opts.SetConnectRetry(true)
 	opts.SetConnectRetryInterval(time.Second << 2)
-	opts.SetKeepAlive(time.Second * 10)
+	opts.SetKeepAlive(time.Second * 15)
 	opts.SetWriteTimeout(time.Minute)
 	opts.SetCleanSession(true)
 	opts.SetOnConnectHandler(func(client mqtt.Client) {
@@ -485,7 +490,26 @@ func setDNSSubscriptions(client mqtt.Client, node *config.Node) {
 	slog.Info("subscribed to DNS sync for node", "node", node.ID, "network", node.Network)
 }
 
-// should only ever use node client configs
+func unzipPayload(data []byte) (resData []byte, err error) {
+	b := bytes.NewBuffer(data)
+
+	var r io.Reader
+	r, err = gzip.NewReader(b)
+	if err != nil {
+		return
+	}
+
+	var resB bytes.Buffer
+	_, err = resB.ReadFrom(r)
+	if err != nil {
+		return
+	}
+
+	resData = resB.Bytes()
+
+	return
+}
+
 func decryptMsg(serverName string, msg []byte) ([]byte, error) {
 	if len(msg) <= 24 { // make sure message is of appropriate length
 		return nil, fmt.Errorf("received invalid message from broker %v", msg)
@@ -508,6 +532,32 @@ func decryptMsg(serverName string, msg []byte) ([]byte, error) {
 	return DeChunk(msg, serverPubKey, diskKey)
 }
 
+func decryptAESGCM(key, ciphertext []byte) ([]byte, error) {
+	// Create AES block cipher
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	// Create GCM (Galois/Counter Mode) cipher
+	aesGCM, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, err
+	}
+
+	// Separate nonce and ciphertext
+	nonceSize := aesGCM.NonceSize()
+	nonce, ciphertext := ciphertext[:nonceSize], ciphertext[nonceSize:]
+
+	// Decrypt the data
+	plaintext, err := aesGCM.Open(nil, nonce, ciphertext, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	return plaintext, nil
+}
+
 func read(network, which string) string {
 	val, isok := messageCache.Load(fmt.Sprintf("%s%s", network, which))
 	if isok {

functions/mqhandlers.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/devilcove/httpclient"
 	mqtt "github.com/eclipse/paho.mqtt.golang"
+	"github.com/gravitl/netclient/auth"
 	"github.com/gravitl/netclient/cache"
 	"github.com/gravitl/netclient/config"
 	"github.com/gravitl/netclient/daemon"
@@ -44,11 +45,22 @@ func NodeUpdate(client mqtt.Client, msg mqtt.Message) {
 	slog.Info("processing node update for network", "network", network)
 	node := config.GetNode(network)
 	server := config.Servers[node.Server]
-	data, err := decryptMsg(server.Name, msg.Payload())
+	data, err := decryptAESGCM(config.Netclient().TrafficKeyPublic[0:32], msg.Payload())
 	if err != nil {
-		slog.Error("error decrypting message", "error", err)
-		return
+		slog.Warn("error decrypting message", "warn", err)
+		data, err = decryptMsg(server.Name, msg.Payload())
+		if err != nil {
+			slog.Error("error decrypting message", "error", err)
+			return
+		}
+	} else {
+		data, err = unzipPayload(data)
+		if err != nil {
+			slog.Error("error unzipping message", "error", err)
+			return
+		}
 	}
+
 	serverNode := models.Node{}
 	if err = json.Unmarshal([]byte(data), &serverNode); err != nil {
 		slog.Error("error unmarshalling node update data", "error", err)
@@ -148,11 +160,22 @@ func HostPeerUpdate(client mqtt.Client, msg mqtt.Message) {
 		return
 	}
 	slog.Info("processing peer update for server", "server", serverName)
-	data, err := decryptMsg(serverName, msg.Payload())
+	data, err := decryptAESGCM(config.Netclient().TrafficKeyPublic[0:32], msg.Payload())
 	if err != nil {
-		slog.Error("error decrypting message", "error", err)
-		return
+		slog.Warn("error decrypting message", "warn", err)
+		data, err = decryptMsg(server.Name, msg.Payload())
+		if err != nil {
+			slog.Error("error decrypting message", "error", err)
+			return
+		}
+	} else {
+		data, err = unzipPayload(data)
+		if err != nil {
+			slog.Error("error unzipping message", "error", err)
+			return
+		}
 	}
+
 	err = json.Unmarshal([]byte(data), &peerUpdate)
 	if err != nil {
 		slog.Error("error unmarshalling peer data", "error", err)
@@ -281,10 +304,20 @@ func HostUpdate(client mqtt.Client, msg mqtt.Message) {
 	if len(msg.Payload()) == 0 {
 		return
 	}
-	data, err := decryptMsg(serverName, msg.Payload())
+	data, err := decryptAESGCM(config.Netclient().TrafficKeyPublic[0:32], msg.Payload())
 	if err != nil {
-		slog.Error("error decrypting message", "error", err)
-		return
+		slog.Warn("error decrypting message", "warn", err)
+		data, err = decryptMsg(server.Name, msg.Payload())
+		if err != nil {
+			slog.Error("error decrypting message", "error", err)
+			return
+		}
+	} else {
+		data, err = unzipPayload(data)
+		if err != nil {
+			slog.Error("error unzipping message", "error", err)
+			return
+		}
 	}
 	err = json.Unmarshal([]byte(data), &hostUpdate)
 	if err != nil {
@@ -606,6 +639,7 @@ func mqFallback(ctx context.Context, wg *sync.WaitGroup) {
 			}
 			// Call netclient http config pull
 			slog.Info("### mqfallback routine execute")
+			auth.CleanJwtToken()
 			response, resetInterface, replacePeers, err := Pull(false)
 			if err != nil {
 				slog.Error("pull failed", "error", err)