From f834516eda22ab8599e9afdac67e3d1dea8ac77a Mon Sep 17 00:00:00 2001 From: ChrispyBacon-dev Date: Tue, 14 Apr 2026 21:44:19 +0200 Subject: [PATCH] better quota logic - KV quota for CF worker not fully working - needs improvements - WIP but I take a break for today.. --- dockflare/app/core/email_manager.py | 19 +++ .../core/worker_templates/inbound_worker.js | 18 +++ dockflare/app/static/js/main.js | 24 +++- dockflare/app/templates/email.html | 3 +- dockflare/app/web/email_routes.py | 89 ++++++++++++++ mail-manager/app/api/routes.py | 7 +- mail-manager/app/api/webhook.py | 115 +++++++++++++++++- mail-manager/app/core/database.py | 3 + 8 files changed, 268 insertions(+), 10 deletions(-) diff --git a/dockflare/app/core/email_manager.py b/dockflare/app/core/email_manager.py index 74bffc6..7628d65 100644 --- a/dockflare/app/core/email_manager.py +++ b/dockflare/app/core/email_manager.py @@ -310,6 +310,25 @@ def scrub_email_dns_records(zone_id, zone_name): errors.append(f"DNS list CNAME: {e}") return errors +def create_kv_namespace(title): + res = cf_api_request('POST', f'/accounts/{config.CF_ACCOUNT_ID}/storage/kv/namespaces', + json_data={'title': title}) + return res.get('result', {}).get('id') + +def update_kv_entry(namespace_id, key, value_dict): + url = f"{config.CF_API_BASE_URL}/accounts/{config.CF_ACCOUNT_ID}/storage/kv/namespaces/{namespace_id}/values/{key}" + headers = {"Authorization": f"Bearer {config.CF_API_TOKEN}", "Content-Type": "text/plain"} + response = requests.put(url, data=json.dumps(value_dict), headers=headers, timeout=10) + response.raise_for_status() + return response.json() + +def delete_kv_entry(namespace_id, key): + try: + cf_api_request('DELETE', + f'/accounts/{config.CF_ACCOUNT_ID}/storage/kv/namespaces/{namespace_id}/values/{key}') + except Exception as e: + logging.warning(f"Could not delete KV entry {key} from {namespace_id}: {e}") + def setup_catchall_routing_rule(zone_id, worker_name): data = { "matchers": [{"type": "all"}], diff --git a/dockflare/app/core/worker_templates/inbound_worker.js b/dockflare/app/core/worker_templates/inbound_worker.js index 6ed6a9f..0d0107c 100644 --- a/dockflare/app/core/worker_templates/inbound_worker.js +++ b/dockflare/app/core/worker_templates/inbound_worker.js @@ -39,6 +39,24 @@ export default { return; } + // Check quota KV before accepting — reject at SMTP level so sender gets a bounce + if (typeof env.QUOTA_KV !== 'undefined') { + try { + const quota = await env.QUOTA_KV.get(message.to, "json"); + if (quota && quota.hard_limit_bytes > 0) { + const currentSize = quota.current_size_bytes || 0; + const msgSize = message.rawSize || 0; + if (currentSize + msgSize > quota.hard_limit_bytes) { + message.setReject("550 5.2.2 Mailbox full"); + return; + } + } + } catch (kvErr) { + // KV unavailable — fall through, webhook safety net handles enforcement + console.warn(`KV quota check failed for ${message.to}: ${kvErr.message}`); + } + } + const messageId = crypto.randomUUID(); const r2Key = `temp_cache/${messageId}.eml`; const receivedAt = new Date().toISOString(); diff --git a/dockflare/app/static/js/main.js b/dockflare/app/static/js/main.js index d14a31d..023a806 100644 --- a/dockflare/app/static/js/main.js +++ b/dockflare/app/static/js/main.js @@ -2331,9 +2331,27 @@ function _fmtBytes(b) { return b.toFixed(1) + '\u00a0' + u[i]; } -function emailUpdateQuotaLabel(labelId, stepIndex) { +function _calcGrace(quotaBytes) { + if (!quotaBytes || quotaBytes <= 0) return 0; + return Math.max(Math.round(quotaBytes * 0.15), 10 * 1024 * 1024); +} + +function emailUpdateQuotaLabel(labelId, stepIndex, graceInfoId) { const el = document.getElementById(labelId); - if (el) el.textContent = QUOTA_STEPS[parseInt(stepIndex)]?.label ?? '10 GB'; + const step = QUOTA_STEPS[parseInt(stepIndex)]; + if (el) el.textContent = step?.label ?? '10 GB'; + if (graceInfoId) { + const graceEl = document.getElementById(graceInfoId); + if (graceEl) { + const quota = step?.bytes ?? 0; + if (quota > 0) { + const grace = _calcGrace(quota); + graceEl.textContent = `Hard limit: ${_fmtBytes(quota + grace)} (includes ${_fmtBytes(grace)} grace buffer)`; + } else { + graceEl.textContent = ''; + } + } + } } function emailLoadStats() { @@ -2983,7 +3001,7 @@ function emailEditQuota(address, domain, currentQuotaBytes) { const stepIndex = QUOTA_STEPS.findIndex(s => s.bytes === currentQuotaBytes); const slider = document.getElementById('editMailboxQuota'); slider.value = stepIndex >= 0 ? stepIndex : 6; - emailUpdateQuotaLabel('editMailboxQuotaLabel', slider.value); + emailUpdateQuotaLabel('editMailboxQuotaLabel', slider.value, 'emailEditQuotaGraceInfo'); const submitBtn = document.getElementById('emailEditQuotaSubmitBtn'); const handler = async () => { submitBtn.removeEventListener('click', handler); diff --git a/dockflare/app/templates/email.html b/dockflare/app/templates/email.html index aa3d13f..a7e854e 100644 --- a/dockflare/app/templates/email.html +++ b/dockflare/app/templates/email.html @@ -460,12 +460,13 @@

- +
100 MB 10 GB Unlimited
+

diff --git a/dockflare/app/web/email_routes.py b/dockflare/app/web/email_routes.py index a6446b6..9342428 100644 --- a/dockflare/app/web/email_routes.py +++ b/dockflare/app/web/email_routes.py @@ -132,6 +132,14 @@ def setup_email_domain(): webmail_hostname = f"mail.{zone_name}" webhook_url = f"https://{webmail_hostname}/api/v1/webhook/inbound" + quota_kv_ns_id = None + try: + quota_kv_ns_id = email_manager.create_kv_namespace( + f"dockflare-quota-{zone_name.replace('.', '-')}" + ) + except Exception as e: + logging.warning(f"Could not create quota KV namespace for {zone_name}: {e}") + inbound_bindings = [ {"type": "r2_bucket", "name": "EMAIL_BUCKET", "bucket_name": bucket_name}, {"type": "plain_text", "name": "WEBHOOK_URL", "text": webhook_url}, @@ -139,6 +147,10 @@ def setup_email_domain(): {"type": "plain_text", "name": "ALLOWED_RECIPIENTS", "text": "[]"}, {"type": "plain_text", "name": "DOMAIN_NAME", "text": zone_name} ] + if quota_kv_ns_id: + inbound_bindings.append( + {"type": "kv_namespace", "name": "QUOTA_KV", "namespace_id": quota_kv_ns_id} + ) email_manager.deploy_worker(inbound_worker_name, _read_worker_template('inbound_worker.js'), inbound_bindings) email_manager.set_worker_cron(inbound_worker_name, ['*/5 * * * *']) @@ -166,6 +178,7 @@ def setup_email_domain(): 'outbound_worker_name': outbound_worker_name, 'outbound_worker_url': outbound_worker_url, 'outbound_auth_secret': outbound_auth_secret, + 'quota_kv_namespace_id': quota_kv_ns_id, 'mailboxes': {} } @@ -370,6 +383,19 @@ def _redeploy_inbound_worker(email_cfg, domain): all_addresses = list(d['mailboxes'].keys()) webmail_hostname = f"mail.{domain}" webhook_url = f"https://{webmail_hostname}/api/v1/webhook/inbound" + + kv_ns_id = d.get('quota_kv_namespace_id') + if not kv_ns_id: + try: + kv_ns_id = email_manager.create_kv_namespace( + f"dockflare-quota-{domain.replace('.', '-')}" + ) + email_cfg['domains'][domain]['quota_kv_namespace_id'] = kv_ns_id + save_email_config(email_cfg) + logging.info(f"Created quota KV namespace for existing domain {domain}: {kv_ns_id}") + except Exception as e: + logging.warning(f"Could not create quota KV namespace for {domain}: {e}") + inbound_bindings = [ {"type": "r2_bucket", "name": "EMAIL_BUCKET", "bucket_name": d['r2_bucket']}, {"type": "plain_text", "name": "WEBHOOK_URL", "text": webhook_url}, @@ -377,6 +403,11 @@ def _redeploy_inbound_worker(email_cfg, domain): {"type": "plain_text", "name": "ALLOWED_RECIPIENTS", "text": json.dumps(all_addresses)}, {"type": "plain_text", "name": "DOMAIN_NAME", "text": domain} ] + if kv_ns_id: + inbound_bindings.append( + {"type": "kv_namespace", "name": "QUOTA_KV", "namespace_id": kv_ns_id} + ) + email_manager.deploy_worker(d['inbound_worker_name'], _read_worker_template('inbound_worker.js'), inbound_bindings) email_manager.set_worker_cron(d['inbound_worker_name'], ['*/5 * * * *']) @@ -464,6 +495,31 @@ def set_mailbox_quota(): ) except Exception: pass + + kv_ns_id = email_cfg['domains'][domain].get('quota_kv_namespace_id') + if kv_ns_id and quota_bytes and quota_bytes > 0 and token: + try: + current_size = 0 + mb_resp = requests.get( + f"{config.MAIL_MANAGER_INTERNAL_URL}/api/v1/mailboxes", + headers={'Authorization': f'Bearer {token}'}, + timeout=5, + ) + if mb_resp.ok: + for mb in mb_resp.json(): + if mb.get('address') == address: + current_size = mb.get('storage_bytes', 0) + break + grace = max(int(quota_bytes * 0.15), 10 * 1024 * 1024) + email_manager.update_kv_entry(kv_ns_id, address, { + "hard_limit_bytes": quota_bytes + grace, + "current_size_bytes": current_size, + }) + except Exception as e: + logging.warning(f"Could not update quota KV for {address}: {e}") + elif kv_ns_id and quota_bytes == 0: + email_manager.delete_kv_entry(kv_ns_id, address) + return jsonify({'success': True}) @email_bp.route('/mailbox-stats', methods=['GET']) @@ -1074,6 +1130,7 @@ def internal_mail_config(): 'webhook_secret': d.get('webhook_secret', ''), 'outbound_worker_url': d.get('outbound_worker_url', ''), 'outbound_auth_secret': d.get('outbound_auth_secret', ''), + 'quota_kv_namespace_id': d.get('quota_kv_namespace_id', ''), 'mailboxes': { addr: {'display_name': m.get('display_name', ''), 'quota_bytes': m.get('quota_bytes', 10737418240)} for addr, m in d.get('mailboxes', {}).items() @@ -1090,6 +1147,38 @@ def internal_mail_config(): 'domains': domains_out }) +@email_bp.route('/internal/quota-kv-sync', methods=['POST']) +def quota_kv_sync(): + if not _check_internal_request(): + return jsonify({'error': 'forbidden'}), 403 + data = request.get_json(force=True, silent=True) or {} + domain = data.get('domain') + address = data.get('address') + current_size_bytes = data.get('current_size_bytes', 0) + if not domain or not address: + return jsonify({'error': 'missing domain or address'}), 400 + cfg = config.EMAIL_CONFIG + if not cfg or domain not in cfg.get('domains', {}): + return jsonify({'status': 'domain_not_found'}), 200 + d = cfg['domains'][domain] + kv_ns_id = d.get('quota_kv_namespace_id') + if not kv_ns_id: + return jsonify({'status': 'no_kv'}), 200 + quota_bytes = d.get('mailboxes', {}).get(address, {}).get('quota_bytes', 0) + if not quota_bytes or quota_bytes <= 0: + return jsonify({'status': 'unlimited'}), 200 + grace = max(int(quota_bytes * 0.15), 10 * 1024 * 1024) + try: + email_manager.update_kv_entry(kv_ns_id, address, { + "hard_limit_bytes": quota_bytes + grace, + "current_size_bytes": current_size_bytes, + }) + except Exception as e: + logging.warning(f"quota-kv-sync failed for {address}: {e}") + return jsonify({'error': str(e)}), 500 + return jsonify({'status': 'ok'}) + + def _restart_mail_container(): try: container = docker_client.containers.get('dockflare-mail-manager') diff --git a/mail-manager/app/api/routes.py b/mail-manager/app/api/routes.py index 7f3c925..0c389e6 100644 --- a/mail-manager/app/api/routes.py +++ b/mail-manager/app/api/routes.py @@ -170,7 +170,7 @@ def get_mailboxes(): """).fetchall()} storage = {r['mailbox_address']: r['bytes'] for r in db.execute(""" SELECT mailbox_address, COALESCE(SUM(size_bytes), 0) as bytes - FROM messages GROUP BY mailbox_address + FROM messages WHERE is_system=0 GROUP BY mailbox_address """).fetchall()} sent = {r['from_address']: r['cnt'] for r in db.execute(""" SELECT from_address, COUNT(*) as cnt FROM send_log @@ -244,7 +244,10 @@ def update_mailbox(address): if not db.execute("SELECT 1 FROM mailboxes WHERE address=?", (address,)).fetchone(): return jsonify({"error": "not found"}), 404 if 'quota_bytes' in data: - db.execute("UPDATE mailboxes SET quota_bytes=?, quota_exceeded_count=0 WHERE address=?", (data['quota_bytes'], address)) + db.execute( + "UPDATE mailboxes SET quota_bytes=?, quota_exceeded_count=0, last_quota_warning_at=NULL WHERE address=?", + (data['quota_bytes'], address) + ) if 'display_name' in data: db.execute("UPDATE mailboxes SET display_name=? WHERE address=?", (data['display_name'], address)) db.commit() diff --git a/mail-manager/app/api/webhook.py b/mail-manager/app/api/webhook.py index fe1e80a..31fa7fd 100644 --- a/mail-manager/app/api/webhook.py +++ b/mail-manager/app/api/webhook.py @@ -3,6 +3,7 @@ import hmac import hashlib import json import os +import shutil import sqlite3 import logging import uuid @@ -20,6 +21,16 @@ log = logging.getLogger(__name__) webhook_bp = Blueprint('webhook', __name__) +def _fmt_bytes(n): + if n >= 1073741824: + return f"{n / 1073741824:.1f} GB" + if n >= 1048576: + return f"{n / 1048576:.1f} MB" + if n >= 1024: + return f"{n / 1024:.1f} KB" + return f"{n} B" + + def _detect_and_log_bounce(eml_bytes, parsed): from_addr = (parsed.get('from_address') or '').lower() headers = {} @@ -282,20 +293,95 @@ def inbound(): _detect_and_log_bounce(eml_bytes, parsed) quota_row = db.execute( - "SELECT quota_bytes FROM mailboxes WHERE address=?", (to_address,) + """SELECT m.quota_bytes, m.last_quota_warning_at, d.grace_buffer_bytes + FROM mailboxes m + LEFT JOIN domain_configs d ON d.domain_name = m.domain + WHERE m.address = ?""", + (to_address,) ).fetchone() + if quota_row and quota_row['quota_bytes'] and quota_row['quota_bytes'] > 0: + quota = quota_row['quota_bytes'] + raw_buffer = quota_row['grace_buffer_bytes'] + grace = raw_buffer if raw_buffer else max(int(quota * 0.15), 10 * 1024 * 1024) + hard_limit = quota + grace + used = db.execute( - "SELECT COALESCE(SUM(size_bytes), 0) FROM messages WHERE mailbox_address=?", + "SELECT COALESCE(SUM(size_bytes), 0) FROM messages WHERE mailbox_address=? AND is_system=0", (to_address,), ).fetchone()[0] - if used > quota_row['quota_bytes']: + + if used > quota: db.execute( "UPDATE mailboxes SET quota_exceeded_count = quota_exceeded_count + 1 WHERE address=?", (to_address,), ) + log.warning("Quota exceeded for %s: %s used / %s limit", to_address, _fmt_bytes(used), _fmt_bytes(quota)) + + if not quota_row['last_quota_warning_at']: + inbox = db.execute( + "SELECT id FROM folders WHERE mailbox_address=? AND name='Inbox'", + (to_address,) + ).fetchone() + if inbox: + db.execute(""" + INSERT INTO messages ( + message_id, mailbox_address, folder_id, + from_address, from_name, to_addresses, + cc_addresses, bcc_addresses, subject, + text_body, html_body, received_at, + is_read, is_starred, is_draft, + in_reply_to, reference_ids, size_bytes, + has_attachments, headers_json, created_at, is_system + ) VALUES (?, ?, ?, 'noreply@dockflare', 'DockFlare System', ?, + '[]', '[]', + 'Action Required: Your mailbox is nearly full', + ?, '', ?, 0, 0, 0, NULL, NULL, 0, 0, '{}', ?, 1) + """, ( + f"quota-warning-{to_address}-{now}", + to_address, + inbox['id'], + f'["{to_address}"]', + ( + f"Your mailbox ({to_address}) has reached its storage quota " + f"({_fmt_bytes(quota)}). You have a grace buffer of " + f"{_fmt_bytes(grace)} before new emails are rejected.\n\n" + f"Current usage: {_fmt_bytes(used)}\n" + f"Soft limit: {_fmt_bytes(quota)}\n" + f"Hard limit: {_fmt_bytes(hard_limit)}\n\n" + f"Please delete old messages or contact your administrator " + f"to increase your quota." + ), + now, now, + )) + db.execute( + "UPDATE mailboxes SET last_quota_warning_at=? WHERE address=?", + (now, to_address) + ) + + elif used < quota * 0.90 and quota_row['last_quota_warning_at']: + db.execute( + "UPDATE mailboxes SET last_quota_warning_at=NULL WHERE address=?", + (to_address,) + ) + + db.commit() + + if used > hard_limit: + att_dir = os.path.join(config.ATTACHMENTS_PATH, str(msg_id)) + if os.path.isdir(att_dir): + shutil.rmtree(att_dir, ignore_errors=True) + db.execute("DELETE FROM messages WHERE id=?", (msg_id,)) db.commit() - log.warning("Quota exceeded for %s: %d bytes used / %d bytes limit", to_address, used, quota_row['quota_bytes']) + log.warning( + "Hard quota exceeded for %s: %s used / %s hard limit — message %d rejected", + to_address, _fmt_bytes(used), _fmt_bytes(hard_limit), msg_id + ) + try: + delete_from_r2(r2_key, domain_cfg) + except Exception: + pass + return jsonify({"status": "rejected", "reason": "over_hard_quota"}), 200 send_push_notifications(to_address, { 'message_id': msg_id, @@ -304,6 +390,27 @@ def inbound(): 'mailbox': to_address, }) _check_and_send_auto_reply(db, to_address, parsed, domain_cfg) + + master_url = os.environ.get('DOCKFLARE_MASTER_URL', '').rstrip('/') + if master_url: + try: + current_size = db.execute( + "SELECT COALESCE(SUM(size_bytes),0) FROM messages WHERE mailbox_address=? AND is_system=0", + (to_address,) + ).fetchone()[0] + _http_requests.post( + f"{master_url}/email/internal/quota-kv-sync", + json={ + 'domain': to_address.split('@')[1], + 'address': to_address, + 'current_size_bytes': int(current_size), + }, + headers={'X-Bootstrap-Token': os.environ.get('INTERNAL_BOOTSTRAP_SECRET', '')}, + timeout=3, + ) + except Exception: + pass + delete_from_r2(r2_key, domain_cfg) log.info("Inbound delivered: message=%s to=%s db_id=%s", diff --git a/mail-manager/app/core/database.py b/mail-manager/app/core/database.py index 76e3940..28d7731 100644 --- a/mail-manager/app/core/database.py +++ b/mail-manager/app/core/database.py @@ -206,6 +206,9 @@ def _migrate(conn): "CREATE INDEX IF NOT EXISTS idx_bounce_log_received_at ON bounce_log(received_at DESC)", "CREATE INDEX IF NOT EXISTS idx_bounce_log_message_id ON bounce_log(original_message_id)", "ALTER TABLE domain_configs ADD COLUMN catch_all_mailbox TEXT DEFAULT NULL", + "ALTER TABLE mailboxes ADD COLUMN last_quota_warning_at TEXT DEFAULT NULL", + "ALTER TABLE domain_configs ADD COLUMN grace_buffer_bytes INTEGER DEFAULT NULL", + "ALTER TABLE messages ADD COLUMN is_system INTEGER NOT NULL DEFAULT 0", ]: try: conn.execute(sql)