Merge pull request #3 from PostHog/feature/rate-limiting

Add connection rate limiting to prevent brute-force attacks

Author: fuziontech

TODO.md

@@ -6,7 +6,7 @@
 - [x] **TLS/SSL Support**: Add encrypted connections for production use
 - [ ] **MD5 Password Authentication**: Support PostgreSQL's MD5-based auth for better security
 - [ ] **SCRAM-SHA-256 Authentication**: Modern PostgreSQL authentication method
-- [ ] **Connection Rate Limiting**: Prevent brute-force attacks
+- [x] **Connection Rate Limiting**: Prevent brute-force attacks
 
 ### Configuration
 - [x] **Config File Support**: Load configuration from YAML/TOML file instead of hardcoding

duckgres.example.yaml

@@ -18,3 +18,14 @@ users:
   postgres: "postgres"
   alice: "alice123"
   bob: "bob123"
+
+# Rate limiting configuration (optional - these are the defaults)
+rate_limit:
+  # Max failed auth attempts before banning an IP
+  max_failed_attempts: 5
+  # Time window for counting failed attempts (e.g., "5m", "1h")
+  failed_attempt_window: "5m"
+  # How long to ban an IP after too many failed attempts
+  ban_duration: "15m"
+  # Max concurrent connections from a single IP (0 = unlimited)
+  max_connections_per_ip: 100

main.go

@@ -8,25 +8,34 @@ import (
 	"os/signal"
 	"strconv"
 	"syscall"
+	"time"
 
 	"github.com/posthog/duckgres/server"
 	"gopkg.in/yaml.v3"
 )
 
 // FileConfig represents the YAML configuration file structure
 type FileConfig struct {
-	Host    string            `yaml:"host"`
-	Port    int               `yaml:"port"`
-	DataDir string            `yaml:"data_dir"`
-	TLS     TLSConfig         `yaml:"tls"`
-	Users   map[string]string `yaml:"users"`
+	Host      string              `yaml:"host"`
+	Port      int                 `yaml:"port"`
+	DataDir   string              `yaml:"data_dir"`
+	TLS       TLSConfig           `yaml:"tls"`
+	Users     map[string]string   `yaml:"users"`
+	RateLimit RateLimitFileConfig `yaml:"rate_limit"`
 }
 
 type TLSConfig struct {
 	Cert string `yaml:"cert"`
 	Key  string `yaml:"key"`
 }
 
+type RateLimitFileConfig struct {
+	MaxFailedAttempts   int    `yaml:"max_failed_attempts"`
+	FailedAttemptWindow string `yaml:"failed_attempt_window"` // e.g., "5m"
+	BanDuration         string `yaml:"ban_duration"`          // e.g., "15m"
+	MaxConnectionsPerIP int    `yaml:"max_connections_per_ip"`
+}
+
 // loadConfigFile loads configuration from a YAML file
 func loadConfigFile(path string) (*FileConfig, error) {
 	data, err := os.ReadFile(path)
@@ -129,6 +138,28 @@ func main() {
 		if len(fileCfg.Users) > 0 {
 			cfg.Users = fileCfg.Users
 		}
+
+		// Apply rate limit config
+		if fileCfg.RateLimit.MaxFailedAttempts > 0 {
+			cfg.RateLimit.MaxFailedAttempts = fileCfg.RateLimit.MaxFailedAttempts
+		}
+		if fileCfg.RateLimit.MaxConnectionsPerIP > 0 {
+			cfg.RateLimit.MaxConnectionsPerIP = fileCfg.RateLimit.MaxConnectionsPerIP
+		}
+		if fileCfg.RateLimit.FailedAttemptWindow != "" {
+			if d, err := time.ParseDuration(fileCfg.RateLimit.FailedAttemptWindow); err == nil {
+				cfg.RateLimit.FailedAttemptWindow = d
+			} else {
+				log.Printf("Warning: invalid failed_attempt_window duration: %v", err)
+			}
+		}
+		if fileCfg.RateLimit.BanDuration != "" {
+			if d, err := time.ParseDuration(fileCfg.RateLimit.BanDuration); err == nil {
+				cfg.RateLimit.BanDuration = d
+			} else {
+				log.Printf("Warning: invalid ban_duration duration: %v", err)
+			}
+		}
 	}
 
 	// Apply environment variables (override config file)

scripts/test_ratelimit.sh

@@ -0,0 +1,92 @@
+#!/bin/bash
+# Test rate limiting functionality
+
+set -e
+
+echo "=== Testing Rate Limiting ==="
+
+# Create a test config with aggressive rate limiting
+cat > /tmp/test-ratelimit.yaml <<EOF
+host: "127.0.0.1"
+port: 35433
+data_dir: "./data"
+tls:
+  cert: "./certs/server.crt"
+  key: "./certs/server.key"
+users:
+  postgres: "postgres"
+rate_limit:
+  max_failed_attempts: 3
+  failed_attempt_window: "1m"
+  ban_duration: "10s"
+  max_connections_per_ip: 5
+EOF
+
+# Kill any existing duckgres
+pkill -f "duckgres.*35433" 2>/dev/null || true
+sleep 1
+
+# Start server
+echo "Starting server with rate limiting config..."
+./duckgres --config /tmp/test-ratelimit.yaml &
+SERVER_PID=$!
+sleep 2
+
+cleanup() {
+    kill $SERVER_PID 2>/dev/null || true
+    rm -f /tmp/test-ratelimit.yaml
+}
+trap cleanup EXIT
+
+PASS_COUNT=0
+FAIL_COUNT=0
+
+echo ""
+echo "Test 1: Successful authentication should work"
+if PGPASSWORD=postgres psql "host=127.0.0.1 port=35433 user=postgres dbname=test sslmode=require" -c "SELECT 'auth works' as result;" 2>&1 | grep -q "auth works"; then
+    echo "  PASS: Successful auth works"
+    ((PASS_COUNT++))
+else
+    echo "  FAIL: Successful auth failed"
+    ((FAIL_COUNT++))
+fi
+
+echo ""
+echo "Test 2: Failed auth attempts should be counted (3 attempts)"
+for i in 1 2 3; do
+    echo "  Attempt $i with wrong password..."
+    PGPASSWORD=wrongpassword psql "host=127.0.0.1 port=35433 user=postgres dbname=test sslmode=require" -c "SELECT 1" 2>&1 | head -1 || true
+done
+
+echo ""
+echo "Test 3: After 3 failures, connection should be rejected"
+RESULT=$(PGPASSWORD=postgres psql "host=127.0.0.1 port=35433 user=postgres dbname=test sslmode=require" -c "SELECT 'should fail'" 2>&1 || true)
+if echo "$RESULT" | grep -q "server closed the connection unexpectedly"; then
+    echo "  PASS: Connection rejected after too many failures"
+    ((PASS_COUNT++))
+else
+    echo "  FAIL: Connection was not rejected as expected"
+    echo "  Result: $RESULT"
+    ((FAIL_COUNT++))
+fi
+
+echo ""
+echo "Test 4: Wait for ban to expire (10s) and try again"
+echo "  Waiting 12 seconds for ban to expire..."
+sleep 12
+
+if PGPASSWORD=postgres psql "host=127.0.0.1 port=35433 user=postgres dbname=test sslmode=require" -c "SELECT 'unbanned' as result;" 2>&1 | grep -q "unbanned"; then
+    echo "  PASS: Connection works after ban expires"
+    ((PASS_COUNT++))
+else
+    echo "  FAIL: Connection still blocked after ban should have expired"
+    ((FAIL_COUNT++))
+fi
+
+echo ""
+echo "=== Rate Limiting Tests Complete ==="
+echo "Results: $PASS_COUNT passed, $FAIL_COUNT failed"
+
+if [ $FAIL_COUNT -gt 0 ]; then
+    exit 1
+fi

server/conn.go

@@ -151,10 +151,18 @@ func (c *clientConn) handleStartup() error {
 	// Validate password
 	expectedPassword, ok := c.server.cfg.Users[c.username]
 	if !ok || expectedPassword != password {
+		// Record failed authentication attempt
+		banned := c.server.rateLimiter.RecordFailedAuth(c.conn.RemoteAddr())
+		if banned {
+			log.Printf("IP %s banned after too many failed auth attempts", c.conn.RemoteAddr())
+		}
 		c.sendError("FATAL", "28P01", "password authentication failed")
 		return fmt.Errorf("authentication failed for user %q", c.username)
 	}
 
+	// Record successful authentication (clears failed attempt counter)
+	c.server.rateLimiter.RecordSuccessfulAuth(c.conn.RemoteAddr())
+
 	// Send auth OK
 	if err := writeAuthOK(c.writer); err != nil {
 		return err