From cc25d17bdc13c31bca170ba675a6236c2853d858 Mon Sep 17 00:00:00 2001
From: Andrey Petrov <andrey.petrov@shazow.net>
Date: Fri, 16 Jan 2015 12:35:57 -0800
Subject: [PATCH] Configurable rate limiting for sshd

---
 cmd.go      |  1 +
 sshd/net.go | 11 ++++++++---
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/cmd.go b/cmd.go
index 9a22209..a763885 100644
--- a/cmd.go
+++ b/cmd.go
@@ -102,6 +102,7 @@ func main() {
 		os.Exit(4)
 	}
 	defer s.Close()
+	s.RateLimit = true
 
 	fmt.Printf("Listening for connections on %v\n", s.Addr().String())
 
diff --git a/sshd/net.go b/sshd/net.go
index 7ded1e7..5e782d8 100644
--- a/sshd/net.go
+++ b/sshd/net.go
@@ -11,7 +11,8 @@ import (
 // Container for the connection and ssh-related configuration
 type SSHListener struct {
 	net.Listener
-	config *ssh.ServerConfig
+	config    *ssh.ServerConfig
+	RateLimit bool
 }
 
 // Make an SSH listener socket
@@ -20,13 +21,17 @@ func ListenSSH(laddr string, config *ssh.ServerConfig) (*SSHListener, error) {
 	if err != nil {
 		return nil, err
 	}
-	l := SSHListener{socket, config}
+	l := SSHListener{Listener: socket, config: config}
 	return &l, nil
 }
 
 func (l *SSHListener) handleConn(conn net.Conn) (*Terminal, error) {
+	if l.RateLimit {
+		// TODO: Configurable Limiter?
+		conn = ReadLimitConn(conn, rateio.NewGracefulLimiter(1000, time.Minute*2, time.Second*3))
+	}
+
 	// 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