From b376abb3869b8a3caf732e144c8f5c95977f801b Mon Sep 17 00:00:00 2001
From: brettlangdon <brett@blangdon.com>
Date: Sat, 14 Nov 2015 09:31:41 -0500
Subject: [PATCH] some improvments in the API

---
 Makefile           |  6 ++++++
 cmd/terse/terse.go |  4 ++++
 server.go          | 28 +++++++++++++++++++++++-----
 url.go             | 20 --------------------
 util.go            | 40 ++++++++++++++++++++++++++++++++++++++++
 5 files changed, 73 insertions(+), 25 deletions(-)
 delete mode 100644 url.go
 create mode 100644 util.go

diff --git a/Makefile b/Makefile
index de263dc..0cd8caf 100644
--- a/Makefile
+++ b/Makefile
@@ -1,5 +1,11 @@
 terse: ./url.go ./server.go ./cmd/terse/terse.go
 	go build ./cmd/terse
 
+run:
+	go run ./cmd/terse/terse.go
+
 clean:
 	rm -f ./terse
+
+
+.PHONY: run clean
diff --git a/cmd/terse/terse.go b/cmd/terse/terse.go
index 64e7524..de3a303 100644
--- a/cmd/terse/terse.go
+++ b/cmd/terse/terse.go
@@ -13,10 +13,14 @@ var args struct {
 }
 
 func main() {
+	// Setup default args
 	args.MaxEntries = 1000
 	args.Bind = "127.0.0.1:5892"
+
+	// Parse args from CLI
 	arg.MustParse(&args)
 
+	// Start the server
 	server := terse.NewServer(args.Bind, args.MaxEntries)
 	log.Fatal(server.ListenAndServe())
 }
diff --git a/server.go b/server.go
index ae99a2f..22529d1 100644
--- a/server.go
+++ b/server.go
@@ -27,8 +27,12 @@ func (handler *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 func (handler *Handler) HandleGet(w http.ResponseWriter, r *http.Request) {
 	url := strings.Trim(r.URL.Path, "/")
+	// If they gave us something like `/<shortcode>/<who knows>`
+	// then jst grab the first part
 	code := strings.SplitN(url, "/", 2)[0]
 	value, ok := handler.cache.Get(code)
+
+	// If the url exists in the cache and is a string, redirect to it
 	if ok {
 		switch url := value.(type) {
 		case string:
@@ -40,15 +44,29 @@ func (handler *Handler) HandleGet(w http.ResponseWriter, r *http.Request) {
 }
 
 func (handler *Handler) HandlePost(w http.ResponseWriter, r *http.Request) {
+	// Parse url from body
 	reader := bufio.NewScanner(r.Body)
 	reader.Scan()
-	url, err := url.ParseRequestURI(reader.Text())
+	rawurl := reader.Text()
+
+	// Ensure url given is a real url
+	cleanUrl, err := CleanURL(rawurl)
 	if err != nil {
-		panic(err)
+		http.Error(w, fmt.Sprintf("Invalid url \"%s\"", reader.Text()), http.StatusBadRequest)
+		return
+	}
+
+	// Generate short code and store in cache
+	code := GetShortCode([]byte(cleanUrl))
+	handler.cache.Add(code, cleanUrl)
+
+	// Generate response url
+	codeUrl := &url.URL{
+		Scheme: "https",
+		Host:   r.Host,
+		Path:   "/" + code,
 	}
-	code := GetShortCode([]byte(url.String()))
-	handler.cache.Add(code, url.String())
-	fmt.Fprintf(w, code)
+	fmt.Fprintf(w, codeUrl.String())
 }
 
 func NewServer(bind string, maxEntries int) *http.Server {
diff --git a/url.go b/url.go
deleted file mode 100644
index 16a7b2b..0000000
--- a/url.go
+++ /dev/null
@@ -1,20 +0,0 @@
-package terse
-
-import (
-	"bytes"
-	"hash/crc32"
-)
-
-const ALPHABET string = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-._~"
-const BASE uint32 = uint32(len(ALPHABET))
-
-func GetShortCode(url []byte) string {
-	var code bytes.Buffer
-	num := crc32.ChecksumIEEE(url)
-	for num > 0 {
-		next := (num % BASE)
-		code.WriteRune(rune(ALPHABET[next]))
-		num = num / 62
-	}
-	return code.String()
-}
diff --git a/util.go b/util.go
new file mode 100644
index 0000000..2951cd8
--- /dev/null
+++ b/util.go
@@ -0,0 +1,40 @@
+package terse
+
+import (
+	"bytes"
+	"fmt"
+	"hash/crc32"
+	"net/url"
+)
+
+// URL safe characters, as definied by section 2.3 of RFC 3986 (http://www.ietf.org/rfc/rfc3986.txt)
+const ALPHABET string = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-._~"
+const BASE uint32 = uint32(len(ALPHABET))
+
+func GetShortCode(url []byte) string {
+	var code bytes.Buffer
+	num := crc32.ChecksumIEEE(url)
+	for num > 0 {
+		next := (num % BASE)
+		code.WriteRune(rune(ALPHABET[next]))
+		num = num / 62
+	}
+	return code.String()
+}
+
+func CleanURL(rawurl string) (string, error) {
+	parsed, err := url.ParseRequestURI(rawurl)
+	if err != nil {
+		return "", fmt.Errorf("Error parsing url \"%s\": %s", rawurl, err)
+	}
+
+	if parsed.Scheme != "http" && parsed.Scheme != "https" {
+		return "", fmt.Errorf("Scheme \"%s\" not allowed for url \"%s\"", parsed.Scheme, rawurl)
+	}
+
+	if parsed.Host == "" {
+		return "", fmt.Errorf("No hostname provided for url \"%s\"", rawurl)
+	}
+
+	return parsed.String(), nil
+}