65 lines
1.4 KiB
Go
65 lines
1.4 KiB
Go
|
package honeypot
|
||
|
|
||
|
import (
|
||
|
"net"
|
||
|
"time"
|
||
|
|
||
|
"github.com/fujiwara/shapeio"
|
||
|
"github.com/google/uuid"
|
||
|
)
|
||
|
|
||
|
type throttledConn struct {
|
||
|
ID uuid.UUID
|
||
|
conn net.Conn
|
||
|
speed float64
|
||
|
CloseCallback func(c *throttledConn)
|
||
|
}
|
||
|
|
||
|
func newThrottledConn(conn net.Conn) *throttledConn {
|
||
|
id := uuid.Must(uuid.NewRandom())
|
||
|
return &throttledConn{ID: id, conn: conn, speed: 1024 * 10}
|
||
|
}
|
||
|
|
||
|
func (sc *throttledConn) SetSpeed(bytesPerSec float64) {
|
||
|
sc.speed = bytesPerSec
|
||
|
}
|
||
|
|
||
|
func (sc *throttledConn) Read(b []byte) (n int, err error) {
|
||
|
slowReader := shapeio.NewReader(sc.conn)
|
||
|
slowReader.SetRateLimit(sc.speed)
|
||
|
return slowReader.Read(b)
|
||
|
}
|
||
|
|
||
|
func (sc *throttledConn) Write(b []byte) (n int, err error) {
|
||
|
slowWriter := shapeio.NewWriter(sc.conn)
|
||
|
slowWriter.SetRateLimit(sc.speed)
|
||
|
return slowWriter.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)
|
||
|
}
|