Fix uint64 overflow issue

Author: mohanson

daze.go

@@ -627,7 +627,7 @@ func (l *Locale) Run() error {
 func NewLocale(listen string, dialer Dialer) *Locale {
 	return &Locale{
 		Dialer: dialer,
-		Limits: rate.NewLimits(math.MaxUint64, time.Second),
+		Limits: rate.NewLimits(math.MaxUint32, time.Second),
 		Listen: listen,
 	}
 }

lib/rate/rate.go

@@ -1,39 +1,50 @@
 package rate
 
 import (
+	"math"
 	"sync"
 	"time"
+
+	"github.com/mohanson/daze/lib/doa"
 )
 
 // Limits represents a rate limiter that controls resource allocation over time.
 type Limits struct {
 	addition uint64
 	capacity uint64
 	last     time.Time
+	loop     uint64
 	mu       sync.Mutex
 	size     uint64
 	step     time.Duration
 }
 
 // Wait ensures there are enough resources (n) available, blocking if necessary.
 func (l *Limits) Wait(n uint64) {
+	doa.Doa(n > 0 && n < math.MaxUint64/2)
 	l.mu.Lock()
 	defer l.mu.Unlock()
-	loop := uint64(time.Since(l.last) / l.step)
-	l.last = l.last.Add(l.step * time.Duration(loop))
-	l.size = l.size + loop*l.addition
-	l.size = min(l.size, l.capacity)
+	l.loop = uint64(time.Since(l.last) / l.step)
+	if l.loop > 0 {
+		l.last = l.last.Add(l.step * time.Duration(l.loop))
+		doa.Doa(l.loop <= math.MaxUint64/l.addition)
+		doa.Doa(l.size <= math.MaxUint64-l.addition*l.loop)
+		l.size = l.size + l.addition*l.loop
+		l.size = min(l.size, l.capacity)
+	}
 	if l.size < n {
-		loop := (n - l.size + l.addition - 1) / l.addition
-		time.Sleep(l.step * time.Duration(loop))
-		l.last = l.last.Add(l.step * time.Duration(loop))
-		l.size = l.size + loop*l.addition
+		l.loop = (n - l.size + l.addition - 1) / l.addition
+		time.Sleep(l.step * time.Duration(l.loop))
+		l.last = l.last.Add(l.step * time.Duration(l.loop))
+		l.size = l.size + l.addition*l.loop
 	}
 	l.size -= n
 }
 
 // NewLimits creates a new rate limiter with rate r over period p.
 func NewLimits(r uint64, p time.Duration) *Limits {
+	doa.Doa(r > 0 && r < math.MaxUint64/2)
+	doa.Doa(p > 0)
 	g := func(a, b uint64) uint64 {
 		t := uint64(0)
 		for b != 0 {

protocol/ashe/engine.go

@@ -247,7 +247,7 @@ func (s *Server) Run() error {
 func NewServer(listen string, cipher string) *Server {
 	return &Server{
 		Cipher: daze.Salt(cipher),
-		Limits: rate.NewLimits(math.MaxUint64, time.Second),
+		Limits: rate.NewLimits(math.MaxUint32, time.Second),
 		Listen: listen,
 	}
 }

protocol/baboon/engine.go

@@ -149,7 +149,7 @@ func (s *Server) Run() error {
 func NewServer(listen string, cipher string) *Server {
 	return &Server{
 		Cipher: daze.Salt(cipher),
-		Limits: rate.NewLimits(math.MaxUint64, time.Second),
+		Limits: rate.NewLimits(math.MaxUint32, time.Second),
 		Listen: listen,
 		Masker: Conf.Masker,
 		NextID: uint32(math.MaxUint32),

protocol/czar/engine.go

@@ -126,7 +126,7 @@ func (s *Server) Run() error {
 func NewServer(listen string, cipher string) *Server {
 	return &Server{
 		Cipher: daze.Salt(cipher),
-		Limits: rate.NewLimits(math.MaxUint64, time.Second),
+		Limits: rate.NewLimits(math.MaxUint32, time.Second),
 		Listen: listen,
 	}
 }

protocol/dahlia/engine.go

@@ -91,7 +91,7 @@ func NewServer(listen string, server string, cipher string) *Server {
 	return &Server{
 		Cipher: daze.Salt(cipher),
 		Listen: listen,
-		Limits: rate.NewLimits(math.MaxUint64, time.Second),
+		Limits: rate.NewLimits(math.MaxUint32, time.Second),
 		Server: server,
 	}
 }
@@ -171,7 +171,7 @@ func (c *Client) Run() error {
 func NewClient(listen string, server string, cipher string) *Client {
 	return &Client{
 		Cipher: daze.Salt(cipher),
-		Limits: rate.NewLimits(math.MaxUint64, time.Second),
+		Limits: rate.NewLimits(math.MaxUint32, time.Second),
 		Listen: listen,
 		Server: server,
 	}