Skip to content
This repository has been archived by the owner on Apr 17, 2024. It is now read-only.

Commit

Permalink
db(flags): change unique index name (#125)
Browse files Browse the repository at this point in the history
  • Loading branch information
@wuhan005
wuhan005 committed May 25, 2021
1 parent 7f9a8de commit 489a3c6
Show file tree
Hide file tree
Showing 3 changed files with 200 additions and 32 deletions.
60 changes: 49 additions & 11 deletions internal/db/actions.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,12 @@ var Actions ActionsStore
type ActionsStore interface {
// Create creates a new action and persists to database.
Create(ctx context.Context, opts CreateActionOptions) error
// Get returns all the actions.
// Get returns the actions according to the given options.
Get(ctx context.Context, opts GetActionOptions) ([]*Action, error)
// SetScore updates the action's score.
SetScore(ctx context.Context, round, gameBoxID uint, score float64, replace ...bool) error
// GetEmptyScore returns the empty score actions in the given round.
GetEmptyScore(ctx context.Context, round uint, actionType ActionType) ([]*Action, error)
// DeleteAll deletes all the actions.
DeleteAll(ctx context.Context) error
}
Expand All @@ -36,8 +40,10 @@ func NewActionsStore(db *gorm.DB) ActionsStore {
type ActionType uint

const (
ActionTypeAttack ActionType = iota
ActionTypeBeenAttack ActionType = iota
ActionTypeCheckDown
ActionTypeAttack
ActionTypeServiceOnline
)

// Action represents the action such as check down or being attacked.
Expand All @@ -50,6 +56,8 @@ type Action struct {
GameBoxID uint `gorm:"uniqueIndex:action_unique_idx"`
AttackerTeamID uint `gorm:"uniqueIndex:action_unique_idx"`
Round uint `gorm:"uniqueIndex:action_unique_idx"`

Score float64
}

type actions struct {
Expand All @@ -58,8 +66,6 @@ type actions struct {

type CreateActionOptions struct {
Type ActionType
TeamID uint
ChallengeID uint
GameBoxID uint
AttackerTeamID uint
Round uint
Expand All @@ -72,13 +78,19 @@ func (db *actions) Create(ctx context.Context, opts CreateActionOptions) error {
opts.AttackerTeamID = 0
}

gameBoxStore := NewGameBoxesStore(db.DB)
gameBox, err := gameBoxStore.GetByID(ctx, opts.GameBoxID)
if err != nil {
return err
}

tx := db.WithContext(ctx).Begin()
var action Action
err := tx.Model(&Action{}).Where(&Action{
err = tx.Model(&Action{}).Where(&Action{
Type: opts.Type,
TeamID: opts.TeamID,
ChallengeID: opts.ChallengeID,
GameBoxID: opts.GameBoxID,
TeamID: gameBox.TeamID,
ChallengeID: gameBox.ChallengeID,
GameBoxID: gameBox.ID,
AttackerTeamID: opts.AttackerTeamID,
Round: opts.Round,
}).First(&action).Error
Expand All @@ -92,9 +104,9 @@ func (db *actions) Create(ctx context.Context, opts CreateActionOptions) error {

err = tx.Create(&Action{
Type: opts.Type,
TeamID: opts.TeamID,
ChallengeID: opts.ChallengeID,
GameBoxID: opts.GameBoxID,
TeamID: gameBox.TeamID,
ChallengeID: gameBox.ChallengeID,
GameBoxID: gameBox.ID,
AttackerTeamID: opts.AttackerTeamID,
Round: opts.Round,
}).Error
Expand Down Expand Up @@ -140,6 +152,32 @@ func (db *actions) Get(ctx context.Context, opts GetActionOptions) ([]*Action, e
}).Find(&actions).Error
}

var ErrActionNotExists = errors.New("action does not exist")

func (db *actions) SetScore(ctx context.Context, round, gameBoxID uint, score float64, replace ...bool) error {
actions, err := db.Get(ctx, GetActionOptions{
GameBoxID: gameBoxID,
Round: round,
})
if err != nil {
return errors.Wrap(err, "get action")
}
if len(actions) == 0 {
return ErrActionNotExists
}

action := actions[0]
if action.Score == 0 || (len(replace) == 1 && replace[0]) {
return db.WithContext(ctx).Model(&Action{}).Where("id = ?", action.ID).Update("score", score).Error
}
return nil
}

func (db *actions) GetEmptyScore(ctx context.Context, round uint, actionType ActionType) ([]*Action, error) {
var actions []*Action
return actions, db.WithContext(ctx).Model(&Action{}).Where("round = ? AND type = ? AND score = 0", round, actionType).Find(&actions).Error
}

func (db *actions) DeleteAll(ctx context.Context) error {
return db.WithContext(ctx).Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(&Action{}).Error
}
166 changes: 148 additions & 18 deletions internal/db/actions_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,14 +21,75 @@ func TestActions(t *testing.T) {
t.Parallel()

db, cleanup := newTestDB(t)
store := NewActionsStore(db)

ctx := context.Background()
challengesStore := NewChallengesStore(db)
_, err := challengesStore.Create(ctx, CreateChallengeOptions{
Title: "Web1",
BaseScore: 1000,
AutoRenewFlag: true,
RenewFlagCommand: "echo {{FLAG}} > /flag",
})
assert.Nil(t, err)
_, err = challengesStore.Create(ctx, CreateChallengeOptions{
Title: "Pwn1",
BaseScore: 1000,
AutoRenewFlag: true,
RenewFlagCommand: "echo {{FLAG}} > /flag",
})
assert.Nil(t, err)

teamsStore := NewTeamsStore(db)
_, err = teamsStore.Create(ctx, CreateTeamOptions{
Name: "Vidar",
Password: "123456",
Logo: "https://vidar.club/logo.png",
})
assert.Nil(t, err)
_, err = teamsStore.Create(ctx, CreateTeamOptions{
Name: "E99p1ant",
Password: "asdfgh",
Logo: "https://github.red/",
})
assert.Nil(t, err)

gameBoxStore := NewGameBoxesStore(db)
_, err = gameBoxStore.Create(ctx, CreateGameBoxOptions{
TeamID: 1,
ChallengeID: 1,
Address: "192.168.1.1",
Description: "Web1 For Vidar",
InternalSSH: SSHConfig{
Port: 22,
User: "root",
Password: "passw0rd",
},
})
assert.Nil(t, err)

_, err = gameBoxStore.Create(ctx, CreateGameBoxOptions{
TeamID: 2,
ChallengeID: 1,
Address: "192.168.2.1",
Description: "Web1 For E99p1ant",
InternalSSH: SSHConfig{
Port: 22,
User: "root",
Password: "s3crets",
},
})
assert.Nil(t, err)

actionsStore := NewActionsStore(db)

for _, tc := range []struct {
name string
test func(t *testing.T, ctx context.Context, db *actions)
}{
{"Create", testActionsCreate},
{"Get", testActionsGet},
{"SetScore", testActionsSetScore},
{"GetEmptyScore", testActionsGetEmptyScore},
{"DeleteAll", testActionsDeleteAll},
} {
t.Run(tc.name, func(t *testing.T) {
Expand All @@ -38,16 +99,14 @@ func TestActions(t *testing.T) {
t.Fatal(err)
}
})
tc.test(t, context.Background(), store.(*actions))
tc.test(t, context.Background(), actionsStore.(*actions))
})
}
}

func testActionsCreate(t *testing.T, ctx context.Context, db *actions) {
err := db.Create(ctx, CreateActionOptions{
Type: ActionTypeAttack,
TeamID: 1,
ChallengeID: 1,
GameBoxID: 1,
AttackerTeamID: 2,
Round: 1,
Expand All @@ -56,20 +115,25 @@ func testActionsCreate(t *testing.T, ctx context.Context, db *actions) {

err = db.Create(ctx, CreateActionOptions{
Type: ActionTypeAttack,
TeamID: 1,
ChallengeID: 1,
GameBoxID: 1,
AttackerTeamID: 2,
Round: 1,
})
assert.Equal(t, ErrDuplicateAction, err)

// Game box not found.
err = db.Create(ctx, CreateActionOptions{
Type: ActionTypeAttack,
GameBoxID: 3,
AttackerTeamID: 2,
Round: 1,
})
assert.Equal(t, ErrGameBoxNotExists, err)
}

func testActionsGet(t *testing.T, ctx context.Context, db *actions) {
err := db.Create(ctx, CreateActionOptions{
Type: ActionTypeAttack,
TeamID: 1,
ChallengeID: 1,
GameBoxID: 1,
AttackerTeamID: 2,
Round: 1,
Expand All @@ -78,8 +142,6 @@ func testActionsGet(t *testing.T, ctx context.Context, db *actions) {

err = db.Create(ctx, CreateActionOptions{
Type: ActionTypeCheckDown,
TeamID: 1,
ChallengeID: 1,
GameBoxID: 1,
AttackerTeamID: 1, // Will be set to 0.
Round: 1,
Expand All @@ -88,8 +150,6 @@ func testActionsGet(t *testing.T, ctx context.Context, db *actions) {

err = db.Create(ctx, CreateActionOptions{
Type: ActionTypeCheckDown,
TeamID: 2,
ChallengeID: 1,
GameBoxID: 2,
AttackerTeamID: 2, // Will be set to 0.
Round: 1,
Expand Down Expand Up @@ -133,11 +193,85 @@ func testActionsGet(t *testing.T, ctx context.Context, db *actions) {
assert.Equal(t, want, got)
}

func testActionsSetScore(t *testing.T, ctx context.Context, db *actions) {
err := db.Create(ctx, CreateActionOptions{
Type: ActionTypeAttack,
GameBoxID: 1,
AttackerTeamID: 2,
Round: 1,
})
assert.Nil(t, err)

err = db.SetScore(ctx, 1, 1, 50)
assert.Nil(t, err)
action, err := db.Get(ctx, GetActionOptions{
GameBoxID: 1,
Round: 1,
})
assert.Nil(t, err)
assert.NotZero(t, action)
assert.Equal(t, float64(50), action[0].Score)

// Replace score
err = db.SetScore(ctx, 1, 1, 60, true)
assert.Nil(t, err)
action, err = db.Get(ctx, GetActionOptions{
GameBoxID: 1,
Round: 1,
})
assert.Nil(t, err)
assert.NotZero(t, action)
assert.Equal(t, float64(60), action[0].Score)

err = db.SetScore(ctx, 2, 1, 50)
assert.Equal(t, ErrActionNotExists, err)
}

func testActionsGetEmptyScore(t *testing.T, ctx context.Context, db *actions) {
err := db.Create(ctx, CreateActionOptions{
Type: ActionTypeAttack,
GameBoxID: 1,
AttackerTeamID: 2,
Round: 1,
})
assert.Nil(t, err)

got, err := db.GetEmptyScore(ctx, 1, ActionTypeAttack)
assert.Nil(t, err)

for _, score := range got {
score.CreatedAt = time.Time{}
score.UpdatedAt = time.Time{}
}

want := []*Action{
{
Model: gorm.Model{
ID: 1,
},
Type: ActionTypeAttack,
TeamID: 1,
ChallengeID: 1,
GameBoxID: 1,
AttackerTeamID: 2,
Round: 1,
Score: 0,
},
}
assert.Equal(t, want, got)

err = db.SetScore(ctx, 1, 1, 60, true)
assert.Nil(t, err)

got, err = db.GetEmptyScore(ctx, 1, ActionTypeAttack)
assert.Nil(t, err)
want = []*Action{}
assert.Equal(t, want, got)
}

func testActionsDeleteAll(t *testing.T, ctx context.Context, db *actions) {
err := db.Create(ctx, CreateActionOptions{
Type: ActionTypeAttack,
TeamID: 1,
ChallengeID: 1,
GameBoxID: 1,
AttackerTeamID: 2,
Round: 1,
Expand All @@ -146,8 +280,6 @@ func testActionsDeleteAll(t *testing.T, ctx context.Context, db *actions) {

err = db.Create(ctx, CreateActionOptions{
Type: ActionTypeCheckDown,
TeamID: 1,
ChallengeID: 1,
GameBoxID: 1,
AttackerTeamID: 1, // Will be set to 0.
Round: 1,
Expand All @@ -156,8 +288,6 @@ func testActionsDeleteAll(t *testing.T, ctx context.Context, db *actions) {

err = db.Create(ctx, CreateActionOptions{
Type: ActionTypeCheckDown,
TeamID: 2,
ChallengeID: 1,
GameBoxID: 2,
AttackerTeamID: 2, // Will be set to 0.
Round: 1,
Expand Down
6 changes: 3 additions & 3 deletions internal/db/flag.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,9 +39,9 @@ func NewFlagsStore(db *gorm.DB) FlagsStore {
type Flag struct {
gorm.Model

TeamID uint `gorm:"uniqueIndex:gamebox_unique_idx"`
ChallengeID uint `gorm:"uniqueIndex:gamebox_unique_idx"`
GameBoxID uint `gorm:"uniqueIndex:gamebox_unique_idx"`
TeamID uint `gorm:"uniqueIndex:flag_unique_idx"`
ChallengeID uint `gorm:"uniqueIndex:flag_unique_idx"`
GameBoxID uint `gorm:"uniqueIndex:flag_unique_idx"`

Round uint
Value string
Expand Down

0 comments on commit 489a3c6

Please sign in to comment.