1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
|
package sshd
import (
"io"
"net"
"time"
"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),
}
}
// Count each read as 1 unless it exceeds some number of bytes.
type inputLimiter struct {
// TODO: Could do all kinds of fancy things here, like be more forgiving of
// connections that have been around for a while.
Amount int
Frequency time.Duration
remaining int
readCap int
numRead int
timeRead time.Time
}
// NewInputLimiter returns a rateio.Limiter with sensible defaults for
// differentiating between humans typing and bots spamming.
func NewInputLimiter() rateio.Limiter {
grace := time.Second * 3
return &inputLimiter{
Amount: 2 << 14, // ~16kb, should be plenty for a high typing rate/copypasta/large key handshakes.
Frequency: time.Minute * 1,
readCap: 128, // Allow up to 128 bytes per read (anecdotally, 1 character = 52 bytes over ssh)
numRead: -1024 * 1024, // Start with a 1mb grace
timeRead: time.Now().Add(grace),
}
}
// Count applies 1 if n<readCap, else n
func (limit *inputLimiter) Count(n int) error {
now := time.Now()
if now.After(limit.timeRead) {
limit.numRead = 0
limit.timeRead = now.Add(limit.Frequency)
}
if n <= limit.readCap {
limit.numRead += 1
} else {
limit.numRead += n
}
if limit.numRead > limit.Amount {
return rateio.ErrRateExceeded
}
return nil
}
|
