release revsuit vBeta0.1.0 (#15)

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

pkg/ftp/ftp.go

@@ -1,6 +1,7 @@
 package ftp
 
 import (
+	"bufio"
 	"bytes"
 	"fmt"
 	"go/types"
@@ -12,11 +13,10 @@ import (
 	"time"
 
 	"github.com/li4n0/revsuit/internal/database"
-	"github.com/li4n0/revsuit/internal/file"
-	"github.com/li4n0/revsuit/internal/qqwry"
 	"github.com/li4n0/revsuit/internal/recycler"
 	"github.com/li4n0/revsuit/internal/rule"
 	"github.com/patrickmn/go-cache"
+	"github.com/pkg/errors"
 	log "unknwon.dev/clog/v2"
 )
 
@@ -24,6 +24,7 @@ type Server struct {
 	Config
 	rules       []*Rule
 	rulesLock   sync.RWMutex
+	livingLock  sync.Mutex
 	dataChannel chan map[string]interface{}
 }
 
@@ -64,7 +65,12 @@ func (s *Server) updateRules() error {
 	db := database.DB.Model(new(Rule))
 	defer s.rulesLock.Unlock()
 	s.rulesLock.Lock()
-	return db.Order("rank desc").Find(&s.rules).Error
+	return errors.Wrap(db.Order("rank desc").Find(&s.rules).Error, "FTP update rules error")
+}
+
+func getClientPasvConnAddress(ip, port string) string {
+	dataPort, _ := strconv.Atoi(port)
+	return fmt.Sprintf("%s:%d", ip, dataPort+1)
 }
 
 func (s *Server) authenticate(user, password string) (_rule *Rule, flag, flagGroup string, vars map[string]string) {
@@ -99,7 +105,28 @@ func (s *Server) getPasvAddressFromCache(ip, pasvAddressTpl string) (pasvAddress
 	return pasvAddress
 }
 
+const (
+	NeedAccount             = "332 Need account for login.\r\n"
+	PasswordPlease          = "331 password please - version check\r\n"
+	PasswordError           = "331 please specify the password\r\n"
+	UserLogged              = "230 User logged in\r\n"
+	NoSuchFile              = "550 %s: No such file or directory.\r\n"
+	CommandNotFound         = "500 '%s': command not understood.\r\n"
+	EnteringPassiveMode     = "227 Entering Passive Mode (%s,%v,%d)\r\n"
+	OpeningBinaryMode       = "150 Opening BINARY mode data connection for '%s' (%d bytes).\r\n"
+	OpeningBinaryModeUpload = "150 Opening BINARY mode data connection for '%s'.\r\n"
+	TransferComplete        = "226 Transfer complete.\r\n"
+	Goodbye                 = "221 Goodbye.\r\n"
+	DirectoryChanged        = "250 Directory successfully changed.\r\n"
+	CurrentDirectory        = "257 \"%s\" is the current directory\r\n"
+)
+
 func (s *Server) handleConnection(conn net.Conn) {
+	defer func() {
+		if err := recover(); err != nil {
+			recycler.Recycle(err)
+		}
+	}()
 	log.Trace("New FTP connection from addr [%s]", conn.RemoteAddr())
 	defer func() {
 		_ = conn.Close()
@@ -117,9 +144,9 @@ func (s *Server) handleConnection(conn net.Conn) {
 	}
 
 	ip, port, _ := net.SplitHostPort(conn.RemoteAddr().String())
-	dataPort, _ := strconv.Atoi(port)
-	clientPasvConnAddress := fmt.Sprintf("%s:%d", ip, dataPort+1)
+	clientPasvConnAddress := getClientPasvConnAddress(ip, port)
 	buf := &bytes.Buffer{}
+	connBuf := bufio.NewWriter(conn)
 
 	var user, password, method, flag, flagGroup, pasvAddress, filename string
 	status := CRASHED
@@ -148,23 +175,23 @@ loop:
 			log.Trace("FTP connection[%s] exec command: %s", conn.RemoteAddr(), strings.TrimRight(buf.String(), "\r\n"))
 
 			if _rule == nil && cmd != "USER" && cmd != "PASS" {
-				_, _ = conn.Write([]byte("332 Need account for login.\r\n"))
+				_, _ = connBuf.WriteString(NeedAccount)
 				break loop
 			}
 
 			switch cmd {
 			case "USER":
 				user = args
-				_, _ = conn.Write([]byte("331 password please - version check\r\n"))
+				_, _ = connBuf.WriteString(PasswordPlease)
 			case "PASS":
 				password = args
 				if _rule, flag, flagGroup, vars = s.authenticate(user, password); _rule == nil {
-					_, _ = conn.Write([]byte("331 please specify the password\r\n"))
+					_, _ = connBuf.WriteString(PasswordError)
 					break loop
 				}
 
 				log.Trace("FTP connection[%s] matched rule[rule_name: %s, flag: %s]", conn.RemoteAddr(), _rule.Name, flag)
-				_, _ = conn.Write([]byte("230 User logged in\r\n"))
+				_, _ = connBuf.WriteString(UserLogged)
 
 				if pasvAddress = s.getPasvAddressFromCache(ip, _rule.PasvAddress); pasvAddress == "" {
 					pasvAddress = fmt.Sprintf("%s:%d", s.PasvIP, s.PasvPort)
@@ -174,13 +201,13 @@ loop:
 			case "SIZE":
 				path += strings.TrimLeft(args, "/")
 				if _rule == nil || isRedirect || len(_rule.Data) == 0 {
-					_, _ = conn.Write([]byte(fmt.Sprintf("550 %s: No such file or directory.\r\n", args)))
+					_, _ = connBuf.WriteString(fmt.Sprintf(NoSuchFile, args))
 					break
 				}
-				_, _ = conn.Write([]byte(fmt.Sprintf("213 %d\r\n", len(_rule.Data))))
+				_, _ = connBuf.WriteString(fmt.Sprintf("213 %d\r\n", len(_rule.Data)))
 			case "EPSV", "EPRT", "PORT":
 				// refuse to use EPSV/EPRT/PORT in order to make the client to use PASV mode.
-				_, _ = conn.Write([]byte(fmt.Sprintf("500 '%s': command not understood.\r\n", cmd)))
+				_, _ = connBuf.WriteString(fmt.Sprintf(CommandNotFound, cmd))
 			case "PASV":
 				//Just so that ide does not prompt that there may be a nil value
 				if _rule != nil {
@@ -188,17 +215,17 @@ loop:
 					pasvAddress := rule.CompileTpl(pasvAddress, vars)
 					pasvIP, pasvPort, err := net.SplitHostPort(pasvAddress)
 					if err != nil {
-						log.Warn("FTP failed to split rule[id%d] pasv_address(%s) :%s", _rule.ID, pasvAddress, err)
+						log.Warn("FTP failed to split rule[id:%d] pasv_address(%s) :%s", _rule.ID, pasvAddress, err)
 						break
 					}
 
 					port, err := strconv.Atoi(pasvPort)
 					if err != nil {
-						log.Warn("FTP failed to convert rule[id%d] pasv_port(%s) :%s", _rule.ID, pasvPort, err)
+						log.Warn("FTP failed to convert rule[id:%d] pasv_port(%s) :%s", _rule.ID, pasvPort, err)
 						break
 					}
-					ret := fmt.Sprintf("227 Entering Passive Mode (%s,%v,%d)\r\n", strings.ReplaceAll(pasvIP, ".", ","), float64(port/256), port%256)
-					_, _ = conn.Write([]byte(ret))
+					ret := fmt.Sprintf(EnteringPassiveMode, strings.ReplaceAll(pasvIP, ".", ","), float64(port/256), port%256)
+					_, _ = connBuf.WriteString(ret)
 					if isRedirect {
 						log.Trace("FTP connection[%s] will be redirect[pasv_address: %s]", conn.RemoteAddr(), pasvAddress)
 					}
@@ -210,16 +237,18 @@ loop:
 					method = DOWNLOAD
 
 					//send data to client
-					_, _ = conn.Write([]byte(fmt.Sprintf("150 Opening BINARY mode data connection for '%s' (%d bytes).\r\n", filename, len(_rule.Data))))
+					_, _ = connBuf.WriteString(fmt.Sprintf(OpeningBinaryMode, filename, len(_rule.Data)))
+					_ = connBuf.Flush()
 					s.dataChannel <- map[string]interface{}{clientPasvConnAddress: []byte(rule.CompileTpl(_rule.Data, vars))}
-					_, _ = conn.Write([]byte("226 Transfer complete.\r\n"))
+					_, _ = connBuf.WriteString(TransferComplete)
 				}
 
 			case "STOR":
 				filename = args
 				method = UPLOAD
 
-				_, _ = conn.Write([]byte(fmt.Sprintf("150 Opening BINARY mode data connection for '%s'.\r\n", filename)))
+				_, _ = connBuf.WriteString(fmt.Sprintf(OpeningBinaryModeUpload, filename))
+				_ = connBuf.Flush()
 				//only could read data send to local pasv server.
 				if !isRedirect {
 					dataChannel := make(chan []byte)
@@ -227,68 +256,36 @@ loop:
 					uploadData = <-dataChannel
 					log.Trace("FTP connection[%s] uploaded %d bytes", conn.RemoteAddr(), len(uploadData))
 				}
-				_, _ = conn.Write([]byte("226 Transfer complete.\r\n"))
+				_, _ = connBuf.WriteString(TransferComplete)
 			case "QUIT":
-				_, _ = conn.Write([]byte("221 Goodbye.\r\n"))
+				_, _ = connBuf.WriteString(Goodbye)
 				status = FINISHED
 				break loop
 			case "CWD":
-				_, _ = conn.Write([]byte("250 Directory successfully changed.\r\n"))
+				_, _ = connBuf.WriteString(DirectoryChanged)
 				path += strings.TrimRight(args, "\r\n") + "/"
 			case "PWD":
-				_, _ = conn.Write([]byte(fmt.Sprintf("257 \"%s\" is the current directory\r\n", path)))
+				_, _ = connBuf.WriteString(fmt.Sprintf(CurrentDirectory, path))
 			default:
 				_, _ = conn.Write([]byte("230 more data please!\r\n"))
 			}
+			_ = connBuf.Flush()
 		}
 		buf = &bytes.Buffer{}
 	}
 
 	if _rule != nil {
-		area := qqwry.Area(ip)
-		var r *Record
-		var err error
-		// create new record
-		ftpFile := &file.FTPFile{}
-		if len(uploadData) != 0 {
-			ftpFile = &file.FTPFile{
-				Name:    filename,
-				Content: uploadData,
-			}
-		}
-
-		r, err = NewRecord(_rule, flag, user, password, method, path, ip, area, ftpFile, status)
-		if err != nil {
-			log.Warn("FTP record[rule_id:%d] created failed :%s", _rule.ID, err)
-			return
-		}
-		log.Info("FTP 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("FTP record[id%d] has been put to client message queue", r.ID)
-				}
-			}
-			r.PushToClient()
-			log.Trace("FTP record[id%d] has been put to client message queue", r.ID)
-		}
-
-		//send notice
-		if _rule.Notice {
-			go func() {
-				r.Notice()
-				log.Trace("FTP record[id%d] notice has been sent", r.ID)
-			}()
-		}
+		createRecord(_rule, flag, flagGroup, user, password, method, path, filename, ip, uploadData, status)
 	}
 }
 
 func (s *Server) handlePasvConnection(conn net.Conn, data map[string]interface{}) {
+	defer func() {
+		if err := recover(); err != nil {
+			recycler.Recycle(err)
+		}
+	}()
+
 	remoteAddress := conn.RemoteAddr().String()
 	switch v := data[remoteAddress].(type) {
 	case types.Nil:
@@ -312,42 +309,97 @@ func (s *Server) handlePasvConnection(conn net.Conn, data map[string]interface{}
 	_ = conn.Close()
 }
 
-func (s *Server) Run() {
-	if err := s.updateRules(); err != nil {
-		log.Fatal(err.Error())
+// run pasv server
+func (s *Server) runPasvServer() (net.Listener, error) {
+	pasvAddress := fmt.Sprintf("%s:%d", strings.Split(s.Addr, ":")[0], s.PasvPort)
+	log.Info("Start to listen FTP PASV port at %v, PasvIP is %v", pasvAddress, s.PasvIP)
+	listener, err := net.Listen("tcp", pasvAddress)
+
+	if err != nil {
+		return nil, errors.Wrap(err, "FTP failed to listen on pasv port")
 	}
-	// run pasv server
+
 	go func() {
-		pasvAddress := fmt.Sprintf("%s:%d", strings.Split(s.Addr, ":")[0], s.PasvPort)
-		log.Info("Start to listen FTP PASV port at %v", pasvAddress)
-		listener, err := net.Listen("tcp", pasvAddress)
-		if err != nil {
-			log.Fatal("FTP failed to listen on pasv port : %v", err)
-		}
 		for data := range s.dataChannel {
 			tcpConn, err := listener.Accept()
 			if err != nil {
-				log.Warn("FTP accept connection error: %v", err)
+				if !strings.Contains(err.Error(), net.ErrClosed.Error()) {
+					log.Warn("FTP accept connection error: %v", err)
+				} else {
+					break
+				}
 				continue
 			}
 			s.handlePasvConnection(tcpConn, data)
 		}
 	}()
+	return listener, nil
+}
+
+func (s *Server) Stop() {
+	log.Info("FTP Server is stopping...")
+	s.Enable = false
+	s.livingLock.Unlock()
+}
+
+func (s *Server) Restart() {
+	s.Stop()
+	time.Sleep(time.Second * 2)
+	go s.Run()
+}
+
+func (s *Server) Run() {
+	s.Enable = true
+	s.livingLock.Lock()
+	defer func() {
+		if s.Enable {
+			log.Error("FTP Server exited unexpectedly")
+		}
+		s.Enable = false
+		s.livingLock.Unlock()
+	}()
+
+	if err := s.updateRules(); err != nil {
+		log.Error(err.Error())
+		return
+	}
+
+	pasvListener, err := s.runPasvServer()
+	if err != nil {
+		log.Error(err.Error())
+	}
+	defer func() {
+		if pasvListener != nil {
+			_ = pasvListener.Close()
+		}
+	}()
 
 	// run ftp server
 	log.Info("Starting FTP Server at %v", s.Addr)
 
 	listener, err := net.Listen("tcp", s.Addr)
 	if err != nil {
-		log.Fatal(err.Error())
+		log.Error(errors.Wrap(err, "FTP failed to start").Error())
+		return
 	}
-	for {
+
+	go func() {
+		s.livingLock.Lock()
+		if !s.Enable {
+			_ = listener.Close()
+		}
+	}()
+
+	for s.Enable {
 		tcpConn, err := listener.Accept()
 		if err != nil {
-			log.Warn("FTP accept connection error: %v", err)
+			if !strings.Contains(err.Error(), net.ErrClosed.Error()) {
+				log.Warn("FTP accept connection error: %v", err)
+			} else {
+				break
+			}
 			continue
 		}
 		go s.handleConnection(tcpConn)
 	}
-
 }