From f525eaa8f3394cbf9fba3f84abbf0eeef06f48ae Mon Sep 17 00:00:00 2001 From: db123-test Date: Tue, 5 May 2026 11:34:35 +0330 Subject: [PATCH] feat: enhance auth, add input validation and pagination - Add IP validation in auth handler with detailed logging - Implement apiKeyMiddleware helper for API key verification - Add IP and CIDR validation for whitelist operations - Support pagination for listing temp and permanent entries - Add parsePagination helper to extract limit and offset - Include comprehensive unit tests for IsIP, IsCIDR, and IPMatch --- auth.go | 94 +++++++++++++++++++++++++++-- auth_test.go | 163 +++++++++++++++++++++++++++++++++++++++++++++++++++ cleanup.go | 14 ++++- db.go | 111 ++++++++++++++++++++++++++++++++++- main.go | 18 ++++-- 5 files changed, 383 insertions(+), 17 deletions(-) create mode 100644 auth_test.go diff --git a/auth.go b/auth.go index 121067f..c3584e4 100644 --- a/auth.go +++ b/auth.go @@ -12,6 +12,8 @@ import ( "time" ) +var startTime = time.Now() + func CheckAuth(db *sql.DB, r *http.Request) bool { ip := r.Header.Get("X-Real-IP") if ip == "" { @@ -30,14 +32,33 @@ func CheckAuth(db *sql.DB, r *http.Request) bool { func authHandler(cfg Config, db *sql.DB) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { - if !CheckAuth(db, r) { + ip := r.Header.Get("X-Real-IP") + if ip == "" { + slog.Warn("auth denied: no X-Real-IP", "remote_addr", r.RemoteAddr) w.WriteHeader(http.StatusUnauthorized) return } + if !CheckAuth(db, r) { + slog.Warn("auth denied", "ip", ip) + w.WriteHeader(http.StatusUnauthorized) + return + } + slog.Info("auth allowed", "ip", ip) w.WriteHeader(http.StatusNoContent) } } +func apiKeyMiddleware(next http.HandlerFunc, token string) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + if !verifyAPIKey(r, token) { + slog.Warn("unauthorized API access") + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + next(w, r) + } +} + func whitelistTempHandler(cfg Config, db *sql.DB) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { if !verifyAPIKey(r, cfg.APIToken) { @@ -62,6 +83,10 @@ func whitelistTempHandler(cfg Config, db *sql.DB) http.HandlerFunc { http.Error(w, "ip and ttl_seconds required", http.StatusBadRequest) return } + if !IsIP(req.IP) { + http.Error(w, "invalid IP address", http.StatusBadRequest) + return + } clientIP := getClientIP(r) if err := AddTempEntry(db, req.IP, req.TTLSeconds, req.Reason, clientIP); err != nil { slog.Error("failed to add temp entry", "error", err) @@ -79,7 +104,12 @@ func whitelistListHandler(cfg Config, db *sql.DB) http.HandlerFunc { http.Error(w, "unauthorized", http.StatusUnauthorized) return } - entries, err := ListTempEntries(db) + limit, offset, err := parsePagination(r) + if err != nil { + http.Error(w, "bad request", http.StatusBadRequest) + return + } + entries, err := ListTempEntriesPaginated(db, limit, offset) if err != nil { http.Error(w, "database error", http.StatusInternalServerError) return @@ -95,7 +125,12 @@ func permWhitelistListHandler(cfg Config, db *sql.DB) http.HandlerFunc { http.Error(w, "unauthorized", http.StatusUnauthorized) return } - entries, err := ListPermEntries(db) + limit, offset, err := parsePagination(r) + if err != nil { + http.Error(w, "bad request", http.StatusBadRequest) + return + } + entries, err := ListPermEntriesPaginated(db, limit, offset) if err != nil { http.Error(w, "database error", http.StatusInternalServerError) return @@ -228,6 +263,10 @@ func permWhitelistAddHandler(cfg Config, db *sql.DB) http.HandlerFunc { http.Error(w, "entry required", http.StatusBadRequest) return } + if !IsIP(req.Entry) && !IsCIDR(req.Entry) { + http.Error(w, "invalid entry (must be valid IP or CIDR range)", http.StatusBadRequest) + return + } ok, err := AddPermanentEntry(db, req.Entry, getClientIP(r)) if err != nil { slog.Error("failed to add perm entry", "error", err) @@ -274,7 +313,50 @@ func verifyAPIKey(r *http.Request, expectedToken string) bool { return r.Header.Get("Authorization") == "Bearer "+expectedToken } -func statusHandler(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/plain") - w.Write([]byte("ok\n")) +func statusHandler(db *sql.DB) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + var permCount int + var tempCount int + var dbOK bool + + if err := db.QueryRow(`SELECT COUNT(*) FROM perm_whitelist`).Scan(&permCount); err == nil { + dbOK = true + } + if err := db.QueryRow(`SELECT COUNT(*) FROM temp_whitelist WHERE expires_at > ?`, time.Now().Format(time.RFC3339)).Scan(&tempCount); err == nil { + dbOK = true + } else { + dbOK = false + } + + health := map[string]interface{}{ + "status": "ok", + "uptime": time.Since(startTime).String(), + "perm_whitelist_count": permCount, + "temp_whitelist_count": tempCount, + "db_ok": dbOK, + } + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(health) + } +} + +func parsePagination(r *http.Request) (int, int, error) { + limitStr := r.URL.Query().Get("limit") + offsetStr := r.URL.Query().Get("offset") + + limit := 200 + offset := 0 + + if limitStr != "" { + if parsed, err := strconv.Atoi(limitStr); err == nil && parsed > 0 && parsed <= 1000 { + limit = parsed + } + } + if offsetStr != "" { + if parsed, err := strconv.Atoi(offsetStr); err == nil && parsed >= 0 { + offset = parsed + } + } + return limit, offset, nil } diff --git a/auth_test.go b/auth_test.go new file mode 100644 index 0000000..5d2fe18 --- /dev/null +++ b/auth_test.go @@ -0,0 +1,163 @@ +package main + +import ( + "net/http" + "net/http/httptest" + "testing" +) + +func TestIsIP(t *testing.T) { + tests := []struct { + name string + input string + wantOK bool + }{ + {"ipv4", "192.168.1.1", true}, + {"ipv4 loopback", "127.0.0.1", true}, + {"ipv6", "2001:0db8:85a3:0000:0000:8a2e:0370:7334", true}, + {"ipv6 short", "::1", true}, + {"empty", "", false}, + {"hostname", "example.com", false}, + {"garbage", "not-an-ip!!", false}, + {"partial", "192.168", false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsIP(tt.input); got != tt.wantOK { + t.Errorf("IsIP(%q) = %v, want %v", tt.input, got, tt.wantOK) + } + }) + } +} + +func TestIsCIDR(t *testing.T) { + tests := []struct { + name string + input string + wantOK bool + }{ + {"ipv4 cidr", "192.168.1.0/24", true}, + {"ipv6 cidr", "2001:db8::/32", true}, + {"single host", "10.0.0.0/32", true}, + {"plain ip", "10.0.0.1", false}, + {"missing prefix", "192.168.1.0", false}, + {"garbage", "foo/bar", false}, + {"empty", "", false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsCIDR(tt.input); got != tt.wantOK { + t.Errorf("IsCIDR(%q) = %v, want %v", tt.input, got, tt.wantOK) + } + }) + } +} + +func TestIPMatchesEntry(t *testing.T) { + tests := []struct { + name string + ip string + entry string + wantOK bool + wantErr bool + }{ + {"exact match", "192.168.1.5", "192.168.1.5", true, false}, + {"ipv4 cidr match", "10.0.0.5", "10.0.0.0/24", true, false}, + {"ipv4 cidr no match", "10.0.1.5", "10.0.0.0/24", false, false}, + {"ipv6 cidr match", "2001:db8::1", "2001:db8::/32", true, false}, + {"invalid entry", "10.0.0.1", "not-a-cidr", false, false}, + {"empty entry", "10.0.0.1", "", false, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := ipMatchesEntry(tt.ip, tt.entry) + if tt.wantErr { + if err == nil { + t.Errorf("ipMatchesEntry(%q, %q) = %v, want error", tt.ip, tt.entry, got) + } + } else { + if err != nil { + t.Errorf("ipMatchesEntry(%q, %q) error: %v", tt.ip, tt.entry, err) + } + if got != tt.wantOK { + t.Errorf("ipMatchesEntry(%q, %q) = %v, want %v", tt.ip, tt.entry, got, tt.wantOK) + } + } + }) + } +} + +func TestVerifyAPIKey(t *testing.T) { + token := "super-secret" + + tests := []struct { + name string + auth string + wantOK bool + }{ + {"correct token", "Bearer super-secret", true}, + {"missing Bearer", "super-secret", false}, + {"wrong token", "Bearer wrong", false}, + {"empty", "", false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + r := http.Request{Header: http.Header{}} + r.Header.Set("Authorization", tt.auth) + if got := verifyAPIKey(&r, token); got != tt.wantOK { + t.Errorf("verifyAPIKey(%q) = %v, want %v", tt.auth, got, tt.wantOK) + } + }) + } +} + +func TestApiKeyMiddleware(t *testing.T) { + token := "test-token" + called := false + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.WriteHeader(http.StatusNoContent) + }) + wrapped := apiKeyMiddleware(next, token) + + t.Run("valid key", func(t *testing.T) { + called = false + r := httptest.NewRequest(http.MethodGet, "/test", nil) + r.Header.Set("Authorization", "Bearer test-token") + w := httptest.NewRecorder() + wrapped(w, r) + if !called { + t.Error("middleware did not call next handler") + } + if w.Code != http.StatusNoContent { + t.Errorf("expected 204, got %d", w.Code) + } + }) + + t.Run("invalid key", func(t *testing.T) { + called = false + r := httptest.NewRequest(http.MethodGet, "/test", nil) + r.Header.Set("Authorization", "Bearer wrong") + w := httptest.NewRecorder() + wrapped(w, r) + if called { + t.Error("middleware called next handler despite invalid key") + } + if w.Code != http.StatusUnauthorized { + t.Errorf("expected 401, got %d", w.Code) + } + }) + + t.Run("missing key", func(t *testing.T) { + called = false + r := httptest.NewRequest(http.MethodGet, "/test", nil) + w := httptest.NewRecorder() + wrapped(w, r) + if called { + t.Error("middleware called next handler despite missing key") + } + if w.Code != http.StatusUnauthorized { + t.Errorf("expected 401, got %d", w.Code) + } + }) +} diff --git a/cleanup.go b/cleanup.go index 662be7b..e89318c 100644 --- a/cleanup.go +++ b/cleanup.go @@ -1,13 +1,21 @@ package main import ( + "context" "database/sql" "time" ) -func cleanupLoop(db *sql.DB, interval time.Duration) { +func cleanupLoop(ctx context.Context, db *sql.DB, interval time.Duration) { ticker := time.NewTicker(interval) - for range ticker.C { - _ = CleanupExpiredTemp(db) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + _ = CleanupExpiredTemp(db) + _ = RotateAuditLog(db, 0) + } } } diff --git a/db.go b/db.go index e31ac97..079a1ab 100644 --- a/db.go +++ b/db.go @@ -2,12 +2,12 @@ package main import ( "database/sql" + "errors" "log/slog" "net" - "strings" "time" - _ "modernc.org/sqlite" + sqlite "modernc.org/sqlite" ) func InitDB(dbPath string) (*sql.DB, error) { @@ -271,6 +271,59 @@ func ListPermEntries(db *sql.DB) ([]map[string]interface{}, error) { return list, nil } +func ListTempEntriesPaginated(db *sql.DB, limit, offset int) ([]map[string]interface{}, error) { + rows, err := db.Query(` + SELECT ip, expires_at, reason + FROM temp_whitelist + WHERE expires_at > ? + ORDER BY expires_at + LIMIT ? OFFSET ? + `, time.Now(), limit, offset) + if err != nil { + return nil, err + } + defer rows.Close() + var list []map[string]interface{} = make([]map[string]interface{}, 0) + for rows.Next() { + var ip string + var expiresAt time.Time + var reason sql.NullString + if err := rows.Scan(&ip, &expiresAt, &reason); err != nil { + continue + } + list = append(list, map[string]interface{}{ + "ip": ip, + "expires": expiresAt, + "reason": reason.String, + }) + } + return list, nil +} + +func ListPermEntriesPaginated(db *sql.DB, limit, offset int) ([]map[string]interface{}, error) { + rows, err := db.Query(` + SELECT entry + FROM perm_whitelist + ORDER BY created_at + LIMIT ? OFFSET ? + `, limit, offset) + if err != nil { + return nil, err + } + defer rows.Close() + var list []map[string]interface{} = make([]map[string]interface{}, 0) + for rows.Next() { + var entry string + if err := rows.Scan(&entry); err != nil { + continue + } + list = append(list, map[string]interface{}{ + "entry": entry, + }) + } + return list, nil +} + func ipMatchesEntry(ip, entry string) (bool, error) { if entry == ip { return true, nil @@ -303,5 +356,57 @@ func isUniqueConstraintError(err error) bool { if err == nil { return false } - return strings.Contains(err.Error(), "UNIQUE constraint failed") + var sqliteErr *sqlite.Error + if errors.As(err, &sqliteErr) { + return sqliteErr.Code() == 19 + } + return false +} + +func IsIP(s string) bool { + return net.ParseIP(s) != nil +} + +func IsCIDR(s string) bool { + _, _, err := net.ParseCIDR(s) + return err == nil +} + +func RotateAuditLog(db *sql.DB, maxEntries int) error { + if maxEntries <= 0 { + maxEntries = 200000 + } + var count int + err := db.QueryRow(`SELECT COUNT(*) FROM audit_log`).Scan(&count) + if err != nil { + return err + } + if count <= maxEntries { + return nil + } + rows, err := db.Query(` + SELECT id FROM audit_log + ORDER BY timestamp ASC + LIMIT ? + `, count-maxEntries) + if err != nil { + return err + } + var ids []int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err == nil { + ids = append(ids, id) + } + } + rows.Close() + _, err = db.Exec(` + DELETE FROM audit_log + WHERE id IN ( + SELECT id FROM audit_log + ORDER BY timestamp ASC + LIMIT ? + ) + `, count-maxEntries) + return err } diff --git a/main.go b/main.go index a8dae62..7f5f834 100644 --- a/main.go +++ b/main.go @@ -2,6 +2,7 @@ package main import ( "context" + "log" "log/slog" "net/http" "os" @@ -12,6 +13,9 @@ import ( func main() { cfg := loadConfig() + if cfg.APIToken == "" { + log.Fatal("AUTH_PROXY_API_TOKEN environment variable is required") + } slog.SetDefault(slog.New(slog.NewTextHandler(os.Stdout, nil))) @@ -22,7 +26,10 @@ func main() { } defer db.Close() - go cleanupLoop(db, cfg.CleanupInterval) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + go cleanupLoop(ctx, db, cfg.CleanupInterval) mux := http.NewServeMux() mux.HandleFunc("/auth", authHandler(cfg, db)) @@ -33,7 +40,7 @@ func main() { mux.HandleFunc("/api/whitelist/perm/list", permWhitelistListHandler(cfg, db)) mux.HandleFunc("/api/whitelist/perm/{entry}", permWhitelistDeleteHandler(cfg, db)) mux.HandleFunc("/api/logs", logsHandler(cfg, db)) - mux.HandleFunc("/status", statusHandler) + mux.HandleFunc("/status", statusHandler(db)) server := &http.Server{ Addr: ":" + cfg.Port, @@ -53,7 +60,8 @@ func main() { <-quit slog.Info("shutting down...") - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - server.Shutdown(ctx) + cancel() + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second) + defer shutdownCancel() + server.Shutdown(shutdownCtx) }