middleware: handle Host routing via Request.Host

middleware/route_headers.go

@@ -1,24 +1,32 @@
 package middleware
 
 import (
+	"net"
 	"net/http"
 	"strings"
 )
 
 // RouteHeaders is a neat little header-based router that allows you to direct
 // the flow of a request through a middleware stack based on a request header.
 //
-// For example, lets say you'd like to setup multiple routers depending on the
-// request Host header, you could then do something as so:
+// Note that the HTTP Host header is promoted to Request.Host by net/http and
+// removed from Request.Header. RouteHeaders handles "Host" specially by reading
+// Request.Host, but for larger host-routing setups the dedicated
+// github.com/go-chi/hostrouter package is often a better fit.
 //
-//	r := chi.NewRouter()
+// For example, let's say you'd like to set up multiple routers depending on the
+// request Host header:
+//
+//	root := chi.NewRouter()
 //	rSubdomain := chi.NewRouter()
-//	r.Use(middleware.RouteHeaders().
-//		Route("Host", "example.com", middleware.New(r)).
+//	hostRouter := chi.NewRouter()
+//	hostRouter.Use(middleware.RouteHeaders().
+//		Route("Host", "example.com", middleware.New(root)).
 //		Route("Host", "*.example.com", middleware.New(rSubdomain)).
 //		Handler)
-//	r.Get("/", h)
+//	root.Get("/", h)
 //	rSubdomain.Get("/", h2)
+//	http.ListenAndServe(":8080", hostRouter)
 //
 // Another example, imagine you want to setup multiple CORS handlers, where for
 // your origin servers you allow authorized requests, but for third-party public
@@ -84,15 +92,23 @@ func (hr HeaderRouter) Handler(next http.Handler) http.Handler {
 
 		// find first matching header route, and continue
 		for header, matchers := range hr {
-			headerValue := r.Header.Get(header)
+			headerValue := getHeaderValue(r, header)
 			if headerValue == "" {
 				continue
 			}
-			headerValue = strings.ToLower(headerValue)
-			for _, matcher := range matchers {
-				if matcher.IsMatch(headerValue) {
-					matcher.Middleware(next).ServeHTTP(w, r)
-					return
+
+			if matcher, ok := matchHeaderRoute(matchers, headerValue); ok {
+				matcher.Middleware(next).ServeHTTP(w, r)
+				return
+			}
+
+			if header == "host" {
+				host, _, err := net.SplitHostPort(headerValue)
+				if err == nil {
+					if matcher, ok := matchHeaderRoute(matchers, host); ok {
+						matcher.Middleware(next).ServeHTTP(w, r)
+						return
+					}
 				}
 			}
 		}
@@ -107,6 +123,23 @@ func (hr HeaderRouter) Handler(next http.Handler) http.Handler {
 	})
 }
 
+func getHeaderValue(r *http.Request, header string) string {
+	if header == "host" {
+		return r.Host
+	}
+	return r.Header.Get(header)
+}
+
+func matchHeaderRoute(matchers []HeaderRoute, value string) (HeaderRoute, bool) {
+	value = strings.ToLower(value)
+	for _, matcher := range matchers {
+		if matcher.IsMatch(value) {
+			return matcher, true
+		}
+	}
+	return HeaderRoute{}, false
+}
+
 type HeaderRoute struct {
 	Middleware func(next http.Handler) http.Handler
 	MatchOne   Pattern

middleware/route_headers_test.go

@@ -51,7 +51,6 @@ func TestRouteHeaders(t *testing.T) {
 
 		req := httptest.NewRequest("GET", "/", nil)
 		req.Host = "example.com"
-		req.Header.Set("Host", "example.com")
 		rec := httptest.NewRecorder()
 
 		handler.ServeHTTP(rec, req)
@@ -77,7 +76,7 @@ func TestRouteHeaders(t *testing.T) {
 		}))
 
 		req := httptest.NewRequest("GET", "/", nil)
-		req.Header.Set("Host", "api.example.com")
+		req.Host = "api.example.com"
 		rec := httptest.NewRecorder()
 
 		handler.ServeHTTP(rec, req)
@@ -108,7 +107,7 @@ func TestRouteHeaders(t *testing.T) {
 		}))
 
 		req := httptest.NewRequest("GET", "/", nil)
-		req.Header.Set("Host", "other.com")
+		req.Host = "other.com"
 		rec := httptest.NewRecorder()
 
 		handler.ServeHTTP(rec, req)
@@ -173,7 +172,7 @@ func TestRouteHeaders(t *testing.T) {
 		}))
 
 		req := httptest.NewRequest("GET", "/", nil)
-		req.Header.Set("Host", "other.com")
+		req.Host = "other.com"
 		rec := httptest.NewRecorder()
 
 		handler.ServeHTTP(rec, req)
@@ -182,6 +181,58 @@ func TestRouteHeaders(t *testing.T) {
 			t.Error("expected next handler to be called when no match and no default")
 		}
 	})
+
+	t.Run("host header with port should match exact hostname pattern", func(t *testing.T) {
+		var matched bool
+
+		hr := RouteHeaders().
+			Route("Host", "example.com", func(next http.Handler) http.Handler {
+				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					matched = true
+					next.ServeHTTP(w, r)
+				})
+			})
+
+		handler := hr.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		}))
+
+		req := httptest.NewRequest("GET", "http://example.com:8080/", nil)
+		req.Host = "example.com:8080"
+		rec := httptest.NewRecorder()
+
+		handler.ServeHTTP(rec, req)
+
+		if !matched {
+			t.Error("expected hostname match to ignore the port when needed")
+		}
+	})
+
+	t.Run("host wildcard with port should match", func(t *testing.T) {
+		var matched bool
+
+		hr := RouteHeaders().
+			Route("Host", "*.example.com", func(next http.Handler) http.Handler {
+				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					matched = true
+					next.ServeHTTP(w, r)
+				})
+			})
+
+		handler := hr.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		}))
+
+		req := httptest.NewRequest("GET", "http://api.example.com:8080/", nil)
+		req.Host = "api.example.com:8080"
+		rec := httptest.NewRecorder()
+
+		handler.ServeHTTP(rec, req)
+
+		if !matched {
+			t.Error("expected wildcard host match to ignore the port when needed")
+		}
+	})
 }
 
 func TestPattern(t *testing.T) {