rushlink/handlers.go
2019-08-29 00:51:40 +02:00

237 lines
5.9 KiB
Go

package main
//go:generate go get github.com/go-bindata/go-bindata
//go:generate go get -u github.com/go-bindata/go-bindata/...
//go:generate go-bindata -pkg $GOPACKAGE assets/
import (
"crypto/rand"
"fmt"
"io"
"log"
"net/http"
"net/url"
"time"
gobmarsh "gitea.hashru.nl/dsprenkels/rushlink/gobmarsh"
mux "github.com/gorilla/mux"
errors "github.com/pkg/errors"
bolt "go.etcd.io/bbolt"
)
type StoredURLState int
type StoredURL struct {
State StoredURLState
RawURL string
Key []byte
TimeCreated time.Time
}
const (
Present StoredURLState = iota
Deleted
)
var base64Alphabet = []byte("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_")
var indexContents = MustAsset("assets/index.txt")
func indexGetHandler(w http.ResponseWriter, r *http.Request) {
_, err := w.Write(indexContents)
if err != nil {
panic(err)
}
}
func indexPostHandler(w http.ResponseWriter, r *http.Request) {
if err := r.ParseMultipartForm(50 * 1000 * 1000); err != nil {
w.WriteHeader(http.StatusInternalServerError)
fmt.Fprintf(w, "Internal server error: %v", err)
return
}
// Determine what kind of post this is, currently only `shorten=...`
if len(r.PostForm) == 0 {
w.WriteHeader(http.StatusBadRequest)
var buf []byte
r.Body.Read(buf)
io.WriteString(w, "empty body in POST request")
return
}
shorten_values, prs := r.PostForm["shorten"]
if !prs {
w.WriteHeader(http.StatusBadRequest)
io.WriteString(w, "no 'shorten' param supplied")
return
}
if len(shorten_values) != 1 {
w.WriteHeader(http.StatusBadRequest)
io.WriteString(w, "only one 'shorten' param is allowed per request")
return
}
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)
if err != nil {
w.WriteHeader(http.StatusBadRequest)
fmt.Fprintf(w, "invalid url (%v): %v", err, rawurl)
return
}
if userURL.Scheme == "" {
w.WriteHeader(http.StatusBadRequest)
fmt.Fprintf(w, "invalid url (unspecified scheme)", rawurl)
return
}
if userURL.Host == "" {
w.WriteHeader(http.StatusBadRequest)
fmt.Fprintf(w, "invalid url (unspecified host)", rawurl)
return
}
var storedURL *StoredURL
if err := db.Update(func(tx *bolt.Tx) error {
u, err := shortenURL(tx, userURL)
storedURL = u
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, "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
//
// 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))
if shortenBucket == nil {
return nil, fmt.Errorf("bucket %v does not exist", BUCKET_SHORTEN)
}
// Generate a key until it is not in the database, this occurs in O(log N),
// where N is the amount of keys stored in the url-shorten database.
epoch := 0
var urlKey []byte
for {
var err error
urlKey, err = generateURLKey(epoch)
if err != nil {
return nil, errors.Wrap(err, "url-key generation failed")
}
found := shortenBucket.Get(urlKey)
if found == nil {
break
}
epoch++
}
// Store the new key
storedURL := StoredURL{
State: Present,
RawURL: userURL.String(),
Key: urlKey,
TimeCreated: time.Now().UTC(),
}
storedBytes, err := gobmarsh.Marshal(storedURL)
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
}
func generateURLKey(epoch int) ([]byte, error) {
urlKey := make([]byte, 4+epoch)
_, err := rand.Read(urlKey)
if err != nil {
return nil, err
}
// Put all the values in the range 0..64 for easier base64-encoding
for i := 0; i < len(urlKey); i++ {
urlKey[i] &= 0x3F
}
// Implement truncate-resistance by forcing the prefix to
// 0b111110xxxxxxxxxx
// ^----- {epoch} ones followed by a single 0
//
// Example when epoch is 1: prefix is 0b10.
i := 0
for i < epoch {
// Set this bit to 1
limb := i / 6
bit := i % 6
urlKey[limb] |= 1 << uint(5-bit)
i++
}
// Finally set the next bit to 0
limb := i / 6
bit := i % 6
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
}