From 061bb23625dc892b1839e4ac1808d9dd954b4916 Mon Sep 17 00:00:00 2001 From: Daan Sprenkels Date: Sun, 25 Aug 2019 21:33:56 +0200 Subject: [PATCH] Initial commit --- .gitignore | 2 + assets/index.txt | 10 +++ db.go | 97 +++++++++++++++++++++++ go.mod | 7 ++ go.sum | 10 +++ gobmarsh/marshal.go | 2 + index.go | 185 ++++++++++++++++++++++++++++++++++++++++++++ rushlink.go | 68 ++++++++++++++++ 8 files changed, 381 insertions(+) create mode 100644 assets/index.txt create mode 100644 db.go create mode 100644 go.sum create mode 100644 index.go create mode 100644 rushlink.go diff --git a/.gitignore b/.gitignore index b899f97..fbf375e 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,5 @@ # Generated code from assets using bindata bindata.go +# Output binary +/rushlink diff --git a/assets/index.txt b/assets/index.txt new file mode 100644 index 0000000..052f227 --- /dev/null +++ b/assets/index.txt @@ -0,0 +1,10 @@ +#RU URL SHORTENER +================= + +Based on https://0x0.st/, this site allows you to easily shorten URLs using +the command line. + +## USAGE + + # Shorten a URL + curl -F'shorten=http://example.com/some/long/url' https://hashru.link diff --git a/db.go b/db.go new file mode 100644 index 0000000..5961853 --- /dev/null +++ b/db.go @@ -0,0 +1,97 @@ +package main + +import ( + "fmt" + "log" + + gobmarsh "gitea.hashru.nl/dsprenkels/rushlink/gobmarsh" + bolt "go.etcd.io/bbolt" +) + +// The current database version +// +// If we alter the database format, we bump this number and write a new +// database migration in migrateDatabase(). +const CURRENT_MIGRATE_VERSION = 1 + +// Bucket storing everything that is not a bulk value. This includes stuff like +// the database version, secret site-wide keys. +const BUCKET_CONF = "conf" + +// The main bucket for URL shortening values +const BUCKET_SHORTEN = "shorten" + +// This value stores the current migration version. If this value is less than +// CURRENT_MIGRATE_VERSION, the database has to be migrated. +const KEY_MIGRATE_VERSION = "migrate_version" + +// Initialize and migrate the database to the current version +func migrateDatabase(tx *bolt.Tx) error { + dbVersion, err := dbVersion(tx) + if err != nil { + return err + } + + // Migrate the database to version 1 + if dbVersion < 1 { + log.Println("migrating database to version 1") + + // Create conf bucket + _, err := tx.CreateBucket([]byte(BUCKET_CONF)) + if err != nil { + return err + } + + // Create URL shortening bucket + _, err = tx.CreateBucket([]byte(BUCKET_SHORTEN)) + if err != nil { + return err + } + + // Update the version number + if err := setDBVersion(tx, 1); err != nil { + return err + } + } + + return nil +} + +// Get the current migrate version from the database +func dbVersion(tx *bolt.Tx) (int, error) { + conf := tx.Bucket([]byte(BUCKET_CONF)) + if conf == nil { + return 0, nil + } + dbVersionBytes := conf.Get([]byte(KEY_MIGRATE_VERSION)) + if dbVersionBytes == nil { + return 0, nil + } + + // Version was already stored + var dbVersion int + if err := gobmarsh.Unmarshal(dbVersionBytes, &dbVersion); err != nil { + return 0, err + } + if dbVersion == 0 { + return 0, fmt.Errorf("database version is invalid (%v)", dbVersion) + } + if dbVersion > CURRENT_MIGRATE_VERSION { + return 0, fmt.Errorf("database version is too recent (%v > %v)", dbVersion, CURRENT_MIGRATE_VERSION) + } + return dbVersion, nil +} + +// Update the current migrate version in the database +func setDBVersion(tx *bolt.Tx, version int) error { + conf, err := tx.CreateBucketIfNotExists([]byte(BUCKET_CONF)) + if err != nil { + return err + } + + versionBytes, err := gobmarsh.Marshal(version) + if err != nil { + return err + } + return conf.Put([]byte(KEY_MIGRATE_VERSION), versionBytes) +} diff --git a/go.mod b/go.mod index e75a316..0705d5d 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,10 @@ module gitea.hashru.nl/dsprenkels/rushlink go 1.12 + +require ( + github.com/go-bindata/go-bindata v3.1.2+incompatible // indirect + github.com/gorilla/mux v1.7.3 + github.com/pkg/errors v0.8.1 + go.etcd.io/bbolt v1.3.3 +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..ae2948b --- /dev/null +++ b/go.sum @@ -0,0 +1,10 @@ +github.com/go-bindata/go-bindata v3.1.2+incompatible h1:5vjJMVhowQdPzjE1LdxyFF7YFTXg5IgGVW4gBr5IbvE= +github.com/go-bindata/go-bindata v3.1.2+incompatible/go.mod h1:xK8Dsgwmeed+BBsSy2XTopBn/8uK2HWuGSnA11C3Joo= +github.com/gorilla/mux v1.7.3 h1:gnP5JzjVOuiZD07fKKToCAOjS0yOpj/qPETTXCCS6hw= +github.com/gorilla/mux v1.7.3/go.mod h1:1lud6UwP+6orDFRuTfBEV8e9/aOM/c4fVVCaMa2zaAs= +github.com/joomcode/errorx v1.0.0 h1:RJAKLTy1Sv2Tszhu14m5RZP4VGRlhXutG/XlL1En5VM= +github.com/joomcode/errorx v1.0.0/go.mod h1:kgco15ekB6cs+4Xjzo7SPeXzx38PbJzBwbnu9qfVNHQ= +github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I= +github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +go.etcd.io/bbolt v1.3.3 h1:MUGmc65QhB3pIlaQ5bB4LwqSj6GIonVJXpZiaKNyaKk= +go.etcd.io/bbolt v1.3.3/go.mod h1:IbVyRI1SCnLcuJnV2u8VeU0CEYM7e686BmAb1XKL+uU= diff --git a/gobmarsh/marshal.go b/gobmarsh/marshal.go index baa0154..329fa4e 100644 --- a/gobmarsh/marshal.go +++ b/gobmarsh/marshal.go @@ -1,5 +1,7 @@ package gobmarsh +// Easier marshalling to and from gob encoding + import ( "bytes" "encoding/gob" diff --git a/index.go b/index.go new file mode 100644 index 0000000..960fe03 --- /dev/null +++ b/index.go @@ -0,0 +1,185 @@ +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" + + 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 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) + fmt.Fprintf(w, "internal server error: %v", err) + return + } + + w.WriteHeader(http.StatusOK) + fmt.Fprintf(w, "new URL saved: %+v", storedURL) +} + +// 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") + } + log.Println(urlKey) + 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 + return urlKey, nil +} diff --git a/rushlink.go b/rushlink.go new file mode 100644 index 0000000..3780778 --- /dev/null +++ b/rushlink.go @@ -0,0 +1,68 @@ +package main + +import ( + "flag" + "log" + "net/http" + "time" + + mux "github.com/gorilla/mux" + errors "github.com/pkg/errors" + bolt "go.etcd.io/bbolt" +) + +type ParsedArguments struct { + databaseName string +} + +var appConfig ParsedArguments +var db *bolt.DB + +func main() { + // Parse the arguments and construct the ParsedArguments + appConfigRef, err := parseArguments() + if err != nil { + log.Fatal(err) + } + appConfig = *appConfigRef + + // Open the bolt database + db, err = bolt.Open(appConfig.databaseName, 0666, &bolt.Options{Timeout: 1 * time.Second}) + if err != nil { + decoratedErr := errors.Wrapf(err, "failed to open database at '%v'", appConfig.databaseName) + log.Fatal(decoratedErr) + } + defer db.Close() + err = db.Update(migrateDatabase) + if err != nil { + decoratedErr := errors.Wrapf(err, "failed to migrate database") + log.Fatal(decoratedErr) + } + + // Initialize Gorilla router + router := mux.NewRouter() + router.HandleFunc("/", indexGetHandler).Methods("GET") + router.HandleFunc("/", indexPostHandler).Methods("POST") + + // Start the server + srv := &http.Server{ + Handler: router, + Addr: "127.0.0.1:8000", + WriteTimeout: 15 * time.Second, + ReadTimeout: 15 * time.Second, + } + log.Fatal(srv.ListenAndServe()) +} + +// Parse the input arguments and return the initialized application config struct +func parseArguments() (*ParsedArguments, error) { + config := ParsedArguments{} + + flag.StringVar(&config.databaseName, "database", "", "Location of the database file") + flag.Parse() + + if config.databaseName == "" { + return nil, errors.New("database not set") + } + return &config, nil +}