tcjlj/nezha
1
0
Fork 0
mirror of https://github.com/nezhahq/nezha.git synced 2025-01-22 12:48:14 -05:00
Code Issues Actions 7 Packages Projects Releases Wiki Activity

prevent writing response to websocket connections (#457)

Browse Source
This commit is contained in:
UUBulb 2024-10-25 21:45:05 +08:00 committed by GitHub
parent d086e98711
commit 380973a200
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 33 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
cmd/dashboard/controller/controller.go
View File

@ -136,6 +136,22 @@ func (ge *gormError) Error() string {
return fmt.Sprintf(ge.msg, ge.a...)
}
type wsError struct {
msg string
a []interface{}
}
func newWsError(format string, args ...interface{}) error {
return &wsError{
msg: format,
a: args,
}
}
func (we *wsError) Error() string {
return fmt.Sprintf(we.msg, we.a...)
}
func commonHandler[T any](handler handlerFunc[T]) func(*gin.Context) {
return func(c *gin.Context) {
data, err := handler(c)
@ -143,11 +159,18 @@ func commonHandler[T any](handler handlerFunc[T]) func(*gin.Context) {
c.JSON(http.StatusOK, model.CommonResponse[T]{Success: true, Data: data})
return
}
if _, ok := err.(*gormError); ok {
switch err.(type) {
case *gormError:
log.Printf("NEZHA>> gorm error: %v", err)
c.JSON(http.StatusOK, newErrorResponse(errors.New("database error")))
return
} else {
case *wsError:
// Connection is upgraded to WebSocket, so c.Writer is no longer usable
if msg := err.Error(); msg != "" {
log.Printf("NEZHA>> websocket error: %v", err)
}
return
default:
c.JSON(http.StatusOK, newErrorResponse(err))
return
}

7
cmd/dashboard/controller/fm.go
View File

@ -21,7 +21,7 @@ import (
// @Description Create an "attached" FM. It is advised to only call this within a terminal session.
// @Tags auth required
// @Accept json
// @Param id path uint true "Server ID"
// @Param id query uint true "Server ID"
// @Produce json
// @Success 200 {object} model.CreateFMResponse
// @Router /file [get]
@ -66,6 +66,7 @@ func createFM(c *gin.Context) (*model.CreateFMResponse, error) {
// @Description Start FM stream
// @Tags auth required
// @Param id path string true "Stream UUID"
// @Success 200 {object} model.CommonResponse[any]
// @Router /ws/file/{id} [get]
func fmStream(c *gin.Context) (any, error) {
streamId := c.Param("id")
@ -92,8 +93,8 @@ func fmStream(c *gin.Context) (any, error) {
}()
if err = rpc.NezhaHandlerSingleton.UserConnected(streamId, conn); err != nil {
return nil, err
return nil, newWsError("%v", err)
}
return nil, rpc.NezhaHandlerSingleton.StartStream(streamId, time.Second*10)
return nil, newWsError("%v", rpc.NezhaHandlerSingleton.StartStream(streamId, time.Second*10))
}

5
cmd/dashboard/controller/terminal.go
View File

@ -66,6 +66,7 @@ func createTerminal(c *gin.Context) (*model.CreateTerminalResponse, error) {
// @Description Terminal stream
// @Tags auth required
// @Param id path string true "Stream UUID"
// @Success 200 {object} model.CommonResponse[any]
// @Router /ws/terminal/{id} [get]
func terminalStream(c *gin.Context) (any, error) {
streamId := c.Param("id")
@ -92,8 +93,8 @@ func terminalStream(c *gin.Context) (any, error) {
}()
if err = rpc.NezhaHandlerSingleton.UserConnected(streamId, conn); err != nil {
return nil, err
return nil, newWsError("%v", err)
}
return nil, rpc.NezhaHandlerSingleton.StartStream(streamId, time.Second*10)
return nil, newWsError("%v", rpc.NezhaHandlerSingleton.StartStream(streamId, time.Second*10))
}

2
cmd/dashboard/controller/ws.go
View File

@ -51,7 +51,7 @@ func serverStream(c *gin.Context) (any, error) {
}
time.Sleep(time.Second * 2)
}
return nil, nil
return nil, newWsError("")
}
var requestGroup singleflight.Group