1 /* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
2 /* vim: set ts=8 sts=2 et sw=2 tw=80: */
3 /* This Source Code Form is subject to the terms of the Mozilla Public
4  * License, v. 2.0. If a copy of the MPL was not distributed with this
5  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
6 
7 #ifndef mozilla_ProcThreadAttributes_h
8 #define mozilla_ProcThreadAttributes_h
9 
10 #include <windows.h>
11 
12 #include <utility>
13 
14 #include "mozilla/Attributes.h"
15 #include "mozilla/Maybe.h"
16 #include "mozilla/UniquePtr.h"
17 #include "mozilla/Vector.h"
18 
19 namespace mozilla {
20 
21 class MOZ_RAII ProcThreadAttributes final {
22   struct ProcThreadAttributeListDeleter {
operatorProcThreadAttributeListDeleter23     void operator()(LPPROC_THREAD_ATTRIBUTE_LIST aList) {
24       ::DeleteProcThreadAttributeList(aList);
25       delete[] reinterpret_cast<char*>(aList);
26     }
27   };
28 
29   using ProcThreadAttributeListPtr =
30       UniquePtr<_PROC_THREAD_ATTRIBUTE_LIST, ProcThreadAttributeListDeleter>;
31 
32  public:
ProcThreadAttributes()33   ProcThreadAttributes() : mMitigationPolicies(0) {}
34 
35   ~ProcThreadAttributes() = default;
36 
37   ProcThreadAttributes(const ProcThreadAttributes&) = delete;
38   ProcThreadAttributes(ProcThreadAttributes&&) = delete;
39   ProcThreadAttributes& operator=(const ProcThreadAttributes&) = delete;
40   ProcThreadAttributes& operator=(ProcThreadAttributes&&) = delete;
41 
AddMitigationPolicy(DWORD64 aPolicy)42   void AddMitigationPolicy(DWORD64 aPolicy) { mMitigationPolicies |= aPolicy; }
43 
AddInheritableHandle(HANDLE aHandle)44   bool AddInheritableHandle(HANDLE aHandle) {
45     DWORD type = ::GetFileType(aHandle);
46     if (type != FILE_TYPE_DISK && type != FILE_TYPE_PIPE) {
47       return false;
48     }
49 
50     if (!::SetHandleInformation(aHandle, HANDLE_FLAG_INHERIT,
51                                 HANDLE_FLAG_INHERIT)) {
52       return false;
53     }
54 
55     return mInheritableHandles.append(aHandle);
56   }
57 
58   template <size_t N>
AddInheritableHandles(HANDLE (& aHandles)[N])59   bool AddInheritableHandles(HANDLE (&aHandles)[N]) {
60     bool ok = true;
61     for (auto handle : aHandles) {
62       ok &= AddInheritableHandle(handle);
63     }
64 
65     return ok;
66   }
67 
HasMitigationPolicies()68   bool HasMitigationPolicies() const { return !!mMitigationPolicies; }
69 
HasInheritableHandles()70   bool HasInheritableHandles() const { return !mInheritableHandles.empty(); }
71 
72   /**
73    * @return false if the STARTUPINFOEXW::lpAttributeList was set to null
74    *               as expected based on the state of |this|;
75    *         true  if the STARTUPINFOEXW::lpAttributeList was set to
76    *               non-null;
77    */
AssignTo(STARTUPINFOEXW & aSiex)78   LauncherResult<bool> AssignTo(STARTUPINFOEXW& aSiex) {
79     ZeroMemory(&aSiex, sizeof(STARTUPINFOEXW));
80 
81     // We'll set the size to sizeof(STARTUPINFOW) until we determine whether the
82     // extended fields will be used.
83     aSiex.StartupInfo.cb = sizeof(STARTUPINFOW);
84 
85     DWORD numAttributes = 0;
86     if (HasMitigationPolicies()) {
87       ++numAttributes;
88     }
89 
90     if (HasInheritableHandles()) {
91       ++numAttributes;
92     }
93 
94     if (!numAttributes) {
95       return false;
96     }
97 
98     SIZE_T listSize = 0;
99     if (!::InitializeProcThreadAttributeList(nullptr, numAttributes, 0,
100                                              &listSize)) {
101       DWORD err = ::GetLastError();
102       if (err != ERROR_INSUFFICIENT_BUFFER) {
103         return LAUNCHER_ERROR_FROM_WIN32(err);
104       }
105     }
106 
107     auto buf = MakeUnique<char[]>(listSize);
108 
109     LPPROC_THREAD_ATTRIBUTE_LIST tmpList =
110         reinterpret_cast<LPPROC_THREAD_ATTRIBUTE_LIST>(buf.get());
111 
112     if (!::InitializeProcThreadAttributeList(tmpList, numAttributes, 0,
113                                              &listSize)) {
114       return LAUNCHER_ERROR_FROM_LAST();
115     }
116 
117     // Transfer buf to a ProcThreadAttributeListPtr - now that the list is
118     // initialized, we are no longer dealing with a plain old char array. We
119     // must now deinitialize the attribute list before deallocating the
120     // underlying buffer.
121     ProcThreadAttributeListPtr attrList(
122         reinterpret_cast<LPPROC_THREAD_ATTRIBUTE_LIST>(buf.release()));
123 
124     if (mMitigationPolicies) {
125       if (!::UpdateProcThreadAttribute(
126               attrList.get(), 0, PROC_THREAD_ATTRIBUTE_MITIGATION_POLICY,
127               &mMitigationPolicies, sizeof(mMitigationPolicies), nullptr,
128               nullptr)) {
129         return LAUNCHER_ERROR_FROM_LAST();
130       }
131     }
132 
133     if (!mInheritableHandles.empty()) {
134       if (!::UpdateProcThreadAttribute(
135               attrList.get(), 0, PROC_THREAD_ATTRIBUTE_HANDLE_LIST,
136               mInheritableHandles.begin(),
137               mInheritableHandles.length() * sizeof(HANDLE), nullptr,
138               nullptr)) {
139         return LAUNCHER_ERROR_FROM_LAST();
140       }
141     }
142 
143     mAttrList = std::move(attrList);
144     aSiex.lpAttributeList = mAttrList.get();
145     aSiex.StartupInfo.cb = sizeof(STARTUPINFOEXW);
146     return true;
147   }
148 
149  private:
150   static const uint32_t kNumInline = 3;  // Inline storage for the std handles
151 
152   DWORD64 mMitigationPolicies;
153   Vector<HANDLE, kNumInline> mInheritableHandles;
154   ProcThreadAttributeListPtr mAttrList;
155 };
156 
157 }  // namespace mozilla
158 
159 #endif  // mozilla_ProcThreadAttributes_h
160