二阶段差不多

2026-06-01 19:46:51 +08:00
parent 9e50d05e71
commit 4e233c82b4
34 changed files with 1631 additions and 67 deletions
@@ -0,0 +1,115 @@
+package store
+
+import (
+	"time"
+
+	"mail_go/internal/db"
+
+	"gorm.io/gorm"
+)
+
+// BanStore defines the interface for IP ban operations.
+type BanStore interface {
+	Create(entry *db.BanEntry) error
+	GetByIP(ip string) (*db.BanEntry, error)
+	Delete(id uint) error
+	List(page, size int) ([]db.BanEntry, int64, error)
+	IsBanned(ip string) (bool, *db.BanEntry)
+	IncrementFail(ip string) (int, error)
+	ResetFail(ip string) error
+	Cleanup() error
+}
+
+// banStoreGorm implements BanStore using GORM.
+type banStoreGorm struct {
+	db *gorm.DB
+}
+
+// newBanStore creates a new GORM-backed BanStore.
+func newBanStore(database *gorm.DB) BanStore {
+	return &banStoreGorm{db: database}
+}
+
+// Create inserts a new ban entry record.
+func (s *banStoreGorm) Create(entry *db.BanEntry) error {
+	return s.db.Create(entry).Error
+}
+
+// GetByIP retrieves the most recent ban entry for a given IP address.
+func (s *banStoreGorm) GetByIP(ip string) (*db.BanEntry, error) {
+	var entry db.BanEntry
+	if err := s.db.Where("ip_address = ?", ip).Order("id DESC").First(&entry).Error; err != nil {
+		return nil, err
+	}
+	return &entry, nil
+}
+
+// Delete removes a ban entry by ID.
+func (s *banStoreGorm) Delete(id uint) error {
+	return s.db.Delete(&db.BanEntry{}, id).Error
+}
+
+// List retrieves a paginated list of ban entries.
+func (s *banStoreGorm) List(page, size int) ([]db.BanEntry, int64, error) {
+	var entries []db.BanEntry
+	var total int64
+
+	if err := s.db.Model(&db.BanEntry{}).Count(&total).Error; err != nil {
+		return nil, 0, err
+	}
+
+	offset := (page - 1) * size
+	if err := s.db.Order("id DESC").Offset(offset).Limit(size).Find(&entries).Error; err != nil {
+		return nil, 0, err
+	}
+	return entries, total, nil
+}
+
+// IsBanned checks whether an IP address is currently banned.
+// An IP is considered banned if there is a record with expires_at in the future.
+func (s *banStoreGorm) IsBanned(ip string) (bool, *db.BanEntry) {
+	var entry db.BanEntry
+	if err := s.db.Where("ip_address = ? AND expires_at > ?", ip, time.Now()).First(&entry).Error; err != nil {
+		return false, nil
+	}
+	return true, &entry
+}
+
+// IncrementFail increments the fail count for an IP address.
+// If no record exists, it creates one with fail_count=1 and a zero expires_at.
+// Returns the updated fail count.
+func (s *banStoreGorm) IncrementFail(ip string) (int, error) {
+	var entry db.BanEntry
+	err := s.db.Where("ip_address = ?", ip).First(&entry).Error
+	if err != nil {
+		// No record exists, create a new one
+		entry = db.BanEntry{
+			IPAddress: ip,
+			FailCount: 1,
+			ExpiresAt: time.Time{}, // Zero time, not yet banned
+		}
+		if createErr := s.db.Create(&entry).Error; createErr != nil {
+			return 0, createErr
+		}
+		return 1, nil
+	}
+
+	// Record exists, increment fail count
+	newCount := entry.FailCount + 1
+	if updateErr := s.db.Model(&entry).Update("fail_count", newCount).Error; updateErr != nil {
+		return 0, updateErr
+	}
+	return newCount, nil
+}
+
+// ResetFail resets the fail count for an IP address by deleting its record.
+func (s *banStoreGorm) ResetFail(ip string) error {
+	return s.db.Where("ip_address = ?", ip).Delete(&db.BanEntry{}).Error
+}
+
+// Cleanup removes expired ban entries.
+// It deletes records where expires_at is in the past and is not zero
+// (preserving records that have fail counts but are not yet banned).
+func (s *banStoreGorm) Cleanup() error {
+	return s.db.Where("expires_at < ? AND expires_at > ?", time.Now(), time.Time{}).Delete(&db.BanEntry{}).Error
+}
@@ -2,6 +2,7 @@ package store

 import (
 	"errors"
+	"time"

 	"mail_go/internal/db"

@@ -28,6 +29,13 @@ type MailStore interface {
 	MoveToFolder(id uint, folder string) error
 	Delete(id uint) error
 	CountUnread(userID uint, folder string) (int64, error)
+	CountByFolder(folder string) (int64, error)
+	CountAll() (int64, error)
+	TotalSizeByFolder(folder string) (int64, error)
+	TotalSize() (int64, error)
+	CountByFolderSince(folder string, since time.Time) (int64, error)
+	ListAll(page, size int) ([]db.Message, int64, error)
+	ListAllByFolder(folder string, page, size int) ([]db.Message, int64, error)
 }

 // mailStoreGorm implements MailStore using GORM.
@@ -123,3 +131,90 @@ func (s *mailStoreGorm) CountByUserAndFolder(userID uint, folder string) (int64,
 	}
 	return count, nil
 }
+
+// CountByFolder returns the total count of messages in a given folder.
+func (s *mailStoreGorm) CountByFolder(folder string) (int64, error) {
+	var count int64
+	if err := s.db.Model(&db.Message{}).Where("folder = ?", folder).Count(&count).Error; err != nil {
+		return 0, err
+	}
+	return count, nil
+}
+
+// CountAll returns the total count of all messages.
+func (s *mailStoreGorm) CountAll() (int64, error) {
+	var count int64
+	if err := s.db.Model(&db.Message{}).Count(&count).Error; err != nil {
+		return 0, err
+	}
+	return count, nil
+}
+
+// TotalSizeByFolder returns the total size (in bytes) of message bodies in a given folder.
+func (s *mailStoreGorm) TotalSizeByFolder(folder string) (int64, error) {
+	var total int64
+	err := s.db.Model(&db.Message{}).
+		Where("folder = ?", folder).
+		Select("COALESCE(SUM(LENGTH(text_body) + LENGTH(html_body)), 0)").
+		Scan(&total).Error
+	if err != nil {
+		return 0, err
+	}
+	return total, nil
+}
+
+// TotalSize returns the total size (in bytes) of all message bodies.
+func (s *mailStoreGorm) TotalSize() (int64, error) {
+	var total int64
+	err := s.db.Model(&db.Message{}).
+		Select("COALESCE(SUM(LENGTH(text_body) + LENGTH(html_body)), 0)").
+		Scan(&total).Error
+	if err != nil {
+		return 0, err
+	}
+	return total, nil
+}
+
+// CountByFolderSince returns the count of messages in a folder since a given time.
+func (s *mailStoreGorm) CountByFolderSince(folder string, since time.Time) (int64, error) {
+	var count int64
+	if err := s.db.Model(&db.Message{}).
+		Where("folder = ? AND created_at >= ?", folder, since).
+		Count(&count).Error; err != nil {
+		return 0, err
+	}
+	return count, nil
+}
+
+// ListAll retrieves a paginated list of all messages across all users.
+func (s *mailStoreGorm) ListAll(page, size int) ([]db.Message, int64, error) {
+	var messages []db.Message
+	var total int64
+
+	if err := s.db.Model(&db.Message{}).Count(&total).Error; err != nil {
+		return nil, 0, err
+	}
+
+	offset := (page - 1) * size
+	if err := s.db.Preload("User").Order("date DESC").Offset(offset).Limit(size).Find(&messages).Error; err != nil {
+		return nil, 0, err
+	}
+	return messages, total, nil
+}
+
+// ListAllByFolder retrieves a paginated list of all messages in a given folder across all users.
+func (s *mailStoreGorm) ListAllByFolder(folder string, page, size int) ([]db.Message, int64, error) {
+	var messages []db.Message
+	var total int64
+
+	query := s.db.Where("folder = ?", folder)
+	if err := query.Model(&db.Message{}).Count(&total).Error; err != nil {
+		return nil, 0, err
+	}
+
+	offset := (page - 1) * size
+	if err := s.db.Preload("User").Where("folder = ?", folder).Order("date DESC").Offset(offset).Limit(size).Find(&messages).Error; err != nil {
+		return nil, 0, err
+	}
+	return messages, total, nil
+}
@@ -12,6 +12,7 @@ type Stores struct {
 	Mails       MailStore
 	Domains     DomainStore
 	Attachments AttachmentStore
+	Bans        BanStore
 }

 // NewStores creates a new Stores instance with all GORM-backed implementations.
@@ -21,6 +22,7 @@ func NewStores(database *gorm.DB) *Stores {
 		Mails:       newMailStore(database),
 		Domains:     newDomainStore(database),
 		Attachments: newAttachmentStore(database),
+		Bans:        newBanStore(database),
 	}
 }

@@ -29,3 +31,4 @@ var _ = db.User{}
 var _ = db.Domain{}
 var _ = db.Message{}
 var _ = db.Attachment{}
+var _ = db.BanEntry{}