From cd1374cc333c9b3fa7feb8abe004a9c2a0698a7c Mon Sep 17 00:00:00 2001
From: Brandon Liu <bdon@bdon.org>
Date: Tue, 1 Nov 2022 17:24:29 +0800
Subject: [PATCH] python: fix magic number in writer, assert current spec
 version and magic number in reader

---
 python/pmtiles/tile.py   | 13 ++++++++++++-
 python/test/test_tile.py |  9 ++++++++-
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/python/pmtiles/tile.py b/python/pmtiles/tile.py
index bb6b5fc..b8f6b00 100644
--- a/python/pmtiles/tile.py
+++ b/python/pmtiles/tile.py
@@ -191,8 +191,19 @@ def serialize_directory(entries):
 
     return gzip.compress(b_io.getvalue())
 
+class SpecVersionUnsupported(Exception):
+    pass
+
+class MagicNumberNotFound(Exception):
+    pass
 
 def deserialize_header(buf):
+    if buf[0:7].decode() != "PMTiles":
+        raise MagicNumberNotFound()
+
+    if buf[7] != 0x3:
+        raise SpecVersionUnsupported()
+
     def read_uint64(pos):
         return int.from_bytes(buf[pos : pos + 8], byteorder="little")
 
@@ -240,7 +251,7 @@ def serialize_header(h):
         b_io.write(i.to_bytes(1, byteorder="little"))
 
     b_io.write("PMTiles".encode())
-    b_io.write(b"3")
+    b_io.write(b"\x03")
     write_uint64(h["root_offset"])
     write_uint64(h["root_length"])
     write_uint64(h["metadata_offset"])
diff --git a/python/test/test_tile.py b/python/test/test_tile.py
index 50314dd..08e5f53 100644
--- a/python/test/test_tile.py
+++ b/python/test/test_tile.py
@@ -3,7 +3,7 @@ from pmtiles.tile import zxy_to_tileid, tileid_to_zxy, Entry
 from pmtiles.tile import read_varint, write_varint
 from pmtiles.tile import Entry, find_tile, Compression, TileType
 from pmtiles.tile import serialize_directory, deserialize_directory
-from pmtiles.tile import serialize_header, deserialize_header
+from pmtiles.tile import serialize_header, deserialize_header, SpecVersionUnsupported, MagicNumberNotFound
 import io
 
 
@@ -160,3 +160,10 @@ class TestHeader(unittest.TestCase):
         self.assertEqual(result["center_zoom"], 3)
         self.assertEqual(result["center_lon_e7"], 3.1 * 10000000)
         self.assertEqual(result["center_lat_e7"], 3.2 * 10000000)
+
+    def test_spec_version(self):
+        with self.assertRaises(SpecVersionUnsupported):
+            result = deserialize_header(b'PMTiles\x04')
+
+        with self.assertRaises(MagicNumberNotFound):
+            result = deserialize_header(b'PM\x00\x02')