Connection-level rate limiting.

diff --git a/sshd/net.go b/sshd/net.go
index bb94432984b47e78eef0100e476224058aad42c4..7ded1e72cf4abde54db83190966c604686c3e43b 100644 (file)
--- a/sshd/net.go
+++ b/sshd/net.go
@@ -2,7 +2,9 @@ package sshd
 
 import (
        "net"
+       "time"
 
+       "github.com/shazow/rateio"
        "golang.org/x/crypto/ssh"
 )
 
@@ -24,6 +26,7 @@ func ListenSSH(laddr string, config *ssh.ServerConfig) (*SSHListener, error) {
 
 func (l *SSHListener) handleConn(conn net.Conn) (*Terminal, error) {
        // Upgrade TCP connection to SSH connection
+       conn = ReadLimitConn(conn, rateio.NewGracefulLimiter(1000, time.Minute*2, time.Second*3))
        sshConn, channels, requests, err := ssh.NewServerConn(conn, l.config)
        if err != nil {
                return nil, err

diff --git a/sshd/ratelimit.go b/sshd/ratelimit.go
new file mode 100644 (file)
index 0000000..c80f0ac
--- /dev/null
+++ b/sshd/ratelimit.go
@@ -0,0 +1,25 @@
+package sshd
+
+import (
+       "io"
+       "net"
+
+       "github.com/shazow/rateio"
+)
+
+type limitedConn struct {
+       net.Conn
+       io.Reader // Our rate-limited io.Reader for net.Conn
+}
+
+func (r *limitedConn) Read(p []byte) (n int, err error) {
+       return r.Reader.Read(p)
+}
+
+// ReadLimitConn returns a net.Conn whose io.Reader interface is rate-limited by limiter.
+func ReadLimitConn(conn net.Conn, limiter rateio.Limiter) net.Conn {
+       return &limitedConn{
+               Conn:   conn,
+               Reader: rateio.NewReader(conn, limiter),
+       }
+}