List active blobs helper script

resources/scripts/list_active_blobs.py

@@ -0,0 +1,156 @@
+#!/usr/bin/env python3
+"""List the S3 object keys of every blob that is still actively linked.
+
+Stalwart reference-counts blobs through link entries kept in the data store
+under the SUBSPACE_BLOB_LINK subspace (the PostgreSQL table named "k"). A blob
+is garbage-collectable once it has no surviving link; this tool reports the
+opposite set: the S3 keys of blobs that are still referenced, so they can be
+diffed against the actual contents of an S3 bucket to find orphans.
+
+Key layout in table "k" (the subspace byte is NOT stored, it only selects the
+table). All integers are big-endian:
+
+  32 bytes  Commit  : <hash:32>                              (blob exists marker)
+  40 bytes  Id link : <hash:32><id:8>
+  41 bytes  Doc link: <hash:32><account_id:4><collection:1><document_id:4>
+  44 bytes  Tmp link: <hash:32><account_id:4><until:8>        (active iff until > now)
+
+A blob is active when it has at least one Id/Doc link, or a Temporary link whose
+"until" (unix seconds) is still in the future. The 32-byte Commit marker alone
+does not keep a blob alive.
+
+The S3 object key is the optional configured key prefix (literal string)
+followed by the custom-base32 encoding of the 32-byte blob hash, matching
+S3Store::build_key in crates/store/src/backend/s3/mod.rs.
+"""
+
+import argparse
+import os
+import sys
+import time
+
+BLOB_HASH_LEN = 32
+ID_LINK = BLOB_HASH_LEN + 8          # 40
+DOC_LINK = BLOB_HASH_LEN + 8 + 1     # 41
+TEMP_LINK = BLOB_HASH_LEN + 4 + 8    # 44
+
+BASE32_ALPHABET = b"abcdefghijklmnopqrstuvwxyz792013"
+
+
+class Base32Writer:
+    """Faithful port of utils::codec::base32_custom::Base32Writer."""
+
+    def __init__(self, prefix=""):
+        self.last_byte = 0
+        self.pos = 0
+        self.out = [prefix] if prefix else []
+
+    def _push_byte(self, byte, is_remainder):
+        p = self.pos % 5
+        if p == 0:
+            ch1 = (byte & 0xF8) >> 3
+            ch2 = 0xFF
+        elif p == 1:
+            ch1 = ((self.last_byte & 0x07) << 2) | ((byte & 0xC0) >> 6)
+            ch2 = (byte & 0x3E) >> 1
+        elif p == 2:
+            ch1 = ((self.last_byte & 0x01) << 4) | ((byte & 0xF0) >> 4)
+            ch2 = 0xFF
+        elif p == 3:
+            ch1 = ((self.last_byte & 0x0F) << 1) | (byte >> 7)
+            ch2 = (byte & 0x7C) >> 2
+        else:
+            ch1 = ((self.last_byte & 0x03) << 3) | ((byte & 0xE0) >> 5)
+            ch2 = byte & 0x1F
+
+        self.out.append(chr(BASE32_ALPHABET[ch1]))
+        if not is_remainder:
+            if ch2 != 0xFF:
+                self.out.append(chr(BASE32_ALPHABET[ch2]))
+            self.last_byte = byte
+            self.pos += 1
+
+    def write(self, data):
+        for byte in data:
+            self._push_byte(byte, False)
+        return self
+
+    def finalize(self):
+        if self.pos % 5 != 0:
+            self._push_byte(0, True)
+        return "".join(self.out)
+
+
+def s3_key(blob_hash, prefix=""):
+    return Base32Writer(prefix).write(blob_hash).finalize()
+
+
+def be_u64(b):
+    return int.from_bytes(b, "big")
+
+
+def collect_active_hashes(rows, now, include_expired_temporary=False):
+    active = set()
+    unknown = 0
+    for (key,) in rows:
+        key = bytes(key)
+        n = len(key)
+        if n == BLOB_HASH_LEN:
+            continue
+        if n in (ID_LINK, DOC_LINK):
+            active.add(key[:BLOB_HASH_LEN])
+        elif n == TEMP_LINK:
+            until = be_u64(key[BLOB_HASH_LEN + 4:BLOB_HASH_LEN + 12])
+            if include_expired_temporary or until > now:
+                active.add(key[:BLOB_HASH_LEN])
+        else:
+            unknown += 1
+    if unknown:
+        print(f"warning: skipped {unknown} key(s) of unexpected length",
+              file=sys.stderr)
+    return active
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__,
+                                 formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--host", default=os.environ.get("PGHOST", "localhost"))
+    ap.add_argument("--port", type=int, default=int(os.environ.get("PGPORT", "5432")))
+    ap.add_argument("--user", default=os.environ.get("PGUSER", "stalwart"))
+    ap.add_argument("--password", default=os.environ.get("PGPASSWORD", "stalwart"))
+    ap.add_argument("--dbname", default=os.environ.get("PGDATABASE", "stalwart"))
+    ap.add_argument("--table", default="k",
+                    help="SUBSPACE_BLOB_LINK table name (default: k)")
+    ap.add_argument("--prefix", default="",
+                    help="S3 key_prefix configured on the blob store (default: none)")
+    ap.add_argument("--now", type=int, default=None,
+                    help="override unix-seconds used to expire temporary links")
+    ap.add_argument("--include-expired-temporary", action="store_true",
+                    help="treat expired temporary links as active too")
+    args = ap.parse_args()
+
+    try:
+        import psycopg2
+        conn = psycopg2.connect(host=args.host, port=args.port, user=args.user,
+                                password=args.password, dbname=args.dbname)
+    except ImportError:
+        import psycopg
+        conn = psycopg.connect(host=args.host, port=args.port, user=args.user,
+                               password=args.password, dbname=args.dbname)
+
+    now = args.now if args.now is not None else int(time.time())
+
+    with conn, conn.cursor() as cur:
+        cur.execute(f'SELECT k FROM "{args.table}"')
+        rows = cur.fetchall()
+    conn.close()
+
+    active = collect_active_hashes(rows, now,
+                                   include_expired_temporary=args.include_expired_temporary)
+
+    for h in sorted(active):
+        print(s3_key(h, args.prefix))
+
+
+if __name__ == "__main__":
+    main()