From 8f6a6484d2a823897478c53080bbbad168e3bc37 Mon Sep 17 00:00:00 2001
From: Paper <37962225+mrpapersonic@users.noreply.github.com>
Date: Sat, 2 Sep 2023 00:52:46 -0400
Subject: [PATCH] *: make ignoring sleep mode optional

---
 src/Context.cc              | 10 ++++++++++
 src/OptionHandlerFactory.cc |  8 ++++++++
 src/main.cc                 |  4 ----
 src/prefs.cc                |  2 ++
 src/prefs.h                 |  2 ++
 src/usage_text.h            |  5 +++++
 6 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/Context.cc b/src/Context.cc
index d12775bb..0930ae14 100644
--- a/src/Context.cc
+++ b/src/Context.cc
@@ -221,6 +221,11 @@ Context::Context(bool standalone, int argc, char** argv, const KeyVals& options)
   if (op->getAsBool(PREF_DISABLE_IPV6)) {
     SocketCore::setProtocolFamily(AF_INET);
   }
+#ifdef _WIN32
+  if (op->getAsBool(PREF_DISABLE_SLEEP)) {
+    SetThreadExecutionState(ES_CONTINUOUS | ES_SYSTEM_REQUIRED);
+  }
+#endif // _WIN32
   SocketCore::setIpDscp(op->getAsInt(PREF_DSCP));
   SocketCore::setSocketRecvBufferSize(
       op->getAsInt(PREF_SOCKET_RECV_BUFFER_SIZE));
@@ -313,6 +318,11 @@ Context::Context(bool standalone, int argc, char** argv, const KeyVals& options)
     reqinfo = std::make_shared<MultiUrlRequestInfo>(std::move(requestGroups),
                                                     op, uriListParser);
   }
+#ifdef _WIN32
+  if (op->getAsBool(PREF_DISABLE_SLEEP)) {
+    SetThreadExecutionState(ES_CONTINUOUS);
+  }
+#endif // _WIN32
 }
 
 Context::~Context() = default;
diff --git a/src/OptionHandlerFactory.cc b/src/OptionHandlerFactory.cc
index 2cd2c43e..1ed62275 100644
--- a/src/OptionHandlerFactory.cc
+++ b/src/OptionHandlerFactory.cc
@@ -738,6 +738,14 @@ std::vector<OptionHandler*> OptionHandlerFactory::createOptionHandlers()
     op->addTag(TAG_ADVANCED);
     handlers.push_back(op);
   }
+#if _WIN32
+  {
+    OptionHandler* op(new BooleanOptionHandler(
+        PREF_DISABLE_SLEEP, TEXT_DISABLE_SLEEP, A2_V_FALSE, OptionHandler::OPT_ARG));
+    op->addTag(TAG_ADVANCED);
+    handlers.push_back(op);
+  }
+#endif // _WIN32
   {
     OptionHandler* op(new BooleanOptionHandler(
         PREF_RPC_ALLOW_ORIGIN_ALL, TEXT_RPC_ALLOW_ORIGIN_ALL, A2_V_FALSE,
diff --git a/src/main.cc b/src/main.cc
index fbaa9195..e38f47f0 100644
--- a/src/main.cc
+++ b/src/main.cc
@@ -69,11 +69,7 @@ error_code::Value main(int argc, char** argv)
     pargv[i] = winArgStrs.back().get();
   }
 
-  SetThreadExecutionState(ES_CONTINUOUS | ES_SYSTEM_REQUIRED);
-
   Context context(true, winArgc, pargv.get(), KeyVals());
-
-  SetThreadExecutionState(ES_CONTINUOUS);
 #else  // !__MINGW32__
   Context context(true, argc, argv, KeyVals());
 #endif // !__MINGW32__
diff --git a/src/prefs.cc b/src/prefs.cc
index 2591b9f0..0e782391 100644
--- a/src/prefs.cc
+++ b/src/prefs.cc
@@ -335,6 +335,8 @@ PrefPtr PREF_STREAM_PIECE_SELECTOR = makePref("stream-piece-selector");
 // value: true | false
 PrefPtr PREF_TRUNCATE_CONSOLE_READOUT = makePref("truncate-console-readout");
 // value: true | false
+PrefPtr PREF_DISABLE_SLEEP = makePref("disable-sleep-until-finished")
+// value: true | false
 PrefPtr PREF_PAUSE = makePref("pause");
 // value: default | full | hide
 PrefPtr PREF_DOWNLOAD_RESULT = makePref("download-result");
diff --git a/src/prefs.h b/src/prefs.h
index 338fd6e6..aa259f39 100644
--- a/src/prefs.h
+++ b/src/prefs.h
@@ -288,6 +288,8 @@ extern PrefPtr PREF_STREAM_PIECE_SELECTOR;
 // value: true | false
 extern PrefPtr PREF_TRUNCATE_CONSOLE_READOUT;
 // value: true | false
+extern PrefPtr PREF_DISABLE_SLEEP;
+// value: true | false
 extern PrefPtr PREF_PAUSE;
 // value: default | full | hide
 extern PrefPtr PREF_DOWNLOAD_RESULT;
diff --git a/src/usage_text.h b/src/usage_text.h
index 0052b798..329daa6e 100644
--- a/src/usage_text.h
+++ b/src/usage_text.h
@@ -884,6 +884,11 @@
 #define TEXT_TRUNCATE_CONSOLE_READOUT                                   \
   _(" --truncate-console-readout[=true|false] Truncate console readout to fit in\n"\
     "                              a single line.")
+#ifdef _WIN32
+#define TEXT_DISABLE_SLEEP
+  _(" --disable-sleep-until-finished[=true|false] Disables sleep mode on Windows\n" \
+    "                              until downloads are finished.")
+#endif // _WIN32
 #define TEXT_PAUSE                              \
   _(" --pause[=true|false]         Pause download after added. This option is\n" \
     "                              effective only when --enable-rpc=true is given.")