apiary/honeypot/conn.go

65 lines
1.4 KiB
Go
Raw Normal View History

2021-04-10 05:58:01 +00:00
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)
}