Updated CLI to use flags

README.md

@@ -20,10 +20,19 @@ export POSTGRES=postgresql://root@localhost:26257/short?sslmode=disable
 ./short
 ```
 
-# ENV
+# Usage
 
-BASEURL: Base URL used in returning the short address.
+```
+$ short -h
+Usage of ./short:
+  -baseurl string
+    	baseurl URL of short links (default "localhost:port")
+  -num int
+    	number of characters in shortened url (default 4)
+  -port int
+    	listen on port (default 8080)
+  -postgres string
+    	postgres string (default "postgresql://root@localhost:26257/short?sslmode=disable")
+```
 
-PORT: Port to listen on
-
-POSTGRES: [Postgres string](https://godoc.org/github.com/lib/pq#hdr-Connection_String_Parameters)
+See [pq docs](https://godoc.org/github.com/lib/pq#hdr-Connection_String_Parameters) for more information on the postgres string. 

generate.go

@@ -13,15 +13,14 @@ var chars = []byte{
 }
 
 // Increase below in the future
-var length = 4
 var seeded = false
 
-func GenerateID() string {
+func GenerateID(n int) string {
 	if !seeded {
 		rand.Seed(time.Now().UnixNano())
 	}
-	id := make([]byte, length)
-	for i := 0; i < length; i++ {
+	id := make([]byte, n)
+	for i := 0; i < n; i++ {
 		id[i] = chars[rand.Intn(len(chars))]
 	}
 	return string(id)

handlers.go

@@ -1,7 +1,9 @@
 package main
 
 import (
+	"database/sql"
 	"fmt"
+	"log"
 	"net/http"
 	"net/url"
 
@@ -18,8 +20,9 @@ func (h *Handlers) Create(w http.ResponseWriter, r *http.Request, _ httprouter.P
 	u := r.FormValue("url")
 
 	// Check that it is a URL
-	_, err := url.ParseRequestURI(u)
-	if err != nil {
+	parsed, err := url.Parse(u)
+	if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") {
+		log.Print("handlers: invalid url: " + u)
 		http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
 		return
 	}
@@ -27,6 +30,7 @@ func (h *Handlers) Create(w http.ResponseWriter, r *http.Request, _ httprouter.P
 	// Create a short URL
 	id, err := h.store.Create(u)
 	if err != nil {
+		log.Print("handlers: " + err.Error())
 		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
 		return
 	}
@@ -39,6 +43,9 @@ func (h *Handlers) Get(w http.ResponseWriter, r *http.Request, p httprouter.Para
 	// Get the short URL
 	url, err := h.store.Get(p.ByName("id"))
 	if err != nil {
+		if err != sql.ErrNoRows {
+			log.Print("handlers: " + err.Error())
+		}
 		http.Error(w, "Not found, or an error occurred, idk.", http.StatusNotFound)
 		return
 	}

main.go

@@ -2,41 +2,47 @@ package main
 
 import (
 	"database/sql"
+	"flag"
+	"fmt"
 	"log"
 	"net/http"
-	"os"
 
 	_ "github.com/lib/pq"
 )
 
+var num int
+var port int
+var baseUrl string
+var postgres string
+
+func init() {
+	flag.IntVar(&num, "num", 4, "number of characters in shortened url")
+	flag.IntVar(&port, "port", 8080, "listen on port")
+	flag.StringVar(&baseUrl, "baseurl", "localhost:port", "baseurl URL of short links")
+	flag.StringVar(&postgres, "postgres", "postgresql://root@localhost:26257/short?sslmode=disable", "postgres string")
+}
+
 func main() {
-	listenPort := os.Getenv("PORT")
-	if len(listenPort) < 1 {
-		listenPort = "8080"
-	}
-	baseUrl := os.Getenv("BASEURL")
-	if len(baseUrl) < 4 {
-		baseUrl = "http://localhost:" + listenPort
-	}
-	postgresString := os.Getenv("POSTGRES")
-	if len(postgresString) < 1 {
-		postgresString = "postgresql://root@localhost:26257/short?sslmode=disable"
+	// Parse commandline flags
+	flag.Parse()
+	if baseUrl == "localhost:port" {
+		baseUrl = fmt.Sprintf("localhost:%d", port)
 	}
 
 	// Open database connection
-	db, err := sql.Open("postgres", postgresString)
+	db, err := sql.Open("postgres", postgres)
 	if err != nil {
 		log.Fatal(err)
 	}
 
 	// Create storage abstraction
-	store := NewStore(db)
+	store := NewStore(db, num)
 	// Setup handlers
 	handlers := NewHandlers(store, baseUrl)
 	// Create router
 	router := NewRouter(handlers)
 
 	// Listen
-	log.Println("Listening on port " + listenPort)
-	log.Fatal(http.ListenAndServe(":"+listenPort, router))
+	log.Println("main: Listening on port", port)
+	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", port), router))
 }

router.go

@@ -6,8 +6,8 @@ import (
 
 func NewRouter(h *Handlers) *httprouter.Router {
 	router := httprouter.New()
-	router.GET("/", h.Index)
 	router.POST("/new", h.Create)
+	router.GET("/", h.Index)
 	router.GET("/:id", h.Get)
 	return router
 }

store.go

@@ -6,17 +6,18 @@ import (
 )
 
 type Store struct {
-	db *sql.DB
+	db  *sql.DB
+	num int
 }
 
 func (s *Store) Create(url string) (id string, err error) {
-	id = GenerateID()
+	id = GenerateID(s.num)
 	var existing string
 	err = s.db.QueryRow(`
 		SELECT url FROM links WHERE id = $1
 	`, id).Scan(&existing)
 	if err == nil {
-		log.Print("Collision occurred on " + id + ", regenerating...")
+		log.Print("store: collision occurred on " + id + ", regenerating...")
 		return s.Create(url)
 	}
 	if err != nil && err != sql.ErrNoRows {
@@ -40,7 +41,7 @@ func (s *Store) Get(id string) (url string, err error) {
 	return url, err
 }
 
-func NewStore(db *sql.DB) *Store {
+func NewStore(db *sql.DB, num int) *Store {
 	// Check for table and initialise if necessary
 	_, err := db.Exec(`
 		CREATE TABLE IF NOT EXISTS links (
@@ -52,7 +53,5 @@ func NewStore(db *sql.DB) *Store {
 		log.Fatal(err)
 	}
 
-	return &Store{
-		db: db,
-	}
+	return &Store{db, num}
 }