refactor: apply traffic control on all cacheable requests by default

Author: RoyXiang

handler/dynamic.go

@@ -7,6 +7,7 @@ import (
 
 func DynamicMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		handler := TrafficMiddleware(next)
 		for dynamicCacheClient != nil {
 			if rangeInHeader := r.Header.Get(headerRange); rangeInHeader != "" {
 				break
@@ -15,11 +16,10 @@ func DynamicMiddleware(next http.Handler) http.Handler {
 				ctx := context.WithValue(context.Background(), cacheClientCtxKey, dynamicCacheClient)
 				r = r.Clone(ctx)
 				r.URL.Query().Set(headerToken, token)
-				CacheMiddleware(next).ServeHTTP(w, r)
-				return
+				handler = CacheMiddleware(handler)
 			}
 			break
 		}
-		next.ServeHTTP(w, r)
+		handler.ServeHTTP(w, r)
 	})
 }

handler/static.go

@@ -7,12 +7,12 @@ import (
 
 func StaticMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		handler := TrafficMiddleware(next)
 		if staticCacheClient != nil {
 			ctx := context.WithValue(context.Background(), cacheClientCtxKey, staticCacheClient)
 			r = r.Clone(ctx)
-			CacheMiddleware(next).ServeHTTP(w, r)
-			return
+			handler = CacheMiddleware(handler)
 		}
-		next.ServeHTTP(w, r)
+		handler.ServeHTTP(w, r)
 	})
 }

handler/user.go

@@ -7,15 +7,15 @@ import (
 
 func UserMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		handler := TrafficMiddleware(next)
 		if userCacheClient != nil {
 			if token := r.Header.Get(headerToken); token != "" {
 				ctx := context.WithValue(context.Background(), cacheClientCtxKey, userCacheClient)
 				r = r.Clone(ctx)
 				r.URL.Query().Set(headerToken, token)
-				CacheMiddleware(next).ServeHTTP(w, r)
-				return
+				handler = CacheMiddleware(handler)
 			}
 		}
-		next.ServeHTTP(w, r)
+		handler.ServeHTTP(w, r)
 	})
 }

main.go

@@ -25,21 +25,18 @@ func newRouter() http.Handler {
 	refreshRouter.Path("/library/sections/{id}/refresh").HandlerFunc(handler.Handler)
 
 	staticRouter := r.Methods(http.MethodGet).Subrouter()
-	staticRouter.Use(handler.TrafficMiddleware)
 	staticRouter.Use(handler.StaticMiddleware)
 	staticRouter.Path("/library/metadata/{key}/art/{id}").HandlerFunc(handler.Handler)
 	staticRouter.Path("/library/metadata/{key}/thumb/{id}").HandlerFunc(handler.Handler)
 	staticRouter.Path("/photo/:/transcode").HandlerFunc(handler.Handler)
 
 	userRouter := r.Methods(http.MethodGet).Subrouter()
-	userRouter.Use(handler.TrafficMiddleware)
 	userRouter.Use(handler.UserMiddleware)
 	userRouter.PathPrefix("/library/collections/").HandlerFunc(handler.Handler)
 	userRouter.PathPrefix("/library/metadata/").HandlerFunc(handler.Handler)
 	userRouter.PathPrefix("/library/sections/").HandlerFunc(handler.Handler)
 
 	dynamicRouter := r.Methods(http.MethodGet).Subrouter()
-	dynamicRouter.Use(handler.TrafficMiddleware)
 	dynamicRouter.Use(handler.DynamicMiddleware)
 	dynamicRouter.PathPrefix("/").HandlerFunc(handler.Handler)