Add KMS context header support (ref bluelabsio#16)

akka-http-aws/src/main/scala/com/bluelabs/akkaaws/AwsHeaders.scala

@@ -2,14 +2,18 @@ package com.bluelabs.akkaaws
 
 import akka.http.scaladsl.model.headers.{ModeledCustomHeaderCompanion, ModeledCustomHeader}
 
+import scala.util.parsing.json.JSON
 import scala.util.{Failure, Success, Try}
-
+import scala.util.parsing.combinator._
 
 object AwsHeaders {
 
   sealed abstract class ServerSideEncryptionAlgorithm(val name: String)
+
   object ServerSideEncryptionAlgorithm {
+
     case object AES256 extends ServerSideEncryptionAlgorithm("AES256")
+
     case object KMS extends ServerSideEncryptionAlgorithm("aws:kms")
 
     def fromString(raw: String): Try[ServerSideEncryptionAlgorithm] = raw match {
@@ -21,6 +25,7 @@ object AwsHeaders {
 
   object `X-Amz-Server-Side-Encryption` extends ModeledCustomHeaderCompanion[`X-Amz-Server-Side-Encryption`] {
     override def name: String = "X-Amz-Server-Side-Encryption"
+
     override def parse(value: String): Try[`X-Amz-Server-Side-Encryption`] =
       ServerSideEncryptionAlgorithm.fromString(value).map(new `X-Amz-Server-Side-Encryption`(_))
   }
@@ -38,6 +43,7 @@ object AwsHeaders {
 
   object `X-Amz-Server-Side-Encryption-Aws-Kms-Key-Id` extends ModeledCustomHeaderCompanion[`X-Amz-Server-Side-Encryption-Aws-Kms-Key-Id`] {
     override def name: String = "X-Amz-Server-Side-Encryption-Aws-Kms-Key-Id"
+
     override def parse(value: String): Try[`X-Amz-Server-Side-Encryption-Aws-Kms-Key-Id`] =
       Success(new `X-Amz-Server-Side-Encryption-Aws-Kms-Key-Id`(value))
   }
@@ -53,6 +59,31 @@ object AwsHeaders {
     override def renderInRequests(): Boolean = true
   }
 
-  // TODO add `x-amz-server-side-encryption-context` header.
+  object `X-Amz-Server-Side-Encryption-Context` extends ModeledCustomHeaderCompanion[`X-Amz-Server-Side-Encryption-Context`] {
+    override def name: String = "X-Amz-Server-Side-Encryption-Context"
+
+    override def parse(value: String): Try[`X-Amz-Server-Side-Encryption-Context`] =
+      JSON.parseFull(value) match {
+        case Some(context: Map[String, Any]) if context.forall(_._2.isInstanceOf[String]) =>
+          Success(new `X-Amz-Server-Side-Encryption-Context`(context.asInstanceOf[Map[String, String]]))
+        case _ =>
+          Failure(new IllegalArgumentException("$value is not a valid AWS KMS context"))
+      }
+  }
+
+  final case class `X-Amz-Server-Side-Encryption-Context`(context: Map[String, String]) extends ModeledCustomHeader[`X-Amz-Server-Side-Encryption-Context`] {
+
+    override def companion: ModeledCustomHeaderCompanion[`X-Amz-Server-Side-Encryption-Context`] = `X-Amz-Server-Side-Encryption-Context`
+
+    override def value(): String =
+      context.map { case (key, value) => s""""$key":"$value"""" }.mkString("{", ",", "}")
+
+    override def renderInResponses(): Boolean = true
+
+    override def renderInRequests(): Boolean = true
+  }
+
 
 }
+
+

akka-http-aws/src/test/scala/com/bluelabs/akkaaws/AwsHeadersSpec.scala

@@ -2,7 +2,7 @@ package com.bluelabs.akkaaws
 
 import akka.http.scaladsl.model.headers.RawHeader
 import com.bluelabs.akkaaws.AwsHeaders.ServerSideEncryptionAlgorithm.AES256
-import com.bluelabs.akkaaws.AwsHeaders.{`X-Amz-Server-Side-Encryption-Aws-Kms-Key-Id`, `X-Amz-Server-Side-Encryption`, ServerSideEncryptionAlgorithm}
+import com.bluelabs.akkaaws.AwsHeaders.{`X-Amz-Server-Side-Encryption-Context`, `X-Amz-Server-Side-Encryption-Aws-Kms-Key-Id`, `X-Amz-Server-Side-Encryption`, ServerSideEncryptionAlgorithm}
 import org.scalatest.{FlatSpec, Matchers}
 
 import scala.util.{Failure, Success}
@@ -44,4 +44,27 @@ class AwsHeadersSpec extends FlatSpec with Matchers {
     value shouldBe "myId"
   }
 
+  "`X-Amz-Server-Side-Encryption-Context`" should "parse context" in {
+    val expectedContext = Map("foo"->"bar", "foo2"->"bar2")
+    val `X-Amz-Server-Side-Encryption-Context`(context) = `X-Amz-Server-Side-Encryption-Context`(expectedContext)
+
+    context shouldBe expectedContext
+  }
+
+  it should "set the X-Amz-Server-Side-Encryption-Context header" in {
+    val RawHeader(key, value) = `X-Amz-Server-Side-Encryption-Context`(Map("foo"->"bar", "foo2"->"bar2"))
+    key shouldBe "X-Amz-Server-Side-Encryption-Context"
+    value shouldBe """{"foo":"bar","foo2":"bar2"}"""
+  }
+
+  it should "parse the raw context" in {
+    val header = `X-Amz-Server-Side-Encryption-Context`.parse("""{"foo":"bar","foo2":"bar2"}""")
+    header shouldBe Success(`X-Amz-Server-Side-Encryption-Context`(Map("foo"->"bar", "foo2"->"bar2")))
+  }
+
+  it should "not parse the raw context if it is not string->string" in {
+    val header = `X-Amz-Server-Side-Encryption-Context`.parse("""{"foo":"bar","foo2":2}""")
+    header shouldBe a[Failure[_]]
+  }
+
 }

s3-stream/src/main/scala/com/bluelabs/s3stream/HttpRequests.scala

@@ -31,12 +31,15 @@ object HttpRequests {
     metadata.serverSideEncryption match {
       case ServerSideEncryption.Aes256 =>
         List(new `X-Amz-Server-Side-Encryption`(AES256))
-      case ServerSideEncryption.Kms(keyId) =>
+      case ServerSideEncryption.Kms(keyId, context) =>
         List(
           new `X-Amz-Server-Side-Encryption`(KMS),
           new `X-Amz-Server-Side-Encryption-Aws-Kms-Key-Id`(keyId)
-          // TODO add `x-amz-server-side-encryption-context` header.
-        )
+        ) ++ (if(context.isEmpty) {
+          List.empty
+        } else {
+          List(`X-Amz-Server-Side-Encryption-Context`(context))
+        })
       case ServerSideEncryption.None =>
         Nil
     }

s3-stream/src/main/scala/com/bluelabs/s3stream/ServerSideEncryption.scala

@@ -9,6 +9,5 @@ object ServerSideEncryption {
 
   case object Aes256 extends ServerSideEncryption
 
-  // TODO add context
-  case class Kms(keyId: String) extends ServerSideEncryption
+  case class Kms(keyId: String, context: Map[String, String] = Map.empty) extends ServerSideEncryption
 }

s3-stream/src/test/scala/com/bluelabs/s3stream/HttpRequestsSpec.scala

@@ -33,10 +33,14 @@ class HttpRequestsSpec extends FlatSpec with Matchers {
     metadataHeaders(Metadata(serverSideEncryption = ServerSideEncryption.Kms("my-id"))) should contain(`X-Amz-Server-Side-Encryption`(KMS))
   }
 
-  it should "add the x-amz-server-side-encryption-kms-id with KMZ header when the server side encryption is KMS" in {
+  it should "add the x-amz-server-side-encryption-kms-id when the server side encryption is KMS" in {
     metadataHeaders(Metadata(serverSideEncryption = ServerSideEncryption.Kms("my-id"))) should contain(`X-Amz-Server-Side-Encryption-Aws-Kms-Key-Id`("my-id"))
   }
 
+  it should "add the x-amz-server-side-encryption-context when the server side encryption is KMS and context is not empty" in {
+    metadataHeaders(Metadata(serverSideEncryption = ServerSideEncryption.Kms("my-id", Map("foo"->"bar")))) should contain(`X-Amz-Server-Side-Encryption-Context`(Map("foo"->"bar")))
+  }
+
   "initiateMultipartUploadRequest" should "set the default entity contentType when it is not specified" in {
     val requestEntity = initiateMultipartUploadRequest(S3Location("test-bucket", "test-key"), Metadata()).entity