1 /*------------------------------------------------------------------------
2  *
3  * Query cancellation support for frontend code
4  *
5  * Assorted utility functions to control query cancellation with signal
6  * handler for SIGINT.
7  *
8  *
9  * Portions Copyright (c) 1996-2021, PostgreSQL Global Development Group
10  * Portions Copyright (c) 1994, Regents of the University of California
11  *
12  * src/fe-utils/cancel.c
13  *
14  *------------------------------------------------------------------------
15  */
16 
17 #include "postgres_fe.h"
18 
19 #include <unistd.h>
20 
21 #include "common/connect.h"
22 #include "fe_utils/cancel.h"
23 #include "fe_utils/string_utils.h"
24 
25 
26 /*
27  * Write a simple string to stderr --- must be safe in a signal handler.
28  * We ignore the write() result since there's not much we could do about it.
29  * Certain compilers make that harder than it ought to be.
30  */
31 #define write_stderr(str) \
32 	do { \
33 		const char *str_ = (str); \
34 		int		rc_; \
35 		rc_ = write(fileno(stderr), str_, strlen(str_)); \
36 		(void) rc_; \
37 	} while (0)
38 
39 /*
40  * Contains all the information needed to cancel a query issued from
41  * a database connection to the backend.
42  */
43 static PGcancel *volatile cancelConn = NULL;
44 
45 /*
46  * CancelRequested is set when we receive SIGINT (or local equivalent).
47  * There is no provision in this module for resetting it; but applications
48  * might choose to clear it after successfully recovering from a cancel.
49  * Note that there is no guarantee that we successfully sent a Cancel request,
50  * or that the request will have any effect if we did send it.
51  */
52 volatile sig_atomic_t CancelRequested = false;
53 
54 #ifdef WIN32
55 static CRITICAL_SECTION cancelConnLock;
56 #endif
57 
58 /*
59  * Additional callback for cancellations.
60  */
61 static void (*cancel_callback) (void) = NULL;
62 
63 
64 /*
65  * SetCancelConn
66  *
67  * Set cancelConn to point to the current database connection.
68  */
69 void
SetCancelConn(PGconn * conn)70 SetCancelConn(PGconn *conn)
71 {
72 	PGcancel   *oldCancelConn;
73 
74 #ifdef WIN32
75 	EnterCriticalSection(&cancelConnLock);
76 #endif
77 
78 	/* Free the old one if we have one */
79 	oldCancelConn = cancelConn;
80 
81 	/* be sure handle_sigint doesn't use pointer while freeing */
82 	cancelConn = NULL;
83 
84 	if (oldCancelConn != NULL)
85 		PQfreeCancel(oldCancelConn);
86 
87 	cancelConn = PQgetCancel(conn);
88 
89 #ifdef WIN32
90 	LeaveCriticalSection(&cancelConnLock);
91 #endif
92 }
93 
94 /*
95  * ResetCancelConn
96  *
97  * Free the current cancel connection, if any, and set to NULL.
98  */
99 void
ResetCancelConn(void)100 ResetCancelConn(void)
101 {
102 	PGcancel   *oldCancelConn;
103 
104 #ifdef WIN32
105 	EnterCriticalSection(&cancelConnLock);
106 #endif
107 
108 	oldCancelConn = cancelConn;
109 
110 	/* be sure handle_sigint doesn't use pointer while freeing */
111 	cancelConn = NULL;
112 
113 	if (oldCancelConn != NULL)
114 		PQfreeCancel(oldCancelConn);
115 
116 #ifdef WIN32
117 	LeaveCriticalSection(&cancelConnLock);
118 #endif
119 }
120 
121 
122 /*
123  * Code to support query cancellation
124  *
125  * Note that sending the cancel directly from the signal handler is safe
126  * because PQcancel() is written to make it so.  We use write() to report
127  * to stderr because it's better to use simple facilities in a signal
128  * handler.
129  *
130  * On Windows, the signal canceling happens on a separate thread, because
131  * that's how SetConsoleCtrlHandler works.  The PQcancel function is safe
132  * for this (unlike PQrequestCancel).  However, a CRITICAL_SECTION is required
133  * to protect the PGcancel structure against being changed while the signal
134  * thread is using it.
135  */
136 
137 #ifndef WIN32
138 
139 /*
140  * handle_sigint
141  *
142  * Handle interrupt signals by canceling the current command, if cancelConn
143  * is set.
144  */
145 static void
handle_sigint(SIGNAL_ARGS)146 handle_sigint(SIGNAL_ARGS)
147 {
148 	int			save_errno = errno;
149 	char		errbuf[256];
150 
151 	CancelRequested = true;
152 
153 	if (cancel_callback != NULL)
154 		cancel_callback();
155 
156 	/* Send QueryCancel if we are processing a database query */
157 	if (cancelConn != NULL)
158 	{
159 		if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
160 		{
161 			write_stderr(_("Cancel request sent\n"));
162 		}
163 		else
164 		{
165 			write_stderr(_("Could not send cancel request: "));
166 			write_stderr(errbuf);
167 		}
168 	}
169 
170 	errno = save_errno;			/* just in case the write changed it */
171 }
172 
173 /*
174  * setup_cancel_handler
175  *
176  * Register query cancellation callback for SIGINT.
177  */
178 void
setup_cancel_handler(void (* callback)(void))179 setup_cancel_handler(void (*callback) (void))
180 {
181 	cancel_callback = callback;
182 	pqsignal(SIGINT, handle_sigint);
183 }
184 
185 #else							/* WIN32 */
186 
187 static BOOL WINAPI
consoleHandler(DWORD dwCtrlType)188 consoleHandler(DWORD dwCtrlType)
189 {
190 	char		errbuf[256];
191 
192 	if (dwCtrlType == CTRL_C_EVENT ||
193 		dwCtrlType == CTRL_BREAK_EVENT)
194 	{
195 		CancelRequested = true;
196 
197 		if (cancel_callback != NULL)
198 			cancel_callback();
199 
200 		/* Send QueryCancel if we are processing a database query */
201 		EnterCriticalSection(&cancelConnLock);
202 		if (cancelConn != NULL)
203 		{
204 			if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
205 			{
206 				write_stderr(_("Cancel request sent\n"));
207 			}
208 			else
209 			{
210 				write_stderr(_("Could not send cancel request: "));
211 				write_stderr(errbuf);
212 			}
213 		}
214 
215 		LeaveCriticalSection(&cancelConnLock);
216 
217 		return TRUE;
218 	}
219 	else
220 		/* Return FALSE for any signals not being handled */
221 		return FALSE;
222 }
223 
224 void
setup_cancel_handler(void (* callback)(void))225 setup_cancel_handler(void (*callback) (void))
226 {
227 	cancel_callback = callback;
228 
229 	InitializeCriticalSection(&cancelConnLock);
230 
231 	SetConsoleCtrlHandler(consoleHandler, TRUE);
232 }
233 
234 #endif							/* WIN32 */
235