release 0.1.3-beta (#19)

Co-authored-by: E99p1ant <[email protected]>

pkg/server/controller.go

@@ -11,7 +11,6 @@ import (
 	"time"
 
 	"github.com/gin-gonic/gin"
-	"github.com/li4n0/revsuit/internal/record"
 	"github.com/li4n0/revsuit/internal/recycler"
 	log "unknwon.dev/clog/v2"
 )
@@ -26,18 +25,15 @@ func ping(c *gin.Context) {
 	c.String(200, "pong")
 }
 
-func events(c *gin.Context) {
-	log.Info("Receive client connection from %v", c.Request.RemoteAddr)
+func (revsuit *Revsuit) events(c *gin.Context) {
+	id := revsuit.addClient(c)
+	log.Info("Receive client[id:%d] connection from %v", id, c.Request.RemoteAddr)
 	c.Stream(func(w io.Writer) bool {
-		select {
-		case <-c.Writer.CloseNotify():
-			return false
-		case r := <-record.Channel():
-			c.SSEvent("message", r.GetFlag())
-		}
-		return true
+		<-c.Writer.CloseNotify()
+		return false
 	})
-	log.Info("Client %s disconnect", c.Request.RemoteAddr)
+	revsuit.removeClient(id)
+	log.Info("Client[id:%d, remote_addr:%s] disconnect", id, c.Request.RemoteAddr)
 }
 
 func recovery(c *gin.Context) {

pkg/server/router.go

@@ -40,7 +40,7 @@ func (revsuit *Revsuit) registerPlatformRouter() {
 
 	//platform routers
 	api.GET("/auth", auth)
-	api.GET("/events", events)
+	api.GET("/events", revsuit.events)
 	api.GET("/ping", ping)
 	api.GET("/version", version)
 }
@@ -58,7 +58,7 @@ func (revsuit *Revsuit) registerHttpRouter() {
 	// init settings router group
 	settingsGroup := revsuit.http.ApiGroup.Group("setting")
 	settingsGroup.GET("/exportRules", exportRules)
-	settingsGroup.POST("/importRules", importRules)
+	settingsGroup.POST("/importRules", revsuit.importRules)
 	settingsGroup.GET("/getPlatformConfig", revsuit.getPlatformConfig)
 	settingsGroup.POST("/updatePlatformConfig", revsuit.updatePlatformConfig)
 	settingsGroup.GET("/getDnsConfig", revsuit.getDnsConfig)

pkg/server/server.go

@@ -1,10 +1,13 @@
 package server
 
 import (
+	"sync"
+
 	"github.com/gin-gonic/gin"
 	"github.com/li4n0/revsuit/internal/database"
 	"github.com/li4n0/revsuit/internal/file"
 	"github.com/li4n0/revsuit/internal/notice"
+	"github.com/li4n0/revsuit/internal/record"
 	"github.com/li4n0/revsuit/pkg/dns"
 	"github.com/li4n0/revsuit/pkg/ftp"
 	"github.com/li4n0/revsuit/pkg/mysql"
@@ -14,7 +17,7 @@ import (
 	log "unknwon.dev/clog/v2"
 )
 
-const VERSION = "0.1.2-beta"
+const VERSION = "0.1.3-beta"
 
 type Revsuit struct {
 	config   *Config
@@ -25,6 +28,29 @@ type Revsuit struct {
 	mysql *mysql.Server
 	rmi   *rmi.Server
 	ftp   *ftp.Server
+
+	clients     map[int]*gin.Context
+	clientID    int
+	clientsLock sync.RWMutex
+	clientsNum  chan struct{}
+}
+
+func (revsuit *Revsuit) addClient(c *gin.Context) int {
+	revsuit.clientsLock.Lock()
+	defer revsuit.clientsLock.Unlock()
+
+	revsuit.clientID++
+	revsuit.clients[revsuit.clientID] = c
+	revsuit.clientsNum <- struct{}{}
+	return revsuit.clientID
+}
+
+func (revsuit *Revsuit) removeClient(id int) {
+	revsuit.clientsLock.Lock()
+	defer revsuit.clientsLock.Unlock()
+
+	delete(revsuit.clients, id)
+	<-revsuit.clientsNum
 }
 
 func initDatabase(dsn string) {
@@ -188,6 +214,8 @@ func New(c *Config) *Revsuit {
 		s.dns.SetServerIP(c.ExternalIP)
 		s.ftp.SetPasvIP(c.ExternalIP)
 	}
+	s.clients = make(map[int]*gin.Context)
+	s.clientsNum = make(chan struct{}, 100)
 	return s
 }
 
@@ -207,6 +235,17 @@ func (revsuit *Revsuit) Run() {
 	if revsuit.ftp != nil && revsuit.ftp.Enable {
 		go revsuit.ftp.Run()
 	}
-
+	go func() {
+		for r := range record.Channel() {
+			<-revsuit.clientsNum
+			revsuit.clientsLock.RLock()
+			for _, client := range revsuit.clients {
+				client.SSEvent("message", r.GetFlag())
+				client.Writer.Flush()
+			}
+			revsuit.clientsNum <- struct{}{}
+			revsuit.clientsLock.RUnlock()
+		}
+	}()
 	revsuit.http.Run()
 }

pkg/server/settings.go

@@ -45,7 +45,7 @@ func exportRules(c *gin.Context) {
 	c.String(200, string(out))
 }
 
-func importRules(c *gin.Context) {
+func (revsuit *Revsuit) importRules(c *gin.Context) {
 	var (
 		db    = database.DB
 		rules Rules
@@ -87,48 +87,58 @@ func importRules(c *gin.Context) {
 	}
 
 	for _, rule := range rules.Http {
-		err := db.Model(&rhttp.Rule{}).Create(&rule).Error
-		if err != nil {
+		if err := db.Model(&rhttp.Rule{}).Create(&rule).Error; err != nil {
 			errs = append(errs, errors.Wrap(err, fmt.Sprintf("http rule[%s]", rule.Name)).Error())
 			continue
 		}
-		count += 1
+		count++
+	}
+	if err := revsuit.http.UpdateRules(); err != nil {
+		errs = append(errs, err.Error())
 	}
 
 	for _, rule := range rules.Dns {
-		err := db.Model(&dns.Rule{}).Create(&rule).Error
-		if err != nil {
+		if err := db.Model(&dns.Rule{}).Create(&rule).Error; err != nil {
 			errs = append(errs, errors.Wrap(err, fmt.Sprintf("dns rule[%s]", rule.Name)).Error())
 			continue
 		}
-		count += 1
+		count++
+	}
+	if err := revsuit.dns.UpdateRules(); err != nil {
+		errs = append(errs, err.Error())
 	}
 
 	for _, rule := range rules.Mysql {
-		err := db.Model(&mysql.Rule{}).Create(&rule).Error
-		if err != nil {
+		if err := db.Model(&mysql.Rule{}).Create(&rule).Error; err != nil {
 			errs = append(errs, errors.Wrap(err, fmt.Sprintf("mysql rule[%s]", rule.Name)).Error())
 			continue
 		}
-		count += 1
+		count++
+	}
+	if err := revsuit.mysql.UpdateRules(); err != nil {
+		errs = append(errs, err.Error())
 	}
 
 	for _, rule := range rules.Rmi {
-		err := db.Model(&rmi.Rule{}).Create(&rule).Error
-		if err != nil {
+		if err := db.Model(&rmi.Rule{}).Create(&rule).Error; err != nil {
 			errs = append(errs, errors.Wrap(err, fmt.Sprintf("rmi rule[%s]", rule.Name)).Error())
 			continue
 		}
-		count += 1
+		count++
+	}
+	if err := revsuit.rmi.UpdateRules(); err != nil {
+		errs = append(errs, err.Error())
 	}
 
 	for _, rule := range rules.Ftp {
-		err := db.Model(&ftp.Rule{}).Create(&rule).Error
-		if err != nil {
+		if err := db.Model(&ftp.Rule{}).Create(&rule).Error; err != nil {
 			errs = append(errs, errors.Wrap(err, fmt.Sprintf("ftp rule[%s]", rule.Name)).Error())
 			continue
 		}
-		count += 1
+		count++
+	}
+	if err := revsuit.ftp.UpdateRules(); err != nil {
+		errs = append(errs, err.Error())
 	}
 
 	c.JSON(200, gin.H{