Skip to content

Commit f20f0f2

Browse files
DarkPhoenix42DarkPhoenix42
committed
feat: migrate to postgres and refactor database models and queries for consistency
Signed-off-by: Praneeth Sarode <praneethsarode@gmail.com>
1 parent f701a12 commit f20f0f2

File tree

10 files changed

+82
-69
lines changed

10 files changed

+82
-69
lines changed

‎api/config.go

-5
Original file line numberDiff line numberDiff line change
@@ -162,11 +162,6 @@ func updateChallengeInfoHandler(c *gin.Context) {
162162
"Name": name,
163163
}
164164

165-
hints, exist := c.GetPostForm("hints")
166-
if exist {
167-
configInfo["Hints"] = hints
168-
}
169-
170165
desc, exist := c.GetPostForm("desc")
171166
if exist {
172167
configInfo["Description"] = desc

‎api/submit.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -277,7 +277,7 @@ func submitFlagHandler(c *gin.Context) {
277277
}
278278

279279
UserChallengesEntry := database.UserChallenges{
280-
CreatedAt: time.Time{},
280+
CreatedAt: time.Now(),
281281
UserID: user.ID,
282282
ChallengeID: challenge.ID,
283283
Solved: true,

‎core/database/challenges.go

+9-7
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ import (
1414
"github.com/sdslabs/beastv4/core"
1515
tools "github.com/sdslabs/beastv4/templates"
1616

17-
_ "gorm.io/driver/sqlite"
1817
"gorm.io/gorm"
1918
"gorm.io/gorm/clause"
2019
)
@@ -63,18 +62,21 @@ type Challenge struct {
6362
MaxPoints uint `gorm:"default:0"`
6463
MinPoints uint `gorm:"default:0"`
6564
Ports []Port
66-
Tags []*Tag `gorm:"many2many:tag_challenges;"`
65+
Tags []*Tag `gorm:"many2many:tag_challenges;constraint:OnUpdate:CASCADE,OnDelete:CASCADE;"`
6766
Users []*User `gorm:"many2many:user_challenges;"`
6867
ServerDeployed string `gorm:"type:varchar(64)"`
6968
}
7069

7170
type UserChallenges struct {
72-
CreatedAt time.Time
73-
UserID uint
71+
CreatedAt time.Time
72+
User User `gorm:"foreignKey:UserID"`
73+
UserID uint
74+
75+
Challenge Challenge `gorm:"foreignKey:ChallengeID"`
7476
ChallengeID uint
75-
Tries uint
76-
Solved bool
77-
Flag string
77+
Tries uint `gorm:"not null;default:0"`
78+
Solved bool `gorm:"not null;default:false;index"`
79+
Flag string `gorm:"type:text"`
7880
}
7981

8082
// The `DynamicFlags` table has the following columns

‎core/database/database.go

+44-31
Original file line numberDiff line numberDiff line change
@@ -87,11 +87,14 @@ func init() {
8787
log.Fatalf("Cannot create related models: %s", err)
8888
}
8989

90-
Db.AutoMigrate(&Challenge{}, &Transaction{}, &Port{}, &User{}, &Tag{}, &Notification{}, &Hint{}, &DynamicFlag{}, &OTP{})
90+
err := Db.AutoMigrate(&Challenge{}, &Transaction{}, &Port{}, &User{}, &Tag{}, &Notification{}, &Hint{}, &DynamicFlag{}, &OTP{})
91+
if err != nil {
92+
log.Fatalf("failed to migrate database with error: %s", err)
93+
}
94+
9195
users, err := QueryUserEntries("email", core.DEFAULT_USER_EMAIL)
9296
if err != nil {
93-
log.Errorf("Error while checking dummy user entry.")
94-
os.Exit(1)
97+
log.Fatalf("Error while checking dummy user entry.")
9598
}
9699

97100
if len(users) == 0 {
@@ -116,9 +119,8 @@ func init() {
116119
}
117120

118121
func BackupAndReset() {
119-
beastRemoteDir := filepath.Join(BEAST_GLOBAL_DIR, core.BEAST_REMOTES_DIR)
120-
beastStagingDir := filepath.Join(BEAST_GLOBAL_DIR, core.BEAST_STAGING_DIR)
121122
LoadDbConfig()
123+
122124
err := BackupDatabase()
123125
if err != nil {
124126
log.Errorf("Error while backing up database: %s", err)
@@ -129,12 +131,33 @@ func BackupAndReset() {
129131
log.Errorf("Error while resetting up database: %s", err)
130132
return
131133
}
132-
err = os.Rename(beastRemoteDir, beastRemoteDir+time.Now().Format("20060102150405")+".bak")
134+
135+
backupPath := filepath.Join(core.BEAST_GLOBAL_DIR, "backup", core.BEAST_REMOTES_DIR)
136+
err = utils.CreateIfNotExistDir(backupPath)
137+
if err != nil {
138+
log.Errorf("Error while creating backup directory: %s", err)
139+
return
140+
}
141+
142+
backupPath = filepath.Join(backupPath, core.BEAST_REMOTES_DIR+time.Now().Format("20060102150405")+".bak")
143+
oldPath := filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_REMOTES_DIR)
144+
err = os.Rename(oldPath, backupPath)
133145
if err != nil {
134146
log.Errorf("Error while backing up remote dir: %s", err)
135147
return
136148
}
137-
err = os.Rename(beastStagingDir, beastStagingDir+time.Now().Format("20060102150405")+".bak")
149+
150+
backupPath = filepath.Join(core.BEAST_GLOBAL_DIR, "backup", core.BEAST_STAGING_DIR)
151+
152+
err = utils.CreateIfNotExistDir(backupPath)
153+
if err != nil {
154+
log.Errorf("Error while creating backup directory: %s", err)
155+
return
156+
}
157+
158+
oldPath = filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_STAGING_DIR)
159+
backupPath = filepath.Join(backupPath, core.BEAST_STAGING_DIR+time.Now().Format("20060102150405")+".bak")
160+
err = os.Rename(oldPath, backupPath)
138161
if err != nil {
139162
log.Errorf("Error while backing up staging dir: %s", err)
140163
return
@@ -145,8 +168,16 @@ func BackupDatabase() error {
145168
if dbConfig == (Config{}) {
146169
LoadDbConfig()
147170
}
171+
172+
backupPath := filepath.Join(core.BEAST_GLOBAL_DIR, "backup", "db")
173+
err := utils.CreateIfNotExistDir(backupPath)
174+
if err != nil {
175+
log.Errorf("Error while creating backup directory: %s", err)
176+
return err
177+
}
178+
148179
backupFile := fmt.Sprintf("%s_%s.bak", dbConfig.PsqlConf.Dbname, time.Now().Format("20060102150405"))
149-
cmd := exec.Command("pg_dump", "-U", dbConfig.PsqlConf.User, "-h", dbConfig.PsqlConf.Host, "-p", dbConfig.PsqlConf.Port, "-F", "c", "-f", filepath.Join(core.BEAST_GLOBAL_DIR, backupFile), dbConfig.PsqlConf.Dbname)
180+
cmd := exec.Command("pg_dump", "-U", dbConfig.PsqlConf.User, "-h", dbConfig.PsqlConf.Host, "-p", dbConfig.PsqlConf.Port, "-F", "c", "-f", filepath.Join(backupPath, backupFile), dbConfig.PsqlConf.Dbname)
150181
cmd.Env = append(os.Environ(), fmt.Sprintf("PGPASSWORD=%s", dbConfig.PsqlConf.Password))
151182
output, err := cmd.CombinedOutput()
152183
if err != nil {
@@ -163,11 +194,11 @@ func ResetDatabase() error {
163194
}
164195
err := TerminateDatabaseConnections()
165196
if err != nil {
166-
log.Errorf("Unable to terminate connections ", err)
197+
log.Errorf("Unable to terminate connections %s", err)
167198
return err
168199
}
169200

170-
dropCmd := exec.Command("psql", "-U", dbConfig.PsqlConf.User, "-h", dbConfig.PsqlConf.Host, "-p", dbConfig.PsqlConf.Port, "-c", "DROP DATABASE IF EXISTS "+dbConfig.PsqlConf.Dbname)
201+
dropCmd := exec.Command("dropdb", "-U", dbConfig.PsqlConf.User, "-h", dbConfig.PsqlConf.Host, "-p", dbConfig.PsqlConf.Port, "--force", dbConfig.PsqlConf.Dbname)
171202
dropCmd.Env = append(os.Environ(), fmt.Sprintf("PGPASSWORD=%s", dbConfig.PsqlConf.Password))
172203

173204
output, err := dropCmd.CombinedOutput()
@@ -176,7 +207,7 @@ func ResetDatabase() error {
176207
return err
177208
}
178209

179-
createCmd := exec.Command("psql", "-U", dbConfig.PsqlConf.User, "-h", dbConfig.PsqlConf.Host, "-p", dbConfig.PsqlConf.Port, "-c", "CREATE DATABASE "+dbConfig.PsqlConf.Dbname)
210+
createCmd := exec.Command("psql", "-U", dbConfig.PsqlConf.User, "-h", dbConfig.PsqlConf.Host, "-p", dbConfig.PsqlConf.Port, "-d", "postgres", "-c", "CREATE DATABASE "+dbConfig.PsqlConf.Dbname+";")
180211
createCmd.Env = append(os.Environ(), fmt.Sprintf("PGPASSWORD=%s", dbConfig.PsqlConf.Password))
181212

182213
output, err = createCmd.CombinedOutput()
@@ -218,10 +249,10 @@ func RestoreDatabase(backupFile string) error {
218249

219250
err := TerminateDatabaseConnections()
220251
if err != nil {
221-
log.Error("Unable to terminate connections ", err)
252+
log.Errorf("Unable to terminate connections: %s ", err)
222253
return err
223254
}
224-
255+
225256
err = utils.ValidateFileExists(backupFile)
226257
if err != nil {
227258
return fmt.Errorf("backup file does not exist: %s", backupFile)
@@ -249,21 +280,3 @@ func RestoreDatabase(backupFile string) error {
249280
log.Println("Database restored successfully from:", backupFile)
250281
return nil
251282
}
252-
253-
// func BackupDatabase() {
254-
// beastDb := filepath.Join(BEAST_GLOBAL_DIR, BEAST_DATABASE)
255-
// beastRemoteDir := filepath.Join(BEAST_GLOBAL_DIR, core.BEAST_REMOTES_DIR)
256-
// beastStagingDir := filepath.Join(BEAST_GLOBAL_DIR, core.BEAST_STAGING_DIR)
257-
// err := utils.CopyFile(beastDb, beastDb+time.Now().Format("20060102150405")+".bak")
258-
// if err != nil {
259-
// log.Errorf("Error while backing up database: %s", err)
260-
// }
261-
// err = utils.CopyDirectory(beastRemoteDir, beastRemoteDir+time.Now().Format("20060102150405")+".bak")
262-
// if err != nil {
263-
// log.Errorf("Error while backing up remote dir: %s", err)
264-
// }
265-
// err = utils.CopyDirectory(beastStagingDir, beastStagingDir+time.Now().Format("20060102150405")+".bak")
266-
// if err != nil {
267-
// log.Errorf("Error while backing up staging dir: %s", err)
268-
// }
269-
// }

‎core/database/hints.go

+9-5
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package database
33
import (
44
"errors"
55
"fmt"
6+
67
"gorm.io/gorm"
78
)
89

@@ -14,11 +15,14 @@ type Hint struct {
1415
}
1516

1617
type UserHint struct {
17-
UserID uint `gorm:"not null"`
18-
ChallengeID uint `gorm:"not null"`
19-
HintID uint `gorm:"not null"`
20-
Hint Hint `gorm:"foreignKey:HintID;references:HintID"`
21-
Challenge Challenge `gorm:"foreignKey:ChallengeID;references:ID"`
18+
UserID uint
19+
User User `gorm:"foreignKey:UserID"`
20+
21+
HintID uint
22+
Hint Hint `gorm:"foreignKey:HintID"`
23+
24+
ChallengeID uint
25+
Challenge Challenge `gorm:"foreignKey:ChallengeID"`
2226
}
2327

2428
func CreateHintEntry(hint *Hint) error {

‎core/database/ports.go

-1
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"fmt"
55

66
"github.com/jinzhu/gorm"
7-
_ "github.com/jinzhu/gorm/dialects/sqlite"
87
)
98

109
type Port struct {

‎core/database/tag.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import (
1111
type Tag struct {
1212
gorm.Model
1313

14-
Challenges []*Challenge `gorm:"many2many:tag_challenges;"`
14+
Challenges []*Challenge `gorm:"many2many:tag_challenges;constraint:OnUpdate:CASCADE,OnDelete:CASCADE;"`
1515
TagName string `gorm:"not null;unique"`
1616
}
1717

‎core/database/user.go

+6-7
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ import (
1919
tools "github.com/sdslabs/beastv4/templates"
2020
log "github.com/sirupsen/logrus"
2121

22-
_ "gorm.io/driver/sqlite"
2322
"gorm.io/gorm"
2423
)
2524

@@ -39,7 +38,7 @@ type User struct {
3938
// Queries all the users entries where the column represented by key
4039
// have the value in value.
4140
func QueryUserEntries(key string, value string) ([]User, error) {
42-
queryKey := fmt.Sprintf("%s == ?", key)
41+
queryKey := fmt.Sprintf("%s = ?", key)
4342
var users []User
4443
DBMux.Lock()
4544
defer DBMux.Unlock()
@@ -92,7 +91,7 @@ func GetUserRank(userID uint, userScore uint, updatedAt time.Time) (rank int64,
9291
DBMux.Lock()
9392
defer DBMux.Unlock()
9493

95-
tx := Db.Where("id != ? AND score >= ? AND role == ? AND status == ?", userID, userScore, core.USER_ROLES["contestant"], 0).Find(&users)
94+
tx := Db.Where("id != ? AND score >= ? AND role = ? AND status = ?", userID, userScore, core.USER_ROLES["contestant"], 0).Find(&users)
9695

9796
for _, user := range users {
9897
if user.Score > userScore {
@@ -171,7 +170,7 @@ func CheckPreviousSubmissions(userId uint, challId uint) (bool, error) {
171170
DBMux.Lock()
172171
defer DBMux.Unlock()
173172

174-
tx := Db.Where("user_id == ? AND challenge_id == ? AND solved == ?", userId, challId, true).Find(&userChallenges).Count(&count)
173+
tx := Db.Where("user_id = ? AND challenge_id = ? AND solved = ?", userId, challId, true).Find(&userChallenges).Count(&count)
175174

176175
if errors.Is(tx.Error, gorm.ErrRecordNotFound) {
177176
return false, nil
@@ -311,7 +310,7 @@ func QueryTopUsersByScore(limit int) ([]User, error) {
311310
DBMux.Lock()
312311
defer DBMux.Unlock()
313312

314-
tx := Db.Where("role == ? AND status == ?", core.USER_ROLES["contestant"], 0).
313+
tx := Db.Where("role = ? AND status = ?", core.USER_ROLES["contestant"], 0).
315314
Order("score desc, updated_at asc").
316315
Limit(limit).
317316
Find(&users)
@@ -329,7 +328,7 @@ func QueryUsersByScoreOffsetLimit(limit, offset int) ([]User, error) {
329328
DBMux.Lock()
330329
defer DBMux.Unlock()
331330

332-
tx := Db.Where("role == ? AND status == ?", core.USER_ROLES["contestant"], 0).
331+
tx := Db.Where("role = ? AND status = ?", core.USER_ROLES["contestant"], 0).
333332
Order("score desc, updated_at asc").
334333
Limit(limit).
335334
Offset(offset).
@@ -346,7 +345,7 @@ func GetUserCount() (int64, error) {
346345
var count int64
347346
DBMux.Lock()
348347
defer DBMux.Unlock()
349-
tx := Db.Model(&User{}).Where("role == ?", "contestant").Count(&count)
348+
tx := Db.Model(&User{}).Where("role = ?", "contestant").Count(&count)
350349
if errors.Is(tx.Error, gorm.ErrRecordNotFound) {
351350
return 0, nil
352351
}

‎core/manager/challenge.go

+8-11
Original file line numberDiff line numberDiff line change
@@ -690,23 +690,20 @@ func undeployChallenge(challengeName string, purge bool) error {
690690
}
691691

692692
log.Info("Subtracting user points")
693-
err = database.SubtractScoreFromSolvers(challenge.ID);
694-
693+
err = database.SubtractScoreFromSolvers(challenge.ID)
695694
if err != nil {
696695
return fmt.Errorf("error while subtracting user points: %s", err)
697696
}
698697

699-
log.Info("Deleting database entry")
700-
if err := coreUtils.DeleteChallengeEntryWithPorts(challenge.Name);
701-
err != nil {
702-
return fmt.Errorf("error while deleting challenge entry: %s", err)
698+
log.Info("Deleting challenge entry from user challenges table")
699+
err = database.DeleteAllUserChallenges(challenge.ID)
700+
if err != nil {
701+
return fmt.Errorf("error while deleting user_challenges entries: %s", err)
703702
}
704703

705-
log.Info("Deleting challenge entry from user challenge database")
706-
err = database.DeleteAllUserChallenges(challenge.ID);
707-
708-
if err != nil {
709-
return fmt.Errorf("error while subtracting user points: %s", err)
704+
log.Info("Deleting database entry")
705+
if err := coreUtils.DeleteChallengeEntryWithPorts(challenge.Name); err != nil {
706+
return fmt.Errorf("error while deleting challenge entry: %s", err)
710707
}
711708

712709
log.Infof("Challenge purge successful")

‎core/manager/health_check.go

+4
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,10 @@ var HEALTH_CHECKER = false
2323
// At the time of writing, Beast deploys assets to localhost only.
2424
// So it will check only on localhost
2525
func CheckStaticChallenge(chall database.Challenge) error {
26+
if chall.Assets == "" {
27+
return nil
28+
}
29+
2630
assets := strings.Split(chall.Assets, core.DELIMITER)
2731
for _, asset := range assets {
2832
filepath := filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_STAGING_DIR, chall.Name, core.BEAST_STATIC_FOLDER, asset)

0 commit comments

Comments
 (0)