|
23 | 23 | import java.sql.ResultSet; |
24 | 24 | import java.util.HashMap; |
25 | 25 | import java.util.Map; |
| 26 | +import java.util.stream.IntStream; |
26 | 27 | import net.snowflake.client.AbstractDriverIT; |
27 | 28 | import net.snowflake.client.api.datasource.SnowflakeDataSource; |
28 | 29 | import net.snowflake.client.api.datasource.SnowflakeDataSourceFactory; |
@@ -260,6 +261,116 @@ public void testInvalidSSOUrl() { |
260 | 261 | assertTrue(ex.getMessage().contains("Invalid SSOUrl found")); |
261 | 262 | } |
262 | 263 |
|
| 264 | + /** |
| 265 | + * SNOW-3704231: the localhost callback server used to assume a single socket read returned the |
| 266 | + * whole HTTP request. Browsers routinely open speculative/preconnect sockets to the callback port |
| 267 | + * and close them without sending any data, so {@code in.read(buf)} returns -1 (EOF) and {@code |
| 268 | + * new String(buf, 0, -1)} threw StringIndexOutOfBoundsException before the real request was ever |
| 269 | + * processed. The server must ignore such empty connections and keep listening. |
| 270 | + * |
| 271 | + * @throws Throwable if any error occurs |
| 272 | + */ |
| 273 | + @Test |
| 274 | + public void testAuthenticateIgnoresEmptyPreconnectSocket() throws Throwable { |
| 275 | + final SFLoginInput loginInput = initMockLoginInput(); |
| 276 | + |
| 277 | + try (MockedStatic<HttpUtil> mockedHttpUtil = mockStatic(HttpUtil.class)) { |
| 278 | + mockSsoUrlResponse(mockedHttpUtil); |
| 279 | + |
| 280 | + // First connection sends no data (browser preconnect closed immediately) -> read() == -1. |
| 281 | + // Second connection carries the real token. |
| 282 | + Socket emptySocket = mockSocket(new ByteArrayInputStream(new byte[0])); |
| 283 | + Socket tokenSocket = |
| 284 | + mockSocket( |
| 285 | + toStream( |
| 286 | + String.format( |
| 287 | + "GET /?token=%s&confirm=1 HTTP/1.1\r\n" + "USER-AGENT: snowflake client", |
| 288 | + FakeSessionUtilExternalBrowser.MOCK_SAML_TOKEN))); |
| 289 | + |
| 290 | + SessionUtilExternalBrowser sub = |
| 291 | + new SequencedFakeSessionUtilExternalBrowser( |
| 292 | + loginInput, mockServerSocket(emptySocket, tokenSocket)); |
| 293 | + sub.authenticate(); |
| 294 | + assertEquals(FakeSessionUtilExternalBrowser.MOCK_SAML_TOKEN, sub.getToken()); |
| 295 | + } |
| 296 | + } |
| 297 | + |
| 298 | + /** |
| 299 | + * SNOW-3704231: large SSO requests (big tokens plus Sec-Fetch and Sec-GPC headers) fill the |
| 300 | + * request line and headers with several kilobytes of data. This confirms the callback server |
| 301 | + * still parses the token out of such a large request. |
| 302 | + * |
| 303 | + * @throws Throwable if any error occurs |
| 304 | + */ |
| 305 | + @Test |
| 306 | + public void testAuthenticateHandlesLargeRequest() throws Throwable { |
| 307 | + final SFLoginInput loginInput = initMockLoginInput(); |
| 308 | + |
| 309 | + try (MockedStatic<HttpUtil> mockedHttpUtil = mockStatic(HttpUtil.class)) { |
| 310 | + mockSsoUrlResponse(mockedHttpUtil); |
| 311 | + |
| 312 | + // A large token (mimicking a real SSO token) plus large privacy headers, ~4 KB total. |
| 313 | + String largeToken = IntStream.range(0, 3000).mapToObj(i -> "A").reduce("", String::concat); |
| 314 | + String request = |
| 315 | + String.format( |
| 316 | + "GET /?token=%s&confirm=1 HTTP/1.1\r\n" |
| 317 | + + "USER-AGENT: snowflake client\r\n" |
| 318 | + + "Sec-Fetch-Site: none\r\n" |
| 319 | + + "Sec-GPC: 1\r\n" |
| 320 | + + "\r\n", |
| 321 | + largeToken); |
| 322 | + |
| 323 | + Socket socket = mockSocket(toStream(request)); |
| 324 | + |
| 325 | + SessionUtilExternalBrowser sub = |
| 326 | + new SequencedFakeSessionUtilExternalBrowser(loginInput, mockServerSocket(socket)); |
| 327 | + sub.authenticate(); |
| 328 | + assertEquals(largeToken, sub.getToken()); |
| 329 | + } |
| 330 | + } |
| 331 | + |
| 332 | + private static void mockSsoUrlResponse(MockedStatic<HttpUtil> mockedHttpUtil) throws Exception { |
| 333 | + mockedHttpUtil |
| 334 | + .when( |
| 335 | + () -> |
| 336 | + HttpUtil.executeGeneralRequestWithContext( |
| 337 | + Mockito.any(HttpRequestBase.class), |
| 338 | + Mockito.anyInt(), |
| 339 | + Mockito.anyInt(), |
| 340 | + Mockito.anyInt(), |
| 341 | + Mockito.anyInt(), |
| 342 | + Mockito.anyInt(), |
| 343 | + Mockito.nullable(HttpClientSettingsKey.class), |
| 344 | + Mockito.nullable(SFBaseSession.class))) |
| 345 | + .thenReturn( |
| 346 | + new HttpResponseWithHeaders( |
| 347 | + "{\"success\":\"true\",\"data\":{\"proofKey\":\"" |
| 348 | + + MOCK_PROOF_KEY |
| 349 | + + "\", \"ssoUrl\":\"" |
| 350 | + + MOCK_SSO_URL |
| 351 | + + "\"}}", |
| 352 | + new HashMap<>())); |
| 353 | + } |
| 354 | + |
| 355 | + private static InputStream toStream(String data) { |
| 356 | + return new ByteArrayInputStream(data.getBytes(StandardCharsets.UTF_8)); |
| 357 | + } |
| 358 | + |
| 359 | + private static Socket mockSocket(InputStream in) throws IOException { |
| 360 | + Socket socket = mock(Socket.class); |
| 361 | + when(socket.getInputStream()).thenReturn(in); |
| 362 | + when(socket.getOutputStream()) |
| 363 | + .thenReturn(new FakeSessionUtilExternalBrowser.NullOutputStream()); |
| 364 | + return socket; |
| 365 | + } |
| 366 | + |
| 367 | + private static ServerSocket mockServerSocket(Socket first, Socket... rest) throws IOException { |
| 368 | + ServerSocket serverSocket = mock(ServerSocket.class); |
| 369 | + when(serverSocket.getLocalPort()).thenReturn(12345); |
| 370 | + when(serverSocket.accept()).thenReturn(first, rest); |
| 371 | + return serverSocket; |
| 372 | + } |
| 373 | + |
263 | 374 | /** |
264 | 375 | * Mock HttpUtil and SFLoginInput |
265 | 376 | * |
@@ -319,3 +430,22 @@ public void testExternalBrowserTimeout() throws Exception { |
319 | 430 | assertTrue(e.getMessage().contains("External browser authentication failed")); |
320 | 431 | } |
321 | 432 | } |
| 433 | + |
| 434 | +/** |
| 435 | + * SessionUtilExternalBrowser whose ServerSocket is supplied by the test, so that the sequence of |
| 436 | + * accepted connections (e.g. an empty preconnect socket followed by the real request) can be |
| 437 | + * controlled. |
| 438 | + */ |
| 439 | +class SequencedFakeSessionUtilExternalBrowser extends SessionUtilExternalBrowser { |
| 440 | + private final ServerSocket serverSocket; |
| 441 | + |
| 442 | + SequencedFakeSessionUtilExternalBrowser(SFLoginInput loginInput, ServerSocket serverSocket) { |
| 443 | + super(loginInput, new MockAuthExternalBrowserHandlers()); |
| 444 | + this.serverSocket = serverSocket; |
| 445 | + } |
| 446 | + |
| 447 | + @Override |
| 448 | + protected ServerSocket getServerSocket() { |
| 449 | + return serverSocket; |
| 450 | + } |
| 451 | +} |
0 commit comments