From 880ac06a930bd2ddfa9d58590659b1a289ae65ca Mon Sep 17 00:00:00 2001 From: Jianhui Zhao Date: Thu, 31 Jul 2025 00:03:57 +0800 Subject: [PATCH] Improves code organization and readability - Encapsulate API routes into APIServer methods (api.go) - Split large handleUserConnection into smaller methods (user.go) Signed-off-by: Jianhui Zhao --- api.go | 491 ++++++++++++++++++++++++++++++-------------------------- http.go | 6 +- user.go | 107 ++++++------ 3 files changed, 324 insertions(+), 280 deletions(-) diff --git a/api.go b/api.go index 10d8911..f818094 100644 --- a/api.go +++ b/api.go @@ -41,8 +41,6 @@ import ( "github.com/rs/zerolog/log" ) -var httpSessions = cache.NewMemCache(cache.WithClearInterval(time.Minute)) - const httpSessionExpire = 30 * time.Minute //go:embed ui/dist @@ -55,6 +53,21 @@ func (srv *RttyServer) ListenAPI() error { r := gin.New() + fs, err := fs.Sub(staticFs, "ui/dist") + if err != nil { + return err + } + + root := http.FS(fs) + + a := &APIServer{ + sessions: cache.NewMemCache(cache.WithClearInterval(time.Minute)), + fh: http.FileServer(root), + srv: srv, + r: r, + root: root, + } + r.Use(func(c *gin.Context) { c.Next() log.Debug().Msgf("%s - \"%s %s %s %d\"", c.ClientIP(), @@ -71,223 +84,26 @@ func (srv *RttyServer) ListenAPI() error { return } - if !httpAuth(cfg, c) { + if !a.auth(c) { c.AbortWithStatus(http.StatusUnauthorized) return } }) - authorized.GET("/connect/:devid", func(c *gin.Context) { - if !callUserHookUrl(cfg, c) { - c.Status(http.StatusForbidden) - return - } + authorized.GET("/connect/:devid", a.handleConnect) + authorized.GET("/counts", a.handleCounts) + authorized.GET("/groups", a.handleGroups) + authorized.GET("/devs", a.handleDevs) + authorized.GET("/dev/:devid", a.handleDev) + authorized.POST("/cmd/:devid", a.handleCmd) + authorized.Any("/web/:devid/:proto/:addr/*path", a.handleWeb) + authorized.Any("/web2/:group/:devid/:proto/:addr/*path", a.handleWeb2) + authorized.GET("/signout", a.handleSignout) - if c.GetHeader("Upgrade") != "websocket" { - group := c.Query("group") - devid := c.Param("devid") - if dev := srv.GetDevice(group, devid); dev == nil { - c.Redirect(http.StatusFound, "/error/offline") - return - } + r.POST("/signin", a.handleSignin) + r.GET("/alive", a.handleAlive) - url := "/rtty/" + devid - - if group != "" { - url += "?group=" + group - } - - c.Redirect(http.StatusFound, url) - } else { - handleUserConnection(srv, c) - } - }) - - authorized.GET("/counts", func(c *gin.Context) { - count := 0 - - srv.groups.Range(func(key, value any) bool { - count += int(value.(*DeviceGroup).count.Load()) - return true - }) - - c.JSON(http.StatusOK, gin.H{"count": count}) - }) - - authorized.GET("/groups", func(c *gin.Context) { - groups := []string{""} - - srv.groups.Range(func(key, value any) bool { - if key != "" { - groups = append(groups, key.(string)) - } - return true - }) - - c.JSON(http.StatusOK, groups) - }) - - authorized.GET("/devs", func(c *gin.Context) { - devs := make([]*DeviceInfo, 0) - g := srv.GetGroup(c.Query("group"), false) - - if g == nil { - c.JSON(http.StatusOK, devs) - return - } - - g.devices.Range(func(key, value any) bool { - dev := value.(*Device) - - devs = append(devs, &DeviceInfo{ - Group: dev.group, - ID: dev.id, - Desc: dev.desc, - Connected: uint32(time.Now().Unix() - dev.timestamp), - Uptime: dev.uptime, - Proto: dev.proto, - IPaddr: dev.conn.RemoteAddr().(*net.TCPAddr).IP.String(), - }) - - return true - }) - - c.JSON(http.StatusOK, devs) - }) - - authorized.GET("/dev/:devid", func(c *gin.Context) { - if dev := srv.GetDevice(c.Query("group"), c.Param("devid")); dev != nil { - info := &DeviceInfo{ - ID: dev.id, - Desc: dev.desc, - Connected: uint32(time.Now().Unix() - dev.timestamp), - Uptime: dev.uptime, - Proto: dev.proto, - IPaddr: dev.conn.RemoteAddr().(*net.TCPAddr).IP.String(), - } - c.JSON(http.StatusOK, info) - } else { - c.Status(http.StatusNotFound) - } - }) - - authorized.POST("/cmd/:devid", func(c *gin.Context) { - if !callUserHookUrl(cfg, c) { - c.Status(http.StatusForbidden) - return - } - - cmdInfo := &CommandReqInfo{} - - err := c.BindJSON(&cmdInfo) - if err != nil || cmdInfo.Cmd == "" || cmdInfo.Username == "" { - cmdErrResp(c, rttyCmdErrInvalid) - return - } - - dev := srv.GetDevice(c.Query("group"), c.Param("devid")) - if dev == nil { - cmdErrResp(c, rttyCmdErrOffline) - return - } - - dev.handleCmdReq(c, cmdInfo) - }) - - authorized.Any("/web/:devid/:proto/:addr/*path", func(c *gin.Context) { - httpProxyRedirect(srv, c, "") - }) - - authorized.Any("/web2/:group/:devid/:proto/:addr/*path", func(c *gin.Context) { - group := c.Param("group") - httpProxyRedirect(srv, c, group) - }) - - authorized.GET("/signout", func(c *gin.Context) { - sid, err := c.Cookie("sid") - if err != nil || !httpSessions.Exists(sid) { - return - } - - httpSessions.Del(sid) - - c.Status(http.StatusOK) - }) - - r.POST("/signin", func(c *gin.Context) { - type credentials struct { - Password string `json:"password"` - } - - creds := credentials{} - - err := c.BindJSON(&creds) - if err != nil { - c.Status(http.StatusBadRequest) - return - } - - if httpLogin(cfg, creds.Password) { - sid := utils.GenUniqueID() - - httpSessions.Set(sid, true, cache.WithEx(httpSessionExpire)) - - c.SetCookie("sid", sid, 0, "", "", false, true) - c.Status(http.StatusOK) - return - } - - c.Status(http.StatusUnauthorized) - }) - - r.GET("/alive", func(c *gin.Context) { - if !httpAuth(cfg, c) { - c.AbortWithStatus(http.StatusUnauthorized) - } else { - c.Status(http.StatusOK) - } - }) - - fs, err := fs.Sub(staticFs, "ui/dist") - if err != nil { - return err - } - - root := http.FS(fs) - fh := http.FileServer(root) - - r.NoRoute(func(c *gin.Context) { - upath := path.Clean(c.Request.URL.Path) - - if strings.HasSuffix(upath, ".js") || strings.HasSuffix(upath, ".css") { - if strings.Contains(c.Request.Header.Get("Accept-Encoding"), "gzip") { - f, err := root.Open(upath + ".gz") - if err == nil { - f.Close() - - c.Request.URL.Path += ".gz" - - if strings.HasSuffix(upath, ".js") { - c.Writer.Header().Set("Content-Type", "application/javascript") - } else if strings.HasSuffix(upath, ".css") { - c.Writer.Header().Set("Content-Type", "text/css") - } - - c.Writer.Header().Set("Content-Encoding", "gzip") - } - } - } else if upath != "/" { - f, err := root.Open(upath) - if err != nil { - c.Request.URL.Path = "/" - r.HandleContext(c) - return - } - defer f.Close() - } - - fh.ServeHTTP(c.Writer, c.Request) - }) + r.NoRoute(a.handleFile) ln, err := net.Listen("tcp", cfg.AddrUser) if err != nil { @@ -300,7 +116,43 @@ func (srv *RttyServer) ListenAPI() error { return r.RunListener(ln) } -func callUserHookUrl(cfg *Config, c *gin.Context) bool { +func isLocalRequest(c *gin.Context) bool { + addr, _ := net.ResolveTCPAddr("tcp", c.Request.RemoteAddr) + return addr.IP.IsLoopback() +} + +type APIServer struct { + srv *RttyServer + sessions cache.ICache + root http.FileSystem + fh http.Handler + r *gin.Engine +} + +func (a *APIServer) auth(c *gin.Context) bool { + cfg := &a.srv.cfg + + if !cfg.LocalAuth && isLocalRequest(c) { + return true + } + + if cfg.Password == "" { + return true + } + + sid, err := c.Cookie("sid") + if err != nil || !a.sessions.Exists(sid) { + return false + } + + a.sessions.Expire(sid, httpSessionExpire) + + return true +} + +func (a *APIServer) callUserHookUrl(c *gin.Context) bool { + cfg := &a.srv.cfg + if cfg.UserHookUrl == "" { return true } @@ -350,30 +202,209 @@ func callUserHookUrl(cfg *Config, c *gin.Context) bool { return true } -func httpLogin(cfg *Config, password string) bool { - return cfg.Password == password -} - -func isLocalRequest(c *gin.Context) bool { - addr, _ := net.ResolveTCPAddr("tcp", c.Request.RemoteAddr) - return addr.IP.IsLoopback() -} - -func httpAuth(cfg *Config, c *gin.Context) bool { - if !cfg.LocalAuth && isLocalRequest(c) { - return true +func (a *APIServer) handleConnect(c *gin.Context) { + if !a.callUserHookUrl(c) { + c.Status(http.StatusForbidden) + return } - if cfg.Password == "" { + if c.GetHeader("Upgrade") != "websocket" { + group := c.Query("group") + devid := c.Param("devid") + if dev := a.srv.GetDevice(group, devid); dev == nil { + c.Redirect(http.StatusFound, "/error/offline") + return + } + + url := "/rtty/" + devid + + if group != "" { + url += "?group=" + group + } + + c.Redirect(http.StatusFound, url) + } else { + handleUserConnection(a.srv, c) + } +} + +func (a *APIServer) handleCounts(c *gin.Context) { + count := 0 + + a.srv.groups.Range(func(key, value any) bool { + count += int(value.(*DeviceGroup).count.Load()) return true + }) + + c.JSON(http.StatusOK, gin.H{"count": count}) +} + +func (a *APIServer) handleGroups(c *gin.Context) { + groups := []string{""} + + a.srv.groups.Range(func(key, value any) bool { + if key != "" { + groups = append(groups, key.(string)) + } + return true + }) + + c.JSON(http.StatusOK, groups) +} + +func (a *APIServer) handleDevs(c *gin.Context) { + devs := make([]*DeviceInfo, 0) + g := a.srv.GetGroup(c.Query("group"), false) + + if g == nil { + c.JSON(http.StatusOK, devs) + return } + g.devices.Range(func(key, value any) bool { + dev := value.(*Device) + + devs = append(devs, &DeviceInfo{ + Group: dev.group, + ID: dev.id, + Desc: dev.desc, + Connected: uint32(time.Now().Unix() - dev.timestamp), + Uptime: dev.uptime, + Proto: dev.proto, + IPaddr: dev.conn.RemoteAddr().(*net.TCPAddr).IP.String(), + }) + + return true + }) + + c.JSON(http.StatusOK, devs) +} + +func (a *APIServer) handleDev(c *gin.Context) { + if dev := a.srv.GetDevice(c.Query("group"), c.Param("devid")); dev != nil { + info := &DeviceInfo{ + ID: dev.id, + Desc: dev.desc, + Connected: uint32(time.Now().Unix() - dev.timestamp), + Uptime: dev.uptime, + Proto: dev.proto, + IPaddr: dev.conn.RemoteAddr().(*net.TCPAddr).IP.String(), + } + c.JSON(http.StatusOK, info) + } else { + c.Status(http.StatusNotFound) + } +} + +func (a *APIServer) handleCmd(c *gin.Context) { + if !a.callUserHookUrl(c) { + c.Status(http.StatusForbidden) + return + } + + cmdInfo := &CommandReqInfo{} + + err := c.BindJSON(&cmdInfo) + if err != nil || cmdInfo.Cmd == "" || cmdInfo.Username == "" { + cmdErrResp(c, rttyCmdErrInvalid) + return + } + + dev := a.srv.GetDevice(c.Query("group"), c.Param("devid")) + if dev == nil { + cmdErrResp(c, rttyCmdErrOffline) + return + } + + dev.handleCmdReq(c, cmdInfo) +} + +func (a *APIServer) handleWeb(c *gin.Context) { + httpProxyRedirect(a, c, "") +} + +func (a *APIServer) handleWeb2(c *gin.Context) { + group := c.Param("group") + httpProxyRedirect(a, c, group) +} + +func (a *APIServer) handleSignout(c *gin.Context) { sid, err := c.Cookie("sid") - if err != nil || !httpSessions.Exists(sid) { - return false + if err != nil || !a.sessions.Exists(sid) { + return } - httpSessions.Expire(sid, httpSessionExpire) + a.sessions.Del(sid) - return true + c.Status(http.StatusOK) +} + +func (a *APIServer) handleSignin(c *gin.Context) { + cfg := &a.srv.cfg + + type credentials struct { + Password string `json:"password"` + } + + creds := credentials{} + + err := c.BindJSON(&creds) + if err != nil { + c.Status(http.StatusBadRequest) + return + } + + if cfg.Password == creds.Password { + sid := utils.GenUniqueID() + + a.sessions.Set(sid, true, cache.WithEx(httpSessionExpire)) + + c.SetCookie("sid", sid, 0, "", "", false, true) + c.Status(http.StatusOK) + return + } + + c.Status(http.StatusUnauthorized) +} + +func (a *APIServer) handleAlive(c *gin.Context) { + if !a.auth(c) { + c.AbortWithStatus(http.StatusUnauthorized) + } else { + c.Status(http.StatusOK) + } +} + +func (a *APIServer) handleFile(c *gin.Context) { + upath := path.Clean(c.Request.URL.Path) + root := a.root + + if strings.HasSuffix(upath, ".js") || strings.HasSuffix(upath, ".css") { + if strings.Contains(c.Request.Header.Get("Accept-Encoding"), "gzip") { + f, err := root.Open(upath + ".gz") + if err == nil { + f.Close() + + c.Request.URL.Path += ".gz" + + if strings.HasSuffix(upath, ".js") { + c.Writer.Header().Set("Content-Type", "application/javascript") + } else if strings.HasSuffix(upath, ".css") { + c.Writer.Header().Set("Content-Type", "text/css") + } + + c.Writer.Header().Set("Content-Encoding", "gzip") + } + } + } else if upath != "/" { + f, err := root.Open(upath) + if err != nil { + c.Request.URL.Path = "/" + a.r.HandleContext(c) + return + } + defer f.Close() + } + + a.fh.ServeHTTP(c.Writer, c.Request) } diff --git a/http.go b/http.go index 22a47fc..253ca27 100644 --- a/http.go +++ b/http.go @@ -201,14 +201,16 @@ func doHttpProxy(srv *RttyServer, c net.Conn) { } } -func httpProxyRedirect(srv *RttyServer, c *gin.Context, group string) { +func httpProxyRedirect(a *APIServer, c *gin.Context, group string) { + srv := a.srv cfg := &srv.cfg + devid := c.Param("devid") proto := c.Param("proto") addr := c.Param("addr") rawPath := c.Param("path") - if !callUserHookUrl(cfg, c) { + if !a.callUserHookUrl(c) { c.Status(http.StatusForbidden) return } diff --git a/user.go b/user.go index 87aa30c..46000ad 100644 --- a/user.go +++ b/user.go @@ -117,18 +117,68 @@ func handleUserConnection(srv *RttyServer, c *gin.Context) { defer cancel() - if !waitForLogin(user, dev, ctx, sid) { + if !user.waitForLogin(dev, ctx, sid) { return } + user.handleMsg() +} + +func (user *User) SendCloseMsg(code int, text string) { + user.conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(code, text), time.Now().Add(time.Second)) +} + +func (user *User) Close() { + user.close.Do(func() { + dev := user.dev + sid := user.sid + + user.closed.Store(true) + + if _, loaded := dev.users.LoadAndDelete(sid); loaded { + dev.WriteMsg(msgTypeLogout, sid, nil) + } + + dev.pending.Delete(sid) + user.conn.Close() + + log.Debug().Msgf("user with session '%s' closed", sid) + }) +} + +func (user *User) WriteMsg(typ int, data []byte) error { + return user.conn.WriteMessage(typ, data) +} + +func (user *User) waitForLogin(dev *Device, ctx context.Context, sid string) bool { for { - msgType, data, err := conn.ReadMessage() + select { + case <-ctx.Done(): + return false + + case ok := <-user.pending: + return ok + + case <-time.After(TermLoginTimeout): + if _, loaded := dev.pending.LoadAndDelete(sid); loaded { + log.Error().Msgf("login timeout for session %s of device %s", sid, dev.id) + user.SendCloseMsg(LoginErrorTimeout, "login timeout") + return false + } + } + } +} + +func (user *User) handleMsg() { + dev := user.dev + sid := user.sid + + for { + msgType, data, err := user.conn.ReadMessage() if err != nil { if !user.closed.Load() { closeError, ok := err.(*websocket.CloseError) - if !ok || (closeError.Code != websocket.CloseGoingAway && - closeError.Code != websocket.CloseAbnormalClosure && - closeError.Code != websocket.CloseNormalClosure) { + if !ok || ignoredWsCloseError(closeError.Code) { log.Error().Msgf("user read fail: %v", err) } } @@ -192,47 +242,8 @@ func handleUserConnection(srv *RttyServer, c *gin.Context) { } } -func (user *User) SendCloseMsg(code int, text string) { - user.conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(code, text), time.Now().Add(time.Second)) -} - -func (user *User) Close() { - user.close.Do(func() { - dev := user.dev - sid := user.sid - - user.closed.Store(true) - - if _, loaded := dev.users.LoadAndDelete(sid); loaded { - dev.WriteMsg(msgTypeLogout, sid, nil) - } - - dev.pending.Delete(sid) - user.conn.Close() - - log.Debug().Msgf("user with session '%s' closed", sid) - }) -} - -func (user *User) WriteMsg(typ int, data []byte) error { - return user.conn.WriteMessage(typ, data) -} - -func waitForLogin(user *User, dev *Device, ctx context.Context, sid string) bool { - for { - select { - case <-ctx.Done(): - return false - - case ok := <-user.pending: - return ok - - case <-time.After(TermLoginTimeout): - if _, loaded := dev.pending.LoadAndDelete(sid); loaded { - log.Error().Msgf("login timeout for session %s of device %s", sid, dev.id) - user.SendCloseMsg(LoginErrorTimeout, "login timeout") - return false - } - } - } +func ignoredWsCloseError(code int) bool { + return code != websocket.CloseGoingAway && + code != websocket.CloseAbnormalClosure && + code != websocket.CloseNormalClosure }