Merge pull request #161 from ing-bank/feature/sts_cache

add cache to sts call

build.sbt

@@ -44,6 +44,8 @@ libraryDependencies ++= Seq(
     "io.dropwizard.metrics"        % "metrics-core"            % metricVersion,
 //    "io.dropwizard.metrics"        % "metrics-jmx"             % metricVersion, // bring back after persistence update
     "com.auth0"                    % "java-jwt"                % "3.9.0",
+    "com.github.cb372"             %% "scalacache-core"        % "0.28.0",
+    "com.github.cb372"             %% "scalacache-caffeine"    % "0.28.0",
     "com.typesafe.akka"            %% "akka-testkit"           % akkaVersion       % Test,
     "com.typesafe.akka"            %% "akka-http-testkit"      % akkaHttpVersion   % Test,
     "org.scalatest"                %% "scalatest"              % "3.2.3"           % "it,test",

src/it/scala/com/ing/wbaa/rokku/proxy/provider/AuthenticationCachedProviderSTSItTest.scala

@@ -0,0 +1,98 @@
+package com.ing.wbaa.rokku.proxy.provider
+
+import akka.actor.ActorSystem
+import akka.stream.ActorMaterializer
+import com.amazonaws.services.securitytoken.model.{AssumeRoleRequest, GetSessionTokenRequest}
+import com.ing.wbaa.rokku.proxy.config.StsSettings
+import com.ing.wbaa.rokku.proxy.data._
+import com.ing.wbaa.testkit.awssdk.StsSdkHelpers
+import com.ing.wbaa.testkit.oauth.OAuth2TokenRequest
+import org.scalatest.Assertion
+import org.scalatest.diagrams.Diagrams
+import org.scalatest.wordspec.AsyncWordSpec
+
+import scala.concurrent.{ExecutionContext, Future}
+
+class AuthenticationCachedProviderSTSItTest extends AsyncWordSpec with Diagrams
+  with AuthenticationCachedProviderSTS
+  with StsSdkHelpers
+  with OAuth2TokenRequest {
+  override implicit val testSystem: ActorSystem = ActorSystem.create("test-system")
+  override implicit val system: ActorSystem = testSystem
+  override implicit val executionContext: ExecutionContext = testSystem.dispatcher
+  override implicit val materializer: ActorMaterializer = ActorMaterializer()(testSystem)
+
+  override val stsSettings: StsSettings = StsSettings(testSystem)
+
+  implicit val requestId: RequestId = RequestId("test")
+
+  private val validKeycloakCredentials = Map(
+    "grant_type" -> "password",
+    "username" -> "testuser",
+    "password" -> "password",
+    "client_id" -> "sts-rokku"
+  )
+  private val userOneKeycloakCredentials = Map(
+    "grant_type" -> "password",
+    "username" -> "userone",
+    "password" -> "password",
+    "client_id" -> "sts-rokku"
+  )
+
+  def withAwsCredentialsValidInSTS(testCode: AwsRequestCredential => Future[Assertion]): Future[Assertion] = {
+    val stsSdk = getAmazonSTSSdk(StsSettings(testSystem).stsBaseUri)
+    retrieveKeycloackToken(validKeycloakCredentials).flatMap { keycloakToken =>
+      val cred = stsSdk.getSessionToken(new GetSessionTokenRequest()
+        .withTokenCode(keycloakToken.access_token))
+        .getCredentials
+
+      testCode(AwsRequestCredential(AwsAccessKey(cred.getAccessKeyId), Some(AwsSessionToken(cred.getSessionToken))))
+    }
+  }
+
+  def withAssumeRoleInSTS(testCode: AwsRequestCredential => Future[Assertion]): Future[Assertion] = {
+    val stsSdk = getAmazonSTSSdk(StsSettings(testSystem).stsBaseUri)
+    retrieveKeycloackToken(userOneKeycloakCredentials).flatMap { keycloakToken =>
+      val assumeRoleReq = new AssumeRoleRequest().withTokenCode(keycloakToken.access_token)
+      assumeRoleReq.setRoleArn("arn:aws:iam::account-id:role/admin")
+      assumeRoleReq.setRoleSessionName("testRole")
+      val cred = stsSdk.assumeRole(assumeRoleReq).getCredentials
+
+      testCode(AwsRequestCredential(AwsAccessKey(cred.getAccessKeyId), Some(AwsSessionToken(cred.getSessionToken))))
+    }
+  }
+
+  "Authentication Provider STS" should {
+    "check authentication" that {
+      "succeeds for valid credentials" in {
+        withAwsCredentialsValidInSTS { awsCredential =>
+          areCredentialsActive(awsCredential).map { userResult =>
+            assert(userResult.map(_.userName).contains(UserName("testuser")))
+            assert(userResult.map(_.userGroups).head.contains(UserGroup("testgroup")))
+            assert(userResult.map(_.userGroups).head.contains(UserGroup("group3")))
+            assert(userResult.map(_.userGroups).head.size == 2)
+            assert(userResult.exists(_.accessKey.value.length == 32))
+            assert(userResult.exists(_.secretKey.value.length == 32))
+          }
+        }
+      }
+
+      "fail when user is not authenticated" in {
+        areCredentialsActive(AwsRequestCredential(AwsAccessKey("notauthenticated"), Some(AwsSessionToken("okSessionToken")))).map { userResult =>
+          assert(userResult.isEmpty)
+        }
+      }
+
+      "succeeds for valid role" in {
+        withAssumeRoleInSTS { awsCredential =>
+          areCredentialsActive(awsCredential).map { roleResult =>
+            assert(roleResult.map(_.userRole).contains(UserAssumeRole("admin")))
+            assert(roleResult.map(_.userGroups).contains(Set()))
+            assert(roleResult.exists(_.accessKey.value.length == 32))
+            assert(roleResult.exists(_.secretKey.value.length == 32))
+          }
+        }
+      }
+    }
+  }
+}

src/main/resources/application.conf

@@ -35,6 +35,7 @@ rokku {
     sts {
         uri = ${?ROKKU_STS_URI}
         encodeSecret = ${?ROKKU_STS_ENCODE_SECRET}
+        cache.ttlInSeconds = ${?ROKKU_STS_CACHE_TTL_IN_SECONDS}
     }
 
     atlas {

src/main/resources/reference.conf

@@ -44,6 +44,7 @@ rokku {
     sts {
         uri = "http://127.0.0.1:12345"
         encodeSecret = "jwtprivatekey"
+        cache.ttlInSeconds = 5
     }
 
     atlas {

src/main/scala/com/ing/wbaa/rokku/proxy/Server.scala

@@ -11,7 +11,7 @@ import com.typesafe.config.ConfigFactory
 
 object Server extends App {
 
-  new RokkuS3Proxy with AuthorizationProviderRanger with RequestHandlerS3 with AuthenticationProviderSTS with LineageProviderAtlas with SignatureProviderAws with KerberosLoginProvider with FilterRecursiveListBucketHandler with MessageProviderKafka with AuditLogProvider with MemoryUserRequestQueue with RequestParser {
+  new RokkuS3Proxy with AuthorizationProviderRanger with RequestHandlerS3 with AuthenticationCachedProviderSTS with LineageProviderAtlas with SignatureProviderAws with KerberosLoginProvider with FilterRecursiveListBucketHandler with MessageProviderKafka with AuditLogProvider with MemoryUserRequestQueue with RequestParser {
 
     override implicit lazy val system: ActorSystem = ActorSystem.create("rokku")
 

src/main/scala/com/ing/wbaa/rokku/proxy/config/StsSettings.scala

@@ -8,6 +8,8 @@ class StsSettings(config: Config) extends Extension {
   private val stsUri: String = config.getString("rokku.sts.uri")
   val stsBaseUri: Uri = Uri(stsUri)
   val encodeSecret: String = config.getString("rokku.sts.encodeSecret")
+  val cacheTTLInSeconds: Int = config.getInt("rokku.sts.cache.ttlInSeconds")
+
 }
 
 object StsSettings extends ExtensionId[StsSettings] with ExtensionIdProvider {

src/main/scala/com/ing/wbaa/rokku/proxy/provider/AuthenticationCachedProviderSTS.scala

@@ -0,0 +1,23 @@
+package com.ing.wbaa.rokku.proxy.provider
+
+import com.github.benmanes.caffeine.cache.Caffeine
+import com.ing.wbaa.rokku.proxy.data.{ AwsRequestCredential, RequestId, User }
+import scalacache._
+import scalacache.caffeine.CaffeineCache
+import scalacache.modes.scalaFuture._
+
+import scala.concurrent.Future
+import scala.concurrent.duration.DurationInt
+
+trait AuthenticationCachedProviderSTS extends AuthenticationProviderSTS {
+
+  private val stsCacheConfig =
+    Caffeine.newBuilder().
+      maximumSize(10000).
+      build[String, Entry[Future[Option[User]]]]
+  private implicit val stsCache: Cache[Future[Option[User]]] = CaffeineCache(stsCacheConfig)
+
+  override protected[this] def areCredentialsActive(awsRequestCredential: AwsRequestCredential)(implicit id: RequestId): Future[Option[User]] = {
+    caching(keyParts = awsRequestCredential)(ttl = Some(stsSettings.cacheTTLInSeconds.second))(super.areCredentialsActive(awsRequestCredential)).flatten
+  }
+}

src/main/scala/com/ing/wbaa/rokku/proxy/provider/AuthenticationProviderSTS.scala

@@ -27,6 +27,7 @@ trait AuthenticationProviderSTS extends JsonProtocols with JwtToken {
   protected[this] def stsSettings: StsSettings
 
   protected[this] def areCredentialsActive(awsRequestCredential: AwsRequestCredential)(implicit id: RequestId): Future[Option[User]] = {
+    logger.debug("no cache - areCredentialsActive invoked")
     val QueryParameters =
       Map("accessKey" -> awsRequestCredential.accessKey.value) ++
         awsRequestCredential.sessionToken.map(s => "sessionToken" -> s.value)