feat(maxconnection): change max connection behaves

README.md

@@ -65,11 +65,11 @@ Usage of FakeSSH:
   -max maxconn
         see maxconn
   -maxconn max:loss_ratio:hard_max
-        max connections in format max:loss_ratio:hard_max, every value is optional means [default, 1.0, default]
+        max unauthenticated connections in format max:loss_ratio:hard_max, optionalable, see README
   -maxsucc maxsuccconn
         see maxsuccconn
-  -maxsuccconn max:loss_ratio:hard_max
-        max success connections in format max:loss_ratio:hard_max, see maxconn
+  -maxsuccconn max:loss_rate:hard_max
+        max success connections in format max:loss_rate:hard_max, see maxconn
   -mc maxconn
         see maxconn
   -msc maxsuccconn
@@ -106,20 +106,20 @@ Usage of FakeSSH:
 
 ### max connections
 
-You can use the commandline option `-maxconn` (or shorter `-mc`) to set the max connections, the `server.max_conn` in configure file does it the same.
+You can use the commandline option `-maxconn` (or shorter `-mc`) to set the max connections __for unauthenticated connections__, the `server.max_conn` in configure file does it the same.
 
-And `-maxsuccconn` (shorter `-msc` or `server.max_succ_conn` in configure file) to set the max success connections, with the same syntax.
+And `-maxsuccconn` (shorter `-msc` or `server.max_succ_conn` in configure file) to set the max __success__ connections, with the same syntax.
 
 The format of `-maxconn` and `-maxsuccconn` is `max:loss_ratio:hard_max`, and the format of configure file is shown in [this file](./conf/config.toml).
 
 It means when the count of connections mathes `max`, it will loss the connection with the ratio. And the ratio will increase literally, and it will be `1.0` when connections equal or larger than `hard_max`.
 
 * `max` is interger, optional means `0`:
   * `max < 0` => unlimited connections, unless `hard_max`.
-  * `max = 0` => use program default value, current is `100` for `maxconn` and `unlimited` for `maxsuccconn`.
+  * `max = 0` => use program default value(current is `100` for `maxconn` and `5` for `maxsuccconn`).
 * `loss_ratio` is float, optional means `0`:
   * `loss_ratio < 0` => not loss connections until it reaches `hard_max`.
-  * `loss_ratio >= 0` => loss connections with the ratio.
+  * `loss_ratio >= 0` => loss connections with the ratio, and it will increase literally until connections reaches `hard_max`.
 * `hard_max` is interger, optional means `0`:
   * `hard_max <= 0` when `max < 0` => unlimited connections.
-  * `hard_max <= 0` when `max >= 0` => it will be the max value of `max * 2` and default value(current if `65535`)
+  * `hard_max <= 0` when `max >= 0` => it will be the __larger__ value of `max * 2` and default value(current is `65535` for `maxconn` and `10` for `maxsuccconn`).

conf/conf.go

@@ -18,8 +18,8 @@ type RateLimitConfig = utils.RateLimitConfig
 const DefaultMaxConnections = 100
 const DefaultHardMaxConnections = 65535
 
-const DefaultMaxSuccessConnections = 0
-const DefaultHardMaxSucessConnections = 65535
+const DefaultMaxSuccessConnections = 5
+const DefaultHardMaxSucessConnections = 10
 
 type AppConfig struct {
 	BaseConfig

conf/config.toml

@@ -23,7 +23,7 @@
 #name = "root"
 #password = "123456"
 
-# How many connections can be connected at the same time
+# How many unauthenticated connections can be connected at the same time
 [server.max_conn]
 # `max` < 0 means unlimited
 # `max` = 0 or optional means default: 100
@@ -38,7 +38,9 @@ loss_ratio = 1.0
 
 # Same as `max_conn`, but for successful connections
 [server.max_succ_conn]
+# default: 5
 max = 0
+# default: 10
 hard_max = 0
 loss_ratio = 1.0
 

conf/flags.go

@@ -106,7 +106,7 @@ func GetArg() (args *FlagArgsStruct, set StringSet, helper func()) {
 
 	StringArrayVar(f, &args.Users, "user", "users in format `user:password`, can set more than one")
 
-	f.StringVar(&args.MaxConns, "maxconn", "", "max connections in format `max:loss_ratio:hard_max`, every value is optional means [default, 1.0, default]")
+	f.StringVar(&args.MaxConns, "maxconn", "", "max unauthenticated connections in format `max:loss_ratio:hard_max`, optionalable, see README")
 	f.StringVar(&args.MaxConns, "max", "", "see `maxconn`")
 	f.StringVar(&args.MaxConns, "mc", "", "see `maxconn`")
 	f.StringVar(&args.MaxSuccConns, "maxsuccconn", "", "max success connections in format `max:loss_rate:hard_max`, see maxconn")

main.go

@@ -11,7 +11,7 @@ import (
 	"hash/fnv"
 	golog "log"
 	"math"
-	"math/rand"
+	"math/rand/v2"
 	"os"
 	"regexp"
 	"strings"
@@ -291,6 +291,10 @@ func authCallback(c *conf.AppConfig) func(conn ssh.ConnMetadata, password []byte
 		log.Infof("[login] Connection from %v using user %s password %s, login: %t",
 			conn.RemoteAddr(), conn.User(), p, succLogin)
 
+		if succLogin {
+			return nil, nil
+		}
+
 		if delay > 0 {
 			m := c.Server.Deviation
 			if m <= 0 {
@@ -301,14 +305,9 @@ func authCallback(c *conf.AppConfig) func(conn ssh.ConnMetadata, password []byte
 				if start < 0 {
 					start = 0
 				}
-				time.Sleep(time.Millisecond * time.Duration(start+rand.Intn(end-start)))
+				time.Sleep(time.Millisecond * time.Duration(start+rand.IntN(end-start)))
 			}
 		}
-
-		if succLogin {
-			return nil, nil
-		}
-
 		return nil, errAuth
 	}
 }

ssh.go

@@ -23,13 +23,14 @@ type Option struct {
 type SSHConnectionContext struct {
 	net.Conn
 
+	Connections      *atomic.Int64
 	SuccConnections  *atomic.Int64
 	MaxSuccConns     int64
 	HardMaxSuccConns int64
 	SuccLossRatio    float64
 }
 
-func (c *SSHConnectionContext) CheckMaxConnections() bool {
+func (c *SSHConnectionContext) CheckMaxSuccussConnections() bool {
 	return checkMaxConnections(c.SuccConnections.Add(1), c.MaxSuccConns, c.HardMaxSuccConns, c.SuccLossRatio)
 }
 
@@ -133,16 +134,20 @@ func StartSSHServer(config *ssh.ServerConfig, opt *Option) {
 		if !checkMaxConnections(connections.Add(1), maxConn, hardMaxConn, lossRatio) {
 			_ = conn.Close()
 			connections.Add(-1)
-			log.Infof("[Disconnect] matches max connections limit, disconnect from: %s", conn.RemoteAddr().String())
+			log.Infof("[Disconnect] reached max connections limit, disconnect from: %s", conn.RemoteAddr().String())
 			continue
 		}
 
 		var ip string
-		addr, ok := conn.RemoteAddr().(*net.TCPAddr)
-		if !ok {
-			ip = conn.RemoteAddr().String()
-		} else {
+		switch addr := conn.RemoteAddr().(type) {
+		case *net.UDPAddr:
+			ip = addr.IP.String()
+		case *net.TCPAddr:
 			ip = addr.IP.String()
+		case *net.IPAddr:
+			ip = addr.IP.String()
+		default:
+			ip = conn.RemoteAddr().String()
 		}
 
 		pass := limiter.Allow(conn.RemoteAddr().String()).OK()
@@ -154,8 +159,9 @@ func StartSSHServer(config *ssh.ServerConfig, opt *Option) {
 		}
 
 		go func() {
-			defer connections.Add(-1)
 			handleConn(&SSHConnectionContext{
+				Conn:             conn,
+				Connections:      &connections,
 				SuccConnections:  &succConnections,
 				MaxSuccConns:     maxSuccConn,
 				HardMaxSuccConns: hardMaxSuccConn,
@@ -178,10 +184,12 @@ func handleConn(sshCtx *SSHConnectionContext, config *ssh.ServerConfig) {
 		return
 	}
 
-	ok := !sshCtx.CheckMaxConnections()
+	ok := !sshCtx.CheckMaxSuccussConnections()
+	// minus 1 for unauthenticated connection count
+	sshCtx.Connections.Add(-1)
 	defer sshCtx.SuccConnections.Add(-1)
 	if !ok {
-		log.Infof("[Disconnect] matches max success connections, disconnect from %s", sshCtx.RemoteAddr().String())
+		log.Infof("[Disconnect] reached max success connections, disconnect from %s", sshCtx.RemoteAddr().String())
 		return
 	}