From b99083ee6eef1ee15026b344874505e17d622dbb Mon Sep 17 00:00:00 2001
From: Andrey Petrov <andrey.petrov@shazow.net>
Date: Fri, 16 Jan 2015 12:30:18 -0800
Subject: [PATCH] Connection-level rate limiting.

---
 sshd/net.go       |  3 +++
 sshd/ratelimit.go | 25 +++++++++++++++++++++++++
 2 files changed, 28 insertions(+)
 create mode 100644 sshd/ratelimit.go

diff --git a/sshd/net.go b/sshd/net.go
index bb94432..7ded1e7 100644
--- 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
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),
+	}
+}