better quota logic - KV quota for CF worker not fully working - needs improvements - WIP but I take a break for today..

This commit is contained in:
ChrispyBacon-dev 2026-04-14 21:44:19 +02:00
parent e022473439
commit f834516eda
8 changed files with 268 additions and 10 deletions

View file

@ -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"}],

View file

@ -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();

View file

@ -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);

View file

@ -460,12 +460,13 @@
<p class="text-sm opacity-60 mb-1" id="emailEditQuotaTarget"></p>
<p class="text-sm mb-4" id="emailEditQuotaUsage"></p>
<div>
<input type="range" id="editMailboxQuota" min="0" max="10" step="1" value="6" class="range range-sm w-full" oninput="emailUpdateQuotaLabel('editMailboxQuotaLabel', this.value)" />
<input type="range" id="editMailboxQuota" min="0" max="10" step="1" value="6" class="range range-sm w-full" oninput="emailUpdateQuotaLabel('editMailboxQuotaLabel', this.value, 'emailEditQuotaGraceInfo')" />
<div class="flex justify-between text-xs mt-1">
<span class="opacity-50">100 MB</span>
<span id="editMailboxQuotaLabel" class="font-semibold">10 GB</span>
<span class="opacity-50">Unlimited</span>
</div>
<p id="emailEditQuotaGraceInfo" class="text-xs opacity-50 mt-1"></p>
</div>
<div class="modal-action">
<form method="dialog"><button class="btn btn-ghost">Cancel</button></form>

View file

@ -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')

View file

@ -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()

View file

@ -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",

View file

@ -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)