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 }