SNOW-652696 S3 FIPS support fix. (#1136)

* S3 FIPS support fix. This has been tested manually with a FIPS deployment.

src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java

@@ -930,10 +930,17 @@ static StageInfo getStageInfo(JsonNode jsonNode) throws SnowflakeSQLException {
           jsonNode.path("data").path("stageInfo").path("isClientSideEncrypted").asBoolean(true);
     }
 
-    // endPoint and storageAccount are only available in Azure stages. Value
-    // will be present but null in other platforms.
+    // endPoint is currently known to be set for Azure stages or S3. For S3 it will be set
+    // specifically
+    // for FIPS or VPCE S3 endpoint. SNOW-652696
     String endPoint = null;
+    if ("AZURE".equalsIgnoreCase(stageLocationType) || "S3".equalsIgnoreCase(stageLocationType)) {
+      endPoint = jsonNode.path("data").path("stageInfo").findValue("endPoint").asText();
+    }
+
     String stgAcct = null;
+    // storageAccount are only available in Azure stages. Value
+    // will be present but null in other platforms.
     if ("AZURE".equalsIgnoreCase(stageLocationType)) {
       // Jackson is doing some very strange things trying to pull the value of
       // the storageAccount node after adding the GCP library dependencies.
@@ -943,7 +950,6 @@ static StageInfo getStageInfo(JsonNode jsonNode) throws SnowflakeSQLException {
       // then comes back with double quotes around it, so we're stripping them
       // off. As long as our JSON doc doesn't add another node that starts with
       // "sto", this should work fine.
-      endPoint = jsonNode.path("data").path("stageInfo").findValue("endPoint").asText();
       Iterator<Entry<String, JsonNode>> fields = jsonNode.path("data").path("stageInfo").fields();
       while (fields.hasNext()) {
         Entry<String, JsonNode> jsonField = fields.next();

src/main/java/net/snowflake/client/jdbc/cloud/storage/SnowflakeS3Client.java

@@ -190,8 +190,11 @@ private void setupSnowflakeS3Client(
               .withClientConfiguration(clientConfig);
     }
 
-    if (stageRegion != null) {
-      Region region = RegionUtils.getRegion(stageRegion);
+    Region region = RegionUtils.getRegion(stageRegion);
+    if (this.stageEndPoint != null && this.stageEndPoint != "" && this.stageEndPoint != "null") {
+      amazonS3Builder.withEndpointConfiguration(
+          new AwsClientBuilder.EndpointConfiguration(this.stageEndPoint, region.getName()));
+    } else {
       if (region != null) {
         if (this.isUseS3RegionalUrl) {
           amazonS3Builder.withEndpointConfiguration(
@@ -204,13 +207,7 @@ private void setupSnowflakeS3Client(
     }
     // Explicitly force to use virtual address style
     amazonS3Builder.withPathStyleAccessEnabled(false);
-
     amazonClient = (AmazonS3) amazonS3Builder.build();
-    if (this.stageEndPoint != null && this.stageEndPoint != "") {
-      // Set the FIPS endpoint if we need it. GS will tell us if we do by
-      // giving us an endpoint to use if required and supported by the region.
-      amazonClient.setEndpoint(this.stageEndPoint);
-    }
   }
 
   // Returns the Max number of retry attempts

src/test/java/net/snowflake/client/jdbc/FileUploaderSessionlessTest.java

@@ -20,6 +20,67 @@ public class FileUploaderSessionlessTest {
   @Rule public TemporaryFolder folder = new TemporaryFolder();
   private ObjectMapper mapper = new ObjectMapper();
 
+  private final String exampleS3JsonStringWithStageEndpoint =
+      "{\n"
+          + "  \"data\": {\n"
+          + "    \"uploadInfo\": {\n"
+          + "      \"locationType\": \"S3\",\n"
+          + "      \"location\": \"example/location\",\n"
+          + "      \"path\": \"tables/19805757505/\",\n"
+          + "      \"region\": \"us-west-2\",\n"
+          + "      \"storageAccount\": null,\n"
+          + "      \"isClientSideEncrypted\": true,\n"
+          + "      \"creds\": {\n"
+          + "        \"AWS_KEY_ID\": \"EXAMPLE_AWS_KEY_ID\",\n"
+          + "        \"AWS_SECRET_KEY\": \"EXAMPLE_AWS_SECRET_KEY\",\n"
+          + "        \"AWS_TOKEN\": \"EXAMPLE_AWS_TOKEN\",\n"
+          + "        \"AWS_ID\": \"EXAMPLE_AWS_ID\",\n"
+          + "        \"AWS_KEY\": \"EXAMPLE_AWS_KEY\"\n"
+          + "      },\n"
+          + "      \"presignedUrl\": null,\n"
+          + "      \"endPoint\": null\n"
+          + "    },\n"
+          + "    \"src_locations\": [\n"
+          + "      \"/tmp/files/orders_100.csv\"\n"
+          + "    ],\n"
+          + "    \"parallel\": 4,\n"
+          + "    \"threshold\": 209715200,\n"
+          + "    \"autoCompress\": true,\n"
+          + "    \"overwrite\": false,\n"
+          + "    \"sourceCompression\": \"auto_detect\",\n"
+          + "    \"clientShowEncryptionParameter\": true,\n"
+          + "    \"queryId\": \"EXAMPLE_QUERY_ID\",\n"
+          + "    \"encryptionMaterial\": {\n"
+          + "      \"queryStageMasterKey\": \"EXAMPLE_QUERY_STAGE_MASTER_KEY\",\n"
+          + "      \"queryId\": \"EXAMPLE_QUERY_ID\",\n"
+          + "      \"smkId\": 123\n"
+          + "    },\n"
+          + "    \"stageInfo\": {\n"
+          + "      \"locationType\": \"S3\",\n"
+          + "      \"location\": \"stage/location/foo/\",\n"
+          + "      \"path\": \"tables/19805757505/\",\n"
+          + "      \"region\": \"us-west-2\",\n"
+          + "      \"storageAccount\": null,\n"
+          + "      \"isClientSideEncrypted\": true,\n"
+          + "      \"creds\": {\n"
+          + "        \"AWS_KEY_ID\": \"EXAMPLE_AWS_KEY_ID\",\n"
+          + "        \"AWS_SECRET_KEY\": \"EXAMPLE_AWS_SECRET_KEY\",\n"
+          + "        \"AWS_TOKEN\": \"EXAMPLE_AWS_TOKEN\",\n"
+          + "        \"AWS_ID\": \"EXAMPLE_AWS_ID\",\n"
+          + "        \"AWS_KEY\": \"EXAMPLE_AWS_KEY\"\n"
+          + "      },\n"
+          + "      \"presignedUrl\": null,\n"
+          + "      \"endPoint\": \"s3-fips.us-east-1.amazonaws.com\"\n"
+          + "    },\n"
+          + "    \"command\": \"UPLOAD\",\n"
+          + "    \"kind\": null,\n"
+          + "    \"operation\": \"Node\"\n"
+          + "  },\n"
+          + "  \"code\": null,\n"
+          + "  \"message\": null,\n"
+          + "  \"success\": true\n"
+          + "}";
+
   private final String exampleS3JsonString =
       "{\n"
           + "  \"data\": {\n"
@@ -184,13 +245,15 @@ public class FileUploaderSessionlessTest {
           + "}";
 
   protected JsonNode exampleS3JsonNode;
+  protected JsonNode exampleS3StageEndpointJsonNode;
   protected JsonNode exampleAzureJsonNode;
   private JsonNode exampleGCSJsonNode;
   private List<JsonNode> exampleNodes;
 
   @Before
   public void setup() throws Exception {
     exampleS3JsonNode = mapper.readTree(exampleS3JsonString);
+    exampleS3StageEndpointJsonNode = mapper.readTree(exampleS3JsonStringWithStageEndpoint);
     exampleAzureJsonNode = mapper.readTree(exampleAzureJsonString);
     exampleGCSJsonNode = mapper.readTree(exampleGCSJsonString);
     exampleNodes = Arrays.asList(exampleS3JsonNode, exampleAzureJsonNode, exampleGCSJsonNode);
@@ -246,7 +309,26 @@ public void testGetS3StageData() throws Exception {
     Assert.assertEquals("stage/location/foo/", stageInfo.getLocation());
     Assert.assertEquals(expectedCreds, stageInfo.getCredentials());
     Assert.assertEquals("us-west-2", stageInfo.getRegion());
-    Assert.assertEquals(null, stageInfo.getEndPoint());
+    Assert.assertEquals("null", stageInfo.getEndPoint());
+    Assert.assertEquals(null, stageInfo.getStorageAccount());
+    Assert.assertEquals(true, stageInfo.getIsClientSideEncrypted());
+  }
+
+  @Test
+  public void testGetS3StageDataWithStageEndpoint() throws Exception {
+    StageInfo stageInfo = SnowflakeFileTransferAgent.getStageInfo(exampleS3StageEndpointJsonNode);
+    Map<String, String> expectedCreds = new HashMap<>();
+    expectedCreds.put("AWS_ID", "EXAMPLE_AWS_ID");
+    expectedCreds.put("AWS_KEY", "EXAMPLE_AWS_KEY");
+    expectedCreds.put("AWS_KEY_ID", "EXAMPLE_AWS_KEY_ID");
+    expectedCreds.put("AWS_SECRET_KEY", "EXAMPLE_AWS_SECRET_KEY");
+    expectedCreds.put("AWS_TOKEN", "EXAMPLE_AWS_TOKEN");
+
+    Assert.assertEquals(StageInfo.StageType.S3, stageInfo.getStageType());
+    Assert.assertEquals("stage/location/foo/", stageInfo.getLocation());
+    Assert.assertEquals(expectedCreds, stageInfo.getCredentials());
+    Assert.assertEquals("us-west-2", stageInfo.getRegion());
+    Assert.assertEquals("s3-fips.us-east-1.amazonaws.com", stageInfo.getEndPoint());
     Assert.assertEquals(null, stageInfo.getStorageAccount());
     Assert.assertEquals(true, stageInfo.getIsClientSideEncrypted());
   }
@@ -303,7 +385,7 @@ public void testGetFileTransferMetadatasS3() throws Exception {
     Assert.assertEquals("stage/location/foo/", stageInfo.getLocation());
     Assert.assertEquals(expectedCreds, stageInfo.getCredentials());
     Assert.assertEquals("us-west-2", stageInfo.getRegion());
-    Assert.assertEquals(null, stageInfo.getEndPoint());
+    Assert.assertEquals("null", stageInfo.getEndPoint());
     Assert.assertEquals(null, stageInfo.getStorageAccount());
     Assert.assertEquals(true, stageInfo.getIsClientSideEncrypted());
 
@@ -347,7 +429,7 @@ public void testGetFileTransferMetadatasS3MissingEncryption() throws Exception {
     Assert.assertEquals("stage/location/foo/", stageInfo.getLocation());
     Assert.assertEquals(expectedCreds, stageInfo.getCredentials());
     Assert.assertEquals("us-west-2", stageInfo.getRegion());
-    Assert.assertEquals(null, stageInfo.getEndPoint());
+    Assert.assertEquals("null", stageInfo.getEndPoint());
     Assert.assertEquals(null, stageInfo.getStorageAccount());
     Assert.assertEquals(true, stageInfo.getIsClientSideEncrypted());