68 lines
1.4 KiB
Go
68 lines
1.4 KiB
Go
package ssh
|
|
|
|
import (
|
|
"net"
|
|
"time"
|
|
|
|
"github.com/fujiwara/shapeio"
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
type throttledConn struct {
|
|
ID uuid.UUID
|
|
conn net.Conn
|
|
writer *shapeio.Writer
|
|
reader *shapeio.Reader
|
|
CloseCallback func(c *throttledConn)
|
|
}
|
|
|
|
func newThrottledConn(conn net.Conn) *throttledConn {
|
|
id := uuid.Must(uuid.NewRandom())
|
|
return &throttledConn{
|
|
ID: id,
|
|
conn: conn,
|
|
writer: shapeio.NewWriter(conn),
|
|
reader: shapeio.NewReader(conn),
|
|
}
|
|
}
|
|
|
|
func (sc *throttledConn) SetSpeed(bytesPerSec float64) {
|
|
sc.writer.SetRateLimit(bytesPerSec)
|
|
sc.reader.SetRateLimit(bytesPerSec)
|
|
}
|
|
|
|
func (sc *throttledConn) Read(b []byte) (n int, err error) {
|
|
return sc.reader.Read(b)
|
|
}
|
|
|
|
func (sc *throttledConn) Write(b []byte) (n int, err error) {
|
|
return sc.writer.Write(b)
|
|
}
|
|
|
|
func (sc *throttledConn) Close() error {
|
|
if sc.CloseCallback != nil {
|
|
sc.CloseCallback(sc)
|
|
}
|
|
return sc.conn.Close()
|
|
}
|
|
|
|
func (sc *throttledConn) LocalAddr() net.Addr {
|
|
return sc.conn.LocalAddr()
|
|
}
|
|
|
|
func (sc *throttledConn) RemoteAddr() net.Addr {
|
|
return sc.conn.RemoteAddr()
|
|
}
|
|
|
|
func (sc *throttledConn) SetDeadline(t time.Time) error {
|
|
return sc.conn.SetDeadline(t)
|
|
}
|
|
|
|
func (sc *throttledConn) SetReadDeadline(t time.Time) error {
|
|
return sc.conn.SetReadDeadline(t)
|
|
}
|
|
|
|
func (sc *throttledConn) SetWriteDeadline(t time.Time) error {
|
|
return sc.conn.SetWriteDeadline(t)
|
|
}
|