Rename main bucket to "pastes" and add /nr

This commit is contained in:
Daan Sprenkels 2019-08-31 00:02:01 +02:00
parent 498c84541b
commit 5a7394ffd3
3 changed files with 67 additions and 39 deletions

8
db.go
View File

@ -18,8 +18,8 @@ const CURRENT_MIGRATE_VERSION = 1
// the database version, secret site-wide keys. // the database version, secret site-wide keys.
const BUCKET_CONF = "conf" const BUCKET_CONF = "conf"
// The main bucket for URL shortening values // The main bucket for paste values and URL redirects
const BUCKET_SHORTEN = "shorten" const BUCKET_PASTES = "pastes"
// This value stores the current migration version. If this value is less than // This value stores the current migration version. If this value is less than
// CURRENT_MIGRATE_VERSION, the database has to be migrated. // CURRENT_MIGRATE_VERSION, the database has to be migrated.
@ -42,8 +42,8 @@ func migrateDatabase(tx *bolt.Tx) error {
return err return err
} }
// Create URL shortening bucket // Create paste bucket
_, err = tx.CreateBucket([]byte(BUCKET_SHORTEN)) _, err = tx.CreateBucket([]byte(BUCKET_PASTES))
if err != nil { if err != nil {
return err return err
} }

View File

@ -19,18 +19,25 @@ import (
bolt "go.etcd.io/bbolt" bolt "go.etcd.io/bbolt"
) )
type StoredURLState int type PasteType int
type PasteState int
type StoredURL struct { type StoredPaste struct {
State StoredURLState Type PasteType
RawURL string State PasteState
Content []byte
Key []byte Key []byte
TimeCreated time.Time TimeCreated time.Time
} }
const ( const (
Present StoredURLState = iota TypePaste PasteType = iota
Deleted TypeRedirect
)
const (
StatePresent PasteState = iota
StateDeleted
) )
var base64Alphabet = []byte("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_") var base64Alphabet = []byte("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_")
@ -73,13 +80,21 @@ func indexPostHandler(w http.ResponseWriter, r *http.Request) {
shortenPostHandler(w, r) 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) vars := mux.Vars(r)
key := vars["key"] key := vars["key"]
var storedURL *StoredURL var storedPaste *StoredPaste
if err := db.View(func(tx *bolt.Tx) error { if err := db.View(func(tx *bolt.Tx) error {
var err error var err error
storedURL, err = getURL(tx, []byte(key)) storedPaste, err = getURL(tx, []byte(key))
return err return err
}); err != nil { }); err != nil {
w.WriteHeader(http.StatusInternalServerError) w.WriteHeader(http.StatusInternalServerError)
@ -87,21 +102,32 @@ func redirectHandler(w http.ResponseWriter, r *http.Request) {
fmt.Fprintf(w, "internal server error: %v", err) fmt.Fprintf(w, "internal server error: %v", err)
return return
} }
if storedURL == nil { if storedPaste == nil {
w.WriteHeader(http.StatusNotFound) w.WriteHeader(http.StatusNotFound)
fmt.Fprintf(w, "url key not found in the database") fmt.Fprintf(w, "url key not found in the database")
return
} }
switch storedURL.State { switch storedPaste.State {
case Present: case StatePresent:
w.Header().Set("Location", storedURL.RawURL) if !noRedirect {
w.WriteHeader(http.StatusTemporaryRedirect) rawurl := string(storedPaste.Content)
case Deleted: 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) w.WriteHeader(http.StatusGone)
fmt.Fprintf(w, "key has been deleted") fmt.Fprintf(w, "key has been deleted")
default: default:
w.WriteHeader(http.StatusInternalServerError) w.WriteHeader(http.StatusInternalServerError)
log.Printf("error: invalid storedURL.State (%v) for key '%v'\n", storedURL.State, storedURL.Key) log.Printf("error: invalid storedPaste.State (%v) for key '%v'\n", storedPaste.State, storedPaste.Key)
fmt.Fprintf(w, "internal server error: invalid storedURL.State (%v)", storedURL.State) 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 return
} }
var storedURL *StoredURL var storedPaste *StoredPaste
if err := db.Update(func(tx *bolt.Tx) error { if err := db.Update(func(tx *bolt.Tx) error {
u, err := shortenURL(tx, userURL) u, err := shortenURL(tx, userURL)
storedURL = u storedPaste = u
return err return err
}); err != nil { }); err != nil {
w.WriteHeader(http.StatusInternalServerError) w.WriteHeader(http.StatusInternalServerError)
@ -137,31 +163,31 @@ func shortenPostHandler(w http.ResponseWriter, r *http.Request) {
} }
w.WriteHeader(http.StatusOK) 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 // Retrieve a URL from the database
func getURL(tx *bolt.Tx, key []byte) (*StoredURL, error) { func getURL(tx *bolt.Tx, key []byte) (*StoredPaste, error) {
shortenBucket := tx.Bucket([]byte(BUCKET_SHORTEN)) shortenBucket := tx.Bucket([]byte(BUCKET_PASTES))
if shortenBucket == nil { 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) storedBytes := shortenBucket.Get(key)
if storedBytes == nil { if storedBytes == nil {
return nil, nil return nil, nil
} }
storedURL := &StoredURL{} storedPaste := &StoredPaste{}
err := gobmarsh.Unmarshal(storedBytes, storedURL) err := gobmarsh.Unmarshal(storedBytes, storedPaste)
return storedURL, err return storedPaste, err
} }
// Add a new URL to the database // Add a new URL to the database
// //
// Returns the new ID if the url was successfully shortened // Returns the new ID if the url was successfully shortened
func shortenURL(tx *bolt.Tx, userURL *url.URL) (*StoredURL, error) { func shortenURL(tx *bolt.Tx, userURL *url.URL) (*StoredPaste, error) {
shortenBucket := tx.Bucket([]byte(BUCKET_SHORTEN)) shortenBucket := tx.Bucket([]byte(BUCKET_PASTES))
if shortenBucket == nil { 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), // 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 // Store the new key
storedURL := StoredURL{ storedPaste := StoredPaste{
State: Present, Type: TypeRedirect,
RawURL: userURL.String(), State: StatePresent,
Content: []byte(userURL.String()),
Key: urlKey, Key: urlKey,
TimeCreated: time.Now().UTC(), TimeCreated: time.Now().UTC(),
} }
storedBytes, err := gobmarsh.Marshal(storedURL) storedBytes, err := gobmarsh.Marshal(storedPaste)
if err != nil { if err != nil {
return nil, errors.Wrap(err, "encoding for database failed") return nil, errors.Wrap(err, "encoding for database failed")
} }
if err := shortenBucket.Put(urlKey, storedBytes); err != nil { if err := shortenBucket.Put(urlKey, storedBytes); err != nil {
return nil, errors.Wrap(err, "database transaction failed") return nil, errors.Wrap(err, "database transaction failed")
} }
return &storedURL, nil return &storedPaste, nil
} }
func generateURLKey(epoch int) ([]byte, error) { func generateURLKey(epoch int) ([]byte, error) {

View File

@ -46,7 +46,8 @@ func main() {
router := mux.NewRouter() router := mux.NewRouter()
router.HandleFunc("/", indexGetHandler).Methods("GET") router.HandleFunc("/", indexGetHandler).Methods("GET")
router.HandleFunc("/", indexPostHandler).Methods("POST") 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 // Start the server
srv := &http.Server{ srv := &http.Server{