Merge pull request #440 from 0xJacky/enhance/handle-ws-error

api/analytic/analytic.go

@@ -3,6 +3,7 @@ package analytic
 import (
 	"fmt"
 	"github.com/0xJacky/Nginx-UI/internal/analytic"
+	"github.com/0xJacky/Nginx-UI/internal/helper"
 	"github.com/0xJacky/Nginx-UI/internal/logger"
 	"github.com/shirou/gopsutil/v3/cpu"
 	"github.com/shirou/gopsutil/v3/host"
@@ -75,10 +76,7 @@ func Analytic(c *gin.Context) {
 
 		// write
 		err = ws.WriteJSON(stat)
-		if err != nil || websocket.IsUnexpectedCloseError(err,
-			websocket.CloseGoingAway,
-			websocket.CloseNoStatusReceived,
-			websocket.CloseNormalClosure) {
+		if helper.IsUnexpectedWebsocketError(err) {
 			logger.Error(err)
 			break
 		}

api/analytic/nodes.go

@@ -2,6 +2,7 @@ package analytic
 
 import (
 	"github.com/0xJacky/Nginx-UI/internal/analytic"
+	"github.com/0xJacky/Nginx-UI/internal/helper"
 	"github.com/0xJacky/Nginx-UI/internal/logger"
 	"github.com/gin-gonic/gin"
 	"github.com/gorilla/websocket"
@@ -27,10 +28,7 @@ func GetNodeStat(c *gin.Context) {
 	for {
 		// write
 		err = ws.WriteJSON(analytic.GetNodeStat())
-		if err != nil || websocket.IsUnexpectedCloseError(err,
-			websocket.CloseGoingAway,
-			websocket.CloseNoStatusReceived,
-			websocket.CloseNormalClosure) {
+		if helper.IsUnexpectedWebsocketError(err) {
 			logger.Error(err)
 			break
 		}
@@ -57,10 +55,7 @@ func GetNodesAnalytic(c *gin.Context) {
 	for {
 		// write
 		err = ws.WriteJSON(analytic.NodeMap)
-		if err != nil || websocket.IsUnexpectedCloseError(err,
-			websocket.CloseGoingAway,
-			websocket.CloseNoStatusReceived,
-			websocket.CloseNormalClosure) {
+		if helper.IsUnexpectedWebsocketError(err) {
 			logger.Error(err)
 			break
 		}

api/certificate/issue.go

@@ -32,14 +32,12 @@ func handleIssueCertLogChan(conn *websocket.Conn, log *cert.Logger, logChan chan
 	}()
 
 	for logString := range logChan {
-
 		log.Info(logString)
 
 		err := conn.WriteJSON(IssueCertResponse{
 			Status:  Info,
 			Message: logString,
 		})
-
 		if err != nil {
 			logger.Error(err)
 			return
@@ -110,7 +108,6 @@ func IssueCert(c *gin.Context) {
 			Status:  Error,
 			Message: err.Error(),
 		})
-
 		if err != nil {
 			logger.Error(err)
 			return
@@ -132,7 +129,7 @@ func IssueCert(c *gin.Context) {
 
 	if err != nil {
 		logger.Error(err)
-		err = ws.WriteJSON(IssueCertResponse{
+		_ = ws.WriteJSON(IssueCertResponse{
 			Status:  Error,
 			Message: err.Error(),
 		})
@@ -149,7 +146,6 @@ func IssueCert(c *gin.Context) {
 		SSLCertificateKey: payload.GetCertificateKeyPath(),
 		KeyType:           payload.GetKeyType(),
 	})
-
 	if err != nil {
 		logger.Error(err)
 		return

api/system/upgrade.go

@@ -41,6 +41,18 @@ func GetCurrentVersion(c *gin.Context) {
 	c.JSON(http.StatusOK, curVer)
 }
 
+const (
+	UpgradeStatusInfo     = "info"
+	UpgradeStatusError    = "error"
+	UpgradeStatusProgress = "progress"
+)
+
+type CoreUpgradeResp struct {
+	Status   string  `json:"status"`
+	Progress float64 `json:"progress"`
+	Message  string  `json:"message"`
+}
+
 func PerformCoreUpgrade(c *gin.Context) {
 	var upGrader = websocket.Upgrader{
 		CheckOrigin: func(r *http.Request) bool {
@@ -67,49 +79,48 @@ func PerformCoreUpgrade(c *gin.Context) {
 		return
 	}
 
-	_ = ws.WriteJSON(gin.H{
-		"status":  "info",
-		"message": "Initialing core upgrader",
+	_ = ws.WriteJSON(CoreUpgradeResp{
+		Status:  UpgradeStatusInfo,
+		Message: "Initialing core upgrader",
 	})
 
 	u, err := upgrader.NewUpgrader(control.Channel)
 
 	if err != nil {
-		_ = ws.WriteJSON(gin.H{
-			"status":  "error",
-			"message": "Initial core upgrader error",
+		_ = ws.WriteJSON(CoreUpgradeResp{
+			Status:  UpgradeStatusError,
+			Message: "Initial core upgrader error",
 		})
-		_ = ws.WriteJSON(gin.H{
-			"status":  "error",
-			"message": err.Error(),
+		_ = ws.WriteJSON(CoreUpgradeResp{
+			Status:  UpgradeStatusError,
+			Message: err.Error(),
 		})
 		logger.Error(err)
 		return
 	}
-	_ = ws.WriteJSON(gin.H{
-		"status":  "info",
-		"message": "Downloading latest release",
+	_ = ws.WriteJSON(CoreUpgradeResp{
+		Status:  UpgradeStatusInfo,
+		Message: "Downloading latest release",
 	})
 	progressChan := make(chan float64)
 	go func() {
 		for progress := range progressChan {
-			_ = ws.WriteJSON(gin.H{
-				"status":   "progress",
-				"progress": progress,
+			_ = ws.WriteJSON(CoreUpgradeResp{
+				Status:   UpgradeStatusProgress,
+				Progress: progress,
 			})
 		}
 	}()
 
 	tarName, err := u.DownloadLatestRelease(progressChan)
-
 	if err != nil {
-		_ = ws.WriteJSON(gin.H{
-			"status":  "error",
-			"message": "Download latest release error",
+		_ = ws.WriteJSON(CoreUpgradeResp{
+			Status:  UpgradeStatusError,
+			Message: "Download latest release error",
 		})
-		_ = ws.WriteJSON(gin.H{
-			"status":  "error",
-			"message": err.Error(),
+		_ = ws.WriteJSON(CoreUpgradeResp{
+			Status:  UpgradeStatusError,
+			Message: err.Error(),
 		})
 		logger.Error(err)
 		return
@@ -119,9 +130,9 @@ func PerformCoreUpgrade(c *gin.Context) {
 		_ = os.Remove(tarName)
 		_ = os.Remove(tarName + ".digest")
 	}()
-	_ = ws.WriteJSON(gin.H{
-		"status":  "info",
-		"message": "Performing core upgrade",
+	_ = ws.WriteJSON(CoreUpgradeResp{
+		Status:  UpgradeStatusInfo,
+		Message: "Performing core upgrade",
 	})
 	// dry run
 	if control.DryRun || settings.ServerSettings.Demo {
@@ -132,13 +143,13 @@ func PerformCoreUpgrade(c *gin.Context) {
 	// bye, overseer will restart nginx-ui
 	err = u.PerformCoreUpgrade(u.ExPath, tarName)
 	if err != nil {
-		_ = ws.WriteJSON(gin.H{
-			"status":  "error",
-			"message": "Perform core upgrade error",
+		_ = ws.WriteJSON(CoreUpgradeResp{
+			Status:  UpgradeStatusError,
+			Message: "Perform core upgrade error",
 		})
-		_ = ws.WriteJSON(gin.H{
-			"status":  "error",
-			"message": err.Error(),
+		_ = ws.WriteJSON(CoreUpgradeResp{
+			Status:  UpgradeStatusError,
+			Message: err.Error(),
 		})
 		logger.Error(err)
 		return

api/upstream/upstream.go

@@ -1,6 +1,7 @@
 package upstream
 
 import (
+	"github.com/0xJacky/Nginx-UI/internal/helper"
 	"github.com/0xJacky/Nginx-UI/internal/logger"
 	"github.com/0xJacky/Nginx-UI/internal/upstream"
 	"github.com/gin-gonic/gin"
@@ -35,11 +36,7 @@ func AvailabilityTest(c *gin.Context) {
 
 	for {
 		err = ws.WriteJSON(upstream.AvailabilityTest(body))
-
-		if err != nil || websocket.IsUnexpectedCloseError(err,
-			websocket.CloseGoingAway,
-			websocket.CloseNoStatusReceived,
-			websocket.CloseNormalClosure) {
+		if helper.IsUnexpectedWebsocketError(err) {
 			logger.Error(err)
 			break
 		}

internal/helper/websocket_error.go

@@ -0,0 +1,22 @@
+package helper
+
+import (
+	"github.com/gorilla/websocket"
+	"github.com/pkg/errors"
+	"syscall"
+)
+
+func IsUnexpectedWebsocketError(err error) bool {
+	// nil error is an expected error
+	if err == nil {
+		return false
+	}
+	// ignore: write: broken pipe
+	if errors.Is(err, syscall.EPIPE) {
+		return false
+	}
+	return websocket.IsUnexpectedCloseError(err,
+		websocket.CloseGoingAway,
+		websocket.CloseNoStatusReceived,
+		websocket.CloseNormalClosure)
+}

internal/helper/websocket_error_test.go

@@ -0,0 +1,38 @@
+package helper
+
+import (
+	"github.com/gorilla/websocket"
+	"github.com/stretchr/testify/assert"
+	"syscall"
+	"testing"
+)
+
+func TestIsUnexpectedWebsocketError(t *testing.T) {
+	var tests = []struct {
+		input  error
+		output bool
+	}{
+		{nil, false},
+		{input: &websocket.CloseError{
+			Code: websocket.CloseGoingAway,
+		}, output: false},
+		{input: &websocket.CloseError{
+			Code: websocket.CloseNoStatusReceived,
+		}, output: false},
+		{input: &websocket.CloseError{
+			Code: websocket.CloseNormalClosure,
+		}, output: false},
+		{input: &websocket.CloseError{
+			Code: websocket.CloseInternalServerErr,
+		}, output: true},
+		{
+			input:  syscall.EPIPE,
+			output: false,
+		},
+	}
+	for _, test := range tests {
+		if !assert.Equal(t, test.output, IsUnexpectedWebsocketError(test.input)) {
+			t.Log(test.input)
+		}
+	}
+}

internal/upgrader/upgrade.go

@@ -5,7 +5,7 @@ import (
 	"fmt"
 	_github "github.com/0xJacky/Nginx-UI/.github"
 	"github.com/0xJacky/Nginx-UI/app"
-	helper2 "github.com/0xJacky/Nginx-UI/internal/helper"
+	"github.com/0xJacky/Nginx-UI/internal/helper"
 	"github.com/0xJacky/Nginx-UI/internal/logger"
 	"github.com/0xJacky/Nginx-UI/settings"
 	"github.com/pkg/errors"
@@ -282,7 +282,6 @@ func (u *Upgrader) DownloadLatestRelease(progressChan chan float64) (tarName str
 
 	// digest
 	digest, ok := assetsMap[fmt.Sprintf("nginx-ui-%s.tar.gz.digest", arch.Name)]
-
 	if !ok || digest.BrowserDownloadUrl == "" {
 		err = errors.New("upgrader core digest is empty")
 		return
@@ -297,7 +296,6 @@ func (u *Upgrader) DownloadLatestRelease(progressChan chan float64) (tarName str
 	}
 
 	resp, err := http.Get(digest.BrowserDownloadUrl)
-
 	if err != nil {
 		err = errors.Wrap(err, "upgrader core download digest fail")
 		return
@@ -324,16 +322,27 @@ func (u *Upgrader) DownloadLatestRelease(progressChan chan float64) (tarName str
 	// check tar digest
 	digestFileBytes, err := io.ReadAll(resp.Body)
 	if err != nil {
-		err = errors.Wrap(err, "digestFileContent read error")
+		err = errors.Wrap(err, "digest file content read error")
 		return
 	}
 
 	digestFileContent := strings.TrimSpace(string(digestFileBytes))
 
-	logger.Debug("DownloadLatestRelease tar digest", helper2.DigestSHA512(tarName))
+	logger.Debug("DownloadLatestRelease tar digest", helper.DigestSHA512(tarName))
 	logger.Debug("DownloadLatestRelease digestFileContent", digestFileContent)
 
-	if digestFileContent != helper2.DigestSHA512(tarName) {
+	if digestFileContent == "" {
+		err = errors.New("digest file content is empty")
+		return
+	}
+
+	exeSHA512 := helper.DigestSHA512(tarName)
+	if exeSHA512 == "" {
+		err = errors.New("executable binary file is empty")
+		return
+	}
+
+	if digestFileContent != exeSHA512 {
 		err = errors.Wrap(err, "digest not equal")
 		return
 	}
@@ -343,7 +352,7 @@ func (u *Upgrader) DownloadLatestRelease(progressChan chan float64) (tarName str
 
 func (u *Upgrader) PerformCoreUpgrade(exPath string, tarPath string) (err error) {
 	dir := filepath.Dir(exPath)
-	err = helper2.UnTar(dir, tarPath)
+	err = helper.UnTar(dir, tarPath)
 	if err != nil {
 		err = errors.Wrap(err, "PerformCoreUpgrade unTar error")
 		return