Compare commits
26 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b7a1ac06d4 | ||
|
|
49f0acb777 | ||
|
|
a602207369 | ||
|
|
8c0cfcdbe7 | ||
|
|
35a1a89baf | ||
|
|
f3387b169a | ||
|
|
5782072f91 | ||
|
|
7918b3efe6 | ||
|
|
f628e24f58 | ||
|
|
7f54db0cd4 | ||
|
|
2e9b307194 | ||
|
|
726594a712 | ||
|
|
b27038a59c | ||
|
|
72025dab7d | ||
|
|
dd131477e2 | ||
|
|
670aec449a | ||
|
|
fc37e641be | ||
|
|
361de03208 | ||
|
|
a2268fda5d | ||
|
|
9dc0b6002e | ||
|
|
4953be3ef6 | ||
|
|
c386a2d6bc | ||
|
|
dea03a6f70 | ||
|
|
11f755c03c | ||
|
|
1c7a1c1778 | ||
|
|
411574cabe |
@@ -30,6 +30,11 @@ var DaemonCmd = &cobra.Command{
|
||||
log.Error("Failed to create request writer", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
reqDb_r, err := storage.NewRequestsRd()
|
||||
if err != nil {
|
||||
log.Error("Failed to create request reader", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
banDb_r, err := storage.NewBanReader()
|
||||
if err != nil {
|
||||
log.Error("Failed to create ban reader", "error", err)
|
||||
@@ -63,7 +68,7 @@ var DaemonCmd = &cobra.Command{
|
||||
log.Error("Failed to load rules", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
j := judge.New(banDb_r, banDb_w, b, resultCh, entryCh)
|
||||
j := judge.New(banDb_r, banDb_w, reqDb_r, b, resultCh, entryCh)
|
||||
j.LoadRules(r)
|
||||
go j.UnbanChecker()
|
||||
go j.Tribunal()
|
||||
|
||||
@@ -16,13 +16,14 @@ var (
|
||||
port int
|
||||
protocol string
|
||||
)
|
||||
|
||||
var UnbanCmd = &cobra.Command{
|
||||
Use: "unban",
|
||||
Short: "Unban IP",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
err := func() error {
|
||||
if len(args) == 0 {
|
||||
fmt.Println("IP can't be empty")
|
||||
os.Exit(1)
|
||||
return fmt.Errorf("IP can't be empty")
|
||||
}
|
||||
if ttl_fw == "" {
|
||||
ttl_fw = "1y"
|
||||
@@ -30,39 +31,38 @@ var UnbanCmd = &cobra.Command{
|
||||
ip := args[0]
|
||||
db, err := storage.NewBanWriter()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
return err
|
||||
}
|
||||
cfg, err := config.LoadConfig()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
return err
|
||||
}
|
||||
fw := cfg.Firewall.Name
|
||||
b := blocker.GetBlocker(fw, cfg.Firewall.Config)
|
||||
if ip == "" {
|
||||
fmt.Println("IP can't be empty")
|
||||
os.Exit(1)
|
||||
return fmt.Errorf("IP can't be empty")
|
||||
}
|
||||
if net.ParseIP(ip) == nil {
|
||||
fmt.Println("Invalid IP")
|
||||
os.Exit(1)
|
||||
return fmt.Errorf("invalid IP")
|
||||
}
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
return err
|
||||
}
|
||||
err = b.Unban(ip)
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
return err
|
||||
}
|
||||
err = db.RemoveBan(ip)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Println("IP unblocked successfully!")
|
||||
return nil
|
||||
}()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println("IP unblocked successfully!")
|
||||
},
|
||||
}
|
||||
|
||||
@@ -70,9 +70,9 @@ var BanCmd = &cobra.Command{
|
||||
Use: "ban",
|
||||
Short: "Ban IP",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
err := func() error {
|
||||
if len(args) == 0 {
|
||||
fmt.Println("IP can't be empty")
|
||||
os.Exit(1)
|
||||
return fmt.Errorf("IP can't be empty")
|
||||
}
|
||||
if ttl_fw == "" {
|
||||
ttl_fw = "1y"
|
||||
@@ -80,39 +80,38 @@ var BanCmd = &cobra.Command{
|
||||
ip := args[0]
|
||||
db, err := storage.NewBanWriter()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
return err
|
||||
}
|
||||
cfg, err := config.LoadConfig()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
return err
|
||||
}
|
||||
fw := cfg.Firewall.Name
|
||||
b := blocker.GetBlocker(fw, cfg.Firewall.Config)
|
||||
if ip == "" {
|
||||
fmt.Println("IP can't be empty")
|
||||
os.Exit(1)
|
||||
return fmt.Errorf("IP can't be empty")
|
||||
}
|
||||
if net.ParseIP(ip) == nil {
|
||||
fmt.Println("Invalid IP")
|
||||
os.Exit(1)
|
||||
return fmt.Errorf("invalid IP")
|
||||
}
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
return err
|
||||
}
|
||||
err = b.Ban(ip)
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
return err
|
||||
}
|
||||
err = db.AddBan(ip, ttl_fw, "manual ban")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Println("IP blocked successfully!")
|
||||
return nil
|
||||
}()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println("IP blocked successfully!")
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -61,11 +61,12 @@ var ListCmd = &cobra.Command{
|
||||
}
|
||||
for _, rule := range r {
|
||||
fmt.Printf(
|
||||
"Name: %s\nService: %s\nPath: %s\nStatus: %s\nMethod: %s\n\n",
|
||||
"Name: %s\nService: %s\nPath: %s\nStatus: %s\n MaxRetry: %d\nMethod: %s\n\n",
|
||||
rule.Name,
|
||||
rule.ServiceName,
|
||||
rule.Path,
|
||||
rule.Status,
|
||||
rule.MaxRetry,
|
||||
rule.Method,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -6,7 +6,7 @@ import (
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
var version = "0.4.3"
|
||||
var version = "0.5.0"
|
||||
|
||||
var VersionCmd = &cobra.Command{
|
||||
Use: "version",
|
||||
|
||||
@@ -40,6 +40,7 @@ Example:
|
||||
service = "nginx"
|
||||
path = ""
|
||||
status = "304"
|
||||
max_retry = 3
|
||||
method = ""
|
||||
ban_time = "1m"
|
||||
```
|
||||
|
||||
@@ -23,6 +23,7 @@ func (f *Firewalld) Ban(ip string) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// #nosec G204 - ip is validated
|
||||
cmd := exec.Command("firewall-cmd", "--zone=drop", "--add-source", ip, "--permanent")
|
||||
output, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
@@ -44,6 +45,7 @@ func (f *Firewalld) Unban(ip string) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// #nosec G204 - ip is validated
|
||||
cmd := exec.Command("firewall-cmd", "--zone=drop", "--remove-source", ip, "--permanent")
|
||||
output, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
@@ -61,7 +63,7 @@ func (f *Firewalld) Unban(ip string) error {
|
||||
}
|
||||
|
||||
func (f *Firewalld) PortOpen(port int, protocol string) error {
|
||||
// #nosec G204 - handle is extracted from nftables output and validated
|
||||
// #nosec G204 - handle is extracted from Firewalld output and validated
|
||||
if port >= 0 && port <= 65535 {
|
||||
if protocol != "tcp" && protocol != "udp" {
|
||||
f.logger.Error("invalid protocol")
|
||||
@@ -94,7 +96,6 @@ func (f *Firewalld) PortClose(port int, protocol string) error {
|
||||
// #nosec G204 - handle is extracted from nftables output and validated
|
||||
if port >= 0 && port <= 65535 {
|
||||
if protocol != "tcp" && protocol != "udp" {
|
||||
f.logger.Error("invalid protocol")
|
||||
return fmt.Errorf("invalid protocol")
|
||||
}
|
||||
s := strconv.Itoa(port)
|
||||
@@ -106,13 +107,11 @@ func (f *Firewalld) PortClose(port int, protocol string) error {
|
||||
)
|
||||
output, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
f.logger.Error(err.Error())
|
||||
return err
|
||||
}
|
||||
f.logger.Info("Remove port " + s + " " + string(output))
|
||||
output, err = exec.Command("firewall-cmd", "--reload").CombinedOutput()
|
||||
if err != nil {
|
||||
f.logger.Error(err.Error())
|
||||
return err
|
||||
}
|
||||
f.logger.Info("Reload " + string(output))
|
||||
|
||||
@@ -28,6 +28,7 @@ func (f *Iptables) Ban(ip string) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// #nosec G204 - f.config is validated above via validateConfigPath()
|
||||
cmd := exec.Command("iptables", "-A", "INPUT", "-s", ip, "-j", "DROP")
|
||||
output, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
@@ -70,6 +71,7 @@ func (f *Iptables) Unban(ip string) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// #nosec G204 - f.config is validated above via validateConfigPath()
|
||||
cmd := exec.Command("iptables", "-D", "INPUT", "-s", ip, "-j", "DROP")
|
||||
output, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
|
||||
@@ -26,7 +26,7 @@ func (n *Nftables) Ban(ip string) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// #nosec G204 - ip is validated
|
||||
cmd := exec.Command("nft", "add", "rule", "inet", "banforge", "banned",
|
||||
"ip", "saddr", ip, "drop")
|
||||
output, err := cmd.CombinedOutput()
|
||||
@@ -113,6 +113,7 @@ func (n *Nftables) Setup(config string) error {
|
||||
}
|
||||
}
|
||||
`
|
||||
// #nosec G204 - config is managed by adminstartor
|
||||
cmd := exec.Command("tee", config)
|
||||
stdin, err := cmd.StdinPipe()
|
||||
if err != nil {
|
||||
@@ -135,7 +136,7 @@ func (n *Nftables) Setup(config string) error {
|
||||
if err = cmd.Wait(); err != nil {
|
||||
return fmt.Errorf("failed to save config: %w", err)
|
||||
}
|
||||
|
||||
// #nosec G204 - config is managed by adminstartor
|
||||
cmd = exec.Command("nft", "-f", config)
|
||||
output, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
@@ -253,7 +254,7 @@ func saveNftablesConfig(configPath string) error {
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get nftables ruleset: %w", err)
|
||||
}
|
||||
|
||||
// #nosec G204 - managed by system adminstartor
|
||||
cmd = exec.Command("tee", configPath)
|
||||
stdin, err := cmd.StdinPipe()
|
||||
if err != nil {
|
||||
|
||||
@@ -23,7 +23,7 @@ func (u *Ufw) Ban(ip string) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// #nosec G204 - ip is validated
|
||||
cmd := exec.Command("ufw", "--force", "deny", "from", ip)
|
||||
output, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
@@ -42,7 +42,7 @@ func (u *Ufw) Unban(ip string) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// #nosec G204 - ip is validated
|
||||
cmd := exec.Command("ufw", "--force", "delete", "deny", "from", ip)
|
||||
output, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
@@ -57,13 +58,9 @@ func NewRule(
|
||||
return err
|
||||
}
|
||||
defer func() {
|
||||
err = file.Close()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
}
|
||||
err = errors.Join(err, file.Close())
|
||||
}()
|
||||
cfg := Rules{Rules: r}
|
||||
|
||||
err = toml.NewEncoder(file).Encode(cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -126,24 +123,24 @@ func EditRule(Name string, ServiceName string, Path string, Status string, Metho
|
||||
}
|
||||
|
||||
func ParseDurationWithYears(s string) (time.Duration, error) {
|
||||
if strings.HasSuffix(s, "y") {
|
||||
years, err := strconv.Atoi(strings.TrimSuffix(s, "y"))
|
||||
if ss, ok := strings.CutSuffix(s, "y"); ok {
|
||||
years, err := strconv.Atoi(ss)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return time.Duration(years) * 365 * 24 * time.Hour, nil
|
||||
}
|
||||
|
||||
if strings.HasSuffix(s, "M") {
|
||||
months, err := strconv.Atoi(strings.TrimSuffix(s, "M"))
|
||||
if ss, ok := strings.CutSuffix(s, "M"); ok {
|
||||
months, err := strconv.Atoi(ss)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return time.Duration(months) * 30 * 24 * time.Hour, nil
|
||||
}
|
||||
|
||||
if strings.HasSuffix(s, "d") {
|
||||
days, err := strconv.Atoi(strings.TrimSuffix(s, "d"))
|
||||
if ss, ok := strings.CutSuffix(s, "d"); ok {
|
||||
days, err := strconv.Atoi(ss)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
@@ -28,5 +28,6 @@ type Rule struct {
|
||||
Path string `toml:"path"`
|
||||
Status string `toml:"status"`
|
||||
Method string `toml:"method"`
|
||||
MaxRetry int `toml:"max_retry"`
|
||||
BanTime string `toml:"ban_time"`
|
||||
}
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
type Judge struct {
|
||||
db_r *storage.BanReader
|
||||
db_w *storage.BanWriter
|
||||
db_rq *storage.RequestReader
|
||||
logger *logger.Logger
|
||||
Blocker blocker.BlockerEngine
|
||||
rulesByService map[string][]config.Rule
|
||||
@@ -24,6 +25,7 @@ type Judge struct {
|
||||
func New(
|
||||
db_r *storage.BanReader,
|
||||
db_w *storage.BanWriter,
|
||||
db_rq *storage.RequestReader,
|
||||
b blocker.BlockerEngine,
|
||||
resultCh chan *storage.LogEntry,
|
||||
entryCh chan *storage.LogEntry,
|
||||
@@ -31,6 +33,7 @@ func New(
|
||||
return &Judge{
|
||||
db_w: db_w,
|
||||
db_r: db_r,
|
||||
db_rq: db_rq,
|
||||
logger: logger.New(false),
|
||||
rulesByService: make(map[string][]config.Rule),
|
||||
Blocker: b,
|
||||
@@ -75,31 +78,28 @@ func (j *Judge) Tribunal() {
|
||||
methodMatch := rule.Method == "" || entry.Method == rule.Method
|
||||
statusMatch := rule.Status == "" || entry.Status == rule.Status
|
||||
pathMatch := matchPath(entry.Path, rule.Path)
|
||||
|
||||
j.logger.Debug(
|
||||
"Testing rule",
|
||||
"rule", rule.Name,
|
||||
"method_match", methodMatch,
|
||||
"status_match", statusMatch,
|
||||
"path_match", pathMatch,
|
||||
)
|
||||
|
||||
if methodMatch && statusMatch && pathMatch {
|
||||
ruleMatched = true
|
||||
j.logger.Info("Rule matched", "rule", rule.Name, "ip", entry.IP)
|
||||
|
||||
j.resultCh <- entry
|
||||
banned, err := j.db_r.IsBanned(entry.IP)
|
||||
if err != nil {
|
||||
j.logger.Error("Failed to check ban status", "ip", entry.IP, "error", err)
|
||||
break
|
||||
}
|
||||
|
||||
if banned {
|
||||
j.logger.Info("IP already banned", "ip", entry.IP)
|
||||
j.resultCh <- entry
|
||||
break
|
||||
}
|
||||
|
||||
exceeded, err := j.db_rq.IsMaxRetryExceeded(entry.IP, rule.MaxRetry)
|
||||
if err != nil {
|
||||
j.logger.Error("Failed to check retry count", "ip", entry.IP, "error", err)
|
||||
break
|
||||
}
|
||||
if !exceeded {
|
||||
j.logger.Info("Max retry not exceeded", "ip", entry.IP)
|
||||
break
|
||||
}
|
||||
err = j.db_w.AddBan(entry.IP, rule.BanTime, rule.Name)
|
||||
if err != nil {
|
||||
j.logger.Error(
|
||||
@@ -127,7 +127,6 @@ func (j *Judge) Tribunal() {
|
||||
"ban_time",
|
||||
rule.BanTime,
|
||||
)
|
||||
j.resultCh <- entry
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
@@ -24,6 +24,7 @@ type Scanner struct {
|
||||
}
|
||||
|
||||
func NewScannerTail(path string) (*Scanner, error) {
|
||||
// #nosec G204 - managed by system adminstartor
|
||||
cmd := exec.Command("tail", "-F", "-n", "10", path)
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
@@ -46,6 +47,7 @@ func NewScannerTail(path string) (*Scanner, error) {
|
||||
}
|
||||
|
||||
func NewScannerJournald(unit string) (*Scanner, error) {
|
||||
// #nosec G204 - managed by system adminstartor
|
||||
cmd := exec.Command("journalctl", "-u", unit, "-f", "-n", "0", "-o", "short", "--no-pager")
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
|
||||
@@ -21,7 +21,7 @@ type BanWriter struct {
|
||||
func NewBanWriter() (*BanWriter, error) {
|
||||
db, err := sql.Open(
|
||||
"sqlite",
|
||||
"/var/lib/banforge/bans.db?_pragma=journal_mode(WAL)&_pragma=busy_timeout(30000)&_pragma=synchronous(NORMAL)",
|
||||
buildSqliteDsn(banDBPath, pragmas),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -175,7 +175,6 @@ func (d *BanReader) IsBanned(ip string) (bool, error) {
|
||||
}
|
||||
|
||||
func (d *BanReader) BanList() error {
|
||||
|
||||
var count int
|
||||
t := table.NewWriter()
|
||||
t.SetOutputMirror(os.Stdout)
|
||||
|
||||
@@ -2,55 +2,60 @@ package storage
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
func CreateTables() error {
|
||||
const (
|
||||
DBDir = "/var/lib/banforge/"
|
||||
ReqDBPath = DBDir + "requests.db"
|
||||
banDBPath = DBDir + "bans.db"
|
||||
)
|
||||
|
||||
var pragmas = map[string]string{
|
||||
`journal_mode`: `wal`,
|
||||
`synchronous`: `normal`,
|
||||
`busy_timeout`: `30000`,
|
||||
// also consider these
|
||||
// `temp_store`: `memory`,
|
||||
// `cache_size`: `1000000000`,
|
||||
}
|
||||
|
||||
func buildSqliteDsn(path string, pragmas map[string]string) string {
|
||||
pragmastrs := make([]string, len(pragmas))
|
||||
i := 0
|
||||
for k, v := range pragmas {
|
||||
pragmastrs[i] = (fmt.Sprintf(`pragma=%s(%s)`, k, v))
|
||||
i++
|
||||
}
|
||||
return path + "?" + "mode=rwc&" + strings.Join(pragmastrs, "&")
|
||||
}
|
||||
|
||||
func initDB(dsn, sqlstr string) (err error) {
|
||||
db, err := sql.Open("sqlite", dsn)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to open %q: %w", dsn, err)
|
||||
}
|
||||
defer func() {
|
||||
closeErr := db.Close()
|
||||
if closeErr != nil {
|
||||
err = errors.Join(err, fmt.Errorf("failed to close %q: %w", dsn, closeErr))
|
||||
}
|
||||
}()
|
||||
_, err = db.Exec(sqlstr)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create table: %w", err)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func CreateTables() (err error) {
|
||||
// Requests DB
|
||||
db_r, err := sql.Open("sqlite",
|
||||
"/var/lib/banforge/requests.db?"+
|
||||
"mode=rwc&"+
|
||||
"_pragma=journal_mode(WAL)&"+
|
||||
"_pragma=busy_timeout(30000)&"+
|
||||
"_pragma=synchronous(NORMAL)")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to open requests db: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
err = db_r.Close()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
}
|
||||
}()
|
||||
err1 := initDB(buildSqliteDsn(ReqDBPath, pragmas), CreateRequestsTable)
|
||||
err2 := initDB(buildSqliteDsn(banDBPath, pragmas), CreateBansTable)
|
||||
|
||||
_, err = db_r.Exec(CreateRequestsTable)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create requests table: %w", err)
|
||||
}
|
||||
|
||||
// Bans DB
|
||||
db_b, err := sql.Open("sqlite",
|
||||
"/var/lib/banforge/bans.db?"+
|
||||
"mode=rwc&"+
|
||||
"_pragma=journal_mode(WAL)&"+
|
||||
"_pragma=busy_timeout(30000)&"+
|
||||
"_pragma=synchronous(FULL)")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to open bans db: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
err = db_b.Close()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
}
|
||||
}()
|
||||
|
||||
_, err = db_b.Exec(CreateBansTable)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create bans table: %w", err)
|
||||
}
|
||||
fmt.Println("Tables created successfully!")
|
||||
return nil
|
||||
return errors.Join(err1, err2)
|
||||
}
|
||||
|
||||
@@ -7,15 +7,15 @@ import (
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
type Request_Writer struct {
|
||||
type RequestWriter struct {
|
||||
logger *logger.Logger
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
func NewRequestsWr() (*Request_Writer, error) {
|
||||
func NewRequestsWr() (*RequestWriter, error) {
|
||||
db, err := sql.Open(
|
||||
"sqlite",
|
||||
"/var/lib/banforge/requests.db?_pragma=journal_mode(WAL)&_pragma=busy_timeout(30000)&_pragma=synchronous(NORMAL)",
|
||||
buildSqliteDsn(ReqDBPath, pragmas),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -23,8 +23,41 @@ func NewRequestsWr() (*Request_Writer, error) {
|
||||
db.SetMaxOpenConns(1)
|
||||
db.SetMaxIdleConns(1)
|
||||
db.SetConnMaxLifetime(0)
|
||||
return &Request_Writer{
|
||||
return &RequestWriter{
|
||||
logger: logger.New(false),
|
||||
db: db,
|
||||
}, nil
|
||||
}
|
||||
|
||||
type RequestReader struct {
|
||||
logger *logger.Logger
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
func NewRequestsRd() (*RequestReader, error) {
|
||||
db, err := sql.Open(
|
||||
"sqlite",
|
||||
buildSqliteDsn(ReqDBPath, pragmas),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
db.SetMaxOpenConns(1)
|
||||
db.SetMaxIdleConns(1)
|
||||
db.SetConnMaxLifetime(0)
|
||||
return &RequestReader{
|
||||
logger: logger.New(false),
|
||||
db: db,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *RequestReader) IsMaxRetryExceeded(ip string, maxRetry int) (bool, error) {
|
||||
var count int
|
||||
err := r.db.QueryRow("SELECT COUNT(*) FROM requests WHERE ip = ?", ip).Scan(&count)
|
||||
if err != nil {
|
||||
r.logger.Error("error query count: " + err.Error())
|
||||
return false, err
|
||||
}
|
||||
r.logger.Info("Current request count for IP", "ip", ip, "count", count, "maxRetry", maxRetry)
|
||||
return count >= maxRetry, nil
|
||||
}
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
func WriteReq(db *Request_Writer, resultCh <-chan *LogEntry) {
|
||||
func WriteReq(db *RequestWriter, resultCh <-chan *LogEntry) {
|
||||
db.logger.Info("Starting log writer")
|
||||
const batchSize = 100
|
||||
const flushInterval = 1 * time.Second
|
||||
@@ -14,29 +17,36 @@ func WriteReq(db *Request_Writer, resultCh <-chan *LogEntry) {
|
||||
defer ticker.Stop()
|
||||
|
||||
flush := func() {
|
||||
defer db.logger.Debug("Flushed batch", "count", len(batch))
|
||||
err := func() (err error) {
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
return nil
|
||||
}
|
||||
|
||||
tx, err := db.db.Begin()
|
||||
if err != nil {
|
||||
db.logger.Error("Failed to begin transaction", "error", err)
|
||||
return
|
||||
return fmt.Errorf("failed to begin transaction: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if rollbackErr := tx.Rollback(); rollbackErr != nil &&
|
||||
!errors.Is(rollbackErr, sql.ErrTxDone) {
|
||||
err = errors.Join(
|
||||
err,
|
||||
fmt.Errorf("failed to rollback transaction: %w", rollbackErr),
|
||||
)
|
||||
}
|
||||
}()
|
||||
|
||||
stmt, err := tx.Prepare(
|
||||
"INSERT INTO requests (service, ip, path, method, status, created_at) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
if err != nil {
|
||||
db.logger.Error("Failed to prepare statement", "error", err)
|
||||
if rollbackErr := tx.Rollback(); rollbackErr != nil {
|
||||
db.logger.Error("Failed to rollback transaction", "error", rollbackErr)
|
||||
}
|
||||
return
|
||||
err = fmt.Errorf("failed to prepare statement: %w", err)
|
||||
return err
|
||||
}
|
||||
defer func() {
|
||||
if closeErr := stmt.Close(); closeErr != nil {
|
||||
db.logger.Error("Failed to close statement", "error", closeErr)
|
||||
err = errors.Join(err, fmt.Errorf("failed to close statement: %w", closeErr))
|
||||
}
|
||||
}()
|
||||
|
||||
@@ -50,17 +60,20 @@ func WriteReq(db *Request_Writer, resultCh <-chan *LogEntry) {
|
||||
time.Now().Format(time.RFC3339),
|
||||
)
|
||||
if err != nil {
|
||||
db.logger.Error("Failed to insert entry", "error", err)
|
||||
db.logger.Error(fmt.Errorf("failed to insert entry: %w", err).Error())
|
||||
}
|
||||
}
|
||||
|
||||
if err := tx.Commit(); err != nil {
|
||||
db.logger.Error("Failed to commit transaction", "error", err)
|
||||
return
|
||||
return fmt.Errorf("failed to commit transaction: %w", err)
|
||||
}
|
||||
|
||||
db.logger.Debug("Flushed batch", "count", len(batch))
|
||||
batch = batch[:0]
|
||||
return err
|
||||
}()
|
||||
if err != nil {
|
||||
db.logger.Error(err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
for {
|
||||
@@ -81,3 +94,13 @@ func WriteReq(db *Request_Writer, resultCh <-chan *LogEntry) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (w *RequestWriter) GetRequestCount() (int, error) {
|
||||
var count int
|
||||
err := w.db.QueryRow("SELECT COUNT(*) FROM requests").Scan(&count)
|
||||
return count, err
|
||||
}
|
||||
|
||||
func (w *RequestWriter) Close() error {
|
||||
return w.db.Close()
|
||||
}
|
||||
|
||||
@@ -277,7 +277,7 @@ func TestWrite_ChannelClosed(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func NewRequestWriterWithDBPath(dbPath string) (*Request_Writer, error) {
|
||||
func NewRequestWriterWithDBPath(dbPath string) (*RequestWriter, error) {
|
||||
db, err := sql.Open("sqlite", dbPath+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(30000)&_pragma=synchronous(NORMAL)")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -285,13 +285,13 @@ func NewRequestWriterWithDBPath(dbPath string) (*Request_Writer, error) {
|
||||
db.SetMaxOpenConns(1)
|
||||
db.SetMaxIdleConns(1)
|
||||
db.SetConnMaxLifetime(0)
|
||||
return &Request_Writer{
|
||||
return &RequestWriter{
|
||||
logger: logger.New(false),
|
||||
db: db,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (w *Request_Writer) CreateTable() error {
|
||||
func (w *RequestWriter) CreateTable() error {
|
||||
_, err := w.db.Exec(CreateRequestsTable)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -299,21 +299,3 @@ func (w *Request_Writer) CreateTable() error {
|
||||
w.logger.Info("Created requests table")
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Request_Writer) Close() error {
|
||||
w.logger.Info("Closing request database connection")
|
||||
err := w.db.Close()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Request_Writer) GetRequestCount() (int, error) {
|
||||
var count int
|
||||
err := w.db.QueryRow("SELECT COUNT(*) FROM requests").Scan(&count)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user