From 6d0cd4e2921bcced9134c8cf09582b3631ab1ba2 Mon Sep 17 00:00:00 2001
From: JT Olds <hello@jtolds.com>
Date: Wed, 11 Jan 2017 19:25:46 -0700
Subject: [PATCH] fix up some redirect helpers

---
 whfatal/fatal.go    |  6 ++---
 whlog/logging.go    |  6 ++---
 whmon/writer.go     |  6 ++---
 whredir/redirect.go | 53 ++++++++++++++++++++++++++++++++++++++-------
 4 files changed, 54 insertions(+), 17 deletions(-)

diff --git a/whfatal/fatal.go b/whfatal/fatal.go
index 0b00b92..48de00f 100644
--- a/whfatal/fatal.go
+++ b/whfatal/fatal.go
@@ -22,9 +22,9 @@ type fatalBehavior func(w http.ResponseWriter, r *http.Request)
 // handler, wherr.HandleWith handlers, and a few other handlers. Otherwise,
 // the wrapper will be one of the things interrupted by Fatal calls.
 func Catch(h http.Handler) http.Handler {
-	return whroute.HandlerFunc(h,
+	return whmon.MonitorResponse(whroute.HandlerFunc(h,
 		func(w http.ResponseWriter, r *http.Request) {
-			rw := whmon.WrapResponseWriter(w)
+			rw := w.(whmon.ResponseWriter)
 			defer func() {
 				rec := recover()
 				if rec == nil {
@@ -48,7 +48,7 @@ func Catch(h http.Handler) http.Handler {
 				}
 			}()
 			h.ServeHTTP(rw, r)
-		})
+		}))
 }
 
 // Redirect is like whredir.Redirect but panics so all additional request
diff --git a/whlog/logging.go b/whlog/logging.go
index da0d2e1..5cb8f08 100644
--- a/whlog/logging.go
+++ b/whlog/logging.go
@@ -23,10 +23,10 @@ var (
 // whmon's ResponseWriter to keep track of activity. whfatal.Catch should be
 // placed *inside* if applicable. whlog.Default makes a good default logger.
 func LogResponses(logger Loggerf, h http.Handler) http.Handler {
-	return whroute.HandlerFunc(h,
+	return whmon.MonitorResponse(whroute.HandlerFunc(h,
 		func(w http.ResponseWriter, r *http.Request) {
 			method, requestURI := r.Method, r.RequestURI
-			rw := whmon.WrapResponseWriter(w)
+			rw := w.(whmon.ResponseWriter)
 			start := time.Now()
 
 			defer func() {
@@ -46,7 +46,7 @@ func LogResponses(logger Loggerf, h http.Handler) http.Handler {
 
 			logger(`%s %#v %d %d %d %v`, method, requestURI, code,
 				r.ContentLength, rw.Written(), time.Since(start))
-		})
+		}))
 }
 
 // LogRequests takes a Handler and makes it log requests (prior to request
diff --git a/whmon/writer.go b/whmon/writer.go
index 6e114d4..fe48c11 100644
--- a/whmon/writer.go
+++ b/whmon/writer.go
@@ -74,15 +74,15 @@ type ResponseWriter interface {
 func MonitorResponse(h http.Handler) http.Handler {
 	return whroute.HandlerFunc(h,
 		func(w http.ResponseWriter, r *http.Request) {
-			h.ServeHTTP(WrapResponseWriter(w), r)
+			h.ServeHTTP(wrapResponseWriter(w), r)
 		})
 }
 
-// WrapResponseWriter will wrap an http.ResponseWriter with the instrumentation
+// wrapResponseWriter will wrap an http.ResponseWriter with the instrumentation
 // to turn it into a whmon.ResponseWriter. An http.ResponseWriter must be
 // turned into a whmon.ResponseWriter before being used. It's much better
 // to use MonitorResponse instead of WrapResponseWriter.
-func WrapResponseWriter(w http.ResponseWriter) ResponseWriter {
+func wrapResponseWriter(w http.ResponseWriter) ResponseWriter {
 	if ww, ok := w.(ResponseWriter); ok {
 		// don't do it if we already have the methods we need
 		return ww
diff --git a/whredir/redirect.go b/whredir/redirect.go
index c4294f8..4186923 100644
--- a/whredir/redirect.go
+++ b/whredir/redirect.go
@@ -7,7 +7,10 @@ package whredir // import "gopkg.in/webhelp.v1/whredir"
 
 import (
 	"net/http"
+	"net/url"
+	"strings"
 
+	"gopkg.in/webhelp.v1/wherr"
 	"gopkg.in/webhelp.v1/whmux"
 	"gopkg.in/webhelp.v1/whroute"
 )
@@ -64,13 +67,17 @@ var _ whroute.Lister = RedirectHandlerFunc(nil)
 func RequireHTTPS(handler http.Handler) http.Handler {
 	return whroute.HandlerFunc(handler,
 		func(w http.ResponseWriter, r *http.Request) {
-			if r.URL.Scheme != "https" {
-				u := *r.URL
-				u.Scheme = "https"
-				Redirect(w, r, u.String())
-			} else {
+			if r.URL.Scheme == "https" {
 				handler.ServeHTTP(w, r)
+				return
 			}
+			u, err := url.ParseRequestURI(r.RequestURI)
+			if err != nil {
+				wherr.Handle(w, r, err)
+				return
+			}
+			u.Scheme = "https"
+			Redirect(w, r, u.String())
 		})
 }
 
@@ -82,9 +89,39 @@ func RequireHost(host string, handler http.Handler) http.Handler {
 	}
 	return whmux.Host{
 		host: handler,
-		"*": RedirectHandlerFunc(func(r *http.Request) string {
-			u := *r.URL
+		"*": http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			u, err := url.ParseRequestURI(r.RequestURI)
+			if err != nil {
+				wherr.Handle(w, r, err)
+				return
+			}
 			u.Host = host
-			return u.String()
+			Redirect(w, r, u.String())
 		})}
 }
+
+// RequireTrailingSlash makes sure all handled paths have a trailing slash.
+// This helps with relative URLs for other resources.
+func RequireTrailingSlash(h http.Handler) http.Handler {
+	return whroute.HandlerFunc(h,
+		func(w http.ResponseWriter, r *http.Request) {
+			if strings.HasSuffix(r.URL.Path, "/") {
+				h.ServeHTTP(w, r)
+				return
+			}
+
+			u, err := url.ParseRequestURI(r.RequestURI)
+			if err != nil {
+				wherr.Handle(w, r, err)
+				return
+			}
+
+			if strings.HasSuffix(u.Path, "/") {
+				h.ServeHTTP(w, r)
+				return
+			}
+
+			u.Path += "/"
+			Redirect(w, r, u.String())
+		})
+}