Flow display and user login and rate limit bug

lib/conn/conn.go

@@ -440,16 +440,16 @@ func CopyWaitGroup(conn1, conn2 net.Conn, crypt bool, snappy bool, rate *rate.Ra
 }
 
 //get crypt or snappy conn
-func GetConn(conn net.Conn, cpt, snappy bool, rate *rate.Rate, isServer bool) (io.ReadWriteCloser) {
+func GetConn(conn net.Conn, cpt, snappy bool, rt *rate.Rate, isServer bool) (io.ReadWriteCloser) {
 	if cpt {
 		if isServer {
-			return crypt.NewTlsServerConn(conn)
+			return rate.NewRateConn(crypt.NewTlsServerConn(conn), rt)
 		}
-		return crypt.NewTlsClientConn(conn)
+		return rate.NewRateConn(crypt.NewTlsClientConn(conn), rt)
 	} else if snappy {
-		return NewSnappyConn(conn, cpt, rate)
+		return NewSnappyConn(conn, cpt, rt)
 	}
-	return conn
+	return rate.NewRateConn(conn, rt)
 }
 
 //read length or id (content length=4)

lib/file/file.go

@@ -256,6 +256,9 @@ func (s *Csv) LoadClientFromCsv() {
 		if post.RateLimit > 0 {
 			post.Rate = rate.NewRate(int64(post.RateLimit * 1024))
 			post.Rate.Start()
+		} else {
+			post.Rate = rate.NewRate(int64(2 << 23))
+			post.Rate.Start()
 		}
 		post.Flow = new(Flow)
 		post.Flow.FlowLimit = int64(common.GetIntNoErrByStr(item[9]))
@@ -382,6 +385,10 @@ reset:
 		isNotSet = true
 		c.VerifyKey = crypt.GetRandomString(16)
 	}
+	if c.RateLimit == 0 {
+		c.Rate = rate.NewRate(int64(2 << 23))
+		c.Rate.Start()
+	}
 	if !s.VerifyVkey(c.VerifyKey, c.id) {
 		if isNotSet {
 			goto reset
@@ -426,6 +433,10 @@ func (s *Csv) GetHostId() int32 {
 
 func (s *Csv) UpdateClient(t *Client) error {
 	s.Clients.Store(t.Id, t)
+	if t.RateLimit == 0 {
+		t.Rate = rate.NewRate(int64(2 << 23))
+		t.Rate.Start()
+	}
 	return nil
 }
 

lib/file/obj.go

@@ -142,6 +142,7 @@ type Health struct {
 	HealthRemoveArr     []string
 	HealthCheckType     string
 	HealthCheckTarget   string
+	sync.RWMutex
 }
 
 func (s *Tunnel) GetRandomTarget() (string, error) {

lib/rate/conn.go

@@ -0,0 +1,38 @@
+package rate
+
+import (
+	"io"
+	"net"
+)
+
+type rateConn struct {
+	conn net.Conn
+	rate *Rate
+}
+
+func NewRateConn(conn net.Conn, rate *Rate) io.ReadWriteCloser {
+	return &rateConn{
+		conn: conn,
+		rate: rate,
+	}
+}
+
+func (s *rateConn) Read(b []byte) (n int, err error) {
+	n, err = s.conn.Read(b)
+	if s.rate != nil {
+		s.rate.Get(int64(n))
+	}
+	return
+}
+
+func (s *rateConn) Write(b []byte) (n int, err error) {
+	n, err = s.conn.Write(b)
+	if s.rate != nil {
+		s.rate.Get(int64(n))
+	}
+	return
+}
+
+func (s *rateConn) Close() error {
+	return s.conn.Close()
+}

lib/rate/rate.go

@@ -10,6 +10,7 @@ type Rate struct {
 	bucketSurplusSize int64     //当前桶中体积
 	bucketAddSize     int64     //每次加水大小
 	stopChan          chan bool //停止
+	NowRate           int64
 }
 
 func NewRate(addSize int64) *Rate {
@@ -26,7 +27,8 @@ func (s *Rate) Start() {
 }
 
 func (s *Rate) add(size int64) {
-	if (s.bucketSize - s.bucketSurplusSize) < s.bucketAddSize {
+	if res := s.bucketSize - s.bucketSurplusSize; res < s.bucketAddSize {
+		atomic.AddInt64(&s.bucketSurplusSize, res)
 		return
 	}
 	atomic.AddInt64(&s.bucketSurplusSize, size)
@@ -65,6 +67,11 @@ func (s *Rate) session() {
 	for {
 		select {
 		case <-ticker.C:
+			if rs := s.bucketAddSize - s.bucketSurplusSize; rs > 0 {
+				s.NowRate = rs
+			} else {
+				s.NowRate = s.bucketSize - s.bucketSurplusSize
+			}
 			s.add(s.bucketAddSize)
 		case <-s.stopChan:
 			ticker.Stop()