feat(rmi): support receive rmi connection (#5)

pkg/dns/dns.go

@@ -82,7 +82,7 @@ func (s *Server) Run() {
 						log.Error("DNS record(rule_id:%s) created failed :%s", _rule.Name, err.Error())
 						return nil, nil
 					}
-					log.Trace("DNS record(id:%d) has been created", r.ID)
+					log.Info("DNS record(id:%d,rule:%s,remote_ip:%s) has been created", r.ID, _rule.Name, ip)
 
 					//only send to client when this connection recorded first time.
 					if _rule.PushToClient {

pkg/dns/record.go

@@ -42,7 +42,7 @@ func newRecord(rule *Rule, flag, domain, remoteIp, ipArea string) (r *Record, er
 	return r, err
 }
 
-func List(c *gin.Context) {
+func ListRecords(c *gin.Context) {
 	var (
 		dnsRecord Record
 		res       []Record
@@ -80,6 +80,11 @@ func List(c *gin.Context) {
 		})
 		return
 	}
+
+	if order != "desc" && order != "asc" {
+		order = "desc"
+	}
+
 	if err := db.Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil {
 		c.JSON(400, gin.H{
 			"status": "failed",

pkg/dns/rule.go

@@ -107,6 +107,10 @@ func ListRules(c *gin.Context) {
 		return
 	}
 
+	if order != "desc" && order != "asc" {
+		order = "desc"
+	}
+
 	if err := db.Order("rank desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil {
 		c.JSON(400, gin.H{
 			"status": "failed",

pkg/mysql/mysql.go

@@ -125,7 +125,7 @@ func (s *Server) ConnectionClosed(c *vmysql.Conn) {
 		log.Error("MySQL record(rule_id:%s) created failed :%s", _rule.Name, err.Error())
 		return
 	}
-	log.Trace("MySQL record(id:%d) has been created", r.ID)
+	log.Info("MySQL record(id:%d,rule:%s,remote_ip:%s) has been created", r.ID, _rule.Name, ip)
 
 	//only send to client when this connection recorded first time.
 	if _rule.PushToClient {
@@ -267,10 +267,10 @@ func (s *Server) Run() {
 	var authServer = &vmysql.AuthServerNone{}
 	var err error
 
-	log.Info("Starting Mysql Server at %s", s.Addr)
+	log.Info("Starting MySQL Server at %s", s.Addr)
 	s.listener, err = vmysql.NewListener("tcp", s.Addr, authServer, s, s.VersionString, 0, 0)
 	if err != nil {
-		log.Error("New Mysql Server failed: %s", err)
+		log.Error("New MySQL Server failed: %s", err)
 		os.Exit(-1)
 	}
 

pkg/mysql/record.go

@@ -49,13 +49,14 @@ func newRecord(rule *Rule, flag, username, clientName, clientOS, remoteIp, ipAre
 	return r, err
 }
 
-func List(c *gin.Context) {
+func ListRecords(c *gin.Context) {
 	var (
 		mysqlRecord Record
 		res         []Record
 		count       int64
 		order       = c.Query("order")
 	)
+
 	if err := c.ShouldBind(&mysqlRecord); err != nil {
 		c.JSON(400, gin.H{
 			"status": "failed",
@@ -90,6 +91,11 @@ func List(c *gin.Context) {
 		})
 		return
 	}
+
+	if order != "desc" && order != "asc" {
+		order = "desc"
+	}
+
 	if err := db.Preload("Files").Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil {
 		c.JSON(400, gin.H{
 			"status": "failed",

pkg/mysql/rule.go

@@ -89,6 +89,11 @@ func ListRules(c *gin.Context) {
 		})
 		return
 	}
+
+	if order != "desc" && order != "asc" {
+		order = "desc"
+	}
+
 	if err := db.Order("rank desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil {
 		c.JSON(400, gin.H{
 			"status": "failed",

pkg/rhttp/http.go

@@ -157,7 +157,7 @@ func (s *Server) Receive(c *gin.Context) {
 			c.String(code, compileTpl(c, _rule.ResponseBody))
 			return
 		}
-		log.Trace("HTTP record(id:%d) has been created", r.ID)
+		log.Info("HTTP record(id:%d,rule:%s,remote_ip:%s) has been created", r.ID, _rule.Name, ip)
 
 		//only send to client when this connection recorded first time.
 		if _rule.PushToClient {

pkg/rhttp/record.go

@@ -89,6 +89,10 @@ func ListRecords(c *gin.Context) {
 		return
 	}
 
+	if order != "desc" && order != "asc" {
+		order = "desc"
+	}
+
 	if err := db.Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil {
 		c.JSON(400, gin.H{
 			"status": "failed",

pkg/rhttp/rule.go

@@ -108,6 +108,10 @@ func ListRules(c *gin.Context) {
 		return
 	}
 
+	if order != "desc" && order != "asc" {
+		order = "desc"
+	}
+
 	if err := db.Order("rank desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil {
 		c.JSON(400, gin.H{
 			"status": "failed",

pkg/rmi/config.go

@@ -0,0 +1,6 @@
+package rmi
+
+type Config struct {
+	Enable bool
+	Addr   string
+}

pkg/rmi/record.go

@@ -0,0 +1,103 @@
+package rmi
+
+import (
+	"strconv"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/li4n0/revsuit/internal/database"
+	"github.com/li4n0/revsuit/internal/notice"
+	"github.com/li4n0/revsuit/internal/record"
+)
+
+var _ record.Record = (*Record)(nil)
+
+type Record struct {
+	Path string `form:"path" json:"path"`
+	record.BaseRecord
+	Rule Rule `gorm:"foreignKey:RuleName;references:Name;constraint:OnUpdate:CASCADE,OnDelete:SET NULL;" form:"-" json:"-" notice:"-"`
+}
+
+func (Record) TableName() string {
+	return "rmi_records"
+}
+
+func (r Record) Notice() {
+	notice.Notice(r)
+}
+
+func NewRecord(rule *Rule, flag, path, ip, area string) (r *Record, err error) {
+	r = &Record{
+		BaseRecord: record.BaseRecord{
+			Flag:        flag,
+			RemoteIP:    ip,
+			IpArea:      area,
+			RequestTime: time.Now(),
+		},
+		Path: path,
+		Rule: *rule,
+	}
+	err = database.DB.Create(r).Error
+	return r, err
+}
+
+func ListRecords(c *gin.Context) {
+	var (
+		rmiRecord Record
+		res       []Record
+		count     int64
+		order     = c.Query("order")
+	)
+
+	if err := c.ShouldBind(&rmiRecord); err != nil {
+		c.JSON(400, gin.H{
+			"status": "failed",
+			"error":  err,
+			"result": nil,
+		})
+		return
+	}
+
+	db := database.DB.Model(&rmiRecord)
+	if rmiRecord.Flag != "" {
+		db.Where("flag = ?", rmiRecord.Flag)
+	}
+	if rmiRecord.Path != "" {
+		db.Where("path like ?", "%"+rmiRecord.Path+"%")
+	}
+	if rmiRecord.RemoteIP != "" {
+		db.Where("remote_ip = ?", rmiRecord.RemoteIP)
+	}
+	if rmiRecord.RuleName != "" {
+		db.Where("rule_name = ?", rmiRecord.RuleName)
+	}
+
+	page, err := strconv.Atoi(c.Query("page"))
+	if err != nil {
+		c.JSON(400, gin.H{
+			"status": "failed",
+			"error":  err.Error(),
+			"result": nil,
+		})
+		return
+	}
+
+	if order != "desc" && order != "asc" {
+		order = "desc"
+	}
+
+	if err := db.Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil {
+		c.JSON(400, gin.H{
+			"status": "failed",
+			"error":  err.Error(),
+			"data":   nil,
+		})
+		return
+	}
+
+	c.JSON(200, gin.H{
+		"status": "succeed",
+		"error":  nil,
+		"result": gin.H{"count": count, "data": res},
+	})
+}

pkg/rmi/rmi.go

@@ -0,0 +1,158 @@
+package rmi
+
+import (
+	"bytes"
+	"encoding/binary"
+	"net"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/li4n0/revsuit/internal/database"
+	"github.com/li4n0/revsuit/internal/qqwry"
+	log "unknwon.dev/clog/v2"
+)
+
+type Server struct {
+	Config
+	rules     []*Rule
+	rulesLock sync.RWMutex
+}
+
+var (
+	server *Server
+	once   sync.Once
+)
+
+func GetServer() *Server {
+	once.Do(func() {
+		server = &Server{rulesLock: sync.RWMutex{}}
+	})
+	return server
+}
+
+func (s *Server) getRules() []*Rule {
+	defer s.rulesLock.RUnlock()
+	s.rulesLock.RLock()
+	return s.rules
+}
+
+func (s *Server) updateRules() error {
+	db := database.DB.Model(new(Rule))
+	s.rulesLock.Lock()
+	db.Order("rank desc").Find(&s.rules)
+	s.rulesLock.Unlock()
+	return nil
+}
+
+func (s *Server) handleConnection(conn net.Conn) {
+	defer conn.Close()
+
+	ip, port, _ := net.SplitHostPort(conn.RemoteAddr().String())
+	if err := conn.SetDeadline(time.Now().Add(time.Second * 30)); err != nil {
+		log.Error("RMI set connection deadline error:%v", err.Error())
+	}
+
+	buf := make([]byte, 1024)
+	_, err := conn.Read(buf)
+	if err != nil {
+		log.Error("RMI read connection error:%v", err.Error())
+	}
+
+	if !bytes.Contains(buf, []byte{0x4a, 0x52, 0x4d, 0x49}) {
+		return
+	}
+
+	send := []byte{0x4e}
+	bs := make([]byte, 8)
+	binary.BigEndian.PutUint16(bs, uint16(len(ip)))
+	send = append(send, bs...)
+	send = append(send, []byte(ip)...)
+	send = append(send, []byte{0x00, 0x00}...)
+	uintPort, _ := strconv.Atoi(port)
+	bs = make([]byte, 8)
+	binary.BigEndian.PutUint16(bs, uint16(uintPort))
+	send = append(send, bs...)
+
+	_, err = conn.Write(send)
+	if err != nil {
+		log.Error("RMI write connection error: %v", err.Error())
+	}
+
+	data := make([]byte, 512)
+
+	for length := 0; length < 50; {
+		n, err := conn.Read(data)
+		if err != nil {
+			log.Error("RMI read connection error: %v", err.Error())
+		}
+		length += n
+	}
+
+	frags := bytes.Split(data, []byte{0xdf, 0x74})
+	path := strings.TrimRight(string(frags[len(frags)-1][2:]), "\x00")
+
+	for _, _rule := range s.getRules() {
+		flag, flagGroup := _rule.Match(path)
+		if flag == "" {
+			continue
+		}
+
+		area := qqwry.Area(ip)
+
+		// create new record
+		r, err := NewRecord(_rule, flag, path, ip, area)
+		if err != nil {
+			log.Error("RMI record(rule_id:%d) created failed :%s", _rule.ID, err.Error())
+			return
+		}
+		log.Info("RMI record(id:%d,rule:%s,remote_ip:%s) has been created", r.ID, _rule.Name, ip)
+
+		//only send to client when this connection recorded first time.
+		if _rule.PushToClient {
+			if flagGroup != "" {
+				var count int64
+				database.DB.Where("rule_name=? and raw like ?", _rule.Name, "%"+flagGroup+"%").Model(&Record{}).Count(&count)
+				if count <= 1 {
+					r.PushToClient()
+					log.Trace("RMI record(id:%d) has been put to client message queue", r.ID)
+				}
+			}
+			r.PushToClient()
+			log.Trace("RMI record(id:%d) has been put to client message queue", r.ID)
+		}
+
+		//send notice
+		if _rule.Notice {
+			go func() {
+				r.Notice()
+				log.Trace("RMI record(id:%d) notice has been sent", r.ID)
+			}()
+		}
+	}
+}
+
+func (s *Server) Run() {
+	if err := s.updateRules(); err != nil {
+		log.Fatal(err.Error())
+	}
+
+	// run server
+	log.Info("Starting RMI Server at %v", s.Addr)
+
+	listener, err := net.Listen("tcp", s.Addr)
+	if err != nil {
+		log.Fatal(err.Error())
+	}
+
+	for {
+		tcpConn, err := listener.Accept()
+		if err != nil {
+			log.Error("RMI accept connection error: %v", err.Error())
+			continue
+		}
+		go s.handleConnection(tcpConn)
+	}
+
+}

pkg/rmi/rule.go

@@ -0,0 +1,193 @@
+package rmi
+
+import (
+	"strconv"
+
+	"github.com/gin-gonic/gin"
+	"github.com/li4n0/revsuit/internal/database"
+	"github.com/li4n0/revsuit/internal/rule"
+	"gorm.io/gorm/clause"
+	log "unknwon.dev/clog/v2"
+)
+
+// Http rule struct
+type Rule struct {
+	rule.BaseRule
+}
+
+func (Rule) TableName() string {
+	return "rmi_rules"
+}
+
+// New rmi rule struct
+func NewRule(name, flagFormat string, pushToClient, notice bool) *Rule {
+	return &Rule{
+		BaseRule: rule.BaseRule{
+			Name:         name,
+			FlagFormat:   flagFormat,
+			PushToClient: pushToClient,
+			Notice:       notice,
+		},
+	}
+}
+
+// Create or update the rmi rule in database and ruleSet
+func (r *Rule) CreateOrUpdate() (err error) {
+	db := database.DB.Model(r)
+	err = db.Clauses(clause.OnConflict{
+		Columns: []clause.Column{{Name: "id"}},
+		DoUpdates: clause.AssignmentColumns(
+			[]string{
+				"name",
+				"flag_format",
+				"rank",
+				"push_to_client",
+				"notice",
+			}),
+	}).Create(r).Error
+	if err != nil {
+		return
+	}
+
+	err = GetServer().updateRules()
+	return err
+}
+
+// Delete the rmi rule in database and ruleSet
+func (r *Rule) Delete() (err error) {
+	db := database.DB.Model(r)
+	err = db.Delete(r).Error
+	if err != nil {
+		return
+	}
+
+	err = GetServer().updateRules()
+	return err
+}
+
+// List all rmi rules those satisfy the filter
+func ListRules(c *gin.Context) {
+	var (
+		rmiRule Rule
+		res     []Rule
+		count   int64
+		order   = c.Query("order")
+	)
+
+	if err := c.ShouldBind(&rmiRule); err != nil {
+		c.JSON(400, gin.H{
+			"status": "failed",
+			"error":  err,
+			"result": nil,
+		})
+		return
+	}
+
+	db := database.DB.Model(&rmiRule)
+	if rmiRule.Name != "" {
+		db.Where("name = ?", rmiRule.Name)
+	}
+	db.Count(&count)
+
+	page, err := strconv.Atoi(c.Query("page"))
+	if err != nil {
+		c.JSON(400, gin.H{
+			"status": "failed",
+			"error":  err.Error(),
+			"result": nil,
+		})
+		return
+	}
+
+	if order != "desc" && order != "asc" {
+		order = "desc"
+	}
+
+	if err := db.Order("rank desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil {
+		c.JSON(400, gin.H{
+			"status": "failed",
+			"error":  err.Error(),
+			"data":   nil,
+		})
+		return
+	}
+
+	c.JSON(200, gin.H{
+		"status": "succeed",
+		"error":  nil,
+		"result": gin.H{"count": count, "data": res},
+	})
+}
+
+// Create or update rmi rule from user submit
+func UpsertRules(c *gin.Context) {
+	var (
+		rmiRule Rule
+		update  bool
+	)
+
+	if err := c.ShouldBind(&rmiRule); err != nil {
+		c.JSON(400, gin.H{
+			"status": "failed",
+			"error":  err.Error(),
+			"data":   nil,
+		})
+		return
+	}
+
+	if rmiRule.ID != 0 {
+		update = true
+	}
+
+	if err := rmiRule.CreateOrUpdate(); err != nil {
+		c.JSON(400, gin.H{
+			"status": "failed",
+			"error":  err.Error(),
+			"result": nil,
+		})
+		return
+	}
+
+	if update {
+		log.Trace("RMI rule(id:%d) has been updated", rmiRule.ID)
+	} else {
+		log.Trace("RMI rule(id:%d) has been created", rmiRule.ID)
+	}
+
+	c.JSON(200, gin.H{
+		"status": "succeed",
+		"error":  nil,
+		"result": nil,
+	})
+}
+
+// Delete rmi rule from user submit
+func DeleteRules(c *gin.Context) {
+	var rmiRule Rule
+
+	if err := c.ShouldBind(&rmiRule); err != nil {
+		c.JSON(400, gin.H{
+			"status": "failed",
+			"error":  err.Error(),
+			"data":   nil,
+		})
+		return
+	}
+
+	if err := rmiRule.Delete(); err != nil {
+		c.JSON(400, gin.H{
+			"status": "failed",
+			"error":  err.Error(),
+			"data":   nil,
+		})
+		return
+	}
+
+	log.Trace("RMI rule(id:%d) has been deleted", rmiRule.ID)
+
+	c.JSON(200, gin.H{
+		"status": "succeed",
+		"error":  nil,
+		"data":   nil,
+	})
+}

pkg/server/config.go

@@ -4,6 +4,7 @@ import (
 	"github.com/li4n0/revsuit/pkg/dns"
 	"github.com/li4n0/revsuit/pkg/mysql"
 	"github.com/li4n0/revsuit/pkg/rhttp"
+	"github.com/li4n0/revsuit/pkg/rmi"
 )
 
 type noticeConfig struct {
@@ -21,5 +22,6 @@ type Config struct {
 	Notice   noticeConfig
 	rhttp.Config
 	DNS   dns.Config
-	Mysql mysql.Config
+	MySQL mysql.Config
+	RMI   rmi.Config
 }

pkg/server/router.go

@@ -9,6 +9,7 @@ import (
 	"github.com/li4n0/revsuit/pkg/dns"
 	"github.com/li4n0/revsuit/pkg/mysql"
 	"github.com/li4n0/revsuit/pkg/rhttp"
+	"github.com/li4n0/revsuit/pkg/rmi"
 	log "unknwon.dev/clog/v2"
 )
 
@@ -52,10 +53,13 @@ func (revsuit *Revsuit) registerHttpRouter() {
 	httpGroup.GET("", rhttp.ListRecords)
 
 	dnsGroup := recordGroup.Group("/dns")
-	dnsGroup.GET("", dns.List)
+	dnsGroup.GET("", dns.ListRecords)
 
 	mysqlGroup := recordGroup.Group("/mysql")
-	mysqlGroup.GET("", mysql.List)
+	mysqlGroup.GET("", mysql.ListRecords)
+
+	rmiGroup := recordGroup.Group("/rmi")
+	rmiGroup.GET("", rmi.ListRecords)
 
 	// init rule router group
 	ruleGroup := revsuit.http.ApiGroup.Group("/rule")
@@ -75,6 +79,11 @@ func (revsuit *Revsuit) registerHttpRouter() {
 	mysqlGroup.POST("", mysql.UpsertRules)
 	mysqlGroup.DELETE("", mysql.DeleteRules)
 
+	rmiGroup = ruleGroup.Group("/rmi")
+	rmiGroup.GET("", rmi.ListRules)
+	rmiGroup.POST("", rmi.UpsertRules)
+	rmiGroup.DELETE("", rmi.DeleteRules)
+
 	// init file router group
 	fileGroup := revsuit.http.ApiGroup.Group("/file")
 	fileGroup.GET("/mysql/:id", mysql.GetFile)

pkg/server/server.go

@@ -7,6 +7,7 @@ import (
 	"github.com/li4n0/revsuit/pkg/dns"
 	"github.com/li4n0/revsuit/pkg/mysql"
 	http "github.com/li4n0/revsuit/pkg/rhttp"
+	"github.com/li4n0/revsuit/pkg/rmi"
 	"gorm.io/gorm/logger"
 	log "unknwon.dev/clog/v2"
 )
@@ -15,6 +16,7 @@ type Revsuit struct {
 	http  *http.Server
 	dns   *dns.Server
 	mysql *mysql.Server
+	rmi   *rmi.Server
 }
 
 func initDatabase(dsn string) {
@@ -51,6 +53,14 @@ func initDatabase(dsn string) {
 	if err != nil {
 		log.Fatal(err.Error())
 	}
+	err = database.DB.AutoMigrate(&rmi.Record{})
+	if err != nil {
+		log.Fatal(err.Error())
+	}
+	err = database.DB.AutoMigrate(&rmi.Rule{})
+	if err != nil {
+		log.Fatal(err.Error())
+	}
 
 }
 
@@ -122,9 +132,13 @@ func New(c *Config) *Revsuit {
 	if c.DNS.Enable {
 		s.dns = dns.GetServer()
 	}
-	if c.Mysql.Enable {
+	if c.MySQL.Enable {
 		s.mysql = mysql.GetServer()
-		s.mysql.Config = c.Mysql
+		s.mysql.Config = c.MySQL
+	}
+	if c.RMI.Enable {
+		s.rmi = rmi.GetServer()
+		s.rmi.Config = c.RMI
 	}
 
 	if c.Addr != "" {
@@ -149,6 +163,9 @@ func (revsuit *Revsuit) Run() {
 	if revsuit.mysql != nil {
 		go revsuit.mysql.Run()
 	}
+	if revsuit.rmi != nil {
+		go revsuit.rmi.Run()
+	}
 
 	revsuit.http.Run()
 }