Skip to content

Commit

Permalink
fix rate limitter not being used
Browse files Browse the repository at this point in the history
  • Loading branch information
n8maninger committed Nov 19, 2024
1 parent 6a7eaf9 commit 753ce67
Showing 1 changed file with 12 additions and 1 deletion.
13 changes: 12 additions & 1 deletion rhp/listener.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package rhp

import (
"context"
"net"

"golang.org/x/time/rate"
Expand All @@ -24,7 +25,7 @@ type (

rhpConn struct {
net.Conn

rl, wl *rate.Limiter
monitor DataMonitor
}

Expand Down Expand Up @@ -66,6 +67,10 @@ func WithDataMonitor(m DataMonitor) Option {
func (c *rhpConn) Read(b []byte) (int, error) {
n, err := c.Conn.Read(b)
c.monitor.ReadBytes(n)
if err != nil {
return n, err
}
c.rl.WaitN(context.Background(), len(b)) // error can be ignored since context will never be cancelled and len(b) should never exceed burst size
return n, err
}

Expand All @@ -74,6 +79,10 @@ func (c *rhpConn) Read(b []byte) (int, error) {
func (c *rhpConn) Write(b []byte) (int, error) {
n, err := c.Conn.Write(b)
c.monitor.WriteBytes(n)
if err != nil {
return n, err
}
c.wl.WaitN(context.Background(), len(b)) // error can be ignored since context will never be cancelled and len(b) should never exceed burst size
return n, err
}

Expand All @@ -84,6 +93,8 @@ func (l *rhpListener) Accept() (net.Conn, error) {
}
return &rhpConn{
Conn: c,
rl: l.readLimiter,
wl: l.writeLimiter,
monitor: l.monitor,
}, nil
}
Expand Down

0 comments on commit 753ce67

Please sign in to comment.