package main import ( "context" "encoding/json" "errors" bolt "go.etcd.io/bbolt" "math/rand" "net/http" "os" "os/signal" "strings" "time" ) var urlDB *bolt.DB var wordList []string type urlStruc struct { date time.Time path string url string uses int } func fec(err error) { // fatal error check if err != nil { println(err.Error()) urlDB.Close() os.Exit(1) } } func openFiles() { b, err := os.ReadFile("wordlist") fec(err) if len(b) == 0 { println("wordlist is empty") os.Exit(1) } wordList = strings.Split(string(b), "\n") var db *bolt.DB db, err = bolt.Open("urls.db", 0600, nil) fec(err) fec(db.Update(func(tx *bolt.Tx) error { _, err = tx.CreateBucketIfNotExists([]byte("urls")) return err })) urlDB = db } func getWord() string { return wordList[rand.Intn(len(wordList))] } func readURL(words string) (urlStruc, error) { if words == "" { return urlStruc{}, errors.New("No words provided") } var url urlStruc var err error err = urlDB.View(func(tx *bolt.Tx) error { b := tx.Bucket([]byte("urls")) return json.Unmarshal(b.Get([]byte(words)), &url) }) return url, err } func checkURL(words string) bool { url, err := readURL(words) if err == nil && url.url != "" { return true } return false } func writeURL(url urlStruc) error { jsonUrl, err := json.Marshal(url) if err != nil { return err } return urlDB.Update(func(tx *bolt.Tx) error { b := tx.Bucket([]byte("urls")) return b.Put([]byte(url.path), jsonUrl) }) } func http404(w http.ResponseWriter) { w.WriteHeader(404) b, err := os.ReadFile("html/404.html") if err != nil { b = []byte("404 not found") w.Header().Set("Content-Type", "text/plain") } else { w.Header().Set("Content-Type", "text/html") } w.Write(b) } func http500(w http.ResponseWriter) { w.WriteHeader(500) b, err := os.ReadFile("html/500.html") if err != nil { b = []byte("500 internal server error") w.Header().Set("Content-Type", "text/plain") } else { w.Header().Set("Content-Type", "text/html") } w.Write(b) } func hec(w http.ResponseWriter, err error, code int) bool { if err == nil { return false } switch code { case 404: http404(w) break case 500: http500(w) break } return true } func intStr(n int) string { if n == 0 { return "0" } var s, sign string if n < 0 { sign = "-" n = -n } for ; n > 0; n /= 10 { s = string(n%10+int('0')) + s } return sign + s } func infoPage(w http.ResponseWriter, url urlStruc) { w.WriteHeader(200) w.Header().Set("Content-Type", "text/html") b, err := os.ReadFile("html/info.html") if hec(w, err, 500) { return } s := strings.ReplaceAll(string(b), "{{url}}", url.url) s = strings.ReplaceAll(s, "{{path}}", url.path) s = strings.ReplaceAll(s, "{{uses}}", intStr(url.uses)) s = strings.ReplaceAll(s, "{{date}}", url.date.Format("2006-01-02 15:04:05")) w.Write([]byte(s)) } func createPage(w http.ResponseWriter, api bool, code int, resp string) { w.WriteHeader(code) if api { w.Header().Set("Content-Type", "text/plain") w.Write([]byte(resp)) return } s := "Error loading HTML\nYour request was still processed.\nRequest response: " + resp + "\n" b, err := os.ReadFile("html/create.html") if err != nil { w.Header().Set("Content-Type", "text/plain") w.Write([]byte(s)) return } s = strings.ReplaceAll(string(b), "{{response}}", s) if code != 200 { s = strings.ReplaceAll(s, "{{title}}", "Your request failed:") } else { s = strings.ReplaceAll(s, "{{title}}", "Your short URL:") } w.Header().Set("Content-Type", "text/html") w.Write([]byte(s)) } func createHandler(w http.ResponseWriter, r *http.Request) { var err error var url urlStruc var api bool = false url.date = time.Now() url.uses = 0 err = r.ParseForm() if hec(w, err, 500) { return } if r.Form.Get("mode") == "api" { api = true } url.url = r.Form.Get("url") if url.url == "" { createPage(w, api, 400, "Error: No URL provided") return } url.path = r.Form.Get("custom") if url.path == "" { url.path = getWord() + "-" + getWord() + "-" for { url.path += getWord() if !checkURL(url.path) { break } url.path += "-" } } else { if checkURL(url.path) { createPage(w, api, 400, "Error: Custom URL already exists") return } } err = writeURL(url) if hec(w, err, 500) { return } createPage(w, api, 200, r.URL.Host+"/"+url.path) } func httpHandler(w http.ResponseWriter, r *http.Request) { if len(r.URL.Path) < 2 { if r.Method == "POST" { createHandler(w, r) return } http.ServeFile(w, r, "html/index.html") return } path := strings.TrimSuffix(r.URL.Path[1:], "/") var err error if path == "favicon.ico" { _, err = os.Stat("favicon.ico") if hec(w, err, 404) { return } http.ServeFile(w, r, "favicon.ico") return } var info bool = false if strings.HasSuffix(path, "/info") { info = true path = strings.TrimSuffix(path, "/info") } var url urlStruc url, err = readURL(path) if err != nil || url.url == "" { http404(w) return } if info { infoPage(w, url) return } http.Redirect(w, r, url.url, http.StatusFound) url.uses++ _ = writeURL(url) } func main() { var ( port string = ":8088" err error server http.Server ) if len(os.Args) == 2 { port = ":" + os.Args[1] } println("Starting server on port " + port) openFiles() http.HandleFunc("/", httpHandler) server = http.Server{ Addr: port, Handler: nil, } go func() { err = server.ListenAndServe() if err != http.ErrServerClosed && err != nil { urlDB.Close() panic(err) } }() quit := make(chan os.Signal, 1) signal.Notify(quit, os.Interrupt) <-quit urlDB.Close() server.Shutdown(context.Background()) }