diff --git a/internal/db/db.go b/internal/db/db.go index 3d403a9..025aa9d 100644 --- a/internal/db/db.go +++ b/internal/db/db.go @@ -32,10 +32,10 @@ type Database = gorm.DB var ( gormLogger logger.Interface = logger.New(log.New(os.Stdout, "\r\n", log.LstdFlags), logger.Config{ SlowThreshold: 50 * time.Millisecond, - LogLevel: logger.Error, + LogLevel: logger.Warn, Colorful: true, }) - gormConfig = gorm.Config{Logger: gormLogger} + gormConfig = gorm.Config{Logger: gormLogger, PrepareStmt: true} ) // OpenDBFromEnvironment tries to open an SQL database, described by diff --git a/internal/db/paste.go b/internal/db/paste.go index 42c136c..91e3882 100644 --- a/internal/db/paste.go +++ b/internal/db/paste.go @@ -232,15 +232,12 @@ func GeneratePasteKey(db *gorm.DB, minimumEntropy int) (string, error) { return "", errors.Wrap(err, "url-key generation failed") } - alreadyInUse := true - var p Paste - err = db.Unscoped().Where("key = ?", []byte(key)).First(&p).Error - if err != nil && err == gorm.ErrRecordNotFound { - alreadyInUse = false - err = nil - } else if err != nil { + var count int64 + db.Unscoped().Model(&Paste{}).Where("key = ?", []byte(key)).Count(&count) + if err != nil { return "", errors.Wrap(err, "failed to check if key already exists") } + alreadyInUse := count != 0 isReserved := false for _, reservedKey := range ReservedPasteKeys { diff --git a/internal/db/paste_test.go b/internal/db/paste_test.go index 7d66aa2..2d266db 100644 --- a/internal/db/paste_test.go +++ b/internal/db/paste_test.go @@ -1,6 +1,20 @@ package db -import "testing" +import ( + "testing" + "testing/quick" + + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +func OpenTemporaryDB() (*Database, error) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gormConfig) + if err != nil { + return nil, err + } + return db, nil +} func TestValidatePasteKey(t *testing.T) { tests := []struct { @@ -36,3 +50,24 @@ func TestValidatePasteKey(t *testing.T) { }) } } + +func TestGeneratedKeysAreValid(t *testing.T) { + db, err := OpenTemporaryDB() + if err != nil { + t.Error(err) + } + var minimumEntropy int + checkGeneratedKeyValid := func() bool { + key, err := GeneratePasteKey(db, minimumEntropy) + if err != nil { + return false + } + return ValidatePasteKey(key) == nil + } + + for minimumEntropy = 0; minimumEntropy <= 80; minimumEntropy++ { + if err = quick.Check(checkGeneratedKeyValid, nil); err != nil { + t.Error(err) + } + } +}