Seregon/ShadPKG

A tool for deriving PKG packet encryption keys for ps4 written in c++

C++/47.3 KB/No license
common/polyfill_thread.h
ShadPKG / common / polyfill_thread.h
1// SPDX-FileCopyrightText: 2022 yuzu Emulator Project
2// SPDX-License-Identifier: GPL-2.0-or-later
3 
4//
5// TODO: remove this file when jthread is supported by all compilation targets
6//
7 
8#pragma once
9 
10#include <version>
11 
12#ifdef __cpp_lib_jthread
13 
14#include <chrono>
15#include <condition_variable>
16#include <stop_token>
17#include <thread>
18#include <utility>
19 
20namespace Common {
21 
22template <typename Condvar, typename Lock, typename Pred>
23void CondvarWait(Condvar& cv, std::unique_lock<Lock>& lk, std::stop_token token, Pred&& pred) {
24 cv.wait(lk, token, std::forward<Pred>(pred));
25}
26 
27template <typename Rep, typename Period>
28bool StoppableTimedWait(std::stop_token token, const std::chrono::duration<Rep, Period>& rel_time) {
29 std::condition_variable_any cv;
30 std::mutex m;
31 
32 // Perform the timed wait.
33 std::unique_lock lk{m};
34 return !cv.wait_for(lk, token, rel_time, [&] { return token.stop_requested(); });
35}
36 
37} // namespace Common
38 
39#else
40 
41#include <atomic>
42#include <chrono>
43#include <condition_variable>
44#include <functional>
45#include <map>
46#include <memory>
47#include <mutex>
48#include <optional>
49#include <thread>
50#include <type_traits>
51#include <utility>
52 
53namespace std {
54namespace polyfill {
55 
56using stop_state_callback = size_t;
57 
58class stop_state {
59public:
60 stop_state() = default;
61 ~stop_state() = default;
62 
63 bool request_stop() {
64 unique_lock lk{m_lock};
65 
66 if (m_stop_requested) {
67 // Already set, nothing to do.
68 return false;
69 }
70 
71 // Mark stop requested.
72 m_stop_requested = true;
73 
74 while (!m_callbacks.empty()) {
75 // Get an iterator to the first element.
76 const auto it = m_callbacks.begin();
77 
78 // Move the callback function out of the map.
79 function<void()> f;
80 swap(it->second, f);
81 
82 // Erase the now-empty map element.
83 m_callbacks.erase(it);
84 
85 // Run the callback.
86 if (f) {
87 f();
88 }
89 }
90 
91 return true;
92 }
93 
94 bool stop_requested() const {
95 unique_lock lk{m_lock};
96 return m_stop_requested;
97 }
98 
99 stop_state_callback insert_callback(function<void()> f) {
100 unique_lock lk{m_lock};
101 
102 if (m_stop_requested) {
103 // Stop already requested. Don't insert anything,
104 // just run the callback synchronously.
105 if (f) {
106 f();
107 }
108 return 0;
109 }
110 
111 // Insert the callback.
112 stop_state_callback ret = ++m_next_callback;
113 m_callbacks.emplace(ret, std::move(f));
114 return ret;
115 }
116 
117 void remove_callback(stop_state_callback cb) {
118 unique_lock lk{m_lock};
119 m_callbacks.erase(cb);
120 }
121 
122private:
123 mutable recursive_mutex m_lock;
124 map<stop_state_callback, function<void()>> m_callbacks;
125 stop_state_callback m_next_callback{0};
126 bool m_stop_requested{false};
127};
128 
129} // namespace polyfill
130 
131class stop_token;
132class stop_source;
133struct nostopstate_t {
134 explicit nostopstate_t() = default;
135};
136inline constexpr nostopstate_t nostopstate{};
137 
138template <class Callback>
139class stop_callback;
140 
141class stop_token {
142public:
143 stop_token() noexcept = default;
144 
145 stop_token(const stop_token&) noexcept = default;
146 stop_token(stop_token&&) noexcept = default;
147 stop_token& operator=(const stop_token&) noexcept = default;
148 stop_token& operator=(stop_token&&) noexcept = default;
149 ~stop_token() = default;
150 
151 void swap(stop_token& other) noexcept {
152 m_stop_state.swap(other.m_stop_state);
153 }
154 
155 [[nodiscard]] bool stop_requested() const noexcept {
156 return m_stop_state && m_stop_state->stop_requested();
157 }
158 [[nodiscard]] bool stop_possible() const noexcept {
159 return m_stop_state != nullptr;
160 }
161 
162private:
163 friend class stop_source;
164 template <typename Callback>
165 friend class stop_callback;
166 stop_token(shared_ptr<polyfill::stop_state> stop_state) : m_stop_state(std::move(stop_state)) {}
167 
168private:
169 shared_ptr<polyfill::stop_state> m_stop_state;
170};
171 
172class stop_source {
173public:
174 stop_source() : m_stop_state(make_shared<polyfill::stop_state>()) {}
175 explicit stop_source(nostopstate_t) noexcept {}
176 
177 stop_source(const stop_source&) noexcept = default;
178 stop_source(stop_source&&) noexcept = default;
179 stop_source& operator=(const stop_source&) noexcept = default;
180 stop_source& operator=(stop_source&&) noexcept = default;
181 ~stop_source() = default;
182 void swap(stop_source& other) noexcept {
183 m_stop_state.swap(other.m_stop_state);
184 }
185 
186 [[nodiscard]] stop_token get_token() const noexcept {
187 return stop_token(m_stop_state);
188 }
189 [[nodiscard]] bool stop_possible() const noexcept {
190 return m_stop_state != nullptr;
191 }
192 [[nodiscard]] bool stop_requested() const noexcept {
193 return m_stop_state && m_stop_state->stop_requested();
194 }
195 bool request_stop() noexcept {
196 return m_stop_state && m_stop_state->request_stop();
197 }
198 
199private:
200 friend class jthread;
201 explicit stop_source(shared_ptr<polyfill::stop_state> stop_state)
202 : m_stop_state(std::move(stop_state)) {}
203 
204private:
205 shared_ptr<polyfill::stop_state> m_stop_state;
206};
207 
208template <typename Callback>
209class stop_callback {
210 static_assert(is_nothrow_destructible_v<Callback>);
211 static_assert(is_invocable_v<Callback>);
212 
213public:
214 using callback_type = Callback;
215 
216 template <typename C>
217 requires constructible_from<Callback, C>
218 explicit stop_callback(const stop_token& st,
219 C&& cb) noexcept(is_nothrow_constructible_v<Callback, C>)
220 : m_stop_state(st.m_stop_state) {
221 if (m_stop_state) {
222 m_callback = m_stop_state->insert_callback(std::move(cb));
223 }
224 }
225 template <typename C>
226 requires constructible_from<Callback, C>
227 explicit stop_callback(stop_token&& st,
228 C&& cb) noexcept(is_nothrow_constructible_v<Callback, C>)
229 : m_stop_state(std::move(st.m_stop_state)) {
230 if (m_stop_state) {
231 m_callback = m_stop_state->insert_callback(std::move(cb));
232 }
233 }
234 ~stop_callback() {
235 if (m_stop_state && m_callback) {
236 m_stop_state->remove_callback(m_callback);
237 }
238 }
239 
240 stop_callback(const stop_callback&) = delete;
241 stop_callback(stop_callback&&) = delete;
242 stop_callback& operator=(const stop_callback&) = delete;
243 stop_callback& operator=(stop_callback&&) = delete;
244 
245private:
246 shared_ptr<polyfill::stop_state> m_stop_state;
247 polyfill::stop_state_callback m_callback;
248};
249 
250template <typename Callback>
251stop_callback(stop_token, Callback) -> stop_callback<Callback>;
252 
253class jthread {
254public:
255 using id = thread::id;
256 using native_handle_type = thread::native_handle_type;
257 
258 jthread() noexcept = default;
259 
260 template <typename F, typename... Args,
261 typename = enable_if_t<!is_same_v<remove_cvref_t<F>, jthread>>>
262 explicit jthread(F&& f, Args&&... args)
263 : m_stop_state(make_shared<polyfill::stop_state>()),
264 m_thread(make_thread(std::forward<F>(f), std::forward<Args>(args)...)) {}
265 
266 ~jthread() {
267 if (joinable()) {
268 request_stop();
269 join();
270 }
271 }
272 
273 jthread(const jthread&) = delete;
274 jthread(jthread&&) noexcept = default;
275 jthread& operator=(const jthread&) = delete;
276 
277 jthread& operator=(jthread&& other) noexcept {
278 m_thread.swap(other.m_thread);
279 m_stop_state.swap(other.m_stop_state);
280 return *this;
281 }
282 
283 void swap(jthread& other) noexcept {
284 m_thread.swap(other.m_thread);
285 m_stop_state.swap(other.m_stop_state);
286 }
287 [[nodiscard]] bool joinable() const noexcept {
288 return m_thread.joinable();
289 }
290 void join() {
291 m_thread.join();
292 }
293 void detach() {
294 m_thread.detach();
295 m_stop_state.reset();
296 }
297 
298 [[nodiscard]] id get_id() const noexcept {
299 return m_thread.get_id();
300 }
301 [[nodiscard]] native_handle_type native_handle() {
302 return m_thread.native_handle();
303 }
304 [[nodiscard]] stop_source get_stop_source() noexcept {
305 return stop_source(m_stop_state);
306 }
307 [[nodiscard]] stop_token get_stop_token() const noexcept {
308 return stop_source(m_stop_state).get_token();
309 }
310 bool request_stop() noexcept {
311 return get_stop_source().request_stop();
312 }
313 [[nodiscard]] static unsigned int hardware_concurrency() noexcept {
314 return thread::hardware_concurrency();
315 }
316 
317private:
318 template <typename F, typename... Args>
319 thread make_thread(F&& f, Args&&... args) {
320 if constexpr (is_invocable_v<decay_t<F>, stop_token, decay_t<Args>...>) {
321 return thread(std::forward<F>(f), get_stop_token(), std::forward<Args>(args)...);
322 } else {
323 return thread(std::forward<F>(f), std::forward<Args>(args)...);
324 }
325 }
326 
327 shared_ptr<polyfill::stop_state> m_stop_state;
328 thread m_thread;
329};
330 
331} // namespace std
332 
333namespace Common {
334 
335template <typename Condvar, typename Lock, typename Pred>
336void CondvarWait(Condvar& cv, std::unique_lock<Lock>& lk, std::stop_token token, Pred pred) {
337 if (token.stop_requested()) {
338 return;
339 }
340 
341 std::stop_callback callback(token, [&] {
342 { std::scoped_lock lk2{*lk.mutex()}; }
343 cv.notify_all();
344 });
345 
346 cv.wait(lk, [&] { return pred() || token.stop_requested(); });
347}
348 
349template <typename Rep, typename Period>
350bool StoppableTimedWait(std::stop_token token, const std::chrono::duration<Rep, Period>& rel_time) {
351 if (token.stop_requested()) {
352 return false;
353 }
354 
355 bool stop_requested = false;
356 std::condition_variable cv;
357 std::mutex m;
358 
359 std::stop_callback cb(token, [&] {
360 // Wake up the waiting thread.
361 {
362 std::scoped_lock lk{m};
363 stop_requested = true;
364 }
365 cv.notify_one();
366 });
367 
368 // Perform the timed wait.
369 std::unique_lock lk{m};
370 return !cv.wait_for(lk, rel_time, [&] { return stop_requested; });
371}
372 
373} // namespace Common
374 
375#endif
376