diff options
Diffstat (limited to 'toolkit/mozapps/defaultagent/ScheduledTask.cpp')
| -rw-r--r-- | toolkit/mozapps/defaultagent/ScheduledTask.cpp | 409 |
1 files changed, 409 insertions, 0 deletions
diff --git a/toolkit/mozapps/defaultagent/ScheduledTask.cpp b/toolkit/mozapps/defaultagent/ScheduledTask.cpp new file mode 100644 index 0000000000..9c77dc15ed --- /dev/null +++ b/toolkit/mozapps/defaultagent/ScheduledTask.cpp @@ -0,0 +1,409 @@ +/* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */ +/* vim:set ts=2 sw=2 sts=2 et cindent: */ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +#include "ScheduledTask.h" + +#include <string> +#include <time.h> + +#include <comutil.h> +#include <taskschd.h> + +#include "readstrings.h" +#include "updatererrors.h" +#include "EventLog.h" +#include "mozilla/RefPtr.h" +#include "mozilla/ScopeExit.h" +#include "mozilla/UniquePtr.h" +#include "mozilla/WinHeaderOnlyUtils.h" +#include "WindowsDefaultBrowser.h" + +#include "DefaultBrowser.h" + +const wchar_t* kTaskVendor = L"" MOZ_APP_VENDOR; +// kTaskName should have the unique token appended before being used. +const wchar_t* kTaskName = L"" MOZ_APP_DISPLAYNAME " Default Browser Agent "; + +// The task scheduler requires its time values to come in the form of a string +// in the format YYYY-MM-DDTHH:MM:SSZ. This format string is used to get that +// out of the C library wcsftime function. +const wchar_t* kTimeFormat = L"%Y-%m-%dT%H:%M:%SZ"; +// The expanded time string should always be this length, for example: +// 2020-02-12T16:59:32Z +const size_t kTimeStrMaxLen = 20; + +#define ENSURE(x) \ + if (FAILED(hr = (x))) { \ + LOG_ERROR(hr); \ + return hr; \ + } + +struct SysFreeStringDeleter { + void operator()(BSTR aPtr) { ::SysFreeString(aPtr); } +}; +using BStrPtr = mozilla::UniquePtr<OLECHAR, SysFreeStringDeleter>; + +bool GetTaskDescription(mozilla::UniquePtr<wchar_t[]>& description) { + mozilla::UniquePtr<wchar_t[]> installPath; + bool success = GetInstallDirectory(installPath); + if (!success) { + LOG_ERROR_MESSAGE(L"Failed to get install directory"); + return false; + } + const wchar_t* iniFormat = L"%s\\defaultagent_localized.ini"; + int bufferSize = _scwprintf(iniFormat, installPath.get()); + ++bufferSize; // Extra character for terminating null + mozilla::UniquePtr<wchar_t[]> iniPath = + mozilla::MakeUnique<wchar_t[]>(bufferSize); + _snwprintf_s(iniPath.get(), bufferSize, _TRUNCATE, iniFormat, + installPath.get()); + + IniReader reader(iniPath.get()); + reader.AddKey("DefaultBrowserAgentTaskDescription", &description); + int status = reader.Read(); + if (status != OK) { + LOG_ERROR_MESSAGE(L"Failed to read task description: %d", status); + return false; + } + return true; +} + +HRESULT RegisterTask(const wchar_t* uniqueToken, + BSTR startTime /* = nullptr */) { + // Do data migration during the task installation. This might seem like it + // belongs in UpdateTask, but we want to be able to call + // RemoveTasks(); + // RegisterTask(); + // and still have data migration happen. Also, UpdateTask calls this function, + // so migration will still get run in that case. + MaybeMigrateCurrentDefault(); + + // Make sure we don't try to register a task that already exists. + RemoveTasks(uniqueToken, WhichTasks::WdbaTaskOnly); + + // If we create a folder and then fail to create the task, we need to + // remember to delete the folder so that whatever set of permissions it ends + // up with doesn't interfere with trying to create the task again later, and + // so that we don't just leave an empty folder behind. + bool createdFolder = false; + + HRESULT hr = S_OK; + RefPtr<ITaskService> scheduler; + ENSURE(CoCreateInstance(CLSID_TaskScheduler, nullptr, CLSCTX_INPROC_SERVER, + IID_ITaskService, getter_AddRefs(scheduler))); + + ENSURE(scheduler->Connect(VARIANT{}, VARIANT{}, VARIANT{}, VARIANT{})); + + RefPtr<ITaskFolder> rootFolder; + BStrPtr rootFolderBStr = BStrPtr(SysAllocString(L"\\")); + ENSURE( + scheduler->GetFolder(rootFolderBStr.get(), getter_AddRefs(rootFolder))); + + RefPtr<ITaskFolder> taskFolder; + BStrPtr vendorBStr = BStrPtr(SysAllocString(kTaskVendor)); + if (FAILED(rootFolder->GetFolder(vendorBStr.get(), + getter_AddRefs(taskFolder)))) { + hr = rootFolder->CreateFolder(vendorBStr.get(), VARIANT{}, + getter_AddRefs(taskFolder)); + if (SUCCEEDED(hr)) { + createdFolder = true; + } else if (hr != HRESULT_FROM_WIN32(ERROR_ALREADY_EXISTS)) { + // The folder already existing isn't an error, but anything else is. + LOG_ERROR(hr); + return hr; + } + } + + auto cleanupFolder = + mozilla::MakeScopeExit([hr, createdFolder, &rootFolder, &vendorBStr] { + if (createdFolder && FAILED(hr)) { + // If this fails, we can't really handle that intelligently, so + // don't even bother to check the return code. + rootFolder->DeleteFolder(vendorBStr.get(), 0); + } + }); + + RefPtr<ITaskDefinition> newTask; + ENSURE(scheduler->NewTask(0, getter_AddRefs(newTask))); + + mozilla::UniquePtr<wchar_t[]> description; + if (!GetTaskDescription(description)) { + return E_FAIL; + } + BStrPtr descriptionBstr = BStrPtr(SysAllocString(description.get())); + + RefPtr<IRegistrationInfo> taskRegistration; + ENSURE(newTask->get_RegistrationInfo(getter_AddRefs(taskRegistration))); + ENSURE(taskRegistration->put_Description(descriptionBstr.get())); + + RefPtr<ITaskSettings> taskSettings; + ENSURE(newTask->get_Settings(getter_AddRefs(taskSettings))); + ENSURE(taskSettings->put_DisallowStartIfOnBatteries(VARIANT_FALSE)); + ENSURE(taskSettings->put_MultipleInstances(TASK_INSTANCES_IGNORE_NEW)); + ENSURE(taskSettings->put_StartWhenAvailable(VARIANT_TRUE)); + ENSURE(taskSettings->put_StopIfGoingOnBatteries(VARIANT_FALSE)); + // This cryptic string means "12 hours 5 minutes". So, if the task runs for + // longer than that, the process will be killed, because that should never + // happen. See + // https://docs.microsoft.com/en-us/windows/win32/taskschd/tasksettings-executiontimelimit + // for a detailed explanation of these strings. + BStrPtr execTimeLimitBStr = BStrPtr(SysAllocString(L"PT12H5M")); + ENSURE(taskSettings->put_ExecutionTimeLimit(execTimeLimitBStr.get())); + + RefPtr<IRegistrationInfo> regInfo; + ENSURE(newTask->get_RegistrationInfo(getter_AddRefs(regInfo))); + + ENSURE(regInfo->put_Author(vendorBStr.get())); + + RefPtr<ITriggerCollection> triggers; + ENSURE(newTask->get_Triggers(getter_AddRefs(triggers))); + + RefPtr<ITrigger> newTrigger; + ENSURE(triggers->Create(TASK_TRIGGER_DAILY, getter_AddRefs(newTrigger))); + + RefPtr<IDailyTrigger> dailyTrigger; + ENSURE(newTrigger->QueryInterface(IID_IDailyTrigger, + getter_AddRefs(dailyTrigger))); + + if (startTime) { + ENSURE(dailyTrigger->put_StartBoundary(startTime)); + } else { + // The time that the task is scheduled to run at every day is taken from the + // time in the trigger's StartBoundary property. We'll set this to the + // current time, on the theory that the time at which we're being installed + // is a time that the computer is likely to be on other days. If our + // theory is wrong and the computer is offline at the scheduled time, then + // because we've set StartWhenAvailable above, the task will run whenever + // it wakes up. Since our task is entirely in the background and doesn't use + // a lot of resources, we're not concerned about it bothering the user if it + // runs while they're actively using this computer. + time_t now_t = time(nullptr); + // Subtract a minute from the current time, to avoid "winning" a potential + // race with the scheduler that might have it start the task immediately + // after we register it, if we finish doing that and then it evaluates the + // trigger during the same second. We haven't seen this happen in practice, + // but there's no documented guarantee that it won't, so let's be sure. + now_t -= 60; + + tm now_tm; + errno_t errno_rv = gmtime_s(&now_tm, &now_t); + if (errno_rv != 0) { + // The C runtime has a (private) function to convert Win32 error codes to + // errno values, but there's nothing that goes the other way, and it + // isn't worth including one here for something that's this unlikely to + // fail anyway. So just return a generic error. + hr = HRESULT_FROM_WIN32(ERROR_INVALID_TIME); + LOG_ERROR(hr); + return hr; + } + + mozilla::UniquePtr<wchar_t[]> timeStr = + mozilla::MakeUnique<wchar_t[]>(kTimeStrMaxLen + 1); + + if (wcsftime(timeStr.get(), kTimeStrMaxLen + 1, kTimeFormat, &now_tm) == + 0) { + hr = E_NOT_SUFFICIENT_BUFFER; + LOG_ERROR(hr); + return hr; + } + + BStrPtr startTimeBStr = BStrPtr(SysAllocString(timeStr.get())); + ENSURE(dailyTrigger->put_StartBoundary(startTimeBStr.get())); + } + + ENSURE(dailyTrigger->put_DaysInterval(1)); + + RefPtr<IActionCollection> actions; + ENSURE(newTask->get_Actions(getter_AddRefs(actions))); + + RefPtr<IAction> action; + ENSURE(actions->Create(TASK_ACTION_EXEC, getter_AddRefs(action))); + + RefPtr<IExecAction> execAction; + ENSURE(action->QueryInterface(IID_IExecAction, getter_AddRefs(execAction))); + + BStrPtr binaryPathBStr = + BStrPtr(SysAllocString(mozilla::GetFullBinaryPath().get())); + ENSURE(execAction->put_Path(binaryPathBStr.get())); + + std::wstring taskArgs = L"do-task \""; + taskArgs += uniqueToken; + taskArgs += L"\""; + BStrPtr argsBStr = BStrPtr(SysAllocString(taskArgs.c_str())); + ENSURE(execAction->put_Arguments(argsBStr.get())); + + std::wstring taskName(kTaskName); + taskName += uniqueToken; + BStrPtr taskNameBStr = BStrPtr(SysAllocString(taskName.c_str())); + + RefPtr<IRegisteredTask> registeredTask; + ENSURE(taskFolder->RegisterTaskDefinition( + taskNameBStr.get(), newTask, TASK_CREATE_OR_UPDATE, VARIANT{}, VARIANT{}, + TASK_LOGON_INTERACTIVE_TOKEN, VARIANT{}, getter_AddRefs(registeredTask))); + + return hr; +} + +HRESULT UpdateTask(const wchar_t* uniqueToken) { + RefPtr<ITaskService> scheduler; + HRESULT hr = S_OK; + ENSURE(CoCreateInstance(CLSID_TaskScheduler, nullptr, CLSCTX_INPROC_SERVER, + IID_ITaskService, getter_AddRefs(scheduler))); + + ENSURE(scheduler->Connect(VARIANT{}, VARIANT{}, VARIANT{}, VARIANT{})); + + RefPtr<ITaskFolder> taskFolder; + BStrPtr folderBStr = BStrPtr(SysAllocString(kTaskVendor)); + + if (FAILED( + scheduler->GetFolder(folderBStr.get(), getter_AddRefs(taskFolder)))) { + // If our folder doesn't exist, create it and the task. + return RegisterTask(uniqueToken); + } + + std::wstring taskName(kTaskName); + taskName += uniqueToken; + BStrPtr taskNameBStr = BStrPtr(SysAllocString(taskName.c_str())); + + RefPtr<IRegisteredTask> task; + if (FAILED(taskFolder->GetTask(taskNameBStr.get(), getter_AddRefs(task)))) { + // If our task doesn't exist at all, just create one. + return RegisterTask(uniqueToken); + } + + // If we have a task registered already, we need to recreate it because + // something might have changed that we need to update. But we don't + // want to restart the schedule from now, because that might mean the + // task never runs at all for e.g. Nightly. So create a new task, but + // first get and preserve the existing trigger. + RefPtr<ITaskDefinition> definition; + if (FAILED(task->get_Definition(getter_AddRefs(definition)))) { + // This task is broken, make a new one. + return RegisterTask(uniqueToken); + } + + RefPtr<ITriggerCollection> triggerList; + if (FAILED(definition->get_Triggers(getter_AddRefs(triggerList)))) { + // This task is broken, make a new one. + return RegisterTask(uniqueToken); + } + + RefPtr<ITrigger> trigger; + if (FAILED(triggerList->get_Item(1, getter_AddRefs(trigger)))) { + // This task is broken, make a new one. + return RegisterTask(uniqueToken); + } + + BSTR startTimeBstr; + if (FAILED(trigger->get_StartBoundary(&startTimeBstr))) { + // This task is broken, make a new one. + return RegisterTask(uniqueToken); + } + BStrPtr startTime(startTimeBstr); + + return RegisterTask(uniqueToken, startTime.get()); +} + +bool EndsWith(const wchar_t* string, const wchar_t* suffix) { + size_t string_len = wcslen(string); + size_t suffix_len = wcslen(suffix); + if (suffix_len > string_len) { + return false; + } + const wchar_t* substring = string + string_len - suffix_len; + return wcscmp(substring, suffix) == 0; +} + +HRESULT RemoveTasks(const wchar_t* uniqueToken, WhichTasks tasksToRemove) { + if (!uniqueToken || wcslen(uniqueToken) == 0) { + return E_INVALIDARG; + } + + RefPtr<ITaskService> scheduler; + HRESULT hr = S_OK; + ENSURE(CoCreateInstance(CLSID_TaskScheduler, nullptr, CLSCTX_INPROC_SERVER, + IID_ITaskService, getter_AddRefs(scheduler))); + + ENSURE(scheduler->Connect(VARIANT{}, VARIANT{}, VARIANT{}, VARIANT{})); + + RefPtr<ITaskFolder> taskFolder; + BStrPtr folderBStr(SysAllocString(kTaskVendor)); + + hr = scheduler->GetFolder(folderBStr.get(), getter_AddRefs(taskFolder)); + if (FAILED(hr)) { + if (hr == HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND)) { + // Don't return an error code if our folder doesn't exist, + // because that just means it's been removed already. + return S_OK; + } else { + return hr; + } + } + + RefPtr<IRegisteredTaskCollection> tasksInFolder; + ENSURE(taskFolder->GetTasks(TASK_ENUM_HIDDEN, getter_AddRefs(tasksInFolder))); + + LONG numTasks = 0; + ENSURE(tasksInFolder->get_Count(&numTasks)); + + std::wstring WdbaTaskName(kTaskName); + WdbaTaskName += uniqueToken; + + // This will be set to the last error that we encounter while deleting tasks. + // This allows us to keep attempting to remove the remaining tasks, even if + // we encounter an error, while still preserving what error we encountered so + // we can return it from this function. + HRESULT deleteResult = S_OK; + // Set to true if we intentionally skip any tasks. + bool tasksSkipped = false; + + for (LONG i = 0; i < numTasks; ++i) { + RefPtr<IRegisteredTask> task; + // IRegisteredTaskCollection's are 1-indexed. + hr = tasksInFolder->get_Item(_variant_t(i + 1), getter_AddRefs(task)); + if (FAILED(hr)) { + deleteResult = hr; + continue; + } + + BSTR taskName; + hr = task->get_Name(&taskName); + if (FAILED(hr)) { + deleteResult = hr; + continue; + } + // Automatically free taskName when we are done with it. + BStrPtr uniqueTaskName(taskName); + + if (tasksToRemove == WhichTasks::WdbaTaskOnly) { + if (WdbaTaskName.compare(taskName) != 0) { + tasksSkipped = true; + continue; + } + } else { // tasksToRemove == WhichTasks::AllTasksForInstallation + if (!EndsWith(taskName, uniqueToken)) { + tasksSkipped = true; + continue; + } + } + + hr = taskFolder->DeleteTask(taskName, 0 /* flags */); + if (FAILED(hr)) { + deleteResult = hr; + } + } + + // If we successfully removed all the tasks, delete the folder too. + if (!tasksSkipped && SUCCEEDED(deleteResult)) { + RefPtr<ITaskFolder> rootFolder; + BStrPtr rootFolderBStr = BStrPtr(SysAllocString(L"\\")); + ENSURE( + scheduler->GetFolder(rootFolderBStr.get(), getter_AddRefs(rootFolder))); + ENSURE(rootFolder->DeleteFolder(folderBStr.get(), 0)); + } + + return deleteResult; +} |