Fix SQS message attribute decoding

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog][chg] and this project adheres to
 
 ## Unreleased
   - Add `messageGroupId` to SQS Attributes.
+  - Fix the type of `messageAttributes` in `SQSEvent`.
   - Fix header parsing to accept `null` values serialised as `{}`
   - The `use-aeson-2.2` flag has been removed in favour of a
     dependency on the `attoparsec-aeson` compatibility package.

src/AWS/Lambda/Events/SQS.hs

@@ -1,4 +1,5 @@
-{-# LANGUAGE ApplicativeDo #-}
+{-# LANGUAGE ApplicativeDo   #-}
+{-# LANGUAGE NamedFieldPuns  #-}
 {-# LANGUAGE RecordWildCards #-}
 
 {-|
@@ -13,13 +14,21 @@ Stability   : stable
 module AWS.Lambda.Events.SQS (
   Records (..),
   Attributes (..),
+  MessageAttribute (..),
+  MessageAttributeValue (..),
   SQSEvent (..)
 ) where
 
-import Data.Aeson   (FromJSON (..), withObject, (.:), (.:?))
-import Data.Map     (Map)
-import Data.Text    (Text)
-import GHC.Generics (Generic)
+import           Data.Aeson             (FromJSON (..), withObject, (.:), (.:?))
+import           Data.ByteString        (ByteString)
+import qualified Data.ByteString.Base64 as B64
+import           Data.Map               (Map)
+import           Data.Scientific        (Scientific)
+import           Data.Text              (Text)
+import qualified Data.Text              as Text
+import qualified Data.Text.Encoding     as TE
+import           GHC.Generics           (Generic)
+import           Text.Read              (readMaybe)
 
 -- | Represents an event from AWS SQS.
 --
@@ -49,16 +58,56 @@ instance FromJSON Attributes where
     messageGroupId <- v .:? "MessageGroupId"
     pure Attributes {..}
 
+-- | An SQS message attribute as it appears in Lambda SQS event payloads.
+--
+-- See the <https://docs.aws.amazon.com/lambda/latest/dg/with-sqs.html#example-standard-queue-message-event AWS Lambda SQS event payload example>
+-- for the JSON shape used under @messageAttributes@.
+data MessageAttribute = MessageAttribute {
+  customTypeLabel :: Maybe Text,
+  value           :: MessageAttributeValue
+} deriving (Show, Eq, Generic)
+
+instance FromJSON MessageAttribute where
+  parseJSON = withObject "MessageAttribute" $ \v -> do
+    dataType <- v .: "dataType"
+    let (baseType, customTypeLabel) = splitDataType dataType
+
+    value <- case baseType of
+      "Binary" -> Binary . decodeBase64Text <$> v .: "binaryValue"
+      "Number" -> Number <$> (v .: "stringValue" >>= parseNumber)
+      "String" -> String <$> v .: "stringValue"
+      _ -> fail $ "Unexpected message attribute dataType: " <> show dataType
+
+    pure MessageAttribute { customTypeLabel, value }
+    where
+      parseNumber = maybe (fail "can't parse stringValue into Scientific") pure . readMaybe
+
+data MessageAttributeValue
+  = Binary ByteString
+  | Number Scientific
+  | String Text
+  deriving (Show, Eq, Generic)
+
 data SQSEvent = SQSEvent {
   messageId         :: Text,
   receiptHandle     :: Text,
   body              :: Text,
   attributes        :: Attributes,
-  messageAttributes :: Map Text Text,
+  messageAttributes :: Map Text MessageAttribute,
   md5OfBody         :: Text,
   eventSource       :: Text,
   eventSourceARN    :: Text,
   awsRegion         :: Text
 } deriving (Show, Eq, Generic)
 
 instance FromJSON SQSEvent
+
+decodeBase64Text :: Text -> ByteString
+decodeBase64Text = B64.decodeLenient . TE.encodeUtf8
+
+splitDataType :: Text -> (Text, Maybe Text)
+splitDataType dataType =
+  case Text.breakOn "." dataType of
+    (baseType, "") -> (baseType, Nothing)
+    (baseType, customTypeLabel) ->
+      (baseType, Just $ Text.drop 1 customTypeLabel)

test/AWS/Lambda/Events/SQS/Spec.hs

@@ -4,16 +4,20 @@ module AWS.Lambda.Events.SQS.Spec where
 
 import           AWS.Lambda.Events.SQS
 import           Data.Aeson            (eitherDecode)
+import           Data.Either           (isLeft)
 import qualified Data.Map              as M
 import           Data.ByteString.Lazy  (ByteString)
-import           Test.Hspec            (Spec, shouldBe, specify)
+import           Test.Hspec            (Spec, shouldBe, shouldSatisfy, specify)
 import           Text.RawString.QQ     (r)
 
 spec :: Spec
-spec =
+spec = do
   specify "read sample payload" $
     eitherDecode samplePayload `shouldBe` Right expectedRecords
 
+  specify "fail on invalid number message attributes" $
+    (eitherDecode invalidNumberPayload :: Either String Records) `shouldSatisfy` isLeft
+
 samplePayload :: ByteString
 samplePayload = [r|
 {
@@ -30,8 +34,32 @@ samplePayload = [r|
         "MessageGroupId": "group-1"
       },
       "messageAttributes": {
-        "attribute1": "value1",
-        "attribute2": "value2"
+        "attribute1": {
+          "stringValue": "value1",
+          "stringListValues": ["ignored-string-list-value"],
+          "binaryListValues": ["aWdub3JlZC1iaW5hcnktbGlzdC12YWx1ZQ=="],
+          "dataType": "String"
+        },
+        "attribute2": {
+          "stringValue": "ignored-string-value",
+          "binaryValue": "dmFsdWUy",
+          "stringListValues": [],
+          "binaryListValues": [],
+          "dataType": "Binary"
+        },
+        "attribute3": {
+          "stringValue": "123.45",
+          "stringListValues": [],
+          "binaryListValues": [],
+          "dataType": "Number"
+        },
+        "attribute4": {
+          "stringValue": "value4",
+          "binaryValue": "aWdub3JlZC1iaW5hcnktdmFsdWU=",
+          "stringListValues": ["ignored-extra-string"],
+          "binaryListValues": ["aWdub3JlZC1leHRyYS1iaW5hcnk="],
+          "dataType": "String.foo"
+        }
       },
       "md5OfBody": "9a0364b9e99bb480dd25e1f0284c8555",
       "eventSource": "aws:sqs",
@@ -57,8 +85,30 @@ expectedRecords = Records
               , messageGroupId = Just "group-1"
               }
           , messageAttributes = M.fromList
-              [ ("attribute1", "value1")
-              , ("attribute2", "value2")
+              [ ( "attribute1"
+                , MessageAttribute
+                    { customTypeLabel = Nothing
+                    , value = String "value1"
+                    }
+                )
+              , ( "attribute2"
+                , MessageAttribute
+                    { customTypeLabel = Nothing
+                    , value = Binary "value2"
+                    }
+                )
+              , ( "attribute3"
+                , MessageAttribute
+                    { customTypeLabel = Nothing
+                    , value = Number 123.45
+                    }
+                )
+              , ( "attribute4"
+                , MessageAttribute
+                    { customTypeLabel = Just "foo"
+                    , value = String "value4"
+                    }
+                )
               ]
           , md5OfBody = "9a0364b9e99bb480dd25e1f0284c8555"
           , eventSource = "aws:sqs"
@@ -67,3 +117,32 @@ expectedRecords = Records
           }
       ]
   }
+
+invalidNumberPayload :: ByteString
+invalidNumberPayload = [r|
+{
+  "Records": [
+    {
+      "messageId": "11111111-2222-3333-4444-555555555555",
+      "receiptHandle": "AQEBwJnKyrHigUMZj6rYigCgxlaS3SLy0a...",
+      "body": "Hello from SQS!",
+      "attributes": {
+        "ApproximateReceiveCount": "1",
+        "SentTimestamp": "1523232000000",
+        "SenderId": "123456789012",
+        "ApproximateFirstReceiveTimestamp": "1523232000001"
+      },
+      "messageAttributes": {
+        "attribute1": {
+          "stringValue": "not-a-number",
+          "dataType": "Number"
+        }
+      },
+      "md5OfBody": "9a0364b9e99bb480dd25e1f0284c8555",
+      "eventSource": "aws:sqs",
+      "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:queue1",
+      "awsRegion": "us-east-1"
+    }
+  ]
+}
+|]