1 /*
2  *  ReactOS Task Manager
3  *
4  *  affinity.c
5  *
6  *  Copyright (C) 1999 - 2001  Brian Palmer  <brianp@reactos.org>
7  *                2005         Klemens Friedl <frik85@reactos.at>
8  *
9  * This library is free software; you can redistribute it and/or
10  * modify it under the terms of the GNU Lesser General Public
11  * License as published by the Free Software Foundation; either
12  * version 2.1 of the License, or (at your option) any later version.
13  *
14  * This library is distributed in the hope that it will be useful,
15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
17  * Lesser General Public License for more details.
18  *
19  * You should have received a copy of the GNU Lesser General Public
20  * License along with this library; if not, write to the Free Software
21  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
22  */
23 
24 #include "precomp.h"
25 
26 HANDLE        hProcessAffinityHandle;
27 
28 static const DWORD dwCpuTable[] = {
29     IDC_CPU0,   IDC_CPU1,   IDC_CPU2,   IDC_CPU3,
30     IDC_CPU4,   IDC_CPU5,   IDC_CPU6,   IDC_CPU7,
31     IDC_CPU8,   IDC_CPU9,   IDC_CPU10,  IDC_CPU11,
32     IDC_CPU12,  IDC_CPU13,  IDC_CPU14,  IDC_CPU15,
33     IDC_CPU16,  IDC_CPU17,  IDC_CPU18,  IDC_CPU19,
34     IDC_CPU20,  IDC_CPU21,  IDC_CPU22,  IDC_CPU23,
35     IDC_CPU24,  IDC_CPU25,  IDC_CPU26,  IDC_CPU27,
36     IDC_CPU28,  IDC_CPU29,  IDC_CPU30,  IDC_CPU31,
37 };
38 
39 static INT_PTR CALLBACK AffinityDialogWndProc(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam);
40 
41 void ProcessPage_OnSetAffinity(void)
42 {
43     DWORD    dwProcessId;
44     WCHAR    strErrorText[260];
45     WCHAR    szTitle[256];
46 
47     dwProcessId = GetSelectedProcessId();
48 
49     if (dwProcessId == 0)
50         return;
51 
52     hProcessAffinityHandle = OpenProcess(PROCESS_QUERY_INFORMATION|PROCESS_SET_INFORMATION, FALSE, dwProcessId);
53     if (!hProcessAffinityHandle) {
54         GetLastErrorText(strErrorText, sizeof(strErrorText) / sizeof(WCHAR));
55         LoadStringW(hInst, IDS_MSG_ACCESSPROCESSAFF, szTitle, sizeof(szTitle) / sizeof(WCHAR));
56         MessageBoxW(hMainWnd, strErrorText, szTitle, MB_OK|MB_ICONSTOP);
57         return;
58     }
59     DialogBoxW(hInst, MAKEINTRESOURCEW(IDD_AFFINITY_DIALOG), hMainWnd, AffinityDialogWndProc);
60     if (hProcessAffinityHandle)    {
61         CloseHandle(hProcessAffinityHandle);
62         hProcessAffinityHandle = NULL;
63     }
64 }
65 
66 INT_PTR CALLBACK
67 AffinityDialogWndProc(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam)
68 {
69     DWORD_PTR  dwProcessAffinityMask = 0;
70     DWORD_PTR  dwSystemAffinityMask = 0;
71     WCHAR  strErrorText[260];
72     WCHAR  szTitle[256];
73     BYTE   nCpu;
74 
75     switch (message) {
76     case WM_INITDIALOG:
77 
78         /*
79          * Get the current affinity mask for the process and
80          * the number of CPUs present in the system
81          */
82         if (!GetProcessAffinityMask(hProcessAffinityHandle, &dwProcessAffinityMask, &dwSystemAffinityMask))    {
83             GetLastErrorText(strErrorText, sizeof(strErrorText) / sizeof(WCHAR));
84             EndDialog(hDlg, 0);
85             LoadStringW(hInst, IDS_MSG_ACCESSPROCESSAFF, szTitle, sizeof(szTitle) / sizeof(WCHAR));
86             MessageBoxW(hMainWnd, strErrorText, szTitle, MB_OK|MB_ICONSTOP);
87         }
88 
89         for (nCpu=0; nCpu<sizeof(dwCpuTable) / sizeof(dwCpuTable[0]); nCpu++) {
90             /*
91              * Enable a checkbox for each processor present in the system
92              */
93             if (dwSystemAffinityMask & (1 << nCpu))
94                 EnableWindow(GetDlgItem(hDlg, dwCpuTable[nCpu]), TRUE);
95             /*
96              * Check each checkbox that the current process
97              * has affinity with
98              */
99             if (dwProcessAffinityMask & (1 << nCpu))
100                 CheckDlgButton(hDlg, dwCpuTable[nCpu], BST_CHECKED);
101         }
102 
103         return TRUE;
104 
105     case WM_COMMAND:
106 
107         /*
108          * If the user has cancelled the dialog box
109          * then just close it
110          */
111         if (LOWORD(wParam) == IDCANCEL) {
112             EndDialog(hDlg, LOWORD(wParam));
113             return TRUE;
114         }
115 
116         /*
117          * The user has clicked OK -- so now we have
118          * to adjust the process affinity mask
119          */
120         if (LOWORD(wParam) == IDOK) {
121             for (nCpu=0; nCpu<sizeof(dwCpuTable) / sizeof(dwCpuTable[0]); nCpu++) {
122                 /*
123                  * First we have to create a mask out of each
124                  * checkbox that the user checked.
125                  */
126                 if (IsDlgButtonChecked(hDlg, dwCpuTable[nCpu]))
127                     dwProcessAffinityMask |= (1 << nCpu);
128             }
129 
130             /*
131              * Make sure they are giving the process affinity
132              * with at least one processor. I'd hate to see a
133              * process that is not in a wait state get deprived
134              * of it's cpu time.
135              */
136             if (!dwProcessAffinityMask) {
137                 LoadStringW(hInst, IDS_MSG_PROCESSONEPRO, strErrorText, sizeof(strErrorText) / sizeof(WCHAR));
138                 LoadStringW(hInst, IDS_MSG_INVALIDOPTION, szTitle, sizeof(szTitle) / sizeof(WCHAR));
139                 MessageBoxW(hDlg, strErrorText, szTitle, MB_OK|MB_ICONSTOP);
140                 return TRUE;
141             }
142 
143             /*
144              * Try to set the process affinity
145              */
146             if (!SetProcessAffinityMask(hProcessAffinityHandle, dwProcessAffinityMask)) {
147                 GetLastErrorText(strErrorText, sizeof(strErrorText) / sizeof(WCHAR));
148                 EndDialog(hDlg, LOWORD(wParam));
149                 LoadStringW(hInst, IDS_MSG_ACCESSPROCESSAFF, szTitle, sizeof(szTitle) / sizeof(WCHAR));
150                 MessageBoxW(hMainWnd, strErrorText, szTitle, MB_OK|MB_ICONSTOP);
151             }
152 
153             EndDialog(hDlg, LOWORD(wParam));
154             return TRUE;
155         }
156 
157         break;
158     }
159 
160     return 0;
161 }
162