From 5b8627b0d8367c012d9864f1847f0c47558bb9c8 Mon Sep 17 00:00:00 2001
From: Brandon Liu <bdon@users.noreply.github.com>
Date: Wed, 7 Feb 2024 11:46:05 +0800
Subject: [PATCH] Implement ETag logic for AWS and Cloudflare [#90] (#354)

---
 serverless/aws/src/index.ts        | 37 +++++++++++++++++++++---------
 serverless/cloudflare/src/index.ts |  8 +++++++
 2 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/serverless/aws/src/index.ts b/serverless/aws/src/index.ts
index ea62071..1f99c7f 100644
--- a/serverless/aws/src/index.ts
+++ b/serverless/aws/src/index.ts
@@ -5,6 +5,7 @@ import {
 } from "aws-lambda";
 import {
   Compression,
+  EtagMismatch,
   PMTiles,
   RangeResponse,
   ResolvedValueCache,
@@ -16,7 +17,11 @@ import { pmtiles_path, tileJSON, tile_path } from "../../shared/index";
 import { createHash } from "crypto";
 import zlib from "zlib";
 
-import { GetObjectCommand, S3Client } from "@aws-sdk/client-s3";
+import {
+  GetObjectCommand,
+  GetObjectCommandOutput,
+  S3Client,
+} from "@aws-sdk/client-s3";
 import { NodeHttpHandler } from "@aws-sdk/node-http-handler";
 
 // the region should default to the same one as the function
@@ -60,16 +65,26 @@ class S3Source implements Source {
     signal?: AbortSignal,
     etag?: string
   ): Promise<RangeResponse> {
-    const resp = await s3client.send(
-      new GetObjectCommand({
-        // biome-ignore lint: aws api
-        Bucket: process.env.BUCKET!,
-        // biome-ignore lint: aws api
-        Key: pmtiles_path(this.archiveName, process.env.PMTILES_PATH),
-        // biome-ignore lint: aws api
-        Range: "bytes=" + offset + "-" + (offset + length - 1),
-      })
-    );
+    let resp: GetObjectCommandOutput;
+    try {
+      resp = await s3client.send(
+        new GetObjectCommand({
+          // biome-ignore lint: aws api
+          Bucket: process.env.BUCKET!,
+          // biome-ignore lint: aws api
+          Key: pmtiles_path(this.archiveName, process.env.PMTILES_PATH),
+          // biome-ignore lint: aws api
+          Range: "bytes=" + offset + "-" + (offset + length - 1),
+          // biome-ignore lint: aws api
+          IfMatch: etag,
+        })
+      );
+    } catch (e: unknown) {
+      if (e instanceof Error && (e as Error).name === "PreconditionFailed") {
+        throw new EtagMismatch();
+      }
+      throw e;
+    }
 
     const arr = await resp.Body?.transformToByteArray();
 
diff --git a/serverless/cloudflare/src/index.ts b/serverless/cloudflare/src/index.ts
index 55219c9..b90626d 100644
--- a/serverless/cloudflare/src/index.ts
+++ b/serverless/cloudflare/src/index.ts
@@ -1,5 +1,6 @@
 import {
   Compression,
+  EtagMismatch,
   PMTiles,
   RangeResponse,
   ResolvedValueCache,
@@ -63,12 +64,19 @@ class R2Source implements Source {
       pmtiles_path(this.archiveName, this.env.PMTILES_PATH),
       {
         range: { offset: offset, length: length },
+        onlyIf: { etagMatches: etag },
       }
     );
     if (!resp) {
       throw new KeyNotFoundError("Archive not found");
     }
+
     const o = resp as R2ObjectBody;
+
+    if (!o.body) {
+      throw new EtagMismatch();
+    }
+
     const a = await o.arrayBuffer();
     return {
       data: a,