[feat]:

- add country middleware
- refactor middleware chain

internal/middleware/country.go

@@ -0,0 +1,31 @@
+package middleware
+
+import (
+	"net/http"
+	"yobble-gateway-go/internal/logger"
+	"yobble-gateway-go/pkg/geoip"
+)
+
+type CountryMW struct {
+	Geo *geoip.GeoIPService
+}
+
+func (m *CountryMW) AddCountryHeaderIPMiddleware(next http.Handler) http.Handler {
+	op := "middleware.AddCountryHeaderIPMiddleware"
+	log := logger.NewLoggerWithOp(op)
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+
+		realIP := GetRealIP(r)
+		log.Debug(op, "realIP", realIP)
+		_, countryCode := m.Geo.IsCountryBlocked(realIP)
+		log.Debug(op, "countryCode", countryCode)
+
+		if countryCode != "" {
+			r.Header.Set("X-Country-Code", countryCode)
+		} else {
+			r.Header.Set("X-Country-Code", "Undefined")
+		}
+
+		next.ServeHTTP(w, r)
+	})
+}

internal/middleware/middleware.go

@@ -56,25 +56,25 @@ func GetRealIP(r *http.Request) string {
 // RemoveTrailingSlashMiddleware перенаправляет запросы с завершающим слэшем (если это не просто '/')
 // на тот же путь без завершающего слэша.
 func RemoveTrailingSlashMiddleware(next http.Handler) http.Handler {
-    op := "middleware.RemoveTrailingSlashMiddleware"
+	op := "middleware.RemoveTrailingSlashMiddleware"
-    log := logger.NewLoggerWithOp(op)
+	log := logger.NewLoggerWithOp(op)
 
-    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-        urlPath := r.URL.Path
+		urlPath := r.URL.Path
 
-        // Skip redirect for Socket.IO and WebSocket upgrade traffic
+		// Skip redirect for Socket.IO and WebSocket upgrade traffic
-        // Socket.IO may rely on a trailing slash in polling endpoints.
+		// Socket.IO may rely on a trailing slash in polling endpoints.
-        if strings.HasPrefix(urlPath, "/socket.io") || strings.HasPrefix(urlPath, "/ws/socket.io") ||
+		if strings.HasPrefix(urlPath, "/socket.io") || strings.HasPrefix(urlPath, "/ws/socket.io") ||
-            strings.EqualFold(r.Header.Get("Upgrade"), "websocket") ||
+			strings.EqualFold(r.Header.Get("Upgrade"), "websocket") ||
-            strings.Contains(strings.ToLower(r.Header.Get("Connection")), "upgrade") {
+			strings.Contains(strings.ToLower(r.Header.Get("Connection")), "upgrade") {
-            next.ServeHTTP(w, r)
+			next.ServeHTTP(w, r)
-            return
+			return
-        }
+		}
 
-        if urlPath != "" && urlPath != "/" && strings.HasSuffix(urlPath, "/") {
+		if urlPath != "" && urlPath != "/" && strings.HasSuffix(urlPath, "/") {
-            newPath := strings.TrimSuffix(urlPath, "/")
+			newPath := strings.TrimSuffix(urlPath, "/")
-            newURL := *r.URL
+			newURL := *r.URL
-            newURL.Path = newPath
+			newURL.Path = newPath
 			log.Debug("redirecting trailing slash", slog.String("old_path", urlPath), slog.String("new_path", newPath))
 			http.Redirect(w, r, newURL.String(), http.StatusMovedPermanently)
 			return

internal/server/server.go

@@ -31,10 +31,14 @@ func NewServer(cfg *config.Settings, geoIPService *geoip.GeoIPService) *Server {
 	// Create the main proxy handler
 	proxyHandler := proxy.NewProxyHandler(cfg, geoIPService)
 
+	//Initialize country middleware
+	countryMW := &middleware.CountryMW{Geo: geoIPService}
 	// Apply middleware chain
-	chain := middleware.RemoveTrailingSlashMiddleware(
+	chain := countryMW.AddCountryHeaderIPMiddleware(
-		middleware.RealIPMiddleware(
+		middleware.RemoveTrailingSlashMiddleware(
-			proxyHandler,
+			middleware.RealIPMiddleware(
+				proxyHandler,
+			),
 		),
 	)