Display Power Off Utility Code

#include "stdafx.h"

#include "Resource.h"

#include "App.h"
#include "Engine.h"
#include "Config.h"
#include "Workstation.h"
#include "Options.h"
#include "Platform.h"
#include "Screensaver.h"
#include "SettingsDialog.h"
#include "NotificationAreaIcon.h"
#include "UserMessage.h"
#include "UserTaskDialog.h"

namespace
{
	bool lastOperationSucceeded = true;
	bool waitForWtsSessionChange;
	bool fromWtsSessionChange;

	constexpr auto succeeded	= 0;
	constexpr auto failed		= 1;

	constexpr auto timerId			= 1;
	constexpr auto lockTimerId		= 100;
	constexpr auto lockTimerElapse	= 3000;

	enum class Action
	{
		None,
		TurnOff,
		Screensaver
	};

	enum class State
	{
		None,
		Ready
	};

	Action	action;
	State	state;

	void Reset(HWND hwnd)
	{
		action	= Action::None;
		state	=  State::None;

		NotificationAreaIcon::OperationEnded(hwnd);
	}

	void ProceedWithDelay(HWND hwnd, DWORD delay, bool ready = false)
	{
		TRACE_BOOL("Ready for action", ready);
		state = ready ? State::Ready : State::None;

		TRACE_UINT("Delay for ms", delay);
		SetTimer(hwnd, timerId, delay, nullptr);
	}

	void ExecuteDelayed(HWND hwnd, Action requestedAction)
	{
		if (Engine::Busy())
		{
			TRACE(TEXT("Engine is busy, ignoring."));
			return;
		}

		NotificationAreaIcon::OperationStarted(hwnd, requestedAction == Action::TurnOff);

		action = requestedAction;
		ProceedWithDelay(hwnd, Config::StartupDelay());
	}

	[[nodiscard]] bool CanExecute(HWND hwnd)
	{
		Config::Read(true);

		if (fromWtsSessionChange)
		{
			fromWtsSessionChange = false;

			const auto lockWorkStationDelay = Config::LockWorkStationDelay();
			ProceedWithDelay(hwnd, lockWorkStationDelay, true);

			return false;
		}

		if (!Options::LockWorkStation() && !Config::AlwaysLockWorkStation())
		{
			TRACE(TEXT("No lock required."));
			return true;
		}

		TRACE(TEXT("Lock required."));
		if (!Workstation::Lock(hwnd))
		{
			Reset(hwnd);
			UserMessage::LockWorkstationFailed();
		}
		else
		{
			waitForWtsSessionChange = true;
			SetTimer(hwnd, lockTimerId, lockTimerElapse, nullptr);
		}

		return false;
	}

	void DisplayOff(HWND hwnd)
	{
		ASSERT(action == Action::TurnOff || action == Action::Screensaver);

		if (action == Action::None)
		{
			return;
		}

		const auto standalone = !Options::NotificationAreaIcon();

		lastOperationSucceeded = true;

		if (action == Action::Screensaver || (standalone && (Options::Screensaver() || Config::AlwaysLaunchScreensaver())))
		{
			TRACE(TEXT("Launching screensaver..."));
			SendMessage(hwnd, WM_SYSCOMMAND, SC_SCREENSAVE, 0);
			TRACE(TEXT("Screensaver is running"));
		}
		else
		{
			TRACE(TEXT("Turning off display..."));
			SendMessage(hwnd, WM_SYSCOMMAND, SC_MONITORPOWER, 2);
			TRACE(TEXT("Display is turned off"));
		}

		if (standalone)
		{
			Platform::CloseWindow(hwnd);
		}
	}

	void OnWtSessionChange(HWND hwnd, int reason)
	{
		TRACE_UINT("Got WM_WTSSESSION_CHANGE with wParam", reason);

		if (reason != WTS_SESSION_LOCK)
		{
			TRACE(TEXT("Not WTS_SESSION_LOCK. Exit."));
			return;
		}

		fromWtsSessionChange = !waitForWtsSessionChange && Config::TriggerOnLockWorkStation();
		TRACE_BOOL("fromWtsSessionChange", fromWtsSessionChange);

		waitForWtsSessionChange = false;
		KillTimer(hwnd, lockTimerId);

		if (fromWtsSessionChange)
		{
			TRACE(TEXT("Engine::ExecuteDefault"));
			Engine::ExecuteDefault(hwnd);
			return;
		}

		const auto lockWorkStationDelay = Config::LockWorkStationDelay();
		ProceedWithDelay(hwnd, lockWorkStationDelay, true);
	}

	LRESULT CALLBACK WndProc(HWND, UINT, WPARAM, LPARAM);

	[[nodiscard]] HWND CreateWorkerWindow()
	{
		constexpr auto szWindowClass	= TEXT("{0032637B-F093-4A8A-BE15-C56863567DDD}");
		constexpr auto nIconId			= IDI_APP_ICON;

#pragma warning(suppress: 26486)
		DECLARE_CBSTRUCT(WNDCLASSEX, wndClassEx);

		wndClassEx.lpfnWndProc		= WndProc;
		wndClassEx.hInstance		= App::GetInstance();
		wndClassEx.lpszClassName	= szWindowClass;
		wndClassEx.hIcon			= Platform::LoadIconSize(App::GetInstance(), nIconId);
		wndClassEx.hIconSm			= Platform::LoadIconSize(App::GetInstance(), nIconId, true);

		ASSERT(wndClassEx.hIcon);
		ASSERT(wndClassEx.hIconSm);

#pragma warning(suppress: 26486)
		VERIFY(RegisterClassEx(&wndClassEx));

		const auto hwnd = CreateWindow(szWindowClass, nullptr, WS_OVERLAPPED, CW_USEDEFAULT, 0, CW_USEDEFAULT, 0, nullptr, nullptr, App::GetInstance(), nullptr);
		ASSERT(hwnd);

		return hwnd;
	}

	BOOL OnCreate(HWND hwnd, LPCREATESTRUCT)
	{
		// Notification area icon.
		if (Options::NotificationAreaIcon())
		{
			if (NotificationAreaIcon::Add(hwnd))
			{
				TRACE(TEXT("Added notification area icon."));
				return TRUE;
			}

			TRACE(TEXT("Notification area icon add failed."));
			UserMessage::NotificationAreaIconFailed();
			return FALSE;
		}

		// Standalone/screensaver.
		Engine::ExecuteDefault(hwnd);
		return TRUE;
	}

	[[nodiscard]] bool IsLockTimer(HWND hwnd, UINT id)
	{
		if (id != lockTimerId || !waitForWtsSessionChange)
		{
			return false;
		}

		waitForWtsSessionChange = false;
		UserMessage::LockWorkstationFailed();

		Reset(hwnd);

		return true;
	}

	void OnTimer(HWND hwnd, UINT id)
	{
		VERIFY(KillTimer(hwnd, id));

		if (IsLockTimer(hwnd, id))
		{
			return;
		}

		const auto canExecute = (state == State::Ready) || CanExecute(hwnd);

		if (canExecute)
		{
			DisplayOff(hwnd);
			Reset(hwnd);
		}
	}

	LRESULT CALLBACK WndProc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
	{
		NotificationAreaIcon::Process(hwnd, uMsg, wParam, lParam);

		HANDLE_APP_MSG(WM_WTSSESSION_CHANGE,	OnWtSessionChange(hwnd, static_cast<int>(wParam)));

#pragma warning (suppress: 26818)
		switch (uMsg)
		{
			HANDLE_MSG(hwnd, WM_CREATE,			OnCreate);
			HANDLE_MSG(hwnd, WM_DESTROY,		[](HWND) { PostQuitMessage(lastOperationSucceeded ? succeeded : failed); });
			HANDLE_MSG(hwnd, WM_TIMER,			OnTimer);
		}

		return DefWindowProc(hwnd, uMsg, wParam, lParam);
	}
}

bool Engine::Busy() noexcept
{
	return action != Action::None;
}

bool Engine::LaunchScreensaverByDefault()
{
	return Config::AlwaysLaunchScreensaver();
}

namespace
{
	[[nodiscard]] bool CanLaunchScreensaver()
	{
		if (!Screensaver::IsSet())
		{
			TRACE(TEXT("Screensaver is not set."));
			UserMessage::ScreensaverNotSet();
			return false;
		}

		if (Screensaver::IsActive() ||
			Config::AlwaysEnableScreensaver() ||
			UserTaskDialog::ShouldActivateScreensaver() && Screensaver::SetActive())
		{
			return true;
		}

		UserMessage::ScreensaverFailed();
		return false;
	}
}

void Engine::LaunchScreensaver(HWND hwnd)
{
	lastOperationSucceeded = false;

	if (CanLaunchScreensaver())
	{
		TRACE(TEXT("Launching screensaver..."));
		ExecuteDelayed(hwnd, Action::Screensaver);
	}
}

void Engine::TurnOffDisplay(HWND hwnd)
{
	lastOperationSucceeded = false;

	TRACE(TEXT("Turning off display..."));
	ExecuteDelayed(hwnd, Action::TurnOff);
}

void Engine::ExecuteDefault(HWND hwnd)
{
	if(!Options::AsScreensaver() && LaunchScreensaverByDefault())
	{
		TRACE(TEXT("LaunchScreensaverByDefault is true."));
		LaunchScreensaver(hwnd);
	}
	else
	{
		TRACE(TEXT("LaunchScreensaverByDefault is false."));
		TurnOffDisplay(hwnd);
	}
}

UINT Engine::Loop()
{
	if (!CreateWorkerWindow())
	{
		TRACE(TEXT("Worker window creation failed."));
		return failed;
	}

	TRACE(TEXT("Message loop started."));

	MSG msg;
	int nRet;

#pragma warning(suppress: 26486)
	while ((nRet = GetMessage(&msg, nullptr, 0, 0)) != 0)
	{
		if (nRet == -1)
		{
			continue;
		}

		if (!SettingsDialog::Process(msg))
		{
			TranslateMessage(&msg);
			DispatchMessage(&msg);
		}
	}

	TRACE(TEXT("Message loop end"));

#pragma warning(suppress: 26472)
	return static_cast<UINT>(msg.wParam);
}