kontenhumas-be/app/middleware/csrf.middleware.go

package middleware

import (
	"fmt"
	"gorm.io/gorm"
	"netidhub-saas-be/app/database/entity"
	"time"
)

type PostgresStorage struct {
	DB *gorm.DB
}

func (s *PostgresStorage) Get(key string) ([]byte, error) {
	//log.Printf("CSRF Storage: Get token %s", key)

	var record entity.CsrfTokenRecords
	result := s.DB.Where("token = ?", key).First(&record)

	if result.Error != nil {
		//log.Printf("CSRF Storage Get error: %v for token: %s", result.Error, key)
		return nil, result.Error
	}

	if record.ExpireAt.Before(time.Now()) {
		//log.Printf("CSRF token %s is expired", key)
		return nil, fmt.Errorf("CSRF token is expired")
	}

	return record.Value, nil
}

func (s *PostgresStorage) Set(key string, value []byte, exp time.Duration) error {
	//log.Printf("CSRF Storage: Setting token %s with expiration %v", key, exp)

	// Calculate expiration time
	expireAt := time.Now().Add(exp)

	// Try to update existing record first
	result := s.DB.Model(&entity.CsrfTokenRecords{}).
		Where("token = ?", key).
		Updates(map[string]interface{}{
			"expire_at": expireAt,
		})

	// If no rows were affected (not found), create a new record
	if result.RowsAffected == 0 {
		record := entity.CsrfTokenRecords{
			Token:     key,
			Value:     value,
			ExpireAt:  expireAt,
			CreatedAt: time.Now(),
		}

		if err := s.DB.Create(&record).Error; err != nil {
			//log.Printf("CSRF Storage: Error saving token: %v", err)
			return err
		}
	} else if result.Error != nil {
		//log.Printf("CSRF Storage: Error updating token: %v", result.Error)
		return result.Error
	}

	//log.Printf("CSRF Storage: Successfully saved/updated token")
	return nil
}

func (s *PostgresStorage) Delete(key string) error {
	return s.DB.Where("token = ?", key).Delete(&entity.CsrfTokenRecords{}).Error
}

func (s *PostgresStorage) Reset() error {
	return s.DB.Where("expire_at < ?", time.Now()).Delete(&entity.CsrfTokenRecords{}).Error
}

func (s *PostgresStorage) Close() error {
	return nil
}