From d1387558068f00ac2c82cd66757a69fd6818d8c9 Mon Sep 17 00:00:00 2001
From: mutantmonkey <mutantmonkey@mutantmonkey.in>
Date: Tue, 13 Oct 2015 19:52:55 -0700
Subject: [PATCH] do a proper same-origin check

String prefix matching is hacky and provides insufficient checking if it
does not end with a /.
---
 csrf.go      | 27 ++++++++++++++++++---------
 fileserve.go |  6 ++++--
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/csrf.go b/csrf.go
index b70215b..9b8b93e 100644
--- a/csrf.go
+++ b/csrf.go
@@ -2,14 +2,19 @@ package main
 
 import (
 	"net/http"
-	"strings"
+	"net/url"
 )
 
+// Do a strict referrer check, matching against both the Origin header (if
+// present) and the Referrer header. If a list of headers is specified, then
+// Referrer checking will be skipped if any of those headers are present.
 func strictReferrerCheck(r *http.Request, prefix string, whitelistHeaders []string) bool {
-	p := strings.TrimSuffix(prefix, "/")
+	p, _ := url.Parse(prefix)
+
+	// if there's an Origin header, check it and skip other checks
 	if origin := r.Header.Get("Origin"); origin != "" {
-		// if there's an Origin header, check it and ignore the rest
-		return strings.HasPrefix(origin, p)
+		u, _ := url.Parse(origin)
+		return sameOrigin(u, p)
 	}
 
 	for _, header := range whitelistHeaders {
@@ -18,9 +23,13 @@ func strictReferrerCheck(r *http.Request, prefix string, whitelistHeaders []stri
 		}
 	}
 
-	if referrer := r.Header.Get("Referer"); !strings.HasPrefix(referrer, p) {
-		return false
-	}
-
-	return true
+	referrer := r.Header.Get("Referer")
+	u, _ := url.Parse(referrer)
+	return sameOrigin(u, p)
+}
+
+// Check if two URLs have the same origin
+func sameOrigin(u1, u2 *url.URL) bool {
+	// host also contains the port if one was specified
+	return (u1.Scheme == u2.Scheme && u1.Host == u2.Host)
 }
diff --git a/fileserve.go b/fileserve.go
index cc682ca..9b0d62f 100644
--- a/fileserve.go
+++ b/fileserve.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"io"
 	"net/http"
+	"net/url"
 	"os"
 	"path"
 	"strings"
@@ -26,8 +27,9 @@ func fileServeHandler(c web.C, w http.ResponseWriter, r *http.Request) {
 
 	if !Config.allowHotlink {
 		referer := r.Header.Get("Referer")
-		prefix := strings.TrimSuffix(Config.siteURL, "/")
-		if referer != "" && !strings.HasPrefix(referer, prefix) {
+		u, _ := url.Parse(referer)
+		p, _ := url.Parse(Config.siteURL)
+		if referer != "" && !sameOrigin(u, p) {
 			w.WriteHeader(403)
 			return
 		}