version 1.6.6: fix slashes in keys in database

lib/db/queries.go

@@ -26,10 +26,9 @@ func SelectUserByToken(token string) (User, error) {
 
 // ObjectKeyExists returns a boolean specifying if the given key is already in
 // use by another object.
-func ObjectKeyExists(key string) (bool, error) {
+func ObjectKeyExists(bucket, key string) (bool, error) {
 	var count int
-	err := DB.QueryRow(countOfObjectByBucketAndRandomKey, "public", key).
-		Scan(&count)
+	err := DB.QueryRow(countOfObjectByBucketAndRandomKey, bucket, key).Scan(&count)
 	if err != nil {
 		return false, err
 	}
@@ -38,7 +37,16 @@ func ObjectKeyExists(key string) (bool, error) {
 
 // InsertShortURL inserts a short URL object into the database.
 func InsertShortURL(bucket, key, destURL string, associatedUser *string) error {
-	result, err := DB.Exec(insertShortURL, fmt.Sprintf("%s/%s", bucket, key), key, destURL, associatedUser)
+	if !strings.HasPrefix(key, "/") {
+		key = "/" + key
+	}
+	result, err := DB.Exec(insertShortURL,
+		bucket+key,
+		bucket,
+		key,
+		key[1:],
+		destURL,
+		associatedUser)
 	if err != nil {
 		return err
 	}
@@ -54,10 +62,14 @@ func InsertShortURL(bucket, key, destURL string, associatedUser *string) error {
 
 // InsertFile inserts a file object into the database.
 func InsertFile(bucket, key, ext, contentType string, contentLength int64, md5Hash string, associatedUser *string) error {
+	if !strings.HasPrefix(key, "/") {
+		key = "/" + key
+	}
 	result, err := DB.Exec(insertFile,
-		fmt.Sprintf("%s/%s", bucket, key+ext),
+		bucket+key,
+		bucket,
 		key+ext,
-		key,
+		key[1:],
 		contentType,
 		contentLength,
 		md5Hash,

lib/db/sql.go

@@ -31,14 +31,14 @@ var insertShortURL = `
 INSERT INTO
 	objects (bucket_key, bucket, key, random_key, dir, type, dest_url, content_type, associated_user)
 VALUES
-	($1, 'public', $2, $2, '/', 1, $3, NULL, $4)
+	($1, $2, $3, $4, '/', 1, $5, NULL, $6)
 `
 
 var insertFile = `
 INSERT INTO
 	objects (bucket_key, bucket, key, random_key, dir, content_type, content_length, md5_hash, associated_user)
 VALUES
-	($1, 'public', $2, $3, '/', $4, $5, $6, $7)
+	($1, $2, $3, $4, '/', $5, $6, $7, $8)
 `
 
 var selectUserByUsernameOrEmail = `

lib/routes/shortenpolr.go

@@ -86,7 +86,8 @@ func ShortenPolr(associateObjectsWithUser bool) func(http.ResponseWriter, *http.
 		}
 
 		// Generate available random key
-		key, err := util.GenerateAvailableKey(5) // maximum of 5 attempts to generate a key
+		bucket := viper.GetString("database.objectBucket")
+		key, err := util.GenerateAvailableKey(bucket, 5)
 		if err != nil {
 			log.Error().Err(err).Msg("failed to generate short URL key")
 			panic(apierrors.InternalServerErrorPolr)
@@ -97,7 +98,6 @@ func ShortenPolr(associateObjectsWithUser bool) func(http.ResponseWriter, *http.
 		if associateObjectsWithUser {
 			associatedUser = &user.ID
 		}
-		bucket := viper.GetString("database.objectBucket")
 		err = db.InsertShortURL(bucket, key, urlString, associatedUser)
 		if err != nil {
 			log.Error().Err(err).Msg("failed to insert short URL object")

lib/routes/uploadpomf.go

@@ -127,7 +127,8 @@ func UploadPomf(associateObjectsWithUser bool) func(http.ResponseWriter, *http.R
 			}
 
 			// Generate a file key
-			key, err := util.GenerateAvailableKey(5)
+			bucket := viper.GetString("database.objectBucket")
+			key, err := util.GenerateAvailableKey(bucket, 5)
 			if err != nil {
 				log.Error().Err(err).Msg("failed to generate available key within 5 attempts for file object")
 				if len(files) == 1 {
@@ -225,7 +226,6 @@ func UploadPomf(associateObjectsWithUser bool) func(http.ResponseWriter, *http.R
 			if associateObjectsWithUser {
 				associatedUser = &user.ID
 			}
-			bucket := viper.GetString("database.objectBucket")
 			err = db.InsertFile(bucket, key, ext, contentType, file.Size, md5Hash, associatedUser)
 			if err != nil {
 				log.Error().Err(err).Msg("failed to create DB object for file upload")

lib/util/random.go

@@ -25,15 +25,15 @@ func GenerateRandomKey() (string, error) {
 }
 
 // GenerateAvailableKey runs GenerateRandomKey and verifies that the generated
-// key isn't already being used.
-func GenerateAvailableKey(maxAttempts int) (string, error) {
+// key isn't already being used in the given bucket.
+func GenerateAvailableKey(bucket string, maxAttempts int) (string, error) {
 	for n := 0; n < maxAttempts; n++ {
 		key, err := GenerateRandomKey()
 		if err != nil {
 			return "", errors.Wrap(err, "failed to generate random key")
 		}
 
-		exists, err := db.ObjectKeyExists(key)
+		exists, err := db.ObjectKeyExists(bucket, key)
 		if err != nil {
 			return "", errors.Wrapf(err, `failed to check if object exists with key "%s"`, key)
 		}