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 <zhaojh329@gmail.com>

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)
 }

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
 	}

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
 }