From 1ecb11c65f391c111814d9802c68f06680aade6e Mon Sep 17 00:00:00 2001
From: Daan Sprenkels <daan@dsprenkels.com>
Date: Thu, 17 Dec 2020 14:49:52 +0100
Subject: [PATCH] db: Do not warn during paste key generation

---
 internal/db/db.go         |  4 ++--
 internal/db/paste.go      | 11 ++++-------
 internal/db/paste_test.go | 37 ++++++++++++++++++++++++++++++++++++-
 3 files changed, 42 insertions(+), 10 deletions(-)

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)
+		}
+	}
+}