diff --git a/bot.go b/bot.go index 2cc861f..9b3b80c 100644 --- a/bot.go +++ b/bot.go @@ -63,6 +63,10 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock) (*Bot, error) { } if err := db.Create(&owner).Error; err != nil { + // If unique constraint is violated, another owner already exists + if strings.Contains(err.Error(), "unique index") { + return nil, fmt.Errorf("an owner already exists for this bot") + } return nil, fmt.Errorf("failed to create owner user: %w", err) } } else if err != nil { @@ -123,6 +127,10 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User } if err := b.db.Create(&user).Error; err != nil { + // Handle unique constraint for owner + if isOwner && strings.Contains(err.Error(), "unique index") { + return User{}, fmt.Errorf("an owner already exists for this bot") + } return User{}, err } } else { @@ -139,11 +147,11 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User return User{}, err } // Promote to owner - user.Role, err = b.getRoleByName("owner") + role, err := b.getRoleByName("owner") if err != nil { return User{}, err } - user.RoleID = user.Role.ID + user.RoleID = role.ID user.IsOwner = true if err := b.db.Save(&user).Error; err != nil { return User{}, err diff --git a/database.go b/database.go index d35a25d..6c7a8f6 100644 --- a/database.go +++ b/database.go @@ -27,14 +27,21 @@ func initDB() (*gorm.DB, error) { return nil, fmt.Errorf("failed to connect to database: %w", err) } - // AutoMigrate with unique constraint for owners + // AutoMigrate the models err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}) if err != nil { return nil, fmt.Errorf("failed to migrate database schema: %w", err) } - // Add unique index for owners per bot - db.SetupJoinTable(&BotModel{}, "Users", &User{}) + // Enforce unique owner per bot using raw SQL + // Note: SQLite doesn't support partial indexes, but we can simulate it by making a unique index on (BotID, IsOwner) + // and ensuring that IsOwner can only be true for one user per BotID. + // This approach allows multiple users with IsOwner=false for the same BotID, + // but only one user can have IsOwner=true per BotID. + err = db.Exec(`CREATE UNIQUE INDEX IF NOT EXISTS idx_bot_owner ON users (bot_id, is_owner) WHERE is_owner = 1;`).Error + if err != nil { + return nil, fmt.Errorf("failed to create unique index for bot owners: %w", err) + } err = createDefaultRoles(db) if err != nil { diff --git a/handlers.go b/handlers.go index 4d01477..149b772 100644 --- a/handlers.go +++ b/handlers.go @@ -84,7 +84,10 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U userMessage := b.createMessage(chatID, userID, username, user.Role.Name, text, true) userMessage.UserRole = string(anthropic.RoleUser) // Convert to string - b.storeMessage(userMessage) + if err := b.storeMessage(userMessage); err != nil { + log.Printf("Error storing user message: %v", err) + return + } chatMemory := b.getOrCreateChatMemory(chatID) b.addMessageToChatMemory(chatMemory, userMessage) @@ -98,10 +101,16 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U response = "I'm sorry, I'm having trouble processing your request right now." } - b.sendResponse(ctx, chatID, response, businessConnectionID) + if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil { + log.Printf("Error sending response: %v", err) + return + } - assistantMessage := b.createMessage(chatID, 0, "", string(anthropic.RoleAssistant), response, false) - b.storeMessage(assistantMessage) + assistantMessage := b.createMessage(chatID, 0, "", "assistant", response, false) + if err := b.storeMessage(assistantMessage); err != nil { + log.Printf("Error storing assistant message: %v", err) + return + } b.addMessageToChatMemory(chatMemory, assistantMessage) } diff --git a/models.go b/models.go index 3c2247f..6ce7b1e 100644 --- a/models.go +++ b/models.go @@ -61,3 +61,8 @@ type User struct { Role Role `gorm:"foreignKey:RoleID"` IsOwner bool `gorm:"default:false"` // Indicates if the user is the owner } + +// Compound unique index to ensure only one owner per bot +func (User) TableName() string { + return "users" +} diff --git a/rate_limiter_test.go b/rate_limiter_test.go index a2871f3..5965a3e 100644 --- a/rate_limiter_test.go +++ b/rate_limiter_test.go @@ -1,8 +1,12 @@ package main import ( + "strings" "testing" "time" + + "gorm.io/driver/sqlite" + "gorm.io/gorm" ) // TestCheckRateLimits tests the checkRateLimits method of the Bot. @@ -80,6 +84,109 @@ func TestCheckRateLimits(t *testing.T) { } } +func TestOwnerAssignment(t *testing.T) { + // Initialize in-memory database for testing + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatalf("Failed to open in-memory database: %v", err) + } + + // Migrate the schema + err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}) + if err != nil { + t.Fatalf("Failed to migrate database schema: %v", err) + } + + // Create default roles + err = createDefaultRoles(db) + if err != nil { + t.Fatalf("Failed to create default roles: %v", err) + } + + // Create a bot configuration + config := BotConfig{ + ID: "test_bot", + TelegramToken: "TEST_TELEGRAM_TOKEN", + MemorySize: 10, + MessagePerHour: 5, + MessagePerDay: 10, + TempBanDuration: "1m", + SystemPrompts: make(map[string]string), + Active: true, + OwnerTelegramID: 111111111, + } + + // Initialize MockClock + mockClock := &MockClock{ + currentTime: time.Now(), + } + + // Create the bot + bot, err := NewBot(db, config, mockClock) + if err != nil { + t.Fatalf("Failed to create bot: %v", err) + } + + // Verify that the owner exists + var owner User + err = db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", config.OwnerTelegramID, bot.botID, true).First(&owner).Error + if err != nil { + t.Fatalf("Owner was not created: %v", err) + } + + // Attempt to create another owner for the same bot + _, err = bot.getOrCreateUser(222222222, "AnotherOwner", true) + if err == nil { + t.Fatalf("Expected error when creating a second owner, but got none") + } + + // Verify that the error message is appropriate + expectedErrorMsg := "an owner already exists for this bot" + if err.Error() != expectedErrorMsg && !strings.Contains(err.Error(), "unique index") { + t.Fatalf("Unexpected error message: %v", err) + } + + // Assign admin role to a new user + adminUser, err := bot.getOrCreateUser(333333333, "AdminUser", false) + if err != nil { + t.Fatalf("Failed to create admin user: %v", err) + } + + if adminUser.Role.Name != "admin" { + t.Fatalf("Expected role 'admin', got '%s'", adminUser.Role.Name) + } + + // Assign owner role to a user from a different bot + otherBotConfig := BotConfig{ + ID: "other_bot", + TelegramToken: "OTHER_TELEGRAM_TOKEN", + MemorySize: 10, + MessagePerHour: 5, + MessagePerDay: 10, + TempBanDuration: "1m", + SystemPrompts: make(map[string]string), + Active: true, + OwnerTelegramID: 444444444, + } + + otherBot, err := NewBot(db, otherBotConfig, mockClock) + if err != nil { + t.Fatalf("Failed to create other bot: %v", err) + } + + _, err = otherBot.getOrCreateUser(config.OwnerTelegramID, "OwnerOfOtherBot", true) + if err != nil { + t.Fatalf("Failed to assign existing owner to another bot: %v", err) + } + + // Verify multiple bots can have the same owner telegram ID + var ownerOfOtherBot User + err = db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", config.OwnerTelegramID, otherBot.botID, true).First(&ownerOfOtherBot).Error + if err != nil { + t.Fatalf("Owner of other bot was not created: %v", err) + } +} + // To ensure thread safety and avoid race conditions during testing, // you can run the tests with the `-race` flag: // go test -race -v