dan9i10/3x-ui
1
0
Fork 0
mirror of https://github.com/MHSanaei/3x-ui.git synced 2026-05-08 14:36:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(websocket): split controller into service + thin controller

Browse Source 
Move per-connection lifecycle out of the controller and into a new
service.WebSocketService. The controller is now HTTP-layer only:
authenticate, validate origin, upgrade, and hand the connection off.

- web/service/websocket.go (new): owns the read/write pumps, hub
  registration, and connection lifetime. Pump constants are prefixed
  (wsWriteWait, wsPongWait, wsPingPeriod, wsClientReadLimit) to avoid
  collisions in the larger service package namespace.
- web/controller/websocket.go: trimmed to the upgrader, same-origin
  check, auth gate, and hand-off to the service.
- web/web.go: wires controller.NewWebSocketController(service.NewWebSocketService(hub)).

The hub package (web/websocket) stays as low-level fan-out
infrastructure. Behavior is unchanged — this is a structural cleanup
to align with the rest of the codebase's controller/service split.

Also includes a small range-int modernization in login_limiter_test.go
that gopls flagged.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
MHSanaei
2026-05-08 00:00:44 +02:00
parent b84b58ef21
commit c394938f01
4 changed files with 133 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
web/controller/login_limiter_test.go
View File

@@ -10,7 +10,7 @@ func TestLoginLimiterBlocksAfterConfiguredFailures(t *testing.T) {
limiter := newLoginLimiter(5, 5*time.Minute, 15*time.Minute)
limiter.now = func() time.Time { return now }
for i := 0; i < 4; i++ {
for i := range 4 {
if _, blocked := limiter.registerFailure("192.0.2.10", "Admin"); blocked {
t.Fatalf("failure %d should not block yet", i+1)
}
@@ -41,7 +41,7 @@ func TestLoginLimiterPrunesOldFailuresAndResetsOnSuccess(t *testing.T) {
limiter := newLoginLimiter(5, 5*time.Minute, 15*time.Minute)
limiter.now = func() time.Time { return now }
for i := 0; i < 4; i++ {
for range 4 {
limiter.registerFailure("192.0.2.10", "admin")
}
now = now.Add(6 * time.Minute)
@@ -50,7 +50,7 @@ func TestLoginLimiterPrunesOldFailuresAndResetsOnSuccess(t *testing.T) {
}
limiter.registerSuccess("192.0.2.10", "admin")
for i := 0; i < 4; i++ {
for i := range 4 {
if _, blocked := limiter.registerFailure("192.0.2.10", "admin"); blocked {
t.Fatalf("success should reset previous failures; failure %d blocked", i+1)
}
@@ -62,7 +62,7 @@ func TestLoginLimiterSeparatesIPAndUsername(t *testing.T) {
limiter := newLoginLimiter(5, 5*time.Minute, 15*time.Minute)
limiter.now = func() time.Time { return now }
for i := 0; i < 5; i++ {
for range 5 {
limiter.registerFailure("192.0.2.10", "admin")
}
if _, ok := limiter.allow("192.0.2.11", "admin"); !ok {

95
web/controller/websocket.go
View File

@@ -5,25 +5,15 @@ import (
"net/http"
"net/url"
"strings"
"time"
"github.com/google/uuid"
"github.com/mhsanaei/3x-ui/v2/logger"
"github.com/mhsanaei/3x-ui/v2/util/common"
"github.com/mhsanaei/3x-ui/v2/web/service"
"github.com/mhsanaei/3x-ui/v2/web/session"
"github.com/mhsanaei/3x-ui/v2/web/websocket"
"github.com/gin-gonic/gin"
ws "github.com/gorilla/websocket"
)
const (
writeWait = 10 * time.Second
pongWait = 60 * time.Second
pingPeriod = (pongWait * 9) / 10
clientReadLimit = 512
)
var upgrader = ws.Upgrader{
ReadBufferSize: 32768,
WriteBufferSize: 32768,
@@ -57,18 +47,21 @@ func checkSameOrigin(r *http.Request) bool {
return strings.EqualFold(u.Hostname(), host)
}
// WebSocketController handles WebSocket connections for real-time updates.
// WebSocketController handles the HTTP→WebSocket upgrade for real-time updates.
// All per-connection lifecycle (pumps, hub registration) lives in
// service.WebSocketService — this controller is HTTP-layer only.
type WebSocketController struct {
BaseController
hub *websocket.Hub
service *service.WebSocketService
}
// NewWebSocketController creates a new WebSocket controller.
func NewWebSocketController(hub *websocket.Hub) *WebSocketController {
return &WebSocketController{hub: hub}
// NewWebSocketController creates a controller wired to the given service.
func NewWebSocketController(svc *service.WebSocketService) *WebSocketController {
return &WebSocketController{service: svc}
}
// HandleWebSocket upgrades the HTTP connection and starts the read/write pumps.
// HandleWebSocket authenticates the request, upgrades the HTTP connection, and
// hands ownership of the connection off to the service.
func (w *WebSocketController) HandleWebSocket(c *gin.Context) {
if !session.IsLogin(c) {
logger.Warningf("Unauthorized WebSocket connection attempt from %s", getRemoteIp(c))
@@ -82,71 +75,5 @@ func (w *WebSocketController) HandleWebSocket(c *gin.Context) {
return
}
client := websocket.NewClient(uuid.New().String())
w.hub.Register(client)
logger.Debugf("WebSocket client %s registered from %s", client.ID, getRemoteIp(c))
go w.writePump(client, conn)
go w.readPump(client, conn)
}
// readPump consumes inbound frames so the gorilla deadline/pong machinery keeps
// running. Clients send no commands today; frames are discarded.
func (w *WebSocketController) readPump(client *websocket.Client, conn *ws.Conn) {
defer func() {
if r := common.Recover("WebSocket readPump panic"); r != nil {
logger.Error("WebSocket readPump panic recovered:", r)
}
w.hub.Unregister(client)
conn.Close()
}()
conn.SetReadLimit(clientReadLimit)
conn.SetReadDeadline(time.Now().Add(pongWait))
conn.SetPongHandler(func(string) error {
return conn.SetReadDeadline(time.Now().Add(pongWait))
})
for {
if _, _, err := conn.ReadMessage(); err != nil {
if ws.IsUnexpectedCloseError(err, ws.CloseGoingAway, ws.CloseAbnormalClosure) {
logger.Debugf("WebSocket read error for client %s: %v", client.ID, err)
}
return
}
}
}
// writePump pushes hub messages to the connection and emits keepalive pings.
func (w *WebSocketController) writePump(client *websocket.Client, conn *ws.Conn) {
ticker := time.NewTicker(pingPeriod)
defer func() {
if r := common.Recover("WebSocket writePump panic"); r != nil {
logger.Error("WebSocket writePump panic recovered:", r)
}
ticker.Stop()
conn.Close()
}()
for {
select {
case msg, ok := <-client.Send:
conn.SetWriteDeadline(time.Now().Add(writeWait))
if !ok {
conn.WriteMessage(ws.CloseMessage, []byte{})
return
}
if err := conn.WriteMessage(ws.TextMessage, msg); err != nil {
logger.Debugf("WebSocket write error for client %s: %v", client.ID, err)
return
}
case <-ticker.C:
conn.SetWriteDeadline(time.Now().Add(writeWait))
if err := conn.WriteMessage(ws.PingMessage, nil); err != nil {
logger.Debugf("WebSocket ping error for client %s: %v", client.ID, err)
return
}
}
}
w.service.HandleConnection(conn, getRemoteIp(c))
}