mirror of
https://tildegit.org/solderpunk/molly-brown
synced 2024-05-11 21:36:02 +02:00
Initial implementation of leaky bucket rate limiting.
This commit is contained in:
parent
c4866d2965
commit
5016f40edb
11
handler.go
11
handler.go
|
@ -36,7 +36,7 @@ func isSubdir(subdir, superdir string) (bool, error) {
|
||||||
return false, nil
|
return false, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func handleGeminiRequest(conn net.Conn, sysConfig SysConfig, config UserConfig, accessLogEntries chan LogEntry, wg *sync.WaitGroup) {
|
func handleGeminiRequest(conn net.Conn, sysConfig SysConfig, config UserConfig, accessLogEntries chan LogEntry, rl *RateLimiter, wg *sync.WaitGroup) {
|
||||||
defer conn.Close()
|
defer conn.Close()
|
||||||
defer wg.Done()
|
defer wg.Done()
|
||||||
var tlsConn (*tls.Conn) = conn.(*tls.Conn)
|
var tlsConn (*tls.Conn) = conn.(*tls.Conn)
|
||||||
|
@ -49,6 +49,15 @@ func handleGeminiRequest(conn net.Conn, sysConfig SysConfig, config UserConfig,
|
||||||
defer func() { accessLogEntries <- logEntry }()
|
defer func() { accessLogEntries <- logEntry }()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Enforce rate limiting
|
||||||
|
noPort := logEntry.RemoteAddr.String()
|
||||||
|
noPort = noPort[0:strings.LastIndex(noPort, ":")]
|
||||||
|
if !rl.Allowed(noPort) {
|
||||||
|
conn.Write([]byte("44 10 second cool down, please!\r\n"))
|
||||||
|
logEntry.Status = 44
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
// Read request
|
// Read request
|
||||||
URL, err := readRequest(conn, &logEntry)
|
URL, err := readRequest(conn, &logEntry)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|
|
@ -159,11 +159,12 @@ func launch(sysConfig SysConfig, userConfig UserConfig, privInfo userInfo) int {
|
||||||
// Infinite serve loop (SIGTERM breaks out)
|
// Infinite serve loop (SIGTERM breaks out)
|
||||||
running := true
|
running := true
|
||||||
var wg sync.WaitGroup
|
var wg sync.WaitGroup
|
||||||
|
rl := newRateLimiter(100, 5)
|
||||||
for running {
|
for running {
|
||||||
conn, err := listener.Accept()
|
conn, err := listener.Accept()
|
||||||
if err == nil {
|
if err == nil {
|
||||||
wg.Add(1)
|
wg.Add(1)
|
||||||
go handleGeminiRequest(conn, sysConfig, userConfig, accessLogEntries, &wg)
|
go handleGeminiRequest(conn, sysConfig, userConfig, accessLogEntries, &rl, &wg)
|
||||||
} else {
|
} else {
|
||||||
select {
|
select {
|
||||||
case <-shutdown:
|
case <-shutdown:
|
||||||
|
|
|
@ -0,0 +1,54 @@
|
||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
type RateLimiter struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
bucket map[string]int
|
||||||
|
capacity int
|
||||||
|
rate int
|
||||||
|
}
|
||||||
|
|
||||||
|
func newRateLimiter(capacity int, rate int) RateLimiter {
|
||||||
|
var rl = new(RateLimiter)
|
||||||
|
rl.bucket = make(map[string]int)
|
||||||
|
rl.capacity = capacity
|
||||||
|
rl.rate = rate
|
||||||
|
// Leak periodically
|
||||||
|
go func () {
|
||||||
|
for(true) {
|
||||||
|
fmt.Println(rl.bucket)
|
||||||
|
rl.mu.Lock()
|
||||||
|
for addr, drips := range rl.bucket {
|
||||||
|
if drips <= rate {
|
||||||
|
delete(rl.bucket, addr)
|
||||||
|
} else {
|
||||||
|
rl.bucket[addr] = drips - rl.rate
|
||||||
|
}
|
||||||
|
}
|
||||||
|
rl.mu.Unlock()
|
||||||
|
time.Sleep(time.Second)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
return *rl
|
||||||
|
}
|
||||||
|
|
||||||
|
func (rl *RateLimiter) Allowed(addr string) bool {
|
||||||
|
rl.mu.Lock()
|
||||||
|
defer rl.mu.Unlock()
|
||||||
|
drips, present := rl.bucket[addr]
|
||||||
|
if !present {
|
||||||
|
rl.bucket[addr] = 1
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
if drips == rl.capacity {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
rl.bucket[addr] = drips + 1
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
Loading…
Reference in New Issue