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(" --expected-scope SCOPE fail if received scopes do not match SCOPE\n");
40 0 : printf(" --expected-uri URI fail if received configuration link does not match URI\n");
41 0 : printf(" --misbehave=MODE have the hook fail required postconditions\n"
42 : " (MODEs: no-hook, fail-async, no-token, no-socket)\n");
43 0 : printf(" --no-hook don't install OAuth hooks\n");
44 0 : printf(" --hang-forever don't ever return a token (combine with connect_timeout)\n");
45 0 : printf(" --token TOKEN use the provided TOKEN value\n");
46 0 : printf(" --stress-async busy-loop on PQconnectPoll rather than polling\n");
47 0 : }
48 :
49 : /* --options */
50 : static bool no_hook = false;
51 : static bool hang_forever = false;
52 : static bool stress_async = false;
53 : static const char *expected_uri = NULL;
54 : static const char *expected_scope = NULL;
55 : static const char *misbehave_mode = NULL;
56 : static char *token = NULL;
57 :
58 : int
59 0 : main(int argc, char *argv[])
60 : {
61 : static const struct option long_options[] = {
62 : {"help", no_argument, NULL, 'h'},
63 :
64 : {"expected-scope", required_argument, NULL, 1000},
65 : {"expected-uri", required_argument, NULL, 1001},
66 : {"no-hook", no_argument, NULL, 1002},
67 : {"token", required_argument, NULL, 1003},
68 : {"hang-forever", no_argument, NULL, 1004},
69 : {"misbehave", required_argument, NULL, 1005},
70 : {"stress-async", no_argument, NULL, 1006},
71 : {0}
72 : };
73 :
74 0 : const char *conninfo;
75 0 : PGconn *conn;
76 0 : int c;
77 :
78 0 : while ((c = getopt_long(argc, argv, "h", long_options, NULL)) != -1)
79 : {
80 0 : switch (c)
81 : {
82 : case 'h':
83 0 : usage(argv);
84 0 : return 0;
85 :
86 : case 1000: /* --expected-scope */
87 0 : expected_scope = optarg;
88 0 : break;
89 :
90 : case 1001: /* --expected-uri */
91 0 : expected_uri = optarg;
92 0 : break;
93 :
94 : case 1002: /* --no-hook */
95 0 : no_hook = true;
96 0 : break;
97 :
98 : case 1003: /* --token */
99 0 : token = optarg;
100 0 : break;
101 :
102 : case 1004: /* --hang-forever */
103 0 : hang_forever = true;
104 0 : break;
105 :
106 : case 1005: /* --misbehave */
107 0 : misbehave_mode = optarg;
108 0 : break;
109 :
110 : case 1006: /* --stress-async */
111 0 : stress_async = true;
112 0 : break;
113 :
114 : default:
115 0 : usage(argv);
116 0 : return 1;
117 : }
118 : }
119 :
120 0 : if (argc != optind + 1)
121 : {
122 0 : usage(argv);
123 0 : return 1;
124 : }
125 :
126 0 : conninfo = argv[optind];
127 :
128 : /* Set up our OAuth hooks. */
129 0 : PQsetAuthDataHook(handle_auth_data);
130 :
131 : /* Connect. (All the actual work is in the hook.) */
132 0 : if (stress_async)
133 : {
134 : /*
135 : * Perform an asynchronous connection, busy-looping on PQconnectPoll()
136 : * without actually waiting on socket events. This stresses code paths
137 : * that rely on asynchronous work to be done before continuing with
138 : * the next step in the flow.
139 : */
140 0 : PostgresPollingStatusType res;
141 :
142 0 : conn = PQconnectStart(conninfo);
143 :
144 0 : do
145 : {
146 0 : res = PQconnectPoll(conn);
147 0 : } while (res != PGRES_POLLING_FAILED && res != PGRES_POLLING_OK);
148 0 : }
149 : else
150 : {
151 : /* Perform a standard synchronous connection. */
152 0 : conn = PQconnectdb(conninfo);
153 : }
154 :
155 0 : if (PQstatus(conn) != CONNECTION_OK)
156 : {
157 0 : fprintf(stderr, "connection to database failed: %s\n",
158 0 : PQerrorMessage(conn));
159 0 : PQfinish(conn);
160 0 : return 1;
161 : }
162 :
163 0 : printf("connection succeeded\n");
164 0 : PQfinish(conn);
165 0 : return 0;
166 0 : }
167 :
168 : /*
169 : * PQauthDataHook implementation. Replaces the default client flow by handling
170 : * PQAUTHDATA_OAUTH_BEARER_TOKEN.
171 : */
172 : static int
173 0 : handle_auth_data(PGauthData type, PGconn *conn, void *data)
174 : {
175 0 : PGoauthBearerRequest *req = data;
176 :
177 0 : if (no_hook || (type != PQAUTHDATA_OAUTH_BEARER_TOKEN))
178 0 : return 0;
179 :
180 0 : if (hang_forever)
181 : {
182 : /* Start asynchronous processing. */
183 0 : req->async = async_cb;
184 0 : return 1;
185 : }
186 :
187 0 : if (misbehave_mode)
188 : {
189 0 : if (strcmp(misbehave_mode, "no-hook") != 0)
190 0 : req->async = misbehave_cb;
191 0 : return 1;
192 : }
193 :
194 0 : if (expected_uri)
195 : {
196 0 : if (!req->openid_configuration)
197 : {
198 0 : fprintf(stderr, "expected URI \"%s\", got NULL\n", expected_uri);
199 0 : return -1;
200 : }
201 :
202 0 : if (strcmp(expected_uri, req->openid_configuration) != 0)
203 : {
204 0 : fprintf(stderr, "expected URI \"%s\", got \"%s\"\n", expected_uri, req->openid_configuration);
205 0 : return -1;
206 : }
207 0 : }
208 :
209 0 : if (expected_scope)
210 : {
211 0 : if (!req->scope)
212 : {
213 0 : fprintf(stderr, "expected scope \"%s\", got NULL\n", expected_scope);
214 0 : return -1;
215 : }
216 :
217 0 : if (strcmp(expected_scope, req->scope) != 0)
218 : {
219 0 : fprintf(stderr, "expected scope \"%s\", got \"%s\"\n", expected_scope, req->scope);
220 0 : return -1;
221 : }
222 0 : }
223 :
224 0 : req->token = token;
225 0 : return 1;
226 0 : }
227 :
228 : static PostgresPollingStatusType
229 0 : async_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock)
230 : {
231 0 : if (hang_forever)
232 : {
233 : /*
234 : * This code tests that nothing is interfering with libpq's handling
235 : * of connect_timeout.
236 : */
237 : static pgsocket sock = PGINVALID_SOCKET;
238 :
239 0 : if (sock == PGINVALID_SOCKET)
240 : {
241 : /* First call. Create an unbound socket to wait on. */
242 : #ifdef WIN32
243 : WSADATA wsaData;
244 : int err;
245 :
246 : err = WSAStartup(MAKEWORD(2, 2), &wsaData);
247 : if (err)
248 : {
249 : perror("WSAStartup failed");
250 : return PGRES_POLLING_FAILED;
251 : }
252 : #endif
253 0 : sock = socket(AF_INET, SOCK_DGRAM, 0);
254 0 : if (sock == PGINVALID_SOCKET)
255 : {
256 0 : perror("failed to create datagram socket");
257 0 : return PGRES_POLLING_FAILED;
258 : }
259 0 : }
260 :
261 : /* Make libpq wait on the (unreadable) socket. */
262 0 : *altsock = sock;
263 0 : return PGRES_POLLING_READING;
264 : }
265 :
266 0 : req->token = token;
267 0 : return PGRES_POLLING_OK;
268 0 : }
269 :
270 : static PostgresPollingStatusType
271 0 : misbehave_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock)
272 : {
273 0 : if (strcmp(misbehave_mode, "fail-async") == 0)
274 : {
275 : /* Just fail "normally". */
276 0 : return PGRES_POLLING_FAILED;
277 : }
278 0 : else if (strcmp(misbehave_mode, "no-token") == 0)
279 : {
280 : /* Callbacks must assign req->token before returning OK. */
281 0 : return PGRES_POLLING_OK;
282 : }
283 0 : else if (strcmp(misbehave_mode, "no-socket") == 0)
284 : {
285 : /* Callbacks must assign *altsock before asking for polling. */
286 0 : return PGRES_POLLING_READING;
287 : }
288 : else
289 : {
290 0 : fprintf(stderr, "unrecognized --misbehave mode: %s\n", misbehave_mode);
291 0 : exit(1);
292 : }
293 0 : }
|
