Make connection throttle configurable

apiary.toml

@@ -22,6 +22,10 @@ LogLevel = "INFO"
 # Address and port to listen to 
 # Default: ":2222"
 ListenAddr = ":2222"
+# Throttle incoming and outgoing data per connection
+# Values are in bytes per second. Empty means no unlimited
+# Default: ""
+ThrottleSpeed = 10240
 
 [Frontend]
 # Log level for SSH Honeypot

config/config.go

@@ -27,6 +27,7 @@ type HoneypotConfig struct {
 	ListenAddr    string  `toml:"ListenAddr"`
 	LogLevel      string  `toml:"LogLevel"`
 	HostKeyPath   string  `toml:"HostKeyPath"`
+	ThrottleSpeed float64 `toml:"ThrottleSpeed"`
 }
 
 type FrontendConfig struct {

honeypot/conn.go

@@ -11,29 +11,32 @@ import (
 type throttledConn struct {
 	ID            uuid.UUID
 	conn          net.Conn
-	speed         float64
+	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, speed: 1024 * 10}
+	return &throttledConn{
+		ID:     id,
+		conn:   conn,
+		writer: shapeio.NewWriter(conn),
+		reader: shapeio.NewReader(conn),
+	}
 }
 
 func (sc *throttledConn) SetSpeed(bytesPerSec float64) {
-	sc.speed = bytesPerSec
+	sc.writer.SetRateLimit(bytesPerSec)
+	sc.reader.SetRateLimit(bytesPerSec)
 }
 
 func (sc *throttledConn) Read(b []byte) (n int, err error) {
-	slowReader := shapeio.NewReader(sc.conn)
-	slowReader.SetRateLimit(sc.speed)
-	return slowReader.Read(b)
+	return sc.reader.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)
+	return sc.writer.Write(b)
 }
 
 func (sc *throttledConn) Close() error {

honeypot/server.go

@@ -19,12 +19,14 @@ import (
 )
 
 type HoneypotServer struct {
-	attemptStore      store.LoginAttemptStore
-	attemptsCallbacks []func(l models.LoginAttempt)
+	Logger *zap.SugaredLogger
 
 	sshServer *ssh.Server
 
-	Logger *zap.SugaredLogger
+	attemptStore      store.LoginAttemptStore
+	attemptsCallbacks []func(l models.LoginAttempt)
+
+	throttleSpeed float64
 }
 
 func NewHoneypotServer(cfg config.HoneypotConfig, store store.LoginAttemptStore) (*HoneypotServer, error) {
@@ -107,8 +109,9 @@ func (hs *HoneypotServer) passwordHandler(ctx ssh.Context, password string) bool
 
 func (s *HoneypotServer) connCallback(ctx ssh.Context, conn net.Conn) net.Conn {
 	throttledConn := newThrottledConn(conn)
+	throttledConn.SetSpeed(s.throttleSpeed)
 	ctx.SetValue("uuid", throttledConn.ID)
-	throttledConn.SetSpeed(2048)
+	throttledConn.SetSpeed(s.throttleSpeed)
 	return throttledConn
 }