2025-06-22 23:43:55

Author: mohanson

cmd/daze/main.go

@@ -14,6 +14,7 @@ import (
 	"github.com/mohanson/daze"
 	"github.com/mohanson/daze/lib/doa"
 	"github.com/mohanson/daze/lib/gracefulexit"
+	"github.com/mohanson/daze/lib/rate"
 	"github.com/mohanson/daze/protocol/ashe"
 	"github.com/mohanson/daze/protocol/baboon"
 	"github.com/mohanson/daze/protocol/czar"
@@ -68,6 +69,7 @@ func main() {
 			flDnserv = flag.String("dns", "", "specifies the DNS, DoT or DoH server")
 			flExtend = flag.String("e", "", "extend data for different protocols")
 			flGpprof = flag.String("g", "", "specify an address to enable net/http/pprof")
+			flLimits = flag.String("b", "", "set the maximum bandwidth in bytes per second, for example, 128k or 1.5m")
 			flCipher = flag.String("k", "daze", "password, should be same with the one specified by client")
 			flListen = flag.String("l", "0.0.0.0:1081", "listen address")
 			flProtoc = flag.String("p", "ashe", "protocol {ashe, baboon, czar, dahlia}")
@@ -90,21 +92,37 @@ func main() {
 		case "ashe":
 			server := ashe.NewServer(*flListen, *flCipher)
 			defer server.Close()
+			if *flLimits != "" {
+				n := daze.SizeParser(*flLimits)
+				server.Limits = rate.NewLimiter(rate.Limit(n), 1024*1024)
+			}
 			doa.Nil(server.Run())
 		case "baboon":
 			server := baboon.NewServer(*flListen, *flCipher)
+			defer server.Close()
 			if *flExtend != "" {
 				server.Masker = *flExtend
 			}
-			defer server.Close()
+			if *flLimits != "" {
+				n := daze.SizeParser(*flLimits)
+				server.Limits = rate.NewLimiter(rate.Limit(n), 1024*1024)
+			}
 			doa.Nil(server.Run())
 		case "czar":
 			server := czar.NewServer(*flListen, *flCipher)
 			defer server.Close()
+			if *flLimits != "" {
+				n := daze.SizeParser(*flLimits)
+				server.Limits = rate.NewLimiter(rate.Limit(n), 1024*1024)
+			}
 			doa.Nil(server.Run())
 		case "dahlia":
 			server := dahlia.NewServer(*flListen, *flExtend, *flCipher)
 			defer server.Close()
+			if *flLimits != "" {
+				n := daze.SizeParser(*flLimits)
+				server.Limits = rate.NewLimiter(rate.Limit(n), 1024*1024)
+			}
 			doa.Nil(server.Run())
 		}
 		if *flGpprof != "" {

daze.go

@@ -30,6 +30,7 @@ import (
 	"github.com/mohanson/daze/lib/doa"
 	"github.com/mohanson/daze/lib/lru"
 	"github.com/mohanson/daze/lib/pretty"
+	"github.com/mohanson/daze/lib/rate"
 )
 
 // ============================================================================
@@ -1042,18 +1043,6 @@ func OpenFile(name string) (io.ReadCloser, error) {
 	return os.Open(name)
 }
 
-// RandomReader is a simple random number generator. Note that it is not cryptographically secure, but for daze, the
-// randomness it provides is enough.
-type RandomReader struct{}
-
-// Read implements io.Reader.
-func (r *RandomReader) Read(p []byte) (int, error) {
-	for i := range len(p) {
-		p[i] = byte(rand.Uint64())
-	}
-	return len(p), nil
-}
-
 // The PrettyReader struct represents a custom reader that keeps track of read bytes and prints progress.
 type PrettyReader struct {
 	E uint64    // Total number of bytes read so far
@@ -1072,12 +1061,66 @@ func (r *PrettyReader) Read(p []byte) (int, error) {
 	return n, err
 }
 
+// RandomReader is a simple random number generator. Note that it is not cryptographically secure, but for daze, the
+// randomness it provides is enough.
+type RandomReader struct{}
+
+// Read implements io.Reader.
+func (r *RandomReader) Read(p []byte) (int, error) {
+	for i := range len(p) {
+		p[i] = byte(rand.Uint64())
+	}
+	return len(p), nil
+}
+
+// RateConn wraps a net.Conn with a per-conn and a rate limiter.
+type RateConn struct {
+	Conn io.ReadWriteCloser
+	Rate *rate.Limiter
+}
+
+// Close closes the connection.
+func (r *RateConn) Close() error {
+	return r.Conn.Close()
+}
+
+// Read reads up to len(p) bytes into p.
+func (r *RateConn) Read(p []byte) (int, error) {
+	n, err := r.Conn.Read(p)
+	r.Rate.WaitN(context.Background(), n)
+	return n, err
+}
+
+// Write writes len(p) bytes from p to the underlying data stream.
+func (r *RateConn) Write(p []byte) (int, error) {
+	n, err := r.Conn.Write(p)
+	r.Rate.WaitN(context.Background(), n)
+	return n, err
+}
+
 // Salt converts the stupid password passed in by the user to 32-sized byte array.
 func Salt(s string) []byte {
 	h := sha256.Sum256([]byte(s))
 	return h[:]
 }
 
+// SizeParser converts a string like "1K", "1M", or "1G" to bytes as uint64. It expects the input string to end with a
+// unit (K, M, or G) and panics if the unit is invalid. The number part can be a float (e.g., "1.5M") and is converted
+// to bytes based on the unit.
+func SizeParser(s string) uint64 {
+	f := doa.Try(strconv.ParseFloat(s[:len(s)-1], 64))
+	u := strings.ToLower(s[len(s)-1:])
+	switch u {
+	case "k":
+		return uint64(f * 1024)
+	case "m":
+		return uint64(f * 1024 * 1024)
+	case "g":
+		return uint64(f * 1024 * 1024 * 1024)
+	}
+	panic("unreachable")
+}
+
 // ============================================================================
 //                 ___           ___           ___           ___
 //                /\  \         /\  \         /\  \         /\  \