middleware/wrap_writer.go

@@ -111,11 +111,15 @@ func (f *httpFancyWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 }
 func (f *httpFancyWriter) ReadFrom(r io.Reader) (int64, error) {
 	if f.basicWriter.tee != nil {
-		return io.Copy(&f.basicWriter, r)
+		n, err := io.Copy(&f.basicWriter, r)
+		f.basicWriter.bytes += int(n)
+		return n, err
 	}
 	rf := f.basicWriter.ResponseWriter.(io.ReaderFrom)
 	f.basicWriter.maybeWriteHeader()
-	return rf.ReadFrom(r)
+	n, err := rf.ReadFrom(r)
+	f.basicWriter.bytes += int(n)
+	return n, err
 }
 
 var _ http.CloseNotifier = &httpFancyWriter{}

tree.go

@@ -760,3 +760,47 @@ type Route struct {
 	Handlers  map[string]http.Handler
 	SubRoutes Routes
 }
+
+// WalkFunc is the type of the function called for each method and route visited by Walk.
+type WalkFunc func(method string, route string, handler http.Handler, middlewares ...func(http.Handler) http.Handler) error
+
+// Walk walks any router tree that implements Routes interface.
+func Walk(r Routes, walkFn WalkFunc) error {
+	return walk(r, walkFn, "")
+}
+
+func walk(r Routes, walkFn WalkFunc, parentRoute string, parentMw ...func(http.Handler) http.Handler) error {
+	for _, route := range r.Routes() {
+		mws := make([]func(http.Handler) http.Handler, len(parentMw))
+		copy(mws, parentMw)
+		mws = append(mws, r.Middlewares()...)
+
+		if route.SubRoutes != nil {
+			if err := walk(route.SubRoutes, walkFn, parentRoute+route.Pattern, mws...); err != nil {
+				return err
+			}
+			continue
+		}
+
+		for method, handler := range route.Handlers {
+			if method == "*" {
+				// Ignore a "catchAll" method, since we pass down all the specific methods for each route.
+				continue
+			}
+
+			fullRoute := parentRoute + route.Pattern
+
+			if chain, ok := handler.(*ChainHandler); ok {
+				if err := walkFn(method, fullRoute, chain.Endpoint, append(mws, chain.Middlewares...)...); err != nil {
+					return err
+				}
+			} else {
+				if err := walkFn(method, fullRoute, handler, mws...); err != nil {
+					return err
+				}
+			}
+		}
+	}
+
+	return nil
+}

tree_test.go

@@ -424,3 +424,16 @@ func BenchmarkTreeGet(b *testing.B) {
 		tr.FindRoute(mctx, mGET, "/ping/123/456")
 	}
 }
+
+func TestWalker(t *testing.T) {
+	r := bigMux()
+
+	// Walk the muxBig router tree.
+	if err := Walk(r, func(method string, route string, handler http.Handler, middlewares ...func(http.Handler) http.Handler) error {
+		t.Logf("%v %v", method, route)
+
+		return nil
+	}); err != nil {
+		t.Error(err)
+	}
+}