vodkas

commit 612e0ef93e9cd999b3bc68257389c9ff73b4da10
parent c7ccc3dce979ede6c1971cfee9d0e780fdf18f15
Author: vh <vetle.haflan@gmail.com>
Date:   Sun, 23 Feb 2020 19:10:15 +0100

Replace map by bbolt database - storage no longer limited to RAM

Diffstat:
M
vodkas.go
 | 
75
+++++++++++++++++++++++++++++++++++++++++++++++++++++++--------------------

1 file changed, 55 insertions(+), 20 deletions(-)
diff --git a/vodkas.go b/vodkas.go
@@ -6,11 +6,13 @@ import (
         "flag"
         "fmt"
         "github.com/gorilla/mux"
+        "github.com/pkg/errors"
+        "go.etcd.io/bbolt"
         "io/ioutil"
         "log"
         "net/http"
         "strings"
-        "sync"
+        "time"
 )
 
 
@@ -25,24 +27,40 @@ As soon as a specific shot has been accessed both the link and the contents
 are removed completely.
 `
 
-type SyncedStorage struct{
-        links map[string][]byte
-        sync.Mutex
-}
-
-var storage SyncedStorage
+var rootBucket = "root"
+var db *bbolt.DB
 
-func (storage *SyncedStorage) push(code string, contents []byte) {
-        storage.Lock()
-        defer storage.Unlock()
-        storage.links[code] = contents
+func push(code string, contents []byte) {
+        err := db.Update(func(tx *bbolt.Tx) error {
+                b := tx.Bucket([]byte(rootBucket))
+                if b == nil {
+                      return errors.New("Failed to open root bucket")
+                }
+                err := b.Put([]byte(code), contents)
+                return err
+        })
+        if err != nil {
+                log.Printf("Push to '%v' failed: %v", code, err)
+        }
 }
 
-func (storage *SyncedStorage) pop(code string) (contents []byte, found bool) {
-        storage.Lock()
-        defer storage.Unlock()
-        contents, found = storage.links[code]
-        delete(storage.links, code)
+func pop(code string) (contents []byte, found bool) {
+        err := db.Update(func(tx *bbolt.Tx) error {
+                b := tx.Bucket([]byte(rootBucket))
+                if b == nil {
+                        return errors.New("Failed to open root bucket")
+                }
+                contents = b.Get([]byte(code))
+                found = contents != nil
+                if found {
+                        fmt.Printf("Found contents for shotcode %v\n", code)
+                        return b.Delete([]byte(code))
+                }
+                return nil
+        })
+        if err != nil {
+                log.Printf("Push to '%v' failed: %v", code, err)
+        }
         return
 }
 
@@ -53,7 +71,7 @@ func MainHandler(res http.ResponseWriter, r *http.Request) {
                 res.WriteHeader(http.StatusInternalServerError)
                 return
         }
-        storage.push("sorandom", b)
+        push("sorandom", b)
         textOnly := r.Header.Get("Simple") != ""    // Should be able to force simple
         textOnly = textOnly || strings.Contains(r.Header.Get("User-Agent"), "curl")
         var responseText string
@@ -74,7 +92,7 @@ func MainHandler(res http.ResponseWriter, r *http.Request) {
 // TODO: Generate random code if no request given (https://flaviocopes.com/go-random/)
 func ShotHandler(res http.ResponseWriter, r *http.Request) {
         code := mux.Vars(r)["shotCode"]
-        contents, found := storage.pop(code)
+        contents, found := pop(code)
         // GET requests only
         /*if !found {
                 res.WriteHeader(http.StatusNotFound)
@@ -93,7 +111,7 @@ func ShotHandler(res http.ResponseWriter, r *http.Request) {
                         res.WriteHeader(http.StatusInternalServerError)
                         log.Panicln("Error when trying to read body")
                 }
-                storage.push(code, b)
+                push(code, b)
                 if _, err := res.Write([]byte(fmt.Sprint("Contents stored in given link\n"))); err != nil {
                         log.Panicln("Error when trying to write response")
                 }
@@ -104,9 +122,26 @@ func ShotHandler(res http.ResponseWriter, r *http.Request) {
 
 /**************** Main ****************/
 func main(){
-        storage.links = make(map[string][]byte)
         port := flag.Int("p", 8080, "Port")
+        dbFile := flag.String("d", "vodka.db", "Database file")
         flag.Parse()
+        var err error // Because ':=' can't be used on the line below without declaring db as a new *local* variable, making the global one nil
+        db, err = bbolt.Open(*dbFile, 0600, &bbolt.Options{Timeout: 1 * time.Second})
+        defer db.Close()
+        if err != nil {
+                panic(err)
+        }
+        err = db.Update(func(tx *bbolt.Tx) error {
+                _, err := tx.CreateBucketIfNotExists([]byte(rootBucket))
+                if err != nil {
+                        return err
+                }
+                return err
+        })
+        if err != nil {
+                panic(err)
+        }
+
         defer fmt.Println("Server shutting down")
         fmt.Println("Server started listening at port", *port)
         router := mux.NewRouter()