2025-09-09 17:50:08

Author: mohanson

daze.go

@@ -592,9 +592,10 @@ func (l *Locale) Serve(ctx *Context, cli io.ReadWriteCloser) error {
 		buf = make([]byte, 1)
 		err error
 	)
-	cli = &RateConn{
-		Conn: cli,
-		Rate: l.Limits,
+	cli = &ReadWriteCloser{
+		Reader: io.TeeReader(cli, rate.NewLimitsWriter(l.Limits)),
+		Writer: io.MultiWriter(cli, rate.NewLimitsWriter(l.Limits)),
+		Closer: cli,
 	}
 	_, err = io.ReadFull(cli, buf)
 	if err != nil {
@@ -606,7 +607,7 @@ func (l *Locale) Serve(ctx *Context, cli io.ReadWriteCloser) error {
 		}
 		return err
 	}
-	cli = ReadWriteCloser{
+	cli = &ReadWriteCloser{
 		Reader: io.MultiReader(bytes.NewReader(buf), cli),
 		Writer: cli,
 		Closer: cli,
@@ -1140,31 +1141,6 @@ func (r *RandomReader) Read(p []byte) (int, error) {
 	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.Limits
-}
-
-// 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.Wait(uint64(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.Wait(uint64(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))

protocol/ashe/engine.go

@@ -234,9 +234,10 @@ func (s *Server) Run() error {
 			idx++
 			ctx := &daze.Context{Cid: idx}
 			log.Printf("conn: %08x accept remote=%s", ctx.Cid, cli.RemoteAddr())
-			rtc := &daze.RateConn{
-				Conn: cli,
-				Rate: s.Limits,
+			rtc := &daze.ReadWriteCloser{
+				Reader: io.TeeReader(cli, rate.NewLimitsWriter(s.Limits)),
+				Writer: io.MultiWriter(cli, rate.NewLimitsWriter(s.Limits)),
+				Closer: cli,
 			}
 			go func() {
 				defer rtc.Close()

protocol/baboon/engine.go

@@ -79,9 +79,10 @@ func (s *Server) ServeDaze(w http.ResponseWriter, r *http.Request) {
 	spy := &ashe.Server{Cipher: s.Cipher}
 	ctx := &daze.Context{Cid: atomic.AddUint32(&s.NextID, 1)}
 	log.Printf("conn: %08x accept remote=%s", ctx.Cid, cc.RemoteAddr())
-	rtc := &daze.RateConn{
-		Conn: cli,
-		Rate: s.Limits,
+	rtc := &daze.ReadWriteCloser{
+		Reader: io.TeeReader(cli, rate.NewLimitsWriter(s.Limits)),
+		Writer: io.MultiWriter(cli, rate.NewLimitsWriter(s.Limits)),
+		Closer: cli,
 	}
 	if err := spy.Serve(ctx, rtc); err != nil {
 		log.Printf("conn: %08x  error %s", ctx.Cid, err)

protocol/czar/engine.go

@@ -108,9 +108,10 @@ func (s *Server) Run() error {
 				}
 				break
 			}
-			rtc := &daze.RateConn{
-				Conn: cli,
-				Rate: s.Limits,
+			rtc := &daze.ReadWriteCloser{
+				Reader: io.TeeReader(cli, rate.NewLimitsWriter(s.Limits)),
+				Writer: io.MultiWriter(cli, rate.NewLimitsWriter(s.Limits)),
+				Closer: cli,
 			}
 			mux := NewMuxServer(rtc)
 			go func() {

protocol/dahlia/engine.go

@@ -70,9 +70,10 @@ func (s *Server) Run() error {
 			idx++
 			ctx := &daze.Context{Cid: idx}
 			log.Printf("conn: %08x accept remote=%s", ctx.Cid, cli.RemoteAddr())
-			rtc := &daze.RateConn{
-				Conn: cli,
-				Rate: s.Limits,
+			rtc := daze.ReadWriteCloser{
+				Reader: io.TeeReader(cli, rate.NewLimitsWriter(s.Limits)),
+				Writer: io.MultiWriter(cli, rate.NewLimitsWriter(s.Limits)),
+				Closer: cli,
 			}
 			go func() {
 				defer rtc.Close()
@@ -148,9 +149,10 @@ func (c *Client) Run() error {
 				}
 				break
 			}
-			rtc := &daze.RateConn{
-				Conn: cli,
-				Rate: c.Limits,
+			rtc := &daze.ReadWriteCloser{
+				Reader: io.TeeReader(cli, rate.NewLimitsWriter(c.Limits)),
+				Writer: io.MultiWriter(cli, rate.NewLimitsWriter(c.Limits)),
+				Closer: cli,
 			}
 			idx++
 			ctx := &daze.Context{Cid: idx}