Branch data Line data Source code
1 : : /*-------------------------------------------------------------------------
2 : : *
3 : : * oauth_hook_client.c
4 : : * Test driver for t/002_client.pl, which verifies OAuth hook
5 : : * functionality in libpq.
6 : : *
7 : : * Portions Copyright (c) 1996-2026, PostgreSQL Global Development Group
8 : : * Portions Copyright (c) 1994, Regents of the University of California
9 : : *
10 : : *
11 : : * IDENTIFICATION
12 : : * src/test/modules/oauth_validator/oauth_hook_client.c
13 : : *
14 : : *-------------------------------------------------------------------------
15 : : */
16 : :
17 : : #include "postgres_fe.h"
18 : :
19 : : #include <sys/socket.h>
20 : :
21 : : #include "getopt_long.h"
22 : : #include "libpq-fe.h"
23 : :
24 : : static int handle_auth_data(PGauthData type, PGconn *conn, void *data);
25 : : static PostgresPollingStatusType async_cb(PGconn *conn,
26 : : PGoauthBearerRequest *req,
27 : : pgsocket *altsock);
28 : : static PostgresPollingStatusType misbehave_cb(PGconn *conn,
29 : : PGoauthBearerRequest *req,
30 : : pgsocket *altsock);
31 : :
32 : : static void
33 : 0 : usage(char *argv[])
34 : : {
35 : 0 : printf("usage: %s [flags] CONNINFO\n\n", argv[0]);
36 : :
37 : 0 : printf("recognized flags:\n");
38 : 0 : printf(" -h, --help show this message\n");
39 : 0 : printf(" -v VERSION select the hook API version (default 2)\n");
40 : 0 : printf(" --expected-scope SCOPE fail if received scopes do not match SCOPE\n");
41 : 0 : printf(" --expected-uri URI fail if received configuration link does not match URI\n");
42 : 0 : printf(" --expected-issuer ISS fail if received issuer does not match ISS (v2 only)\n");
43 : 0 : printf(" --misbehave=MODE have the hook fail required postconditions\n"
44 : : " (MODEs: no-hook, fail-async, no-token, no-socket)\n");
45 : 0 : printf(" --no-hook don't install OAuth hooks\n");
46 : 0 : printf(" --hang-forever don't ever return a token (combine with connect_timeout)\n");
47 : 0 : printf(" --token TOKEN use the provided TOKEN value\n");
48 : 0 : printf(" --error ERRMSG fail instead, with the given ERRMSG (v2 only)\n");
49 : 0 : printf(" --stress-async busy-loop on PQconnectPoll rather than polling\n");
50 : 0 : }
51 : :
52 : : /* --options */
53 : : static bool no_hook = false;
54 : : static bool hang_forever = false;
55 : : static bool stress_async = false;
56 : : static const char *expected_uri = NULL;
57 : : static const char *expected_issuer = NULL;
58 : : static const char *expected_scope = NULL;
59 : : static const char *misbehave_mode = NULL;
60 : : static char *token = NULL;
61 : : static char *errmsg = NULL;
62 : : static int hook_version = PQAUTHDATA_OAUTH_BEARER_TOKEN_V2;
63 : :
64 : : int
65 : 0 : main(int argc, char *argv[])
66 : : {
67 : : static const struct option long_options[] = {
68 : : {"help", no_argument, NULL, 'h'},
69 : :
70 : : {"expected-scope", required_argument, NULL, 1000},
71 : : {"expected-uri", required_argument, NULL, 1001},
72 : : {"no-hook", no_argument, NULL, 1002},
73 : : {"token", required_argument, NULL, 1003},
74 : : {"hang-forever", no_argument, NULL, 1004},
75 : : {"misbehave", required_argument, NULL, 1005},
76 : : {"stress-async", no_argument, NULL, 1006},
77 : : {"expected-issuer", required_argument, NULL, 1007},
78 : : {"error", required_argument, NULL, 1008},
79 : : {0}
80 : : };
81 : :
82 : : const char *conninfo;
83 : : PGconn *conn;
84 : : int c;
85 : :
86 [ # # ]: 0 : while ((c = getopt_long(argc, argv, "hv:", long_options, NULL)) != -1)
87 : : {
88 [ # # # # : 0 : switch (c)
# # # # #
# # # ]
89 : : {
90 : 0 : case 'h':
91 : 0 : usage(argv);
92 : 0 : return 0;
93 : :
94 : 0 : case 'v':
95 [ # # ]: 0 : if (strcmp(optarg, "1") == 0)
96 : 0 : hook_version = PQAUTHDATA_OAUTH_BEARER_TOKEN;
97 [ # # ]: 0 : else if (strcmp(optarg, "2") == 0)
98 : 0 : hook_version = PQAUTHDATA_OAUTH_BEARER_TOKEN_V2;
99 : : else
100 : : {
101 : 0 : usage(argv);
102 : 0 : return 1;
103 : : }
104 : 0 : break;
105 : :
106 : 0 : case 1000: /* --expected-scope */
107 : 0 : expected_scope = optarg;
108 : 0 : break;
109 : :
110 : 0 : case 1001: /* --expected-uri */
111 : 0 : expected_uri = optarg;
112 : 0 : break;
113 : :
114 : 0 : case 1002: /* --no-hook */
115 : 0 : no_hook = true;
116 : 0 : break;
117 : :
118 : 0 : case 1003: /* --token */
119 : 0 : token = optarg;
120 : 0 : break;
121 : :
122 : 0 : case 1004: /* --hang-forever */
123 : 0 : hang_forever = true;
124 : 0 : break;
125 : :
126 : 0 : case 1005: /* --misbehave */
127 : 0 : misbehave_mode = optarg;
128 : 0 : break;
129 : :
130 : 0 : case 1006: /* --stress-async */
131 : 0 : stress_async = true;
132 : 0 : break;
133 : :
134 : 0 : case 1007: /* --expected-issuer */
135 : 0 : expected_issuer = optarg;
136 : 0 : break;
137 : :
138 : 0 : case 1008: /* --error */
139 : 0 : errmsg = optarg;
140 : 0 : break;
141 : :
142 : 0 : default:
143 : 0 : usage(argv);
144 : 0 : return 1;
145 : : }
146 : : }
147 : :
148 [ # # ]: 0 : if (argc != optind + 1)
149 : : {
150 : 0 : usage(argv);
151 : 0 : return 1;
152 : : }
153 : :
154 : 0 : conninfo = argv[optind];
155 : :
156 : : /* Set up our OAuth hooks. */
157 : 0 : PQsetAuthDataHook(handle_auth_data);
158 : :
159 : : /* Connect. (All the actual work is in the hook.) */
160 [ # # ]: 0 : if (stress_async)
161 : : {
162 : : /*
163 : : * Perform an asynchronous connection, busy-looping on PQconnectPoll()
164 : : * without actually waiting on socket events. This stresses code paths
165 : : * that rely on asynchronous work to be done before continuing with
166 : : * the next step in the flow.
167 : : */
168 : : PostgresPollingStatusType res;
169 : :
170 : 0 : conn = PQconnectStart(conninfo);
171 : :
172 : : do
173 : : {
174 : 0 : res = PQconnectPoll(conn);
175 [ # # # # ]: 0 : } while (res != PGRES_POLLING_FAILED && res != PGRES_POLLING_OK);
176 : : }
177 : : else
178 : : {
179 : : /* Perform a standard synchronous connection. */
180 : 0 : conn = PQconnectdb(conninfo);
181 : : }
182 : :
183 [ # # ]: 0 : if (PQstatus(conn) != CONNECTION_OK)
184 : : {
185 : 0 : fprintf(stderr, "connection to database failed: %s\n",
186 : : PQerrorMessage(conn));
187 : 0 : PQfinish(conn);
188 : 0 : return 1;
189 : : }
190 : :
191 : 0 : printf("connection succeeded\n");
192 : 0 : PQfinish(conn);
193 : 0 : return 0;
194 : : }
195 : :
196 : : /*
197 : : * PQauthDataHook implementation. Replaces the default client flow by handling
198 : : * PQAUTHDATA_OAUTH_BEARER_TOKEN[_V2].
199 : : */
200 : : static int
201 : 0 : handle_auth_data(PGauthData type, PGconn *conn, void *data)
202 : : {
203 : : PGoauthBearerRequest *req;
204 : 0 : PGoauthBearerRequestV2 *req2 = NULL;
205 : :
206 : : Assert(hook_version == PQAUTHDATA_OAUTH_BEARER_TOKEN ||
207 : : hook_version == PQAUTHDATA_OAUTH_BEARER_TOKEN_V2);
208 : :
209 [ # # # # ]: 0 : if (no_hook || type != hook_version)
210 : 0 : return 0;
211 : :
212 : 0 : req = data;
213 [ # # ]: 0 : if (type == PQAUTHDATA_OAUTH_BEARER_TOKEN_V2)
214 : 0 : req2 = data;
215 : :
216 [ # # ]: 0 : if (hang_forever)
217 : : {
218 : : /* Start asynchronous processing. */
219 : 0 : req->async = async_cb;
220 : 0 : return 1;
221 : : }
222 : :
223 [ # # ]: 0 : if (misbehave_mode)
224 : : {
225 [ # # ]: 0 : if (strcmp(misbehave_mode, "no-hook") != 0)
226 : 0 : req->async = misbehave_cb;
227 : 0 : return 1;
228 : : }
229 : :
230 [ # # ]: 0 : if (expected_uri)
231 : : {
232 [ # # ]: 0 : if (!req->openid_configuration)
233 : : {
234 : 0 : fprintf(stderr, "expected URI \"%s\", got NULL\n", expected_uri);
235 : 0 : return -1;
236 : : }
237 : :
238 [ # # ]: 0 : if (strcmp(expected_uri, req->openid_configuration) != 0)
239 : : {
240 : 0 : fprintf(stderr, "expected URI \"%s\", got \"%s\"\n", expected_uri, req->openid_configuration);
241 : 0 : return -1;
242 : : }
243 : : }
244 : :
245 [ # # ]: 0 : if (expected_scope)
246 : : {
247 [ # # ]: 0 : if (!req->scope)
248 : : {
249 : 0 : fprintf(stderr, "expected scope \"%s\", got NULL\n", expected_scope);
250 : 0 : return -1;
251 : : }
252 : :
253 [ # # ]: 0 : if (strcmp(expected_scope, req->scope) != 0)
254 : : {
255 : 0 : fprintf(stderr, "expected scope \"%s\", got \"%s\"\n", expected_scope, req->scope);
256 : 0 : return -1;
257 : : }
258 : : }
259 : :
260 [ # # ]: 0 : if (expected_issuer)
261 : : {
262 [ # # ]: 0 : if (!req2)
263 : : {
264 : 0 : fprintf(stderr, "--expected-issuer cannot be combined with -v1\n");
265 : 0 : return -1;
266 : : }
267 : :
268 [ # # ]: 0 : if (!req2->issuer)
269 : : {
270 : 0 : fprintf(stderr, "expected issuer \"%s\", got NULL\n", expected_issuer);
271 : 0 : return -1;
272 : : }
273 : :
274 [ # # ]: 0 : if (strcmp(expected_issuer, req2->issuer) != 0)
275 : : {
276 : 0 : fprintf(stderr, "expected issuer \"%s\", got \"%s\"\n", expected_issuer, req2->issuer);
277 : 0 : return -1;
278 : : }
279 : : }
280 : :
281 [ # # ]: 0 : if (errmsg)
282 : : {
283 [ # # ]: 0 : if (token)
284 : : {
285 : 0 : fprintf(stderr, "--error cannot be combined with --token\n");
286 : 0 : return -1;
287 : : }
288 [ # # ]: 0 : else if (!req2)
289 : : {
290 : 0 : fprintf(stderr, "--error cannot be combined with -v1\n");
291 : 0 : return -1;
292 : : }
293 : :
294 : 0 : req2->error = errmsg;
295 : 0 : return -1;
296 : : }
297 : :
298 : 0 : req->token = token;
299 : 0 : return 1;
300 : : }
301 : :
302 : : static PostgresPollingStatusType
303 : 0 : async_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock)
304 : : {
305 [ # # ]: 0 : if (hang_forever)
306 : : {
307 : : /*
308 : : * This code tests that nothing is interfering with libpq's handling
309 : : * of connect_timeout.
310 : : */
311 : : static pgsocket sock = PGINVALID_SOCKET;
312 : :
313 [ # # ]: 0 : if (sock == PGINVALID_SOCKET)
314 : : {
315 : : /* First call. Create an unbound socket to wait on. */
316 : : #ifdef WIN32
317 : : WSADATA wsaData;
318 : : int err;
319 : :
320 : : err = WSAStartup(MAKEWORD(2, 2), &wsaData);
321 : : if (err)
322 : : {
323 : : perror("WSAStartup failed");
324 : : return PGRES_POLLING_FAILED;
325 : : }
326 : : #endif
327 : 0 : sock = socket(AF_INET, SOCK_DGRAM, 0);
328 [ # # ]: 0 : if (sock == PGINVALID_SOCKET)
329 : : {
330 : 0 : perror("failed to create datagram socket");
331 : 0 : return PGRES_POLLING_FAILED;
332 : : }
333 : : }
334 : :
335 : : /* Make libpq wait on the (unreadable) socket. */
336 : 0 : *altsock = sock;
337 : 0 : return PGRES_POLLING_READING;
338 : : }
339 : :
340 : 0 : req->token = token;
341 : 0 : return PGRES_POLLING_OK;
342 : : }
343 : :
344 : : static PostgresPollingStatusType
345 : 0 : misbehave_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock)
346 : : {
347 [ # # ]: 0 : if (strcmp(misbehave_mode, "fail-async") == 0)
348 : : {
349 : : /* Just fail "normally". */
350 [ # # ]: 0 : if (errmsg)
351 : : {
352 : : PGoauthBearerRequestV2 *req2;
353 : :
354 [ # # ]: 0 : if (hook_version == PQAUTHDATA_OAUTH_BEARER_TOKEN)
355 : : {
356 : 0 : fprintf(stderr, "--error cannot be combined with -v1\n");
357 : 0 : exit(1);
358 : : }
359 : :
360 : 0 : req2 = (PGoauthBearerRequestV2 *) req;
361 : 0 : req2->error = errmsg;
362 : : }
363 : :
364 : 0 : return PGRES_POLLING_FAILED;
365 : : }
366 [ # # ]: 0 : else if (strcmp(misbehave_mode, "no-token") == 0)
367 : : {
368 : : /* Callbacks must assign req->token before returning OK. */
369 : 0 : return PGRES_POLLING_OK;
370 : : }
371 [ # # ]: 0 : else if (strcmp(misbehave_mode, "no-socket") == 0)
372 : : {
373 : : /* Callbacks must assign *altsock before asking for polling. */
374 : 0 : return PGRES_POLLING_READING;
375 : : }
376 : : else
377 : : {
378 : 0 : fprintf(stderr, "unrecognized --misbehave mode: %s\n", misbehave_mode);
379 : 0 : exit(1);
380 : : }
381 : : }
|
