Merge pull request #629 from metacret/ISSUE-628

AWSCredentialsProvider for s3-extentions

s3-extensions/pom.xml

@@ -41,6 +41,10 @@
             <groupId>net.java.dev.jets3t</groupId>
             <artifactId>jets3t</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-java-sdk</artifactId>
+        </dependency>
         <!-- override httpclient / httpcore version from jets3t -->
         <dependency>
             <groupId>org.apache.httpcomponents</groupId>
@@ -65,6 +69,11 @@
             <artifactId>junit</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.easymock</groupId>
+            <artifactId>easymock</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
     <build>
         <plugins>

s3-extensions/src/main/java/io/druid/storage/s3/AWSCredentialsConfig.java

@@ -31,6 +31,9 @@ public class AWSCredentialsConfig
   @JsonProperty
   private String secretKey = "";
 
+  @JsonProperty
+  private String fileSessionCredentials = "";
+
   public String getAccessKey()
   {
     return accessKey;
@@ -40,4 +43,6 @@ public class AWSCredentialsConfig
   {
     return secretKey;
   }
+
+  public String getFileSessionCredentials() { return fileSessionCredentials; }
 }

s3-extensions/src/main/java/io/druid/storage/s3/AWSSessionCredentialsAdapter.java

+/*
+ * Druid - a distributed column store.
+ * Copyright (C) 2012, 2013  Metamarkets Group Inc.
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+
+package io.druid.storage.s3;
+
+import com.amazonaws.auth.AWSCredentialsProvider;
+import org.jets3t.service.security.AWSSessionCredentials;
+
+public class AWSSessionCredentialsAdapter extends AWSSessionCredentials {
+  private final AWSCredentialsProvider provider;
+
+  public AWSSessionCredentialsAdapter(AWSCredentialsProvider provider) {
+    super(null, null, null);
+    if(provider.getCredentials() instanceof com.amazonaws.auth.AWSSessionCredentials)
+      this.provider = provider;
+    else
+      throw new IllegalArgumentException("provider does not contain session credentials");
+  }
+
+  @Override
+  protected String getTypeName() {
+    return "AWSSessionCredentialsAdapter";
+  }
+
+  @Override
+  public String getVersionPrefix() {
+    return "AWSSessionCredentialsAdapter, version: ";
+  }
+
+  @Override
+  public String getAccessKey() {
+    return provider.getCredentials().getAWSAccessKeyId();
+  }
+
+  @Override
+  public String getSecretKey() {
+    return provider.getCredentials().getAWSSecretKey();
+  }
+
+  public String getSessionToken() {
+    com.amazonaws.auth.AWSSessionCredentials sessionCredentials =
+        (com.amazonaws.auth.AWSSessionCredentials) provider.getCredentials();
+    return sessionCredentials.getSessionToken();
+  }
+}

s3-extensions/src/main/java/io/druid/storage/s3/FileSessionCredentialsProvider.java

+/*
+ * Druid - a distributed column store.
+ * Copyright (C) 2012, 2013  Metamarkets Group Inc.
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+
+package io.druid.storage.s3;
+
+import com.amazonaws.auth.AWSCredentials;
+import com.amazonaws.auth.AWSCredentialsProvider;
+import com.amazonaws.auth.AWSSessionCredentials;
+import com.google.common.base.Charsets;
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Properties;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+public class FileSessionCredentialsProvider implements AWSCredentialsProvider {
+  private final String sessionCredentials;
+  private volatile String sessionToken;
+  private volatile String accessKey;
+  private volatile String secretKey;
+
+  private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(
+      new ThreadFactoryBuilder().setNameFormat("FileSessionCredentialsProviderRefresh-%d")
+          .setDaemon(true).build()
+  );
+
+  public FileSessionCredentialsProvider(String sessionCredentials) {
+    this.sessionCredentials = sessionCredentials;
+    refresh();
+
+    scheduler.scheduleAtFixedRate(new Runnable() {
+      @Override
+      public void run() {
+        refresh();
+      }
+    }, 1, 1, TimeUnit.HOURS); // refresh every hour
+  }
+
+  @Override
+  public AWSCredentials getCredentials() {
+    return new AWSSessionCredentials() {
+      @Override
+      public String getSessionToken() {
+        return sessionToken;
+      }
+
+      @Override
+      public String getAWSAccessKeyId() {
+        return accessKey;
+      }
+
+      @Override
+      public String getAWSSecretKey() {
+        return secretKey;
+      }
+    };
+  }
+
+  @Override
+  public void refresh() {
+    try {
+      Properties props = new Properties();
+      InputStream is = new FileInputStream(new File(sessionCredentials));
+      props.load(is);
+      is.close();
+
+      sessionToken = props.getProperty("sessionToken");
+      accessKey = props.getProperty("accessKey");
+      secretKey = props.getProperty("secretKey");
+    } catch (IOException e) {
+      throw new RuntimeException("cannot refresh AWS credentials", e);
+    }
+  }
+}

s3-extensions/src/main/java/io/druid/storage/s3/S3StorageDruidModule.java

@@ -19,16 +19,16 @@
 
 package io.druid.storage.s3;
 
+import com.amazonaws.auth.AWSCredentialsProvider;
 import com.fasterxml.jackson.databind.Module;
+import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.inject.Binder;
 import com.google.inject.Provides;
-import com.google.inject.ProvisionException;
 import io.druid.guice.Binders;
 import io.druid.guice.JsonConfigProvider;
 import io.druid.guice.LazySingleton;
 import io.druid.initialization.DruidModule;
-import org.jets3t.service.S3ServiceException;
 import org.jets3t.service.impl.rest.httpclient.RestS3Service;
 import org.jets3t.service.security.AWSCredentials;
 
@@ -64,15 +64,44 @@ public class S3StorageDruidModule implements DruidModule
 
   @Provides
   @LazySingleton
-  public AWSCredentials getJets3tAWSCredentials(AWSCredentialsConfig config)
+  public AWSCredentialsProvider getAWSCredentialsProvider(final AWSCredentialsConfig config)
   {
-    return new AWSCredentials(config.getAccessKey(), config.getSecretKey());
+    if (!Strings.isNullOrEmpty(config.getAccessKey()) && !Strings.isNullOrEmpty(config.getSecretKey())) {
+      return new AWSCredentialsProvider() {
+        @Override
+        public com.amazonaws.auth.AWSCredentials getCredentials() {
+          return new com.amazonaws.auth.AWSCredentials() {
+            @Override
+            public String getAWSAccessKeyId() {
+              return config.getAccessKey();
+            }
+
+            @Override
+            public String getAWSSecretKey() {
+              return config.getSecretKey();
+            }
+          };
+        }
+
+        @Override
+        public void refresh() {}
+      };
+    } else {
+      return new FileSessionCredentialsProvider(config.getFileSessionCredentials());
+    }
   }
 
   @Provides
   @LazySingleton
-  public RestS3Service getRestS3Service(AWSCredentials credentials)
+  public RestS3Service getRestS3Service(AWSCredentialsProvider provider)
   {
-      return new RestS3Service(credentials);
+    if(provider.getCredentials() instanceof com.amazonaws.auth.AWSSessionCredentials) {
+      return new RestS3Service(new AWSSessionCredentialsAdapter(provider));
+    } else {
+      return new RestS3Service(new AWSCredentials(
+          provider.getCredentials().getAWSAccessKeyId(),
+          provider.getCredentials().getAWSSecretKey()
+      ));
+    }
   }
 }

s3-extensions/src/test/java/io/druid/storage/s3/TestAWSCredentialsProvider.java

+/*
+ * Druid - a distributed column store.
+ * Copyright (C) 2012, 2013  Metamarkets Group Inc.
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+
+package io.druid.storage.s3;
+
+import com.amazonaws.auth.AWSCredentials;
+import com.amazonaws.auth.AWSCredentialsProvider;
+import com.amazonaws.auth.AWSSessionCredentials;
+import org.easymock.EasyMock;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.PrintWriter;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class TestAWSCredentialsProvider {
+  @Test
+  public void testWithFixedAWSKeys() {
+    S3StorageDruidModule module = new S3StorageDruidModule();
+
+    AWSCredentialsConfig config = EasyMock.createMock(AWSCredentialsConfig.class);
+    EasyMock.expect(config.getAccessKey()).andReturn("accessKeySample").atLeastOnce();
+    EasyMock.expect(config.getSecretKey()).andReturn("secretKeySample").atLeastOnce();
+    EasyMock.replay(config);
+
+    AWSCredentialsProvider provider = module.getAWSCredentialsProvider(config);
+    AWSCredentials credentials = provider.getCredentials();
+    assertEquals(credentials.getAWSAccessKeyId(), "accessKeySample");
+    assertEquals(credentials.getAWSSecretKey(), "secretKeySample");
+
+    // try to create
+    module.getRestS3Service(provider);
+  }
+
+  @Rule
+  public TemporaryFolder folder = new TemporaryFolder();
+
+  @Test
+  public void testWithFileSessionCredentials() throws IOException {
+    S3StorageDruidModule module = new S3StorageDruidModule();
+
+    AWSCredentialsConfig config = EasyMock.createMock(AWSCredentialsConfig.class);
+    EasyMock.expect(config.getAccessKey()).andReturn("");
+    EasyMock.expect(config.getSecretKey()).andReturn("");
+    File file = folder.newFile();
+    PrintWriter out = new PrintWriter(file.getAbsolutePath());
+    out.println("sessionToken=sessionTokenSample\nsecretKey=secretKeySample\naccessKey=accessKeySample");
+    out.close();
+    EasyMock.expect(config.getFileSessionCredentials()).andReturn(file.getAbsolutePath()).atLeastOnce();
+    EasyMock.replay(config);
+
+    AWSCredentialsProvider provider = module.getAWSCredentialsProvider(config);
+    AWSCredentials credentials = provider.getCredentials();
+    assertTrue(credentials instanceof AWSSessionCredentials);
+    AWSSessionCredentials sessionCredentials = (AWSSessionCredentials) credentials;
+    assertEquals(sessionCredentials.getAWSAccessKeyId(), "accessKeySample");
+    assertEquals(sessionCredentials.getAWSSecretKey(), "secretKeySample");
+    assertEquals(sessionCredentials.getSessionToken(), "sessionTokenSample");
+
+    // try to create
+    module.getRestS3Service(provider);
+  }
+}

s3-extensions/src/test/java/io/druid/storage/s3/TestFileSessionCredentialsProvider.java

+/*
+ * Druid - a distributed column store.
+ * Copyright (C) 2012, 2013  Metamarkets Group Inc.
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+
+package io.druid.storage.s3;
+
+import com.amazonaws.auth.AWSSessionCredentials;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.PrintWriter;
+
+import static org.junit.Assert.assertEquals;
+
+public class TestFileSessionCredentialsProvider {
+  @Rule
+  public TemporaryFolder folder = new TemporaryFolder();
+
+  @Test
+  public void test() throws IOException {
+    File file = folder.newFile();
+    PrintWriter out = new PrintWriter(file.getAbsolutePath());
+    out.println("sessionToken=sessionTokenSample\nsecretKey=secretKeySample\naccessKey=accessKeySample");
+    out.close();
+
+    FileSessionCredentialsProvider provider = new FileSessionCredentialsProvider(file.getAbsolutePath());
+    AWSSessionCredentials sessionCredentials = (AWSSessionCredentials) provider.getCredentials();
+    assertEquals(sessionCredentials.getSessionToken(), "sessionTokenSample");
+    assertEquals(sessionCredentials.getAWSAccessKeyId(), "accessKeySample");
+    assertEquals(sessionCredentials.getAWSSecretKey(), "secretKeySample");
+  }
+}