commit 9807b63a307658e6299a18f32ed65459dd826144
parent 29c5f02178b91624f72d00afb73c84b531b17b9b
Author: Vetle Haflan <vetle@haflan.dev>
Date: Tue, 3 Mar 2020 20:32:57 +0100
gofmt
Diffstat:
| M | vodkas.go | | | 421 | ++++++++++++++++++++++++++++++++++++++++--------------------------------------- |
1 file changed, 212 insertions(+), 209 deletions(-)
diff --git a/vodkas.go b/vodkas.go
@@ -2,24 +2,23 @@
package main
import (
- "crypto/rand"
- "encoding/hex"
- "flag"
- "fmt"
- "github.com/gorilla/mux"
- "github.com/pkg/errors"
- "go.etcd.io/bbolt"
- "io"
- "io/ioutil"
- "log"
- "mime/multipart"
- "net/http"
- "strings"
- "strconv"
- "time"
+ "crypto/rand"
+ "encoding/hex"
+ "flag"
+ "fmt"
+ "github.com/gorilla/mux"
+ "github.com/pkg/errors"
+ "go.etcd.io/bbolt"
+ "io"
+ "io/ioutil"
+ "log"
+ "mime/multipart"
+ "net/http"
+ "strconv"
+ "strings"
+ "time"
)
-
/**************** Handler Functions ****************/
const FormNameFile = "file"
@@ -42,36 +41,36 @@ var db *bbolt.DB
// Pops data from database. Will probably be replaced by shot(), and support more
// than a single download (although that will still be the default)
func pop(key 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(key))
- found = contents != nil
- if found {
- fmt.Printf("Found contents for shotkey %v\n", key)
- return b.Delete([]byte(key))
- }
- return nil
- })
- if err != nil {
- log.Printf("Push to '%v' failed: %v", key, err)
- }
- return
+ 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(key))
+ found = contents != nil
+ if found {
+ fmt.Printf("Found contents for shotkey %v\n", key)
+ return b.Delete([]byte(key))
+ }
+ return nil
+ })
+ if err != nil {
+ log.Printf("Push to '%v' failed: %v", key, err)
+ }
+ return
}
// Check if the key is taken without touching the contents
func smell(shotKey string) (found bool) {
- _ = db.View(func(tx *bbolt.Tx) error {
- b := tx.Bucket([]byte(rootBucket))
- if b == nil {
- log.Fatal("Failed to open root bucket")
- }
- found = b.Get([]byte(shotKey)) != nil
- return nil
- })
- return
+ _ = db.View(func(tx *bbolt.Tx) error {
+ b := tx.Bucket([]byte(rootBucket))
+ if b == nil {
+ log.Fatal("Failed to open root bucket")
+ }
+ found = b.Get([]byte(shotKey)) != nil
+ return nil
+ })
+ return
}
/*
@@ -81,60 +80,64 @@ func shot(shotKey string) (contents []byte, error err) {
*/
func pour(shotKey string, r *http.Request) (err error) {
- var contents []byte
- var numshots int
- // Dumps can be both x-www-urlencoded and multipart/form-data.
- // Try multipart first, then x-www-urlencoded if no mpReader is returned
- mpReader, _ := r.MultipartReader()
- if mpReader != nil {
- contents, numshots, err = extractMultipart(mpReader)
- } else {
- numshots = 1
- contents, err = ioutil.ReadAll(r.Body)
- }
- if err != nil {
- return err
- }
- fmt.Printf("Number of shots: %v", numshots)
- err = db.Update(func(tx *bbolt.Tx) error {
- b := tx.Bucket([]byte(rootBucket))
- if b == nil {
- return errors.New("Failed to open root bucket")
- }
- return b.Put([]byte(shotKey), contents)
- })
- return err
+ var contents []byte
+ var numshots int
+ // Dumps can be both x-www-urlencoded and multipart/form-data.
+ // Try multipart first, then x-www-urlencoded if no mpReader is returned
+ mpReader, _ := r.MultipartReader()
+ if mpReader != nil {
+ contents, numshots, err = extractMultipart(mpReader)
+ } else {
+ numshots = 1
+ contents, err = ioutil.ReadAll(r.Body)
+ }
+ if err != nil {
+ return err
+ }
+ fmt.Printf("Number of shots: %v", numshots)
+ err = db.Update(func(tx *bbolt.Tx) error {
+ b := tx.Bucket([]byte(rootBucket))
+ if b == nil {
+ return errors.New("Failed to open root bucket")
+ }
+ return b.Put([]byte(shotKey), contents)
+ })
+ return err
}
func extractMultipart(mr *multipart.Reader) (contents []byte, num int, err error) {
- for {
- var part *multipart.Part
- part, err = mr.NextPart()
- if err == io.EOF {
- err = nil
- break
- }
- formName := part.FormName()
- if err != nil {
- log.Fatal(err)
- }
- if formName == FormNameText || formName == FormNameFile {
- contents, err = ioutil.ReadAll(part)
- if err != nil {
- return
- }
- continue
- }
- if formName == FormNameNumShots {
- var numShotsRaw []byte
- numShotsRaw, err = ioutil.ReadAll(part)
- if err != nil { return }
- num, err = strconv.Atoi(string(numShotsRaw))
- if err != nil { return }
- }
- }
- err = nil
- return
+ for {
+ var part *multipart.Part
+ part, err = mr.NextPart()
+ if err == io.EOF {
+ err = nil
+ break
+ }
+ formName := part.FormName()
+ if err != nil {
+ log.Fatal(err)
+ }
+ if formName == FormNameText || formName == FormNameFile {
+ contents, err = ioutil.ReadAll(part)
+ if err != nil {
+ return
+ }
+ continue
+ }
+ if formName == FormNameNumShots {
+ var numShotsRaw []byte
+ numShotsRaw, err = ioutil.ReadAll(part)
+ if err != nil {
+ return
+ }
+ num, err = strconv.Atoi(string(numShotsRaw))
+ if err != nil {
+ return
+ }
+ }
+ }
+ err = nil
+ return
}
// TODO: Make function that handles responses based on mode?? like
@@ -142,132 +145,132 @@ func extractMultipart(mr *multipart.Reader) (contents []byte, num int, err error
// where textOnly and responseKey maps to response messages or templates
func RootHandler(res http.ResponseWriter, r *http.Request) {
- // Detect whether Simple mode (text only) is active
- textOnly := r.Header.Get("Simple") != "" // for forcing textOnly mode
- textOnly = textOnly || strings.Contains(r.Header.Get("User-Agent"), "curl")
- if r.Method == http.MethodGet {
- if textOnly {
- if _, err := res.Write(InfoMessage); err != nil {
- log.Panicln("Error when trying to write response body")
- }
- } else {
- templateData := struct { ShotKey string }{ "" }
- uploadPageTemplate.Execute(res, templateData)
- }
- } else if r.Method == http.MethodPost {
- // Generate random shot key
- random := make([]byte, 16)
- rand.Read(random)
- shotKey := hex.EncodeToString(random)
- // Try to pour
- if err := pour(shotKey, r); err != nil {
- log.Println(err)
- res.WriteHeader(http.StatusInternalServerError)
- }
- if /*textOnly*/ true {
- response := r.Host + "/" + shotKey
- if _, err := res.Write([]byte(response)); err != nil {
- log.Panicln("Error when trying to write response body")
- }
- }
- }
+ // Detect whether Simple mode (text only) is active
+ textOnly := r.Header.Get("Simple") != "" // for forcing textOnly mode
+ textOnly = textOnly || strings.Contains(r.Header.Get("User-Agent"), "curl")
+ if r.Method == http.MethodGet {
+ if textOnly {
+ if _, err := res.Write(InfoMessage); err != nil {
+ log.Panicln("Error when trying to write response body")
+ }
+ } else {
+ templateData := struct{ ShotKey string }{""}
+ uploadPageTemplate.Execute(res, templateData)
+ }
+ } else if r.Method == http.MethodPost {
+ // Generate random shot key
+ random := make([]byte, 16)
+ rand.Read(random)
+ shotKey := hex.EncodeToString(random)
+ // Try to pour
+ if err := pour(shotKey, r); err != nil {
+ log.Println(err)
+ res.WriteHeader(http.StatusInternalServerError)
+ }
+ if /*textOnly*/ true {
+ response := r.Host + "/" + shotKey
+ if _, err := res.Write([]byte(response)); err != nil {
+ log.Panicln("Error when trying to write response body")
+ }
+ }
+ }
}
func KeyHandler(res http.ResponseWriter, r *http.Request) {
- key := mux.Vars(r)["shotKey"]
- textOnly := r.Header.Get("Simple") != "" // for forcing textOnly mode
- textOnly = textOnly || strings.Contains(r.Header.Get("User-Agent"), "curl")
- if r.Method == http.MethodGet {
- //contents, err := shot(key)
- } else if r.Method == http.MethodPost {
- if smell(key) {
- // POSTs to taken shouldn't happen often, so use textOnly always
- if _, err := res.Write(KeyTakenMessage); err != nil {
- res.WriteHeader(http.StatusInternalServerError)
- }
- } else {
- if err := pour(key, r); err != nil {
- res.WriteHeader(http.StatusInternalServerError)
- }
- }
- }
- contents, found := pop(key)
- // GET requests only
- /*if !found {
- res.WriteHeader(http.StatusNotFound)
- if _, err := res.Write([]byte(fmt.Sprint("404 no shot here\n"))); err != nil {
- log.Panicln("Error when trying to write response")
- }
- return
- if found {
- // For POST requests to specific key, this should actually return 'link taken' or something
- if _, err := res.Write([]byte(contents)); err != nil {
- log.Panicln("Error when trying to write response body")
- }
- } else {
- pour(key, r)
- if _, err := res.Write([]byte(fmt.Sprint("Contents stored in given link\n"))); err != nil {
- log.Panicln("Error when trying to write response")
- }
- }*/
- fmt.Printf("Request from client: %v\n", r.Header.Get("User-Agent"))
- return
+ key := mux.Vars(r)["shotKey"]
+ textOnly := r.Header.Get("Simple") != "" // for forcing textOnly mode
+ textOnly = textOnly || strings.Contains(r.Header.Get("User-Agent"), "curl")
+ if r.Method == http.MethodGet {
+ //contents, err := shot(key)
+ } else if r.Method == http.MethodPost {
+ if smell(key) {
+ // POSTs to taken shouldn't happen often, so use textOnly always
+ if _, err := res.Write(KeyTakenMessage); err != nil {
+ res.WriteHeader(http.StatusInternalServerError)
+ }
+ } else {
+ if err := pour(key, r); err != nil {
+ res.WriteHeader(http.StatusInternalServerError)
+ }
+ }
+ }
+ contents, found := pop(key)
+ // GET requests only
+ /*if !found {
+ res.WriteHeader(http.StatusNotFound)
+ if _, err := res.Write([]byte(fmt.Sprint("404 no shot here\n"))); err != nil {
+ log.Panicln("Error when trying to write response")
+ }
+ return
+ if found {
+ // For POST requests to specific key, this should actually return 'link taken' or something
+ if _, err := res.Write([]byte(contents)); err != nil {
+ log.Panicln("Error when trying to write response body")
+ }
+ } else {
+ pour(key, r)
+ if _, err := res.Write([]byte(fmt.Sprint("Contents stored in given link\n"))); err != nil {
+ log.Panicln("Error when trying to write response")
+ }
+ }*/
+ fmt.Printf("Request from client: %v\n", r.Header.Get("User-Agent"))
+ return
}
// Prints all keys in the database along with size of the contents
func statDB() {
- err := db.View(func(tx *bbolt.Tx) error {
- root := tx.Bucket([]byte(rootBucket))
- if root == nil {
- return errors.New("Failed to open root bucket")
- }
- number := 0
- fmt.Println("Elements in database:")
- err := root.ForEach(func (k, v []byte) error {
- fmt.Printf("%v %v\n", string(k), len(v))
- number++
- return nil
- })
- fmt.Printf("\n%v elements \n", number)
- return err
- })
- if err != nil {
- log.Fatal(err)
- }
+ err := db.View(func(tx *bbolt.Tx) error {
+ root := tx.Bucket([]byte(rootBucket))
+ if root == nil {
+ return errors.New("Failed to open root bucket")
+ }
+ number := 0
+ fmt.Println("Elements in database:")
+ err := root.ForEach(func(k, v []byte) error {
+ fmt.Printf("%v %v\n", string(k), len(v))
+ number++
+ return nil
+ })
+ fmt.Printf("\n%v elements \n", number)
+ return err
+ })
+ if err != nil {
+ log.Fatal(err)
+ }
}
/**************** Main ****************/
-func main(){
- port := flag.Int("p", 8080, "Port")
- dbFile := flag.String("d", "vodka.db", "Database file")
- stat := flag.Bool("s", false, "View database keys and size of associated contents")
- 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)
- }
- if *stat {
- statDB()
- return
- }
+func main() {
+ port := flag.Int("p", 8080, "Port")
+ dbFile := flag.String("d", "vodka.db", "Database file")
+ stat := flag.Bool("s", false, "View database keys and size of associated contents")
+ 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)
+ }
+ if *stat {
+ statDB()
+ return
+ }
- fmt.Println("Server started listening at port", *port)
- router := mux.NewRouter()
- router.HandleFunc("/", RootHandler)
- router.HandleFunc("/{shotKey}", KeyHandler)
- http.Handle("/", router)
+ fmt.Println("Server started listening at port", *port)
+ router := mux.NewRouter()
+ router.HandleFunc("/", RootHandler)
+ router.HandleFunc("/{shotKey}", KeyHandler)
+ http.Handle("/", router)
- log.Fatal(http.ListenAndServe(fmt.Sprintf(":%v", *port), nil))
+ log.Fatal(http.ListenAndServe(fmt.Sprintf(":%v", *port), nil))
}
