diff --git a/db.go b/db.go index 7e8249b..38a54a7 100644 --- a/db.go +++ b/db.go @@ -18,8 +18,8 @@ const CURRENT_MIGRATE_VERSION = 1 // the database version, secret site-wide keys. const BUCKET_CONF = "conf" -// The main bucket for URL shortening values -const BUCKET_SHORTEN = "shorten" +// The main bucket for paste values and URL redirects +const BUCKET_PASTES = "pastes" // This value stores the current migration version. If this value is less than // CURRENT_MIGRATE_VERSION, the database has to be migrated. @@ -42,8 +42,8 @@ func migrateDatabase(tx *bolt.Tx) error { return err } - // Create URL shortening bucket - _, err = tx.CreateBucket([]byte(BUCKET_SHORTEN)) + // Create paste bucket + _, err = tx.CreateBucket([]byte(BUCKET_PASTES)) if err != nil { return err } diff --git a/handlers.go b/handlers.go index 21d3056..020db28 100644 --- a/handlers.go +++ b/handlers.go @@ -19,18 +19,25 @@ import ( bolt "go.etcd.io/bbolt" ) -type StoredURLState int +type PasteType int +type PasteState int -type StoredURL struct { - State StoredURLState - RawURL string +type StoredPaste struct { + Type PasteType + State PasteState + Content []byte Key []byte TimeCreated time.Time } const ( - Present StoredURLState = iota - Deleted + TypePaste PasteType = iota + TypeRedirect +) + +const ( + StatePresent PasteState = iota + StateDeleted ) var base64Alphabet = []byte("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_") @@ -73,13 +80,21 @@ func indexPostHandler(w http.ResponseWriter, r *http.Request) { shortenPostHandler(w, r) } -func redirectHandler(w http.ResponseWriter, r *http.Request) { +func pasteGetHandler(w http.ResponseWriter, r *http.Request) { + pasteGetHandlerInner(w, r, false) +} + +func pasteGetHandlerNoRedirect(w http.ResponseWriter, r *http.Request) { + pasteGetHandlerInner(w, r, true) +} + +func pasteGetHandlerInner(w http.ResponseWriter, r *http.Request, noRedirect bool) { vars := mux.Vars(r) key := vars["key"] - var storedURL *StoredURL + var storedPaste *StoredPaste if err := db.View(func(tx *bolt.Tx) error { var err error - storedURL, err = getURL(tx, []byte(key)) + storedPaste, err = getURL(tx, []byte(key)) return err }); err != nil { w.WriteHeader(http.StatusInternalServerError) @@ -87,21 +102,32 @@ func redirectHandler(w http.ResponseWriter, r *http.Request) { fmt.Fprintf(w, "internal server error: %v", err) return } - if storedURL == nil { + if storedPaste == nil { w.WriteHeader(http.StatusNotFound) fmt.Fprintf(w, "url key not found in the database") + return } - switch storedURL.State { - case Present: - w.Header().Set("Location", storedURL.RawURL) - w.WriteHeader(http.StatusTemporaryRedirect) - case Deleted: + switch storedPaste.State { + case StatePresent: + if !noRedirect { + rawurl := string(storedPaste.Content) + urlParse, err := url.Parse(rawurl) + if err != nil { + w.WriteHeader(http.StatusInternalServerError) + log.Printf("error: invalid URL ('%v') in database for key '%v': %v\n", rawurl, storedPaste.Key, err) + fmt.Fprintf(w, "internal server error: invalid url in database") + return + } + http.Redirect(w, r, urlParse.String(), http.StatusSeeOther) + } + w.Write(storedPaste.Content) + case StateDeleted: 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) + log.Printf("error: invalid storedPaste.State (%v) for key '%v'\n", storedPaste.State, storedPaste.Key) + fmt.Fprintf(w, "internal server error: invalid storedPaste.State (%v)", storedPaste.State) } } @@ -124,10 +150,10 @@ func shortenPostHandler(w http.ResponseWriter, r *http.Request) { return } - var storedURL *StoredURL + var storedPaste *StoredPaste if err := db.Update(func(tx *bolt.Tx) error { u, err := shortenURL(tx, userURL) - storedURL = u + storedPaste = u return err }); err != nil { w.WriteHeader(http.StatusInternalServerError) @@ -137,31 +163,31 @@ func shortenPostHandler(w http.ResponseWriter, r *http.Request) { } w.WriteHeader(http.StatusOK) - fmt.Fprintf(w, "URL saved at /%v", string(storedURL.Key)) + fmt.Fprintf(w, "URL saved at /%v", string(storedPaste.Key)) } // Retrieve a URL from the database -func getURL(tx *bolt.Tx, key []byte) (*StoredURL, error) { - shortenBucket := tx.Bucket([]byte(BUCKET_SHORTEN)) +func getURL(tx *bolt.Tx, key []byte) (*StoredPaste, error) { + shortenBucket := tx.Bucket([]byte(BUCKET_PASTES)) if shortenBucket == nil { - return nil, fmt.Errorf("bucket %v does not exist", BUCKET_SHORTEN) + return nil, fmt.Errorf("bucket %v does not exist", BUCKET_PASTES) } storedBytes := shortenBucket.Get(key) if storedBytes == nil { return nil, nil } - storedURL := &StoredURL{} - err := gobmarsh.Unmarshal(storedBytes, storedURL) - return storedURL, err + storedPaste := &StoredPaste{} + err := gobmarsh.Unmarshal(storedBytes, storedPaste) + return storedPaste, err } // Add a new URL to the database // // Returns the new ID if the url was successfully shortened -func shortenURL(tx *bolt.Tx, userURL *url.URL) (*StoredURL, error) { - shortenBucket := tx.Bucket([]byte(BUCKET_SHORTEN)) +func shortenURL(tx *bolt.Tx, userURL *url.URL) (*StoredPaste, error) { + shortenBucket := tx.Bucket([]byte(BUCKET_PASTES)) if shortenBucket == nil { - return nil, fmt.Errorf("bucket %v does not exist", BUCKET_SHORTEN) + return nil, fmt.Errorf("bucket %v does not exist", BUCKET_PASTES) } // Generate a key until it is not in the database, this occurs in O(log N), @@ -182,20 +208,21 @@ func shortenURL(tx *bolt.Tx, userURL *url.URL) (*StoredURL, error) { } // Store the new key - storedURL := StoredURL{ - State: Present, - RawURL: userURL.String(), + storedPaste := StoredPaste{ + Type: TypeRedirect, + State: StatePresent, + Content: []byte(userURL.String()), Key: urlKey, TimeCreated: time.Now().UTC(), } - storedBytes, err := gobmarsh.Marshal(storedURL) + storedBytes, err := gobmarsh.Marshal(storedPaste) if err != nil { return nil, errors.Wrap(err, "encoding for database failed") } if err := shortenBucket.Put(urlKey, storedBytes); err != nil { return nil, errors.Wrap(err, "database transaction failed") } - return &storedURL, nil + return &storedPaste, nil } func generateURLKey(epoch int) ([]byte, error) { diff --git a/rushlink.go b/rushlink.go index f72bdd4..ea444ba 100644 --- a/rushlink.go +++ b/rushlink.go @@ -46,7 +46,8 @@ 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") + router.HandleFunc("/{key:[A-Za-z0-9-_]{4,}}", pasteGetHandler).Methods("GET") + router.HandleFunc("/{key:[A-Za-z0-9-_]{4,}}/nr", pasteGetHandlerNoRedirect).Methods("GET") // Start the server srv := &http.Server{