From 84f02bb7cb727b14815936cdbee5394ca60b8f27 Mon Sep 17 00:00:00 2001 From: Mario Nachbaur Date: Sat, 11 Oct 2025 02:41:08 +0200 Subject: [PATCH] add basic key-value storage --- Makefile | 2 +- src/atproto/__init__.py | 81 +++++++++++++++++++++++------------------ src/db.py | 40 ++++++++++++++++++++ src/main.py | 55 +++++++++++++++------------- src/oauth.py | 18 ++++++--- src/schema.sql | 8 +++- 6 files changed, 136 insertions(+), 68 deletions(-) diff --git a/Makefile b/Makefile index a7c8892..ecc3b02 100644 --- a/Makefile +++ b/Makefile @@ -7,4 +7,4 @@ debug: .PHONY: run run: - uv run -- gunicorn --bind ':$(PORT)' 'src.main:app' + uv run -- dotenv run -- gunicorn -w 4 --bind ':$(PORT)' 'src.main:app' diff --git a/src/atproto/__init__.py b/src/atproto/__init__.py index f2d1905..c2f12a1 100644 --- a/src/atproto/__init__.py +++ b/src/atproto/__init__.py @@ -3,6 +3,8 @@ from re import match as regex_match from typing import Any import httpx +from ..db import KV + from .validator import is_valid_authserver_meta from ..security import is_safe_url @@ -15,10 +17,6 @@ AuthserverUrl = str PdsUrl = str DID = str -authservers: dict[PdsUrl, AuthserverUrl] = {} -dids: dict[str, DID] = {} -pdss: dict[DID, PdsUrl] = {} - def is_valid_handle(handle: str) -> bool: return regex_match(HANDLE_REGEX, handle) is not None @@ -28,12 +26,12 @@ def is_valid_did(did: str) -> bool: return regex_match(DID_REGEX, did) is not None -def resolve_identity(query: str) -> tuple[str, str, dict[str, Any]] | None: +def resolve_identity(query: str, didkv: KV) -> tuple[str, str, dict[str, Any]] | None: """Resolves an identity to a DID, handle and DID document, verifies handles bi directionally.""" if is_valid_handle(query): handle = query - did = resolve_did_from_handle(handle) + did = resolve_did_from_handle(handle, didkv) if not did: return None doc = resolve_doc_from_did(did) @@ -52,7 +50,7 @@ def resolve_identity(query: str) -> tuple[str, str, dict[str, Any]] | None: handle = handle_from_doc(doc) if not handle: return None - if resolve_did_from_handle(handle) != did: + if resolve_did_from_handle(handle, didkv) != did: return None return (did, handle, doc) @@ -79,12 +77,16 @@ def handle_from_doc(doc: dict[str, list[str]]) -> str | None: return None -def resolve_did_from_handle(handle: str, reload: bool = False) -> str | None: +def resolve_did_from_handle(handle: str, kv: KV, reload: bool = False) -> str | None: """Returns the DID for a given handle""" - if handle in dids and not reload: + if not is_valid_handle(handle): + return None + + did = kv.get(handle) + if did is not None and not reload: print(f"returning cached did for {handle}") - return dids[handle] + return did answer = resolve_dns(f"_atproto.{handle}", "TXT") for record in answer: @@ -92,6 +94,8 @@ def resolve_did_from_handle(handle: str, reload: bool = False) -> str | None: if value.startswith("did="): did = value[4:] if is_valid_did(did): + print(f"caching did {did} for {handle}") + kv.set(handle, value=did) return did return None @@ -106,17 +110,20 @@ def pds_endpoint_from_doc(doc: dict[str, list[dict[str, str]]]) -> str | None: return None -def resolve_pds_from_did(did: DID, reload: bool = False) -> PdsUrl | None: - if did in pdss and not reload: +def resolve_pds_from_did(did: DID, kv: KV, reload: bool = False) -> PdsUrl | None: + pds = kv.get(did) + if pds is not None and not reload: print(f"returning cached pds for {did}") - return pdss[did] + return pds doc = resolve_doc_from_did(did) if doc is None: return None pds = doc["service"][0]["serviceEndpoint"] - pdss[did] = pds + if pds is None: + return None print(f"caching pds {pds} for {did}") + kv.set(did, value=pds) return pds @@ -139,13 +146,15 @@ def resolve_doc_from_did( def resolve_authserver_from_pds( pds_url: PdsUrl, + kv: KV, reload: bool = False, ) -> AuthserverUrl | None: """Returns the authserver URL for the PDS.""" - if pds_url in authservers and not reload: + authserver_url = kv.get(pds_url) + if authserver_url is not None and not reload: print(f"returning cached authserver for PDS {pds_url}") - return authservers[pds_url] + return authserver_url assert is_safe_url(pds_url) endpoint = f"{pds_url}/.well-known/oauth-protected-resource" @@ -155,7 +164,7 @@ def resolve_authserver_from_pds( parsed: dict[str, list[str]] = response.json() authserver_url = parsed["authorization_servers"][0] print(f"caching authserver {authserver_url} for PDS {pds_url}") - authservers[pds_url] = authserver_url + kv.set(pds_url, value=authserver_url) return authserver_url @@ -163,27 +172,29 @@ def fetch_authserver_meta(authserver_url: str) -> dict[str, str] | None: """Returns metadata from the authserver""" assert is_safe_url(authserver_url) endpoint = f"{authserver_url}/.well-known/oauth-authorization-server" - meta = http_get_json(endpoint) + response = httpx.get(endpoint) + if not response.is_success: + return None + meta: dict[str, Any] = response.json() assert is_valid_authserver_meta(meta, authserver_url) return meta -def get_record(pds: str, repo: str, collection: str, record: str) -> str | None: - response = http_get( +def get_record( + pds: str, + repo: str, + collection: str, + record: str, +) -> dict[str, Any] | None: + """Retrieve record from PDS. Verifies type is the same as collection name.""" + response = httpx.get( f"{pds}/xrpc/com.atproto.repo.getRecord?repo={repo}&collection={collection}&rkey={record}" ) - return response - - -def http_get_json(url: str) -> Any | None: - response = httpx.get(url) - if response.is_success: - return response.json() - return None - - -def http_get(url: str) -> str | None: - response = httpx.get(url) - if response.is_success: - return response.text - return None + if not response.is_success: + return None + parsed = response.json() + value: dict[str, Any] = parsed["value"] + if value["$type"] != collection: + return None + del value["$type"] + return value diff --git a/src/db.py b/src/db.py index e269fa1..f34ba17 100644 --- a/src/db.py +++ b/src/db.py @@ -1,6 +1,46 @@ +from abc import ABC, abstractmethod +from typing import override from flask import Flask, g import sqlite3 +from sqlite3 import Connection + + +class KV(ABC): + @abstractmethod + def get(self, key: str) -> str | None: + pass + + @abstractmethod + def set(self, key: str, value: str): + pass + + +class Keyval(KV): + db: Connection + prefix: str + + def __init__(self, app: Flask, prefix: str): + self.db = get_db(app) + self.prefix = prefix + + @override + def get(self, key: str) -> str | None: + cursor = self.db.cursor() + row = cursor.execute( + "select value from keyval where prefix = ? and key = ?", + (self.prefix, key), + ).fetchone() + return None if row is None else row["value"] + + @override + def set(self, key: str, value: str): + cursor = self.db.cursor() + _ = cursor.execute( + "insert or replace into keyval (prefix, key, value) values (?, ?, ?)", + (self.prefix, key, value), + ) + self.db.commit() def get_db(app: Flask) -> sqlite3.Connection: diff --git a/src/main.py b/src/main.py index 8d1ef5d..31c3dd1 100644 --- a/src/main.py +++ b/src/main.py @@ -10,7 +10,7 @@ from .atproto import ( resolve_pds_from_did, ) from .atproto.oauth import pds_authed_req -from .db import close_db_connection, init_db +from .db import Keyval, close_db_connection, init_db from .oauth import get_auth_session, oauth, save_auth_session from .types import OAuthSession @@ -19,9 +19,6 @@ _ = app.config.from_prefixed_env() app.register_blueprint(oauth) init_db(app) -links: dict[str, list[dict[str, str]]] = {} -profiles: dict[str, tuple[str, str]] = {} - SCHEMA = "at.ligo" @@ -55,15 +52,16 @@ def page_profile_with_did(did: str): @app.get("/@") def page_profile_with_handle(handle: str): reload = request.args.get("reload") is not None - - did = resolve_did_from_handle(handle, reload=reload) + kv = Keyval(app, "did_from_handle") + did = resolve_did_from_handle(handle, kv, reload=reload) if did is None: return "did not found", 404 return page_profile(did, reload=reload) def page_profile(did: str, reload: bool = False): - pds = resolve_pds_from_did(did, reload=reload) + kv = Keyval(app, "pds_from_did") + pds = resolve_pds_from_did(did, kv, reload=reload) if pds is None: return "pds not found", 404 profile, _ = load_profile(pds, did, reload=reload) @@ -195,41 +193,46 @@ def page_terms(): def load_links(pds: str, did: str, reload: bool = False) -> list[dict[str, str]] | None: - if did in links and not reload: + kv = Keyval(app, "links_from_did") + links = kv.get(did) + + if links is not None and not reload: app.logger.debug(f"returning cached links for {did}") - return links[did] + return json.loads(links) - response = get_record(pds, did, f"{SCHEMA}.actor.links", "self") - if response is None: + record = get_record(pds, did, f"{SCHEMA}.actor.links", "self") + if record is None: return None - record = json.loads(response) - links_ = record["value"]["links"] + links = record["links"] app.logger.debug(f"caching links for {did}") - links[did] = links_ - return links_ + kv.set(did, value=json.dumps(links)) + return links def load_profile( - pds: str, did: str, reload: bool = False + pds: str, + did: str, + reload: bool = False, ) -> tuple[tuple[str, str] | None, bool]: - if did in profiles and not reload: + kv = Keyval(app, "profile_from_did") + profile = kv.get(did) + + if profile is not None and not reload: app.logger.debug(f"returning cached profile for {did}") - return profiles[did], False + return json.loads(profile), False from_bluesky = False - response = get_record(pds, did, f"{SCHEMA}.actor.profile", "self") - if response is None: - response = get_record(pds, did, "app.bsky.actor.profile", "self") + record = get_record(pds, did, f"{SCHEMA}.actor.profile", "self") + if record is None: + record = get_record(pds, did, "app.bsky.actor.profile", "self") from_bluesky = True - if response is None: + if record is None: return None, False - record = json.loads(response) - value: dict[str, str] = record["value"] - profile = (value["displayName"], value["description"]) + profile = (record["displayName"], record["description"]) app.logger.debug(f"caching profile for {did}") - profiles[did] = profile + kv.set(did, value=json.dumps(profile)) return profile, from_bluesky diff --git a/src/oauth.py b/src/oauth.py index c8f4224..74b4fcd 100644 --- a/src/oauth.py +++ b/src/oauth.py @@ -6,6 +6,8 @@ from urllib.parse import urlencode import json +from .db import Keyval + from .atproto import ( is_valid_did, is_valid_handle, @@ -28,9 +30,12 @@ def oauth_start(): if not username: return redirect(url_for("page_login"), 303) + pdskv = Keyval(current_app, "authserver_from_pds") + if is_valid_handle(username) or is_valid_did(username): login_hint = username - identity = resolve_identity(username) + kv = Keyval(current_app, "did_from_handle") + identity = resolve_identity(username, didkv=kv) if identity is None: return "couldnt resolve identity", 500 did, handle, doc = identity @@ -38,14 +43,14 @@ def oauth_start(): if not pds_url: return "pds not found", 404 current_app.logger.debug(f"account PDS: {pds_url}") - authserver_url = resolve_authserver_from_pds(pds_url) + authserver_url = resolve_authserver_from_pds(pds_url, pdskv) if not authserver_url: return "authserver not found", 404 elif username.startswith("https://") and is_safe_url(username): did, handle, pds_url = None, None, None login_hint = None - authserver_url = resolve_authserver_from_pds(username) or username + authserver_url = resolve_authserver_from_pds(username, pdskv) or username else: return "not a valid handle, did or auth server", 400 @@ -134,6 +139,9 @@ def oauth_callback(): row = auth_request + didkv = Keyval(current_app, "did_from_handle") + authserverkv = Keyval(current_app, "authserver_from_pds") + if row.did: # If we started with an account identifier, this is simple did, handle, pds_url = row.did, row.handle, row.pds_url @@ -141,14 +149,14 @@ def oauth_callback(): else: did = tokens.sub assert is_valid_did(did) - identity = resolve_identity(did) + identity = resolve_identity(did, didkv=didkv) if not identity: return "could not resolve identity", 500 did, handle, did_doc = identity pds_url = pds_endpoint_from_doc(did_doc) if not pds_url: return "could not resolve pds", 500 - authserver_url = resolve_authserver_from_pds(pds_url) + authserver_url = resolve_authserver_from_pds(pds_url, authserverkv) assert authserver_url == authserver_iss assert row.scope == tokens.scope diff --git a/src/schema.sql b/src/schema.sql index 12b9f4b..423852b 100644 --- a/src/schema.sql +++ b/src/schema.sql @@ -1 +1,7 @@ --- empty for now +drop table if exists keyval; +create table if not exists keyval ( + prefix text not null, + key text not null, + value text, + primary key (prefix, key) +) strict, without rowid; -- 2.54.0