Skip to content

Commit ff1399e

Browse files
0xJacky0xJacky
committed
enhance: handle websocket error
1 parent 7a9aa3a commit ff1399e

File tree

8 files changed

+126
-60
lines changed

8 files changed

+126
-60
lines changed

‎api/analytic/analytic.go‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package analytic
33
import (
44
"fmt"
55
"github.com/0xJacky/Nginx-UI/internal/analytic"
6+
"github.com/0xJacky/Nginx-UI/internal/helper"
67
"github.com/0xJacky/Nginx-UI/internal/logger"
78
"github.com/shirou/gopsutil/v3/cpu"
89
"github.com/shirou/gopsutil/v3/host"
@@ -75,10 +76,7 @@ func Analytic(c *gin.Context) {
7576

7677
// write
7778
err = ws.WriteJSON(stat)
78-
if err != nil || websocket.IsUnexpectedCloseError(err,
79-
websocket.CloseGoingAway,
80-
websocket.CloseNoStatusReceived,
81-
websocket.CloseNormalClosure) {
79+
if helper.IsUnexpectedWebsocketError(err) {
8280
logger.Error(err)
8381
break
8482
}

‎api/analytic/nodes.go‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package analytic
22

33
import (
44
"github.com/0xJacky/Nginx-UI/internal/analytic"
5+
"github.com/0xJacky/Nginx-UI/internal/helper"
56
"github.com/0xJacky/Nginx-UI/internal/logger"
67
"github.com/gin-gonic/gin"
78
"github.com/gorilla/websocket"
@@ -27,10 +28,7 @@ func GetNodeStat(c *gin.Context) {
2728
for {
2829
// write
2930
err = ws.WriteJSON(analytic.GetNodeStat())
30-
if err != nil || websocket.IsUnexpectedCloseError(err,
31-
websocket.CloseGoingAway,
32-
websocket.CloseNoStatusReceived,
33-
websocket.CloseNormalClosure) {
31+
if helper.IsUnexpectedWebsocketError(err) {
3432
logger.Error(err)
3533
break
3634
}
@@ -57,10 +55,7 @@ func GetNodesAnalytic(c *gin.Context) {
5755
for {
5856
// write
5957
err = ws.WriteJSON(analytic.NodeMap)
60-
if err != nil || websocket.IsUnexpectedCloseError(err,
61-
websocket.CloseGoingAway,
62-
websocket.CloseNoStatusReceived,
63-
websocket.CloseNormalClosure) {
58+
if helper.IsUnexpectedWebsocketError(err) {
6459
logger.Error(err)
6560
break
6661
}

‎api/certificate/issue.go‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -32,14 +32,12 @@ func handleIssueCertLogChan(conn *websocket.Conn, log *cert.Logger, logChan chan
3232
}()
3333

3434
for logString := range logChan {
35-
3635
log.Info(logString)
3736

3837
err := conn.WriteJSON(IssueCertResponse{
3938
Status: Info,
4039
Message: logString,
4140
})
42-
4341
if err != nil {
4442
logger.Error(err)
4543
return
@@ -110,7 +108,6 @@ func IssueCert(c *gin.Context) {
110108
Status: Error,
111109
Message: err.Error(),
112110
})
113-
114111
if err != nil {
115112
logger.Error(err)
116113
return
@@ -132,7 +129,7 @@ func IssueCert(c *gin.Context) {
132129

133130
if err != nil {
134131
logger.Error(err)
135-
err = ws.WriteJSON(IssueCertResponse{
132+
_ = ws.WriteJSON(IssueCertResponse{
136133
Status: Error,
137134
Message: err.Error(),
138135
})
@@ -149,7 +146,6 @@ func IssueCert(c *gin.Context) {
149146
SSLCertificateKey: payload.GetCertificateKeyPath(),
150147
KeyType: payload.GetKeyType(),
151148
})
152-
153149
if err != nil {
154150
logger.Error(err)
155151
return

‎api/system/upgrade.go‎

Lines changed: 42 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,18 @@ func GetCurrentVersion(c *gin.Context) {
4141
c.JSON(http.StatusOK, curVer)
4242
}
4343

44+
const (
45+
UpgradeStatusInfo = "info"
46+
UpgradeStatusError = "error"
47+
UpgradeStatusProgress = "progress"
48+
)
49+
50+
type CoreUpgradeResp struct {
51+
Status string `json:"status"`
52+
Progress float64 `json:"progress"`
53+
Message string `json:"message"`
54+
}
55+
4456
func PerformCoreUpgrade(c *gin.Context) {
4557
var upGrader = websocket.Upgrader{
4658
CheckOrigin: func(r *http.Request) bool {
@@ -67,49 +79,48 @@ func PerformCoreUpgrade(c *gin.Context) {
6779
return
6880
}
6981

70-
_ = ws.WriteJSON(gin.H{
71-
"status": "info",
72-
"message": "Initialing core upgrader",
82+
_ = ws.WriteJSON(CoreUpgradeResp{
83+
Status: UpgradeStatusInfo,
84+
Message: "Initialing core upgrader",
7385
})
7486

7587
u, err := upgrader.NewUpgrader(control.Channel)
7688

7789
if err != nil {
78-
_ = ws.WriteJSON(gin.H{
79-
"status": "error",
80-
"message": "Initial core upgrader error",
90+
_ = ws.WriteJSON(CoreUpgradeResp{
91+
Status: UpgradeStatusError,
92+
Message: "Initial core upgrader error",
8193
})
82-
_ = ws.WriteJSON(gin.H{
83-
"status": "error",
84-
"message": err.Error(),
94+
_ = ws.WriteJSON(CoreUpgradeResp{
95+
Status: UpgradeStatusError,
96+
Message: err.Error(),
8597
})
8698
logger.Error(err)
8799
return
88100
}
89-
_ = ws.WriteJSON(gin.H{
90-
"status": "info",
91-
"message": "Downloading latest release",
101+
_ = ws.WriteJSON(CoreUpgradeResp{
102+
Status: UpgradeStatusInfo,
103+
Message: "Downloading latest release",
92104
})
93105
progressChan := make(chan float64)
94106
go func() {
95107
for progress := range progressChan {
96-
_ = ws.WriteJSON(gin.H{
97-
"status": "progress",
98-
"progress": progress,
108+
_ = ws.WriteJSON(CoreUpgradeResp{
109+
Status: UpgradeStatusProgress,
110+
Progress: progress,
99111
})
100112
}
101113
}()
102114

103115
tarName, err := u.DownloadLatestRelease(progressChan)
104-
105116
if err != nil {
106-
_ = ws.WriteJSON(gin.H{
107-
"status": "error",
108-
"message": "Download latest release error",
117+
_ = ws.WriteJSON(CoreUpgradeResp{
118+
Status: UpgradeStatusError,
119+
Message: "Download latest release error",
109120
})
110-
_ = ws.WriteJSON(gin.H{
111-
"status": "error",
112-
"message": err.Error(),
121+
_ = ws.WriteJSON(CoreUpgradeResp{
122+
Status: UpgradeStatusError,
123+
Message: err.Error(),
113124
})
114125
logger.Error(err)
115126
return
@@ -119,9 +130,9 @@ func PerformCoreUpgrade(c *gin.Context) {
119130
_ = os.Remove(tarName)
120131
_ = os.Remove(tarName + ".digest")
121132
}()
122-
_ = ws.WriteJSON(gin.H{
123-
"status": "info",
124-
"message": "Performing core upgrade",
133+
_ = ws.WriteJSON(CoreUpgradeResp{
134+
Status: UpgradeStatusInfo,
135+
Message: "Performing core upgrade",
125136
})
126137
// dry run
127138
if control.DryRun || settings.ServerSettings.Demo {
@@ -132,13 +143,13 @@ func PerformCoreUpgrade(c *gin.Context) {
132143
// bye, overseer will restart nginx-ui
133144
err = u.PerformCoreUpgrade(u.ExPath, tarName)
134145
if err != nil {
135-
_ = ws.WriteJSON(gin.H{
136-
"status": "error",
137-
"message": "Perform core upgrade error",
146+
_ = ws.WriteJSON(CoreUpgradeResp{
147+
Status: UpgradeStatusError,
148+
Message: "Perform core upgrade error",
138149
})
139-
_ = ws.WriteJSON(gin.H{
140-
"status": "error",
141-
"message": err.Error(),
150+
_ = ws.WriteJSON(CoreUpgradeResp{
151+
Status: UpgradeStatusError,
152+
Message: err.Error(),
142153
})
143154
logger.Error(err)
144155
return

‎api/upstream/upstream.go‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package upstream
22

33
import (
4+
"github.com/0xJacky/Nginx-UI/internal/helper"
45
"github.com/0xJacky/Nginx-UI/internal/logger"
56
"github.com/0xJacky/Nginx-UI/internal/upstream"
67
"github.com/gin-gonic/gin"
@@ -35,11 +36,7 @@ func AvailabilityTest(c *gin.Context) {
3536

3637
for {
3738
err = ws.WriteJSON(upstream.AvailabilityTest(body))
38-
39-
if err != nil || websocket.IsUnexpectedCloseError(err,
40-
websocket.CloseGoingAway,
41-
websocket.CloseNoStatusReceived,
42-
websocket.CloseNormalClosure) {
39+
if helper.IsUnexpectedWebsocketError(err) {
4340
logger.Error(err)
4441
break
4542
}

‎internal/helper/websocket_error.go‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
package helper
2+
3+
import (
4+
"github.com/gorilla/websocket"
5+
"github.com/pkg/errors"
6+
"syscall"
7+
)
8+
9+
func IsUnexpectedWebsocketError(err error) bool {
10+
// nil error is an expected error
11+
if err == nil {
12+
return false
13+
}
14+
// ignore: write: broken pipe
15+
if errors.Is(err, syscall.EPIPE) {
16+
return false
17+
}
18+
return websocket.IsUnexpectedCloseError(err,
19+
websocket.CloseGoingAway,
20+
websocket.CloseNoStatusReceived,
21+
websocket.CloseNormalClosure)
22+
}

‎internal/helper/websocket_error_test.go‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package helper
2+
3+
import (
4+
"github.com/gorilla/websocket"
5+
"github.com/stretchr/testify/assert"
6+
"syscall"
7+
"testing"
8+
)
9+
10+
func TestIsUnexpectedWebsocketError(t *testing.T) {
11+
var tests = []struct {
12+
input error
13+
output bool
14+
}{
15+
{nil, false},
16+
{input: &websocket.CloseError{
17+
Code: websocket.CloseGoingAway,
18+
}, output: false},
19+
{input: &websocket.CloseError{
20+
Code: websocket.CloseNoStatusReceived,
21+
}, output: false},
22+
{input: &websocket.CloseError{
23+
Code: websocket.CloseNormalClosure,
24+
}, output: false},
25+
{input: &websocket.CloseError{
26+
Code: websocket.CloseInternalServerErr,
27+
}, output: true},
28+
{
29+
input: syscall.EPIPE,
30+
output: false,
31+
},
32+
}
33+
for _, test := range tests {
34+
if !assert.Equal(t, test.output, IsUnexpectedWebsocketError(test.input)) {
35+
t.Log(test.input)
36+
}
37+
}
38+
}

‎internal/upgrader/upgrade.go‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import (
55
"fmt"
66
_github "github.com/0xJacky/Nginx-UI/.github"
77
"github.com/0xJacky/Nginx-UI/app"
8-
helper2 "github.com/0xJacky/Nginx-UI/internal/helper"
8+
"github.com/0xJacky/Nginx-UI/internal/helper"
99
"github.com/0xJacky/Nginx-UI/internal/logger"
1010
"github.com/0xJacky/Nginx-UI/settings"
1111
"github.com/pkg/errors"
@@ -282,7 +282,6 @@ func (u *Upgrader) DownloadLatestRelease(progressChan chan float64) (tarName str
282282

283283
// digest
284284
digest, ok := assetsMap[fmt.Sprintf("nginx-ui-%s.tar.gz.digest", arch.Name)]
285-
286285
if !ok || digest.BrowserDownloadUrl == "" {
287286
err = errors.New("upgrader core digest is empty")
288287
return
@@ -297,7 +296,6 @@ func (u *Upgrader) DownloadLatestRelease(progressChan chan float64) (tarName str
297296
}
298297

299298
resp, err := http.Get(digest.BrowserDownloadUrl)
300-
301299
if err != nil {
302300
err = errors.Wrap(err, "upgrader core download digest fail")
303301
return
@@ -324,16 +322,27 @@ func (u *Upgrader) DownloadLatestRelease(progressChan chan float64) (tarName str
324322
// check tar digest
325323
digestFileBytes, err := io.ReadAll(resp.Body)
326324
if err != nil {
327-
err = errors.Wrap(err, "digestFileContent read error")
325+
err = errors.Wrap(err, "digest file content read error")
328326
return
329327
}
330328

331329
digestFileContent := strings.TrimSpace(string(digestFileBytes))
332330

333-
logger.Debug("DownloadLatestRelease tar digest", helper2.DigestSHA512(tarName))
331+
logger.Debug("DownloadLatestRelease tar digest", helper.DigestSHA512(tarName))
334332
logger.Debug("DownloadLatestRelease digestFileContent", digestFileContent)
335333

336-
if digestFileContent != helper2.DigestSHA512(tarName) {
334+
if digestFileContent == "" {
335+
err = errors.New("digest file content is empty")
336+
return
337+
}
338+
339+
exeSHA512 := helper.DigestSHA512(tarName)
340+
if exeSHA512 == "" {
341+
err = errors.New("executable binary file is empty")
342+
return
343+
}
344+
345+
if digestFileContent != exeSHA512 {
337346
err = errors.Wrap(err, "digest not equal")
338347
return
339348
}
@@ -343,7 +352,7 @@ func (u *Upgrader) DownloadLatestRelease(progressChan chan float64) (tarName str
343352

344353
func (u *Upgrader) PerformCoreUpgrade(exPath string, tarPath string) (err error) {
345354
dir := filepath.Dir(exPath)
346-
err = helper2.UnTar(dir, tarPath)
355+
err = helper.UnTar(dir, tarPath)
347356
if err != nil {
348357
err = errors.Wrap(err, "PerformCoreUpgrade unTar error")
349358
return

0 commit comments

Comments
 (0)