medol-be/app/middleware/register.middleware.go

package middleware

import (
	"github.com/gofiber/fiber/v2/middleware/csrf"
	"github.com/gofiber/fiber/v2/middleware/session"
	"go-humas-be/app/database"
	"go-humas-be/config/config"
	utilsSvc "go-humas-be/utils"
	"log"
	"time"

	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/middleware/compress"
	"github.com/gofiber/fiber/v2/middleware/cors"
	"github.com/gofiber/fiber/v2/middleware/limiter"
	"github.com/gofiber/fiber/v2/middleware/monitor"
	"github.com/gofiber/fiber/v2/middleware/pprof"
	"github.com/gofiber/fiber/v2/middleware/recover"
	"github.com/gofiber/fiber/v2/utils"
)

// Middleware is a struct that contains all the middleware functions
type Middleware struct {
	App *fiber.App
	Cfg *config.Config
}

func NewMiddleware(app *fiber.App, cfg *config.Config) *Middleware {
	return &Middleware{
		App: app,
		Cfg: cfg,
	}
}

// Register registers all the middleware functions
func (m *Middleware) Register(db *database.Database) {
	// Add Extra Middlewares

	m.App.Use(limiter.New(limiter.Config{
		Next:       utilsSvc.IsEnabled(m.Cfg.Middleware.Limiter.Enable),
		Max:        m.Cfg.Middleware.Limiter.Max,
		Expiration: m.Cfg.Middleware.Limiter.Expiration * time.Second,
	}))

	m.App.Use(compress.New(compress.Config{
		Next:  utilsSvc.IsEnabled(m.Cfg.Middleware.Compress.Enable),
		Level: m.Cfg.Middleware.Compress.Level,
	}))

	m.App.Use(recover.New(recover.Config{
		Next: utilsSvc.IsEnabled(m.Cfg.Middleware.Recover.Enable),
	}))

	m.App.Use(pprof.New(pprof.Config{
		Next: utilsSvc.IsEnabled(m.Cfg.Middleware.Pprof.Enable),
	}))

	m.App.Use(cors.New(cors.Config{
		Next:             utilsSvc.IsEnabled(m.Cfg.Middleware.Cors.Enable),
		AllowOrigins:     "*",
		AllowMethods:     "HEAD, GET, POST, PUT, DELETE, OPTION, PATCH",
		AllowHeaders:     "Origin, Content-Type, Accept, Accept-Language, Authorization, X-Requested-With, Access-Control-Request-Method, Access-Control-Request-Headers",
		ExposeHeaders:    "Content-Length, Content-Type",
		AllowCredentials: false,
		MaxAge:           12,
	}))

	//===============================
	// CSRF CONFIG
	//===============================

	// Custom storage for CSRF
	csrfSessionStorage := &PostgresStorage{
		DB: db.DB,
	}

	// Store initialization for session
	store := session.New(session.Config{
		Storage: csrfSessionStorage,
	})

	m.App.Use(func(c *fiber.Ctx) error {
		sess, err := store.Get(c)
		if err != nil {
			return err
		}
		c.Locals("session", sess)
		return c.Next()
	})

	// Cleanup the expired token
	go func() {
		ticker := time.NewTicker(1 * time.Hour)
		defer ticker.Stop()

		for range ticker.C {
			if err := csrfSessionStorage.Reset(); err != nil {
				log.Printf("Error cleaning up expired CSRF tokens: %v", err)
			}
		}
	}()

	m.App.Use(csrf.New(csrf.Config{
		Next:              utilsSvc.IsEnabled(m.Cfg.Middleware.Csrf.Enable),
		KeyLookup:         "header:" + csrf.HeaderName,
		CookieName:        "csrf_",
		CookieSameSite:    "Lax",
		CookieSecure:      false,
		CookieSessionOnly: true,
		CookieHTTPOnly:    true,
		Expiration:        1 * time.Hour,
		KeyGenerator:      utils.UUIDv4,
		ContextKey:        "csrf",
		ErrorHandler: func(c *fiber.Ctx, err error) error {
			return utilsSvc.CsrfErrorHandler(c, err)
		},
		Extractor:  csrf.CsrfFromHeader(csrf.HeaderName),
		Session:    store,
		SessionKey: "fiber.csrf.token",
	}))

	//===============================

	m.App.Use(AuditTrailsMiddleware(db.DB))
	StartAuditTrailCleanup(db.DB, m.Cfg.Middleware.AuditTrails.Retention)

	//m.App.Use(filesystem.New(filesystem.Config{
	//	Next:   utils.IsEnabled(m.Cfg.Middleware.FileSystem.Enable),
	//	Root:   http.Dir(m.Cfg.Middleware.FileSystem.Root),
	//	Browse: m.Cfg.Middleware.FileSystem.Browse,
	//	MaxAge: m.Cfg.Middleware.FileSystem.MaxAge,
	//}))

	// ==================================================

	m.App.Get(m.Cfg.Middleware.Monitor.Path, monitor.New(monitor.Config{
		Next: utilsSvc.IsEnabled(m.Cfg.Middleware.Monitor.Enable),
	}))

	// Route for generate CSRF token
	m.App.Get("/csrf-token", func(c *fiber.Ctx) error {
		// Retrieve CSRF token from Fiber's middleware context
		token, ok := c.Locals("csrf").(string)

		//c.Context().VisitUserValues(func(key []byte, value interface{}) {
		//	log.Printf("Local Key: %s, Value: %v", key, value)
		//})

		if !ok || token == "" {
			return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
				"success":  false,
				"code":     500,
				"messages": []string{"Failed to retrieve CSRF token"},
			})
		}

		return c.JSON(fiber.Map{
			"success":    true,
			"csrf_token": token,
		})
	})
}