Age Owner Branch data TLA 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-2025, 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
198 dgustafsson@postgres 33 :UBC 0 : usage(char *argv[])
34 : : {
35 : 0 : printf("usage: %s [flags] CONNINFO\n\n", argv[0]);
36 : :
37 : 0 : printf("recognized flags:\n");
125 peter@eisentraut.org 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");
198 dgustafsson@postgres 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
198 dgustafsson@postgres 59 :CBC 7 : 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 : : const char *conninfo;
75 : : PGconn *conn;
76 : : int c;
77 : :
78 [ + + ]: 17 : while ((c = getopt_long(argc, argv, "h", long_options, NULL)) != -1)
79 : : {
80 [ - + + + : 10 : switch (c)
+ + + +
- ]
81 : : {
198 dgustafsson@postgres 82 :UBC 0 : case 'h':
83 : 0 : usage(argv);
84 : 0 : return 0;
85 : :
198 dgustafsson@postgres 86 :CBC 1 : case 1000: /* --expected-scope */
87 : 1 : expected_scope = optarg;
88 : 1 : break;
89 : :
90 : 1 : case 1001: /* --expected-uri */
91 : 1 : expected_uri = optarg;
92 : 1 : break;
93 : :
94 : 1 : case 1002: /* --no-hook */
95 : 1 : no_hook = true;
96 : 1 : break;
97 : :
98 : 1 : case 1003: /* --token */
99 : 1 : token = optarg;
100 : 1 : break;
101 : :
102 : 1 : case 1004: /* --hang-forever */
103 : 1 : hang_forever = true;
104 : 1 : break;
105 : :
106 : 4 : case 1005: /* --misbehave */
107 : 4 : misbehave_mode = optarg;
108 : 4 : break;
109 : :
110 : 1 : case 1006: /* --stress-async */
111 : 1 : stress_async = true;
112 : 1 : break;
113 : :
198 dgustafsson@postgres 114 :UBC 0 : default:
115 : 0 : usage(argv);
116 : 0 : return 1;
117 : : }
118 : : }
119 : :
198 dgustafsson@postgres 120 [ - + ]:CBC 7 : if (argc != optind + 1)
121 : : {
198 dgustafsson@postgres 122 :UBC 0 : usage(argv);
123 : 0 : return 1;
124 : : }
125 : :
198 dgustafsson@postgres 126 :CBC 7 : conninfo = argv[optind];
127 : :
128 : : /* Set up our OAuth hooks. */
129 : 7 : PQsetAuthDataHook(handle_auth_data);
130 : :
131 : : /* Connect. (All the actual work is in the hook.) */
132 [ + + ]: 7 : 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 : : PostgresPollingStatusType res;
141 : :
142 : 1 : conn = PQconnectStart(conninfo);
143 : :
144 : : do
145 : : {
146 : 517162 : res = PQconnectPoll(conn);
147 [ + - + + ]: 517162 : } while (res != PGRES_POLLING_FAILED && res != PGRES_POLLING_OK);
148 : : }
149 : : else
150 : : {
151 : : /* Perform a standard synchronous connection. */
152 : 6 : conn = PQconnectdb(conninfo);
153 : : }
154 : :
155 [ + + ]: 7 : if (PQstatus(conn) != CONNECTION_OK)
156 : : {
157 : 5 : fprintf(stderr, "connection to database failed: %s\n",
158 : : PQerrorMessage(conn));
159 : 5 : PQfinish(conn);
160 : 5 : return 1;
161 : : }
162 : :
163 : 2 : printf("connection succeeded\n");
164 : 2 : PQfinish(conn);
165 : 2 : return 0;
166 : : }
167 : :
168 : : /*
169 : : * PQauthDataHook implementation. Replaces the default client flow by handling
170 : : * PQAUTHDATA_OAUTH_BEARER_TOKEN.
171 : : */
172 : : static int
173 : 8 : handle_auth_data(PGauthData type, PGconn *conn, void *data)
174 : : {
175 : 8 : PGoauthBearerRequest *req = data;
176 : :
177 [ + + - + ]: 8 : if (no_hook || (type != PQAUTHDATA_OAUTH_BEARER_TOKEN))
178 : 2 : return 0;
179 : :
180 [ + + ]: 6 : if (hang_forever)
181 : : {
182 : : /* Start asynchronous processing. */
183 : 1 : req->async = async_cb;
184 : 1 : return 1;
185 : : }
186 : :
187 [ + + ]: 5 : if (misbehave_mode)
188 : : {
189 [ + + ]: 4 : if (strcmp(misbehave_mode, "no-hook") != 0)
190 : 3 : req->async = misbehave_cb;
191 : 4 : return 1;
192 : : }
193 : :
194 [ + - ]: 1 : if (expected_uri)
195 : : {
196 [ - + ]: 1 : if (!req->openid_configuration)
197 : : {
198 dgustafsson@postgres 198 :UBC 0 : fprintf(stderr, "expected URI \"%s\", got NULL\n", expected_uri);
199 : 0 : return -1;
200 : : }
201 : :
198 dgustafsson@postgres 202 [ - + ]:CBC 1 : if (strcmp(expected_uri, req->openid_configuration) != 0)
203 : : {
198 dgustafsson@postgres 204 :UBC 0 : fprintf(stderr, "expected URI \"%s\", got \"%s\"\n", expected_uri, req->openid_configuration);
205 : 0 : return -1;
206 : : }
207 : : }
208 : :
198 dgustafsson@postgres 209 [ + - ]:CBC 1 : if (expected_scope)
210 : : {
211 [ - + ]: 1 : if (!req->scope)
212 : : {
198 dgustafsson@postgres 213 :UBC 0 : fprintf(stderr, "expected scope \"%s\", got NULL\n", expected_scope);
214 : 0 : return -1;
215 : : }
216 : :
198 dgustafsson@postgres 217 [ - + ]:CBC 1 : if (strcmp(expected_scope, req->scope) != 0)
218 : : {
198 dgustafsson@postgres 219 :UBC 0 : fprintf(stderr, "expected scope \"%s\", got \"%s\"\n", expected_scope, req->scope);
220 : 0 : return -1;
221 : : }
222 : : }
223 : :
198 dgustafsson@postgres 224 :CBC 1 : req->token = token;
225 : 1 : return 1;
226 : : }
227 : :
228 : : static PostgresPollingStatusType
229 : 1 : async_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock)
230 : : {
231 [ + - ]: 1 : 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 [ + - ]: 1 : 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 : 1 : sock = socket(AF_INET, SOCK_DGRAM, 0);
254 [ - + ]: 1 : if (sock == PGINVALID_SOCKET)
255 : : {
198 dgustafsson@postgres 256 :UBC 0 : perror("failed to create datagram socket");
257 : 0 : return PGRES_POLLING_FAILED;
258 : : }
259 : : }
260 : :
261 : : /* Make libpq wait on the (unreadable) socket. */
198 dgustafsson@postgres 262 :CBC 1 : *altsock = sock;
263 : 1 : return PGRES_POLLING_READING;
264 : : }
265 : :
198 dgustafsson@postgres 266 :UBC 0 : req->token = token;
267 : 0 : return PGRES_POLLING_OK;
268 : : }
269 : :
270 : : static PostgresPollingStatusType
198 dgustafsson@postgres 271 :CBC 3 : misbehave_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock)
272 : : {
273 [ + + ]: 3 : if (strcmp(misbehave_mode, "fail-async") == 0)
274 : : {
275 : : /* Just fail "normally". */
276 : 1 : return PGRES_POLLING_FAILED;
277 : : }
278 [ + + ]: 2 : else if (strcmp(misbehave_mode, "no-token") == 0)
279 : : {
280 : : /* Callbacks must assign req->token before returning OK. */
281 : 1 : return PGRES_POLLING_OK;
282 : : }
283 [ + - ]: 1 : else if (strcmp(misbehave_mode, "no-socket") == 0)
284 : : {
285 : : /* Callbacks must assign *altsock before asking for polling. */
286 : 1 : return PGRES_POLLING_READING;
287 : : }
288 : : else
289 : : {
198 dgustafsson@postgres 290 :UBC 0 : fprintf(stderr, "unrecognized --misbehave mode: %s\n", misbehave_mode);
291 : 0 : exit(1);
292 : : }
293 : : }
|