1 files changed, 24 insertions, 9 deletions
diff --git a/main.go b/main.go
index 11d7370..a43b510 100644
--- a/main.go
+++ b/main.go
@@ -51,13 +51,36 @@ func redirect(w http.ResponseWriter, r *http.Request) {
 	original, ok := Table[url]
 
 	if ok {
+		log.Println("Found:", original)
+
 		w.Header().Set("location", original)
 		http.Error(w, http.StatusText(301), 301)
 	} else {
+		log.Println("Unable to find website for", url)
 		http.Error(w, http.StatusText(404), 404)
 	}
 }
 
+func FindOrGenerate(url string) string {
+	for key, value := range Table {
+		if value == url {
+			return key
+		}
+	}
+
+	short := GenerateString(4)
+	_, ok := Table[short]
+
+	for ok {
+		short = GenerateString(4)
+		_, ok = Table[short]
+	}
+
+	Table[short] = url
+
+	return short
+}
+
 func shorten(w http.ResponseWriter, r *http.Request) {
 	log.Println("Request to Shorten")
 	log_header(r)
@@ -67,15 +90,7 @@ func shorten(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, http.StatusText(500), 500)
 	}
 
-	url := GenerateString(4)
-	_, ok := Table[url]
-
-	for ok {
-		url = GenerateString(4)
-		_, ok = Table[url]
-	}
-
-	Table[url] = request.URL
+	url := FindOrGenerate(request.URL)
 
 	response := Response{version, "ok", request.URL, url}