store OAuth session in client side cookie

src/atproto/oauth.py

···
-import sqlite3
-from typing import Any, NamedTuple
+from typing import Any, Callable, NamedTuple
 import time
 import json
 from authlib.jose import JsonWebKey, Key
···
     refresh_token: str
     scope: str
     sub: str
+    # only for parsing
+    token_type: str | None
+    expires_in: int | None
 
 
 # Prepares and sends a pushed auth request (PAR) via HTTP POST to the Authorization Server.
···
 
     resp.raise_for_status()
     token_body = resp.json()
-    try:
-        tokens = OAuthTokens(**token_body)
-    except TypeError:
-        raise Exception("invalid token body")
+    tokens = OAuthTokens(**token_body)
 
     return tokens, dpop_authserver_nonce
 
···
 
     resp.raise_for_status()
     token_body = resp.json()
-    try:
-        tokens = OAuthTokens(**token_body)
-    except TypeError:
-        raise Exception("invalid token body")
+    tokens = OAuthTokens(**token_body)
 
     return tokens, dpop_authserver_nonce
 
···
     method: str,
     url: str,
     user: OAuthSession,
-    db: sqlite3.Connection,
+    update_dpop_pds_nonce: Callable[[str], None],
     body: dict[str, Any] | None = None,
 ) -> Response | None:
     dpop_private_jwk = JsonWebKey.import_key(json.loads(user.dpop_private_jwk))
···
             response.status_code in [400, 401]
             and response.json()["error"] == "use_dpop_nonce"
         ):
-            # print(resp.headers)
             dpop_pds_nonce = response.headers["DPoP-Nonce"]
             print(f"retrying with new PDS DPoP nonce: {dpop_pds_nonce}")
-            # update session database with new nonce
-            cur = db.cursor()
-            _ = cur.execute(
-                "UPDATE oauth_session SET dpop_pds_nonce = ? WHERE did = ?;",
-                [dpop_pds_nonce, user.did],
-            )
-            db.commit()
-            cur.close()
+            update_dpop_pds_nonce(dpop_pds_nonce)
             continue
         break
 

src/main.py

···
 
 from .atproto import PdsUrl, get_record, resolve_did_from_handle, resolve_pds_from_did
 from .atproto.oauth import pds_authed_req
-from .db import close_db_connection, get_db, init_db
-from .oauth import oauth
+from .db import close_db_connection, init_db
+from .oauth import get_auth_session, oauth, save_auth_session
 from .types import OAuthSession
 
 app = Flask(__name__)
···
 
 @app.before_request
 def load_user_to_context():
-    user: OAuthSession | None = None
-    did: str | None = session.get("user_did")
-    if did is not None:
-        db = get_db(app)
-        row = db.execute(
-            "select * from oauth_session where did = ?",
-            (did,),
-        ).fetchone()
-        user = OAuthSession(**row)
-    g.user = user
+    g.user = get_auth_session(session)
 
 
 def get_user() -> OAuthSession | None:
···
     if not username:
         return redirect(url_for("page_login"), 303)
     return redirect(url_for("oauth.oauth_start", username=username), 303)
+
+
[e-mail]("/auth/logout")
+def auth_logout():
+    session.clear()
+    return redirect("/", 303)
 
 
 @app.get("/editor")
···
     return redirect("/editor", 303)
 
 
[e-mail]("/terms")
+def page_terms():
+    return "come back soon"
+
+
 def load_links(pds: str, did: str, reload: bool = False) -> list[dict[str, str]] | None:
     if did in links and not reload:
         app.logger.debug(f"returning cached links for {did}")
···
         "rkey": rkey,
         "record": record,
     }
+
+    def update_dpop_pds_nonce(nonce: str):
+        session_ = user._replace(dpop_pds_nonce=nonce)
+        save_auth_session(session, session_)
+
     response = pds_authed_req(
         method="POST",
         url=endpoint,
         body=body,
         user=user,
-        db=get_db(app),
+        update_dpop_pds_nonce=update_dpop_pds_nonce,
     )
     if not response or not response.ok:
         app.logger.warning("PDS HTTP ERROR")
-
-
-# AUTH
-
-
[e-mail]("/auth/logout")
-def auth_logout():
-    user = get_user()
-    if user is not None:
-        db = get_db(app)
-        cursor = db.cursor()
-        _ = cursor.execute("delete from oauth_session where did = ?", (user.did,))
-        db.commit()
-        cursor.close()
-    session.clear()
-    return redirect("/", 303)

src/oauth.py

···
+from typing import NamedTuple
 from authlib.jose import JsonWebKey, Key
 from flask import Blueprint, current_app, jsonify, redirect, request, session, url_for
+from flask.sessions import SessionMixin
 from urllib.parse import urlencode
 
 import json
···
 )
 from .atproto.oauth import initial_token_request, send_par_auth_request
 from .security import is_safe_url
-from .types import OAuthAuthRequest
-from .db import get_db
+from .types import OAuthAuthRequest, OAuthSession
 
 oauth = Blueprint("oauth", __name__, url_prefix="/oauth")
 
···
     par_request_uri: str = resp.json()["request_uri"]
     current_app.logger.debug(f"saving oauth_auth_request to DB  state={state}")
 
-    db = get_db(current_app)
-    cursor = db.cursor()
-    _ = cursor.execute(
-        "insert or replace into oauth_auth_requests values (?, ?, ?, ?, ?, ?, ?, ?, ?)",
-        (
-            state,
-            authserver_meta["issuer"],
-            did,
-            handle,
-            pds_url,
-            pkce_verifier,
-            scope,
-            dpop_authserver_nonce,
-            dpop_private_jwk.as_json(is_private=True),
-        ),
+    oauth_request = OAuthAuthRequest(
+        state,
+        authserver_meta["issuer"],
+        did,
+        handle,
+        pds_url,
+        pkce_verifier,
+        scope,
+        dpop_authserver_nonce,
+        dpop_private_jwk.as_json(is_private=True),
     )
-    db.commit()
-    cursor.close()
+    save_auth_request(session, oauth_request)
 
     auth_endpoint = authserver_meta["authorization_endpoint"]
     assert is_safe_url(auth_endpoint)
···
     authserver_iss = request.args["iss"]
     authorization_code = request.args["code"]
 
-    db = get_db(current_app)
-    cursor = db.cursor()
-
-    row = cursor.execute(
-        "select * from oauth_auth_requests where state = ?", (state,)
-    ).fetchone()
-    try:
-        auth_request = OAuthAuthRequest(**row)
-    except TypeError:
+    auth_request = get_auth_request(session)
+    if not auth_request:
         return redirect(url_for("page_login"), 303)
 
     current_app.logger.debug(f"Deleting auth request for state={state}")
-    _ = cursor.execute("delete from oauth_auth_requests where state = ?", (state,))
-    db.commit()
+    delete_auth_request(session)
 
     assert auth_request.authserver_iss == authserver_iss
     assert auth_request.state == state
···
 
     row = auth_request
 
-    did = auth_request.did
     if row.did:
         # If we started with an account identifier, this is simple
         did, handle, pds_url = row.did, row.handle, row.pds_url
···
         assert authserver_url == authserver_iss
 
     assert row.scope == tokens.scope
+    assert pds_url is not None
 
-    current_app.logger.debug("storing user did and handle")
-    db = get_db(current_app)
-    cursor = db.cursor()
-    _ = cursor.execute(
-        "insert or replace into oauth_session values (?, ?, ?, ?, ?, ?, ?, ?, ?)",
-        (
-            did,
-            handle,
-            pds_url,
-            authserver_iss,
-            tokens.access_token,
-            tokens.refresh_token,
-            dpop_authserver_nonce,
-            None,
-            auth_request.dpop_private_jwk,
-        ),
+    current_app.logger.debug("storing user oauth session")
+    oauth_session = OAuthSession(
+        did,
+        handle,
+        pds_url,
+        authserver_iss,
+        tokens.access_token,
+        tokens.refresh_token,
+        dpop_authserver_nonce,
+        None,
+        auth_request.dpop_private_jwk,
     )
-    db.commit()
-    cursor.close()
-
-    session["user_did"] = did
-    session["user_handle"] = auth_request.handle
+    save_auth_session(session, oauth_session)
 
     return redirect(url_for("page_login"))
 
···
     CLIENT_SECRET_JWK = JsonWebKey.import_key(current_app.config["CLIENT_SECRET_JWK"])
     CLIENT_PUB_JWK = json.loads(CLIENT_SECRET_JWK.as_json(is_private=False))
     return jsonify({"keys": [CLIENT_PUB_JWK]})
+
+
+# Session storage
+
+
+def save_auth_request(session: SessionMixin, request: OAuthAuthRequest):
+    return _set_into_session(session, "oauth_auth_request", request)
+
+
+def save_auth_session(session: SessionMixin, auth_session: OAuthSession):
+    return _set_into_session(session, "oauth_auth_session", auth_session)
+
+
+def delete_auth_request(session: SessionMixin):
+    return _delete_from_session(session, "oauth_auth_request")
+
+
+def delete_auth_session(session: SessionMixin):
+    return _delete_from_session(session, "oauth_auth_session")
+
+
+def get_auth_request(session: SessionMixin) -> OAuthAuthRequest | None:
+    try:
+        return OAuthAuthRequest(**session["oauth_auth_request"])
+    except TypeError as exception:
+        current_app.logger.debug("unable to load oauth_auth_request")
+        current_app.logger.debug(exception)
+        return None
+
+
+def get_auth_session(session: SessionMixin) -> OAuthSession | None:
+    try:
+        return OAuthSession(**session["oauth_auth_session"])
+    except TypeError as exception:
+        current_app.logger.debug("unable to load oauth_auth_session")
+        current_app.logger.debug(exception)
+        return None
+
+
+def _set_into_session(session: SessionMixin, key: str, value: NamedTuple):
+    session[key] = value._asdict()
+
+
+def _delete_from_session(session: SessionMixin, key: str):
+    del session[key]

src/schema.sql

···
-create table if not exists oauth_auth_requests (
-    state text not null primary key,
-    authserver_iss text not null,
-    did text,
-    handle text,
-    pds_url text,
-    pkce_verifier text not null,
-    scope text not null,
-    dpop_authserver_nonce text not null,
-    dpop_private_jwk text not null
-) strict, without rowid;
-
-create table if not exists oauth_sessions (
-    did text not null primary key,
-    handle text,
-    pds_url text not null,
-    authserver_iss text not null,
-    access_token text,
-    refresh_token text,
-    dpop_authserver_nonce text not null,
-    dpop_pds_nonce text,
-    dpop_private_jwk text not null
-) strict, without rowid;
+-- empty for now