From 3c4d0ca60eca09ffd34c1b9427e370e1257d5b8e Mon Sep 17 00:00:00 2001
From: Daan Sprenkels <hello@dsprenkels.com>
Date: Thu, 29 Aug 2019 00:50:26 +0200
Subject: [PATCH] Implement retrieval of redirects

---
 index.go => handlers.go | 59 ++++++++++++++++++++++++++++++++++++++---
 rushlink.go             |  1 +
 2 files changed, 56 insertions(+), 4 deletions(-)
 rename index.go => handlers.go (72%)

diff --git a/index.go b/handlers.go
similarity index 72%
rename from index.go
rename to handlers.go
index 960fe03..24f31e6 100644
--- a/index.go
+++ b/handlers.go
@@ -15,6 +15,7 @@ import (
 
 	gobmarsh "gitea.hashru.nl/dsprenkels/rushlink/gobmarsh"
 
+	mux "github.com/gorilla/mux"
 	errors "github.com/pkg/errors"
 
 	bolt "go.etcd.io/bbolt"
@@ -25,7 +26,7 @@ type StoredURLState int
 type StoredURL struct {
 	State       StoredURLState
 	RawURL      string
-	Key			[]byte
+	Key         []byte
 	TimeCreated time.Time
 }
 
@@ -74,6 +75,38 @@ func indexPostHandler(w http.ResponseWriter, r *http.Request) {
 	shortenPostHandler(w, r)
 }
 
+func redirectHandler(w http.ResponseWriter, r *http.Request) {
+	vars := mux.Vars(r)
+	key := vars["key"]
+	var storedURL *StoredURL
+	if err := db.View(func(tx *bolt.Tx) error {
+		var err error
+		storedURL, err = getURL(tx, []byte(key))
+		return err
+	}); err != nil {
+		w.WriteHeader(http.StatusInternalServerError)
+		log.Printf("error: %v\n", err)
+		fmt.Fprintf(w, "internal server error: %v", err)
+		return
+	}
+	if storedURL == nil {
+		w.WriteHeader(http.StatusNotFound)
+		fmt.Fprintf(w, "url key not found in the database")
+	}
+	switch storedURL.State {
+	case Present:
+		w.Header().Set("Location", storedURL.RawURL)
+		w.WriteHeader(http.StatusTemporaryRedirect)
+	case Deleted:
+		w.WriteHeader(http.StatusGone)
+		fmt.Fprintf(w, "key has been deleted")
+	default:
+		w.WriteHeader(http.StatusInternalServerError)
+		log.Printf("error: invalid storedURL.State (%v) for key '%v'\n", storedURL.State, storedURL.Key)
+		fmt.Fprintf(w, "internal server error: invalid storedURL.State (%v)", storedURL.State)
+	}
+}
+
 func shortenPostHandler(w http.ResponseWriter, r *http.Request) {
 	rawurl := r.PostForm.Get("shorten")
 	userURL, err := url.ParseRequestURI(rawurl)
@@ -100,12 +133,28 @@ func shortenPostHandler(w http.ResponseWriter, r *http.Request) {
 		return err
 	}); err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
+		log.Printf("error: %v\n", err)
 		fmt.Fprintf(w, "internal server error: %v", err)
 		return
 	}
 
 	w.WriteHeader(http.StatusOK)
-	fmt.Fprintf(w, "new URL saved: %+v", storedURL)
+	fmt.Fprintf(w, "URL saved at /%v", string(storedURL.Key))
+}
+
+// Retrieve a URL from the database
+func getURL(tx *bolt.Tx, key []byte) (*StoredURL, error) {
+	shortenBucket := tx.Bucket([]byte(BUCKET_SHORTEN))
+	if shortenBucket == nil {
+		return nil, fmt.Errorf("bucket %v does not exist", BUCKET_SHORTEN)
+	}
+	storedBytes := shortenBucket.Get(key)
+	if storedBytes == nil {
+		return nil, nil
+	}
+	storedURL := &StoredURL{}
+	err := gobmarsh.Unmarshal(storedBytes, storedURL)
+	return storedURL, err
 }
 
 // Add a new URL to the database
@@ -138,14 +187,13 @@ func shortenURL(tx *bolt.Tx, userURL *url.URL) (*StoredURL, error) {
 	storedURL := StoredURL{
 		State:       Present,
 		RawURL:      userURL.String(),
-		Key:    urlKey,
+		Key:         urlKey,
 		TimeCreated: time.Now().UTC(),
 	}
 	storedBytes, err := gobmarsh.Marshal(storedURL)
 	if err != nil {
 		return nil, errors.Wrap(err, "encoding for database failed")
 	}
-	log.Println(urlKey)
 	if err := shortenBucket.Put(urlKey, storedBytes); err != nil {
 		return nil, errors.Wrap(err, "database transaction failed")
 	}
@@ -181,5 +229,8 @@ func generateURLKey(epoch int) ([]byte, error) {
 	urlKey[limb] &= ^(1 << uint(5-bit))
 
 	// Convert this ID to a canonical base64 notation
+	for i := range urlKey {
+		urlKey[i] = base64Alphabet[urlKey[i]]
+	}
 	return urlKey, nil
 }
diff --git a/rushlink.go b/rushlink.go
index 3780778..037bcd4 100644
--- a/rushlink.go
+++ b/rushlink.go
@@ -43,6 +43,7 @@ func main() {
 	router := mux.NewRouter()
 	router.HandleFunc("/", indexGetHandler).Methods("GET")
 	router.HandleFunc("/", indexPostHandler).Methods("POST")
+	router.HandleFunc("/{key:[A-Za-z0-9-_]{4,}}", redirectHandler).Methods("GET")
 
 	// Start the server
 	srv := &http.Server{