From f3f9ebb5af753fb07626a3e05965346fc8677006 Mon Sep 17 00:00:00 2001
From: LittleSheep <littlesheep.code@hotmail.com>
Date: Sat, 29 Mar 2025 00:26:05 +0800
Subject: [PATCH] :sparkles: IP Blacklist

---
 pkg/internal/web/blacklist.go | 34 ++++++++++++++++++++++++++++++++++
 pkg/internal/web/server.go    | 10 +++++++++-
 pkg/main.go                   |  2 ++
 settings.toml                 |  1 +
 4 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100644 pkg/internal/web/blacklist.go

diff --git a/pkg/internal/web/blacklist.go b/pkg/internal/web/blacklist.go
new file mode 100644
index 0000000..f2c1708
--- /dev/null
+++ b/pkg/internal/web/blacklist.go
@@ -0,0 +1,34 @@
+package web
+
+import (
+	"bufio"
+	"os"
+
+	"github.com/rs/zerolog/log"
+)
+
+var ipBlocklist []string
+
+func AddBlockIP(ip string) {
+	ipBlocklist = append(ipBlocklist, ip)
+}
+
+func ParseBlockIPList(path string) error {
+	file, err := os.Open(path)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		ipBlocklist = append(ipBlocklist, scanner.Text())
+	}
+
+	if err := scanner.Err(); err != nil {
+		return err
+	}
+
+	log.Info().Int("count", len(ipBlocklist)).Msg("Block IP list parsed successfully.")
+	return nil
+}
diff --git a/pkg/internal/web/server.go b/pkg/internal/web/server.go
index 50c402a..795779e 100644
--- a/pkg/internal/web/server.go
+++ b/pkg/internal/web/server.go
@@ -51,7 +51,13 @@ func NewServer() *WebApp {
 		},
 	}))
 
-	app.Use(auth.ContextMiddleware)
+	app.Use(func(c *fiber.Ctx) error {
+		if lo.Contains(ipBlocklist, c.IP()) {
+			return fiber.NewError(fiber.StatusForbidden, "your ip has been listed in the blacklist")
+		}
+		return c.Next()
+	})
+
 	app.Use(limiter.New(limiter.Config{
 		Max:               viper.GetInt("rate_limit"),
 		Expiration:        60 * time.Second,
@@ -69,6 +75,8 @@ func NewServer() *WebApp {
 		},
 	}))
 
+	app.Use(auth.ContextMiddleware)
+
 	api.MapControllers(app)
 
 	return &WebApp{app}
diff --git a/pkg/main.go b/pkg/main.go
index aae8569..2c6983d 100644
--- a/pkg/main.go
+++ b/pkg/main.go
@@ -50,6 +50,8 @@ func main() {
 	// Load settings
 	if err := viper.ReadInConfig(); err != nil {
 		log.Panic().Err(err).Msg("An error occurred when loading settings.")
+	} else if err := web.ParseBlockIPList(viper.GetString("ip_block_path")); err != nil {
+		log.Error().Err(err).Msg("An error occurred when parsing block IP list.")
 	}
 
 	// Connect to kv (etcd)
diff --git a/settings.toml b/settings.toml
index d47eff3..533ebf5 100644
--- a/settings.toml
+++ b/settings.toml
@@ -3,6 +3,7 @@ grpc_bind = "0.0.0.0:7001"
 domain = "localhost"
 
 templates_dir = "./templates"
+ip_block_path = "./ip_block.list"
 
 rate_limit = 120
 rate_limit_advance = 60