132 lines
4.5 KiB
Python
132 lines
4.5 KiB
Python
from os import getenv, urandom
|
|
import os
|
|
import time
|
|
import json
|
|
|
|
from nacl.bindings import crypto_sign_seed_keypair
|
|
from tonsdk.utils import Address
|
|
|
|
from app.core._blockchain.ton.wallet_v3cr3 import WalletV3CR3
|
|
from app.core.logger import make_log
|
|
from sqlalchemy import create_engine, inspect
|
|
from sqlalchemy.orm import Session
|
|
from typing import Optional
|
|
from app.core.models._config import ServiceConfigValue
|
|
|
|
|
|
def _load_seed_from_env_or_generate() -> bytes:
|
|
seed_hex = os.getenv("TON_INIT_HOT_SEED")
|
|
if seed_hex:
|
|
make_log("HotWallet", "Loaded seed from env")
|
|
return bytes.fromhex(seed_hex)
|
|
make_log("HotWallet", "No seed provided; generating ephemeral seed", level='info')
|
|
return urandom(32)
|
|
|
|
|
|
def _init_seed_via_db() -> bytes:
|
|
"""Store and read hot seed from PostgreSQL service_config (key='private_key').
|
|
Primary node writes it once; workers wait until it appears.
|
|
"""
|
|
from app.core._config import DATABASE_URL
|
|
|
|
engine = create_engine(DATABASE_URL, pool_pre_ping=True)
|
|
role = os.getenv("NODE_ROLE", "worker").lower()
|
|
|
|
# Best-effort: ensure service_config table exists before waiting on it.
|
|
# This complements the synchronous init in app.__main__ and protects
|
|
# against ordering issues where _secrets is imported before that init runs.
|
|
try:
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
|
|
with engine.begin() as conn:
|
|
inspector = inspect(conn)
|
|
if not inspector.has_table('service_config'):
|
|
ServiceConfigValue.__table__.create(bind=conn, checkfirst=True)
|
|
except SQLAlchemyError as exc:
|
|
make_log("HotWallet", f"Failed to ensure service_config table: {exc}", level="error")
|
|
except Exception:
|
|
# Avoid failing hard here; the fallback waiter below may still succeed
|
|
pass
|
|
|
|
def db_ready(conn) -> bool:
|
|
try:
|
|
inspector = inspect(conn)
|
|
return inspector.has_table('service_config')
|
|
except Exception:
|
|
return False
|
|
|
|
# Wait for table to exist
|
|
start = time.time()
|
|
# Wait for table existence, reconnecting to avoid stale transactions
|
|
while True:
|
|
with engine.connect() as conn:
|
|
if db_ready(conn):
|
|
break
|
|
time.sleep(0.5)
|
|
if time.time() - start > 120:
|
|
raise TimeoutError("service_config table not available")
|
|
|
|
def read_seed() -> Optional[bytes]:
|
|
# Use a fresh connection/session per read to avoid snapshot staleness
|
|
try:
|
|
with engine.connect() as rconn:
|
|
with Session(bind=rconn) as s:
|
|
row = s.query(ServiceConfigValue).filter(ServiceConfigValue.key == 'private_key').first()
|
|
if not row:
|
|
return None
|
|
packed = row.packed_value or {}
|
|
if isinstance(packed, str):
|
|
packed = json.loads(packed)
|
|
seed_hex = packed.get('value')
|
|
return bytes.fromhex(seed_hex) if seed_hex else None
|
|
except Exception:
|
|
return None
|
|
|
|
seed = read_seed()
|
|
if seed:
|
|
return seed
|
|
|
|
if role == "primary":
|
|
seed = _load_seed_from_env_or_generate()
|
|
# Try insert; if another primary raced, ignore
|
|
try:
|
|
with engine.connect() as wconn:
|
|
with Session(bind=wconn) as s:
|
|
s.add(ServiceConfigValue(key='private_key', packed_value={"value": seed.hex()}))
|
|
s.commit()
|
|
make_log("HotWallet", "Seed saved in service_config by primary", level='info')
|
|
return seed
|
|
except Exception:
|
|
# Read again in case of race
|
|
seed2 = read_seed()
|
|
if seed2:
|
|
return seed2
|
|
raise
|
|
else:
|
|
make_log("HotWallet", "Worker waiting for seed in service_config...", level='info')
|
|
while True:
|
|
seed = read_seed()
|
|
if seed:
|
|
return seed
|
|
time.sleep(0.5)
|
|
|
|
|
|
_extra_ton_wallet_options = {}
|
|
if getenv('TON_CUSTOM_WALLET_ADDRESS'):
|
|
_extra_ton_wallet_options['address'] = Address(getenv('TON_CUSTOM_WALLET_ADDRESS'))
|
|
|
|
|
|
def _init_wallet():
|
|
# Primary writes to DB; workers wait and read from DB
|
|
hot_seed_bytes = _init_seed_via_db()
|
|
pub, priv = crypto_sign_seed_keypair(hot_seed_bytes)
|
|
wallet = WalletV3CR3(
|
|
private_key=priv,
|
|
public_key=pub,
|
|
**_extra_ton_wallet_options
|
|
)
|
|
return hot_seed_bytes, pub, priv, wallet
|
|
|
|
|
|
hot_seed, hot_pubkey, hot_privkey, service_wallet = _init_wallet()
|