strip quotes from shard iterators before parsing base64 (#689)

* strip quotes from shard iterators before parsing base64 * Add CLI based test script * remove created stream after CLI tests --------- Co-authored-by: Eric Meisel <[email protected]>
etspaceman · Dec 18, 2023 · f3f86c6 · f3f86c6
1 parent 2a2dfc7
commit f3f86c6
Show file tree

Hide file tree

Showing 4 changed files with 91 additions and 3 deletions.
diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml
@@ -41,8 +41,10 @@ services:
       - "AWS_ACCESS_KEY_ID=foo"
       - "AWS_SECRET_ACCESS_KEY=bar"
       - "AWS_DEFAULT_REGION=us-east-1"
-    entrypoint: "bash -c"
-    command: "'aws --endpoint-url 'https://kinesis-mock:4567/' --no-verify-ssl kinesis list-streams && exit 0 || exit 1'"
+    entrypoint: "bash"
+    command: "/awscli-tests.sh"
+    volumes:
+      - "./scripts/awscli-tests.sh:/awscli-tests.sh"
     depends_on:
       kinesis-mock:
         condition: service_healthy

diff --git a/docker/scripts/awscli-tests.sh b/docker/scripts/awscli-tests.sh
@@ -0,0 +1,20 @@
+#!/bin/bash
+set -euo pipefail
+
+aws --endpoint-url 'https://kinesis-mock:4567/' --no-verify-ssl kinesis list-streams
+
+KINESIS_STREAM_NAME=test-stream
+
+aws --endpoint-url 'https://kinesis-mock:4567/' --no-verify-ssl kinesis create-stream --stream-name $KINESIS_STREAM_NAME --shard-count 1 --stream-mode-details StreamMode=PROVISIONED
+
+base64_content=$(echo '{"test": "data"}' | base64)
+
+aws --endpoint-url 'https://kinesis-mock:4567/' --no-verify-ssl kinesis put-record \
+  --stream-name $KINESIS_STREAM_NAME \
+  --data  ${base64_content} --partition-key id
+
+SHARD_ITERATOR=$(aws --endpoint-url 'https://kinesis-mock:4567/' --no-verify-ssl kinesis get-shard-iterator --shard-id shardId-000000000000 --shard-iterator-type TRIM_HORIZON --stream-name $KINESIS_STREAM_NAME --query 'ShardIterator');
+
+aws --endpoint-url 'https://kinesis-mock:4567/' --no-verify-ssl kinesis get-records --shard-iterator $SHARD_ITERATOR
+
+aws --endpoint-url 'https://kinesis-mock:4567/' --no-verify-ssl kinesis delete-stream --stream-name $KINESIS_STREAM_NAME
diff --git a/integration-tests/src/test/scala/kinesis/mock/GetRecordsTests.scala b/integration-tests/src/test/scala/kinesis/mock/GetRecordsTests.scala
@@ -79,4 +79,70 @@ class GetRecordsTests extends AwsFunctionalTests {
       s"$res\n$recordRequests"
     )
   }
+
+  fixture.test("It should get records with quotes around shard iterator") { resources =>
+    for {
+      recordRequests <- IO(
+        putRecordRequestArb.arbitrary
+          .take(1)
+          .toVector
+          .map(
+            _.copy(streamName = Some(resources.streamName), streamArn = None)
+          )
+          .map(x =>
+            PutRecordRequest
+              .builder()
+              .partitionKey(x.partitionKey)
+              .streamName(resources.streamName.streamName)
+              .data(SdkBytes.fromByteArray(x.data))
+              .maybeTransform(x.explicitHashKey)(_.explicitHashKey(_))
+              .maybeTransform(x.sequenceNumberForOrdering)((req, sequenceNum) =>
+                req.sequenceNumberForOrdering(sequenceNum.value)
+              )
+              .build()
+          )
+      )
+      _ <- recordRequests.traverse(x =>
+        resources.kinesisClient.putRecord(x).toIO
+      )
+      shards <- resources.kinesisClient
+        .listShards(
+          ListShardsRequest
+            .builder()
+            .streamName(resources.streamName.streamName)
+            .build()
+        )
+        .toIO
+        .map(_.shards().asScala.toVector)
+      shardIterators <- shards.traverse(shard =>
+        resources.kinesisClient
+          .getShardIterator(
+            GetShardIteratorRequest
+              .builder()
+              .shardId(shard.shardId())
+              .streamName(resources.streamName.streamName)
+              .shardIteratorType(ShardIteratorType.TRIM_HORIZON)
+              .build()
+          )
+          .toIO
+          .map(_.shardIterator())
+      )
+      gets <- shardIterators.traverse(shardIterator =>
+        resources.kinesisClient
+          .getRecords(
+            GetRecordsRequest.builder().shardIterator(s"\"$shardIterator\"").build()
+          )
+          .toIO
+      )
+      res = gets.flatMap(_.records().asScala.toVector)
+    } yield assert(
+      res.length == 1 && res.forall(rec =>
+        recordRequests.exists(req =>
+          req.data.asByteArray.sameElements(rec.data.asByteArray)
+            && req.partitionKey == rec.partitionKey
+        )
+      ),
+      s"$res\n$recordRequests"
+    )
+  }
 }
diff --git a/kinesis-mock/src/main/scala/kinesis/mock/models/ShardIterator.scala b/kinesis-mock/src/main/scala/kinesis/mock/models/ShardIterator.scala
@@ -30,7 +30,7 @@ import kinesis.mock.validations.CommonValidations
 
 final case class ShardIterator(value: String) {
   def parse(now: Instant): Response[ShardIteratorParts] = {
-    val decoded = Base64.getDecoder.decode(value)
+    val decoded = Base64.getDecoder.decode(value.replaceAll("^\"|\"$", ""))
 
     val decrypted = new String(
       AES.decrypt(