From 3bc6cc1e63afbb08bf78f9dc8785d67dffd1be60 Mon Sep 17 00:00:00 2001
From: Victor Zverovich <viz@meta.com>
Date: Sat, 2 Mar 2024 09:13:29 -0800
Subject: [PATCH] Protect against locking formatters

---
 include/fmt/base.h       | 27 ++++++++++++++++++++++----
 include/fmt/format-inl.h | 10 ++++++----
 test/format-test.cc      | 41 +++++++++++++++++++++++++++++++++++++++-
 3 files changed, 69 insertions(+), 9 deletions(-)

diff --git a/include/fmt/base.h b/include/fmt/base.h
index c2ec12c7415d..9bbf88fdfd5a 100644
--- a/include/fmt/base.h
+++ b/include/fmt/base.h
@@ -1162,6 +1162,20 @@ using appender = basic_appender<char>;
 
 namespace detail {
 
+template <typename T, typename Enable = void>
+struct locking : std::true_type {};
+template <typename T>
+struct locking<T, void_t<typename formatter<remove_cvref_t<T>>::nonlocking>>
+    : std::false_type {};
+
+template <typename T = int> FMT_CONSTEXPR inline auto is_locking() -> bool {
+  return locking<T>::value;
+}
+template <typename T1, typename T2, typename... Tail>
+FMT_CONSTEXPR inline auto is_locking() -> bool {
+  return locking<T1>::value || is_locking<T2, Tail...>();
+}
+
 // An optimized version of std::copy with the output value type (T).
 template <typename T, typename InputIt>
 auto copy(InputIt begin, InputIt end, appender out) -> appender {
@@ -2796,6 +2810,8 @@ struct formatter<T, Char,
   detail::dynamic_format_specs<Char> specs_;
 
  public:
+  using nonlocking = void;
+
   template <typename ParseContext>
   FMT_CONSTEXPR auto parse(ParseContext& ctx) -> const Char* {
     if (ctx.begin() == ctx.end() || *ctx.begin() == '}') return ctx.begin();
@@ -2978,6 +2994,7 @@ FMT_NODISCARD FMT_INLINE auto formatted_size(format_string<T...> fmt,
 
 FMT_API void vprint(string_view fmt, format_args args);
 FMT_API void vprint(FILE* f, string_view fmt, format_args args);
+FMT_API void vprint_locked(FILE* f, string_view fmt, format_args args);
 FMT_API void vprintln(FILE* f, string_view fmt, format_args args);
 
 /**
@@ -2993,8 +3010,9 @@ FMT_API void vprintln(FILE* f, string_view fmt, format_args args);
 template <typename... T>
 FMT_INLINE void print(format_string<T...> fmt, T&&... args) {
   const auto& vargs = fmt::make_format_args(args...);
-  return detail::is_utf8() ? vprint(fmt, vargs)
-                           : detail::vprint_mojibake(stdout, fmt, vargs);
+  if (!detail::is_utf8()) return detail::vprint_mojibake(stdout, fmt, vargs);
+  return detail::is_locking<T...>() ? vprint(fmt, vargs)
+                                    : vprint_locked(stdout, fmt, vargs);
 }
 
 /**
@@ -3010,8 +3028,9 @@ FMT_INLINE void print(format_string<T...> fmt, T&&... args) {
 template <typename... T>
 FMT_INLINE void print(FILE* f, format_string<T...> fmt, T&&... args) {
   const auto& vargs = fmt::make_format_args(args...);
-  return detail::is_utf8() ? vprint(f, fmt, vargs)
-                           : detail::vprint_mojibake(f, fmt, vargs);
+  if (!detail::is_utf8()) return detail::vprint_mojibake(f, fmt, vargs);
+  return detail::is_locking<T...>() ? vprint(f, fmt, vargs)
+                                    : vprint_locked(f, fmt, vargs);
 }
 
 /**
diff --git a/include/fmt/format-inl.h b/include/fmt/format-inl.h
index 872aa9802df1..c5259ba4c0d1 100644
--- a/include/fmt/format-inl.h
+++ b/include/fmt/format-inl.h
@@ -1677,15 +1677,17 @@ FMT_FUNC void print(std::FILE* f, string_view text) {
 }  // namespace detail
 
 FMT_FUNC void vprint(std::FILE* f, string_view fmt, format_args args) {
-  if (detail::file_ref(f).is_buffered()) {
-    auto&& buffer = detail::file_print_buffer(f);
-    return detail::vformat_to(buffer, fmt, args);
-  }
   auto buffer = memory_buffer();
   detail::vformat_to(buffer, fmt, args);
   detail::print(f, {buffer.data(), buffer.size()});
 }
 
+FMT_FUNC void vprint_locked(std::FILE* f, string_view fmt, format_args args) {
+  if (!detail::file_ref(f).is_buffered()) return vprint(f, fmt, args);
+  auto&& buffer = detail::file_print_buffer(f);
+  return detail::vformat_to(buffer, fmt, args);
+}
+
 FMT_FUNC void vprintln(std::FILE* f, string_view fmt, format_args args) {
   auto buffer = memory_buffer();
   detail::vformat_to(buffer, fmt, args);
diff --git a/test/format-test.cc b/test/format-test.cc
index 67d0631fc2fa..eec50456473c 100644
--- a/test/format-test.cc
+++ b/test/format-test.cc
@@ -242,7 +242,7 @@ TEST(util_test, format_system_error) {
     throws_on_alloc = true;
   }
   if (!throws_on_alloc) {
-    fmt::print("warning: std::allocator allocates {} chars\n", max_size);
+    fmt::print(stderr, "warning: std::allocator allocates {} chars\n", max_size);
     return;
   }
 }
@@ -1785,6 +1785,45 @@ TEST(format_test, line_buffering) {
 }
 #endif
 
+struct deadlockable {
+  int value = 0;
+  mutable std::mutex mutex;
+};
+
+FMT_BEGIN_NAMESPACE
+template <> struct formatter<deadlockable> {
+  FMT_CONSTEXPR auto parse(fmt::format_parse_context& ctx)
+      -> decltype(ctx.begin()) {
+    return ctx.begin();
+  }
+
+  auto format(const deadlockable& d, fmt::format_context& ctx) const
+      -> decltype(ctx.out()) {
+    std::lock_guard<std::mutex> lock(d.mutex);
+    return fmt::format_to(ctx.out(), "{}", d.value);
+  }
+};
+FMT_END_NAMESPACE
+
+TEST(format_test, locking_formatter) {
+  auto f = fmt::buffered_file();
+  try {
+    f = fmt::buffered_file("/dev/null", "w");
+  } catch (const std::system_error&) {
+    fmt::print(stderr, "warning: /dev/null is not supported\n");
+    return;
+  }
+  deadlockable d;
+  auto t = std::thread([&]() {
+    fmt::print(f.get(), "start t\n");
+    std::lock_guard<std::mutex> lock(d.mutex);
+    for (int i = 0; i < 1000000; ++i) d.value += 10;
+    fmt::print(f.get(), "done\n");
+  });
+  for (int i = 0; i < 100; ++i) fmt::print(f.get(), "{}", d);
+  t.join();
+}
+
 TEST(format_test, variadic) {
   EXPECT_EQ(fmt::format("{}c{}", "ab", 1), "abc1");
 }