|
35 | 35 | import java.net.InetAddress; |
36 | 36 | import java.net.ServerSocket; |
37 | 37 | import java.nio.charset.StandardCharsets; |
| 38 | +import java.util.Base64; |
38 | 39 | import java.sql.ResultSet; |
39 | 40 | import java.sql.SQLException; |
40 | 41 | import java.sql.Statement; |
@@ -137,9 +138,17 @@ public void cleanUpIdpEnv() { |
137 | 138 | idpContainer.stop(); |
138 | 139 | idpContainer = null; |
139 | 140 | } |
140 | | - if (miniHS2 != null) { |
| 141 | + if (miniHS2 != null && miniHS2.isStarted()) { |
141 | 142 | miniHS2.stop(); |
142 | 143 | } |
| 144 | + HiveSamlAuthTokenGenerator.shutdown(); |
| 145 | + } |
| 146 | + |
| 147 | + private static ISAMLAuthTokenGenerator createTokenGenerator(String tokenTtl) { |
| 148 | + HiveSamlAuthTokenGenerator.shutdown(); |
| 149 | + HiveConf conf = new HiveConf(); |
| 150 | + conf.setVar(ConfVars.HIVE_SERVER2_SAML_CALLBACK_TOKEN_TTL, tokenTtl); |
| 151 | + return HiveSamlAuthTokenGenerator.get(conf); |
143 | 152 | } |
144 | 153 |
|
145 | 154 | private void setupIDP(boolean useSignedAssertions, String authMode) throws Exception { |
@@ -554,6 +563,86 @@ public void testTokenReuse() throws Exception { |
554 | 563 | } |
555 | 564 | } |
556 | 565 |
|
| 566 | + @Test |
| 567 | + public void testValidTokenRoundTrip() throws Exception { |
| 568 | + ISAMLAuthTokenGenerator tokenGenerator = createTokenGenerator("30s"); |
| 569 | + String token = tokenGenerator.get("alice", "relay-state-1"); |
| 570 | + assertEquals("alice", tokenGenerator.validate(token)); |
| 571 | + } |
| 572 | + |
| 573 | + @Test |
| 574 | + public void testForgedSignatureRejected() throws Exception { |
| 575 | + ISAMLAuthTokenGenerator tokenGenerator = createTokenGenerator("30s"); |
| 576 | + String forgedPayload = "u=alice;id=1337;time=" + System.currentTimeMillis() |
| 577 | + + ";rs=deadbeef;sg=bogus"; |
| 578 | + try { |
| 579 | + String forgedToken = Base64.getEncoder().encodeToString(forgedPayload.getBytes(StandardCharsets.UTF_8)); |
| 580 | + tokenGenerator.validate(forgedToken); |
| 581 | + fail("Expected forged token to be rejected"); |
| 582 | + } catch (HttpSamlAuthenticationException e) { |
| 583 | + assertEquals("Token could not be verified", e.getMessage()); |
| 584 | + } |
| 585 | + } |
| 586 | + |
| 587 | + @Test |
| 588 | + public void testInvalidTokenRejected() throws Exception { |
| 589 | + ISAMLAuthTokenGenerator tokenGenerator = createTokenGenerator("30s"); |
| 590 | + try { |
| 591 | + tokenGenerator.validate("notAValidToken"); |
| 592 | + fail("Expected malformed base64 token to be rejected"); |
| 593 | + } catch (HttpSamlAuthenticationException e) { |
| 594 | + assertEquals("Invalid token", e.getMessage()); |
| 595 | + } |
| 596 | + String invalidStructure = Base64.getEncoder().encodeToString("foo".getBytes()); |
| 597 | + try { |
| 598 | + tokenGenerator.validate(invalidStructure); |
| 599 | + fail("Expected invalid token structure to be rejected"); |
| 600 | + } catch (HttpSamlAuthenticationException e) { |
| 601 | + assertEquals("Invalid token", e.getMessage()); |
| 602 | + } |
| 603 | + } |
| 604 | + |
| 605 | + @Test |
| 606 | + public void testExpiredTokenRejected() throws Exception { |
| 607 | + ISAMLAuthTokenGenerator tokenGenerator = createTokenGenerator("1s"); |
| 608 | + String token = tokenGenerator.get("alice", "relay-state-1"); |
| 609 | + Thread.sleep(1100); |
| 610 | + try { |
| 611 | + tokenGenerator.validate(token); |
| 612 | + fail("Expected expired token to be rejected"); |
| 613 | + } catch (HttpSamlAuthenticationException e) { |
| 614 | + assertEquals("Token is expired", e.getMessage()); |
| 615 | + } |
| 616 | + } |
| 617 | + |
| 618 | + @Test |
| 619 | + public void testParseHandlesBase64PaddingInSignature() { |
| 620 | + Map<String, String> kv = new HashMap<>(); |
| 621 | + String token = "u=alice;id=1;time=1000;rs=rs1;sg=YWJjZA=="; |
| 622 | + assertTrue(HiveSamlAuthTokenGenerator.parse(token, kv)); |
| 623 | + assertEquals("alice", kv.get("u")); |
| 624 | + assertEquals("YWJjZA==", kv.get("sg")); |
| 625 | + } |
| 626 | + |
| 627 | + @Test |
| 628 | + public void testParseRejectsEncodedBearerToken() { |
| 629 | + Map<String, String> kv = new HashMap<>(); |
| 630 | + String encoded = Base64.getEncoder().encodeToString( |
| 631 | + "u=alice;id=1;time=1000;rs=rs1;sg=abc".getBytes()); |
| 632 | + assertFalse(HiveSamlAuthTokenGenerator.parse(encoded, kv)); |
| 633 | + } |
| 634 | + |
| 635 | + @Test |
| 636 | + public void testParseDecodedTokenFromGenerator() throws Exception { |
| 637 | + ISAMLAuthTokenGenerator tokenGenerator = createTokenGenerator("30s"); |
| 638 | + String encoded = tokenGenerator.get("bob", "relay-42"); |
| 639 | + String decoded = new String(Base64.getDecoder().decode(encoded), StandardCharsets.UTF_8); |
| 640 | + Map<String, String> kv = new HashMap<>(); |
| 641 | + assertTrue(HiveSamlAuthTokenGenerator.parse(decoded, kv)); |
| 642 | + assertEquals("bob", kv.get("u")); |
| 643 | + assertEquals("relay-42", kv.get(HiveSamlAuthTokenGenerator.RELAY_STATE)); |
| 644 | + } |
| 645 | + |
557 | 646 | private static void assertLoggedInUser(HiveConnection connection, String expectedUser) |
558 | 647 | throws SQLException { |
559 | 648 | Statement stmt = connection.createStatement(); |
|
0 commit comments