medium / chi rate limiter

chi_base.go

func main() {
	router := chi.NewRouter()

	router.Get("/info", infoHandler)
	router.Group(func(router chi.Router) {
		router.Use(authMiddleware)
		router.Get("/profile", profileHandler)
		router.Get("/stat", statHandler)
	})

	http.ListenAndServe(":8080", router)
}

func infoHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("info."))
}

func profileHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("profile."))
}

func statHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("stat."))
}

func authMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		userID := r.Header.Get("UserID")
		if userID == "" {
			w.WriteHeader(http.StatusUnauthorized)
			w.Write([]byte("unauthorized"))
			return
		}
		ctx := context.WithValue(r.Context(), "userID", userID)
		r = r.WithContext(ctx)
		next.ServeHTTP(w, r)
	})
}