安全与可靠性: 修复配额并发、投票事务、密码授权与导出完整性

internal/db/conversation.go

@@ -6,22 +6,30 @@ package db
 
 import (
 	"context"
+	"time"
 
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"github.com/thanhpk/randstr"
 	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
 
 	"github.com/tamakyi/TamaBox/internal/dbutil"
 )
 
 var Conversations ConversationStore
 
+var (
+	ErrDailyMessageLimit = errors.New("已达到每日私聊上限")
+	ErrConversationLimit = errors.New("会话已达到消息上限或已关闭")
+)
+
 type ConversationStore interface {
 	Create(ctx context.Context, opts CreateConversationOptions) (*Conversation, error)
 	// CreateWithFirstMessage 事务化地建会话 + 首条消息 + 消息计数。
 	// 单独写首条消息（用于「发第一条私聊」这类一次性动作），避免失败时留下空会话。
 	CreateWithFirstMessage(ctx context.Context, opts CreateConversationOptions, senderID uint, content string) (*Conversation, *Message, error)
+	SendMessage(ctx context.Context, opts SendConversationMessageOptions) (*Message, error)
 	GetByID(ctx context.Context, id uint) (*Conversation, error)
 	GetByToken(ctx context.Context, token string) (*Conversation, error)
 	GetByIDAndOwner(ctx context.Context, id, ownerID uint) (*Conversation, error)
@@ -74,6 +82,39 @@ type CreateConversationOptions struct {
 	MaxMessages int
 }
 
+type SendConversationMessageOptions struct {
+	ConversationID uint
+	SenderID       uint
+	Content        string
+	Image          *CreateMessageImageOptions
+}
+
+func reserveDailyMessageTx(tx *gorm.DB, ownerID, senderID uint) error {
+	var sender User
+	if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("id").First(&sender, senderID).Error; err != nil {
+		return errors.Wrap(err, "lock message sender")
+	}
+	var owner User
+	if err := tx.Select("daily_message_limit").First(&owner, ownerID).Error; err != nil {
+		return errors.Wrap(err, "get owner daily message limit")
+	}
+	if owner.DailyMessageLimit <= 0 {
+		return nil
+	}
+	now := time.Now()
+	start := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location())
+	var count int64
+	if err := tx.Table("messages").Joins("JOIN conversations ON conversations.id = messages.conversation_id").
+		Where("conversations.owner_id = ? AND messages.sender_id = ? AND messages.created_at >= ?", ownerID, senderID, start).
+		Count(&count).Error; err != nil {
+		return errors.Wrap(err, "count today messages")
+	}
+	if count >= int64(owner.DailyMessageLimit) {
+		return ErrDailyMessageLimit
+	}
+	return nil
+}
+
 func (db *conversations) Create(ctx context.Context, opts CreateConversationOptions) (*Conversation, error) {
 	conversation := &Conversation{
 		OwnerID:     opts.OwnerID,
@@ -99,6 +140,11 @@ func (db *conversations) CreateWithFirstMessage(
 	var msg *Message
 
 	err := db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		if senderID == opts.AskerID {
+			if err := reserveDailyMessageTx(tx, opts.OwnerID, senderID); err != nil {
+				return err
+			}
+		}
 		conv = &Conversation{
 			OwnerID:     opts.OwnerID,
 			AskerID:     opts.AskerID,
@@ -139,6 +185,59 @@ func (db *conversations) CreateWithFirstMessage(
 	return conv, msg, nil
 }
 
+func (db *conversations) SendMessage(ctx context.Context, opts SendConversationMessageOptions) (*Message, error) {
+	msg := &Message{ConversationID: opts.ConversationID, SenderID: opts.SenderID, Content: opts.Content}
+	err := db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		// 先锁发送者，再锁会话，保证与 CreateWithFirstMessage 的锁顺序一致，
+		// 避免并发下两条路径互相等待造成配额检查窗口。
+		var sender User
+		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("id").First(&sender, opts.SenderID).Error; err != nil {
+			return errors.Wrap(err, "lock message sender")
+		}
+
+		var conv Conversation
+		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&conv, opts.ConversationID).Error; err != nil {
+			return errors.Wrap(err, "lock conversation")
+		}
+
+		// 发送者必须是会话双方之一
+		if conv.OwnerID != opts.SenderID && conv.AskerID != opts.SenderID {
+			return errors.New("sender is not part of the conversation")
+		}
+
+		// 只有 asker 发送才消耗 owner 的每日配额
+		if conv.AskerID == opts.SenderID {
+			if err := reserveDailyMessageTx(tx, conv.OwnerID, opts.SenderID); err != nil {
+				return err
+			}
+		}
+
+		res := tx.Model(&Conversation{}).Where("id = ? AND message_count < max_messages AND is_closed = ?", conv.ID, false).
+			UpdateColumn("message_count", gorm.Expr("message_count + 1"))
+		if res.Error != nil {
+			return errors.Wrap(res.Error, "increment message count")
+		}
+		if res.RowsAffected == 0 {
+			return ErrConversationLimit
+		}
+		if err := tx.Create(msg).Error; err != nil {
+			return errors.Wrap(err, "create message")
+		}
+		if opts.Image != nil {
+			image := *opts.Image
+			image.MessageID = msg.ID
+			if err := tx.Create(&MessageImage{MessageID: image.MessageID, URL: image.URL, ObjectKey: image.ObjectKey, FileName: image.FileName, FileSize: image.FileSize}).Error; err != nil {
+				return errors.Wrap(err, "create message image")
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return msg, nil
+}
+
 func (db *conversations) GetByID(ctx context.Context, id uint) (*Conversation, error) {
 	var conv Conversation
 	if err := db.WithContext(ctx).First(&conv, id).Error; err != nil {