From c05a90d9c4bde30e59f0b117c786b79858eae035 Mon Sep 17 00:00:00 2001
From: Jan-Lukas Else <mail@jlelse.de>
Date: Sun, 29 Mar 2020 11:38:47 +0200
Subject: [PATCH] Allow setting custom slug

---
 main.go | 28 +++++++++++++++++-----------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/main.go b/main.go
index 8fc4dbd..bc86260 100644
--- a/main.go
+++ b/main.go
@@ -98,21 +98,27 @@ func ShortenHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	var slug string
-	err := db.QueryRow("SELECT slug FROM redirect WHERE url = ?", requestUrl).Scan(&slug)
-	if err == nil {
-		writeShortenedUrl(w, slug)
-		return
+	slug := r.URL.Query().Get("slug")
+	if slug == "" {
+		_ = db.QueryRow("SELECT slug FROM redirect WHERE url = ?", requestUrl).Scan(&slug)
 	}
 
-	var exists = true
-	for exists == true {
-		slug = generateSlug()
-		err, exists = slugExists(slug)
-		if err != nil {
-			http.Error(w, err.Error(), http.StatusInternalServerError)
+	if slug != "" {
+		if _, e := slugExists(slug); e {
+			writeShortenedUrl(w, slug)
 			return
 		}
+	} else {
+		var exists = true
+		for exists == true {
+			slug = generateSlug()
+			var err error
+			err, exists = slugExists(slug)
+			if err != nil {
+				http.Error(w, err.Error(), http.StatusInternalServerError)
+				return
+			}
+		}
 	}
 
 	stmt, err := db.Prepare("INSERT INTO redirect (slug, url, hits) VALUES (?, ?, ?)")