diff --git a/internal/middleware/country.go b/internal/middleware/country.go new file mode 100644 index 0000000..7b27bcf --- /dev/null +++ b/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) + }) +} diff --git a/internal/middleware/middleware.go b/internal/middleware/middleware.go index 5f7dbbf..246e3e4 100644 --- a/internal/middleware/middleware.go +++ b/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" - log := logger.NewLoggerWithOp(op) + op := "middleware.RemoveTrailingSlashMiddleware" + log := logger.NewLoggerWithOp(op) - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - urlPath := r.URL.Path + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + urlPath := r.URL.Path - // Skip redirect for Socket.IO and WebSocket upgrade traffic - // Socket.IO may rely on a trailing slash in polling endpoints. - if strings.HasPrefix(urlPath, "/socket.io") || strings.HasPrefix(urlPath, "/ws/socket.io") || - strings.EqualFold(r.Header.Get("Upgrade"), "websocket") || - strings.Contains(strings.ToLower(r.Header.Get("Connection")), "upgrade") { - next.ServeHTTP(w, r) - return - } + // Skip redirect for Socket.IO and WebSocket upgrade traffic + // Socket.IO may rely on a trailing slash in polling endpoints. + if strings.HasPrefix(urlPath, "/socket.io") || strings.HasPrefix(urlPath, "/ws/socket.io") || + strings.EqualFold(r.Header.Get("Upgrade"), "websocket") || + strings.Contains(strings.ToLower(r.Header.Get("Connection")), "upgrade") { + next.ServeHTTP(w, r) + return + } - if urlPath != "" && urlPath != "/" && strings.HasSuffix(urlPath, "/") { - newPath := strings.TrimSuffix(urlPath, "/") - newURL := *r.URL - newURL.Path = newPath + if urlPath != "" && urlPath != "/" && strings.HasSuffix(urlPath, "/") { + newPath := strings.TrimSuffix(urlPath, "/") + newURL := *r.URL + 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 diff --git a/internal/server/server.go b/internal/server/server.go index 0088101..ad596fb 100644 --- a/internal/server/server.go +++ b/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( - middleware.RealIPMiddleware( - proxyHandler, + chain := countryMW.AddCountryHeaderIPMiddleware( + middleware.RemoveTrailingSlashMiddleware( + middleware.RealIPMiddleware( + proxyHandler, + ), ), )