diff --git a/models/campaign.go b/models/campaign.go index 6e1b7b36..9b6d9ecb 100644 --- a/models/campaign.go +++ b/models/campaign.go @@ -414,7 +414,11 @@ func PostCampaign(c *Campaign, uid int64) error { // Insert a result for each target in the group for _, t := range g.Targets { r := &Result{Email: t.Email, Position: t.Position, Status: STATUS_SENDING, CampaignId: c.Id, UserId: c.UserId, FirstName: t.FirstName, LastName: t.LastName} - r.GenerateId() + err = r.GenerateId() + if err != nil { + Logger.Println(err) + continue + } err = db.Save(r).Error if err != nil { Logger.Printf("Error adding result record for target %s\n", t.Email) diff --git a/models/models_test.go b/models/models_test.go index 13bea6b0..1ac56402 100644 --- a/models/models_test.go +++ b/models/models_test.go @@ -1,6 +1,7 @@ package models import ( + "regexp" "strings" "testing" @@ -321,3 +322,11 @@ func (s *ModelsSuite) TestPostPage(c *check.C) { c.Assert(ok, check.Equals, false) }) } + +func (s *ModelsSuite) TestGenerateResultId(c *check.C) { + r := Result{} + r.GenerateId() + match, err := regexp.Match("[a-zA-Z0-9]{7}", []byte(r.RId)) + c.Assert(err, check.Equals, nil) + c.Assert(match, check.Equals, true) +} diff --git a/models/result.go b/models/result.go index 7b7178d8..6e8e9301 100644 --- a/models/result.go +++ b/models/result.go @@ -2,9 +2,8 @@ package models import ( "crypto/rand" - "fmt" - "io" "log" + "math/big" "net" "github.com/jinzhu/gorm" @@ -68,17 +67,25 @@ func (r *Result) UpdateGeo(addr string) error { // GenerateId generates a unique key to represent the result // in the database -func (r *Result) GenerateId() { +func (r *Result) GenerateId() error { // Keep trying until we generate a unique key (shouldn't take more than one or two iterations) - k := make([]byte, 32) + const alphaNum = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" + k := make([]byte, 7) for { - io.ReadFull(rand.Reader, k) - r.RId = fmt.Sprintf("%x", k) + for i := range k { + idx, err := rand.Int(rand.Reader, big.NewInt(int64(len(alphaNum)))) + if err != nil { + return err + } + k[i] = alphaNum[idx.Int64()] + } + r.RId = string(k) err := db.Table("results").Where("r_id=?", r.RId).First(&Result{}).Error if err == gorm.ErrRecordNotFound { break } } + return nil } // GetResult returns the Result object from the database