go-paste/src/routes/shorturl.go

package routes

import (
	"net/http"
	"net/url"
	"regexp"

	"github.com/gorilla/mux"
	"github.com/timshannon/badgerhold"

	"sour.is/x/toolbox/httpsrv"
	"sour.is/x/toolbox/log"
	"sour.is/x/toolbox/uuid"
)

type logger struct{}

func (logger) Errorf(s string, o ...interface{})   { log.Errorf(s, o...) }
func (logger) Warningf(s string, o ...interface{}) { log.Warningf(s, o...) }
func (logger) Infof(s string, o ...interface{})    { log.Infof(s, o...) }
func (logger) Debugf(s string, o ...interface{})   { log.Debugf(s, o...) }

func init() {
	s := &shortDB{}
	httpsrv.RegisterModule("short", s.config)

	httpsrv.HttpRegister("short", httpsrv.HttpRoutes{
		{Name: "get", Method: "GET", Pattern: "/s/{id}", HandlerFunc: s.get},
		{Name: "put", Method: "PUT", Pattern: "/s/{id}", HandlerFunc: s.put},
	})
}

type shortDB struct {
	options badgerhold.Options
}

func (s *shortDB) config(config map[string]string) {
	s.options = badgerhold.DefaultOptions
	s.options.Options = s.options.WithLogger(logger{})

	var ok bool
	if s.options.Dir, ok = config["index"]; !ok {
		s.options.Dir = "data"
	}

	if s.options.ValueDir, ok = config["value"]; !ok {
		s.options.ValueDir = "data"
	}

	store, err := badgerhold.Open(s.options)
	defer func() {
		if err = store.Close(); err != nil {
			log.Fatal(err)
		}
	}()
	if err != nil {
		log.Fatal(err)
	}
}
func (s *shortDB) get(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	id := vars["id"]

	store, err := badgerhold.Open(s.options)
	defer func() {
		if err = store.Close(); err != nil {
			log.Fatal(err)
		}
	}()
	if err != nil {
		httpsrv.WriteText(w, 500, err.Error())
		return
	}

	var url shortURL
	err = store.Get(id, &url)
	if err != nil && err != badgerhold.ErrNotFound {
		httpsrv.WriteText(w, 500, err.Error())
		return
	}
	if err == badgerhold.ErrNotFound {
		httpsrv.WriteText(w, 404, "not found")
		return
	}

	w.Header().Set("Location", url.URL)
	w.WriteHeader(http.StatusFound)
}
func (s *shortDB) put(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	id := vars["id"]
	secret := r.FormValue("secret")
	u, err := url.Parse(r.FormValue("url"))
	if err != nil {
		httpsrv.WriteText(w, 400, "bad url")
		return
	}

	store, err := badgerhold.Open(s.options)
	defer func() {
		if err = store.Close(); err != nil {
			log.Fatal(err)
		}
	}()
	if err != nil {
		httpsrv.WriteText(w, 500, err.Error())
		return
	}

	url := &shortURL{}
	err = store.Get(id, url)
	if err != nil && err != badgerhold.ErrNotFound {
		httpsrv.WriteText(w, 500, err.Error())
		return
	}

	if err == badgerhold.ErrNotFound {
		url = newShortURL(id, secret, u.String())

		err := store.Insert(url.ID, url)
		if err != nil {
			httpsrv.WriteText(w, 500, err.Error())
			return
		}

		httpsrv.WriteObject(w, 200, url)
		return
	}

	if secret == "" {
		httpsrv.WriteError(w, 401, "no auth")
		return
	}

	if secret != url.Secret {
		httpsrv.WriteError(w, 403, "forbidden")
		return
	}

	err = store.Upsert(url.ID, url)
	if err != nil {
		httpsrv.WriteText(w, 500, err.Error())
		return
	}

	httpsrv.WriteObject(w, 200, url)
}

type shortURL struct {
	ID     string `badgerhold:"unique"`
	URL    string
	Secret string
}

func newShortURL(id, secret, u string) *shortURL {
	m, err := regexp.MatchString("[a-z-]{1,64}", id)
	if id == "" || !m || err != nil {
		id = uuid.V4()
	}
	m, err = regexp.MatchString("[a-z-]{1,64}", secret)
	if secret == "" || !m || err != nil {
		secret = uuid.V4()
	}
	return &shortURL{ID: id, Secret: secret, URL: u}
}