[Web][Dovecot] Rework imapsync: sync sources, OAuth2 (XOAUTH2), worker pool, Python runner

data/Dockerfiles/dovecot/imapsync_runner.py

@@ -0,0 +1,274 @@
+#!/usr/bin/python3
+
+import fcntl
+import json
+import os
+import re
+import signal
+import subprocess
+import sys
+import tempfile
+import time
+import MySQLdb
+import MySQLdb.cursors
+
+run_dir = '/tmp'
+# Max size of imapsync output stored in returned_text (the tail holds the summary)
+max_output = 16 * 1024 * 1024
+lock_file = run_dir + '/imapsync_busy.lock'
+
+# Next due job: higher prio first, never-run jobs first, then oldest run
+next_sql = """SELECT i.id, i.user1, i.user2, i.password1, i.exclude, i.delete2duplicates, i.maxage,
+  i.subfolder2, i.delete1, i.delete2, i.automap, i.skipcrossduplicates, i.maxbytespersecond, i.custom_params,
+  i.subscribeall, i.timeout1, i.timeout2, i.dry, i.source_id,
+  s.host1, s.port1, s.enc1, s.auth_type, s.oauth_flow, s.oauth_access_token, s.oauth_token_expires
+    FROM imapsync i JOIN imapsync_source s ON i.source_id = s.id
+      WHERE i.active = 1
+        AND s.active = 1
+        AND i.is_running = 0
+        AND (UNIX_TIMESTAMP(NOW()) - UNIX_TIMESTAMP(i.last_run) > i.mins_interval * 60 OR i.last_run IS NULL)
+  ORDER BY i.prio DESC, i.last_run IS NOT NULL, i.last_run ASC, i.id ASC
+  LIMIT 1"""
+# Atomic claim, a job is only started when this flips is_running 0 -> 1
+claim_sql = """UPDATE imapsync SET is_running = 1, success = NULL, exit_status = NULL
+  WHERE id = %s AND is_running = 0 AND active = 1"""
+# Token of authorization_code sources: source + local mailbox (user2) + remote user (user1)
+token_sql = """SELECT access_token, token_expires FROM imapsync_source_oauth_token
+  WHERE source_id = %s AND owner = %s AND username = %s"""
+
+conn = None
+# Popen => job of running workers
+running = {}
+stopping = False
+
+def num(value):
+  try:
+    return int(value)
+  except (TypeError, ValueError):
+    return 0
+
+def unlink(files):
+  for path in files:
+    try:
+      os.unlink(path)
+    except OSError:
+      pass
+
+def imapsync_running():
+  for pid in filter(str.isdigit, os.listdir('/proc')):
+    try:
+      with open('/proc/' + pid + '/cmdline', 'rb') as f:
+        cmdline = f.read().replace(b'\0', b' ')
+    except OSError:
+      continue
+    if re.search(rb'imapsync\s', cmdline, re.I):
+      return True
+  return False
+
+def connect():
+  return MySQLdb.connect(unix_socket='/var/run/mysqld/mysqld.sock', user=os.environ['DBUSER'],
+    passwd=os.environ['DBPASS'], db=os.environ['DBNAME'], charset='utf8mb4', autocommit=True)
+
+# Reconnect when the server dropped the connection (syncs can outlast wait_timeout)
+def query(sql, args=None):
+  global conn
+  try:
+    conn.ping()
+  except MySQLdb.Error:
+    conn = connect()
+  cur = conn.cursor(MySQLdb.cursors.DictCursor)
+  cur.execute(sql, args)
+  return cur
+
+# Store a result and release the job
+def finish_job(job_id, text, success, exit_status):
+  try:
+    query("UPDATE imapsync SET returned_text = %s, success = %s, exit_status = %s, last_run = NOW(), is_running = 0 WHERE id = %s",
+      (text, success, exit_status, job_id))
+  except MySQLdb.Error:
+    query("UPDATE imapsync SET success = 0, last_run = NOW(), is_running = 0 WHERE id = %s", (job_id,))
+
+# Max concurrent imapsync processes and global per-process bandwidth cap in bytes/s (0 = off),
+# re-read before every claim round so changes apply to long runs
+def read_settings():
+  settings = {r['name']: r['value'] for r in query("SELECT name, value FROM imapsync_settings")}
+  return max(num(settings.get('max_parallel')), 1), max(num(settings.get('max_bytes_per_second')), 0)
+
+# Build temp files and argv for a claimed job. Returns None when the job was finished without running imapsync.
+def prepare_job(j, global_bw, files):
+  def tmpfile(content):
+    fd, path = tempfile.mkstemp(prefix='imapsync.', dir=run_dir)
+    files.append(path)
+    with os.fdopen(fd, 'w', encoding='utf-8') as f:
+      f.write(content)
+    return path
+
+  auth_type = j['auth_type'] or 'PLAIN'
+  if auth_type == 'XOAUTH2':
+    # authorization_code: token of this local mailbox + remote user, client_credentials: shared source token
+    token, expires = j['oauth_access_token'], j['oauth_token_expires']
+    if j['oauth_flow'] == 'authorization_code':
+      row = query(token_sql, (j['source_id'], j['user2'], j['user1'])).fetchone() or {}
+      token, expires = row.get('access_token'), row.get('token_expires')
+    # Skip sync if the cached token is missing or about to expire (<2 min)
+    if not token or expires is None or expires - time.time() < 120:
+      finish_job(j['id'], "OAuth access token for the configured source is missing or expires too soon; refresh cron will retry.",
+        0, 'OAUTH_TOKEN_NOT_READY')
+      return None
+    auth_args = ['--authmech1', 'XOAUTH2', '--oauthaccesstoken1', tmpfile(token + '\n')]
+  else:
+    auth_args = ['--passfile1', tmpfile(j['password1'] + '\n')]
+    # imapsync auto-picks PLAIN, only pass --authmech1 for LOGIN/CRAM-MD5
+    if auth_type != 'PLAIN':
+      auth_args += ['--authmech1', auth_type]
+
+  # Custom params: JSON array of {o,v}, validated by the backend on write. One argv token per pair,
+  # "--opt" or "--opt=value" (value glued to its option, no shell)
+  custom_params = []
+  try:
+    pairs = json.loads(j['custom_params'] or '[]')
+  except ValueError:
+    pairs = []
+  for p in pairs if isinstance(pairs, list) else []:
+    if not isinstance(p, dict) or not isinstance(p.get('o'), str) or not re.fullmatch(r'[a-z0-9]+', p['o']):
+      continue
+    value = '' if p.get('v') is None else str(p['v'])
+    custom_params.append('--' + p['o'] if value == '' else '--' + p['o'] + '=' + value)
+
+  # Effective per-process bandwidth cap (bytes/s): per-job value, capped by the global limit
+  eff_bw = num(j['maxbytespersecond'])
+  if global_bw > 0 and (eff_bw <= 0 or eff_bw > global_bw):
+    eff_bw = global_bw
+
+  cmd = ['/usr/local/bin/imapsync', '--tmpdir', '/tmp', '--nofoldersizes', '--addheader']
+  if num(j['timeout1']) > 0:
+    cmd += ['--timeout1', str(j['timeout1'])]
+  if num(j['timeout2']) > 0:
+    cmd += ['--timeout2', str(j['timeout2'])]
+  if j['exclude']:
+    cmd += ['--exclude', j['exclude']]
+  if j['subfolder2']:
+    cmd += ['--subfolder2', j['subfolder2']]
+  if num(j['maxage']) != 0:
+    cmd += ['--maxage', str(j['maxage'])]
+  if eff_bw > 0:
+    cmd += ['--maxbytespersecond', str(eff_bw)]
+  for col, opt in (('delete2duplicates', '--delete2duplicates'), ('subscribeall', '--subscribeall'),
+      ('delete1', '--delete'), ('delete2', '--delete2'), ('automap', '--automap'),
+      ('skipcrossduplicates', '--skipcrossduplicates')):
+    if j[col] == 1:
+      cmd.append(opt)
+  enc1 = {'TLS': '--tls1', 'SSL': '--ssl1'}.get(j['enc1'])
+  if enc1:
+    cmd.append(enc1)
+  cmd += ['--host1', j['host1'], '--user1', j['user1']] + auth_args + ['--port1', str(j['port1'])]
+  cmd += ['--host2', 'localhost', '--user2', j['user2'] + '*' + master_user, '--passfile2', tmpfile(master_pass + '\n')]
+  if j['dry'] == 1:
+    cmd.append('--dry')
+  cmd += ['--no-modulesversion', '--noreleasecheck'] + custom_params
+
+  return {'id': j['id'], 'files': files, 'outfile': tmpfile(''), 'cmd': cmd}
+
+# Worker: imapsync with output to a file, no inherited fds (DB socket, lock)
+def start_job(job):
+  with open(job['outfile'], 'wb') as out:
+    return subprocess.Popen(job['cmd'], stdin=subprocess.DEVNULL, stdout=out, stderr=subprocess.STDOUT)
+
+# Store the result of a finished worker and remove its temp files
+def reap_job(job):
+  stdout = b''
+  try:
+    with open(job['outfile'], 'rb') as f:
+      size = os.fstat(f.fileno()).st_size
+      if size > max_output:
+        f.seek(size - max_output)
+        stdout = b'[output truncated, last %d bytes]\n' % max_output
+      stdout += f.read()
+  except OSError:
+    pass
+  unlink(job['files'])
+  text = stdout.decode('utf-8', 'replace')
+  # Last match: the echoed command line at the top may contain user controlled values
+  m = re.findall(r'Exiting\swith\sreturn\svalue\s(\d+)\s\(([^:)]+)', text)
+  exit_code, exit_status = m[-1] if m else (None, None)
+  if exit_code is None:
+    text = 'Could not start or finish imapsync\n' + text
+  finish_job(job['id'], text, 1 if exit_code is not None and int(exit_code) == 0 else 0, exit_status)
+
+# Stop claiming new jobs, pass the signal on to the running imapsync processes
+def stop(signum, frame):
+  global stopping
+  stopping = True
+  for proc in list(running):
+    proc.send_signal(signal.SIGTERM)
+
+# Released by the kernel when the process ends, a crashed run leaves no stale lock
+lock_fd = os.open(lock_file, os.O_RDWR | os.O_CREAT, 0o600)
+try:
+  fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+except OSError:
+  sys.exit("can't lock " + lock_file)
+
+if imapsync_running():
+  print('imapsync is active, exiting...')
+  sys.exit(0)
+
+for sig in (signal.SIGINT, signal.SIGTERM, signal.SIGQUIT):
+  signal.signal(sig, stop)
+
+conn = connect()
+# Stale flags from crashed runs, safe only here: under the lock and with no imapsync running
+query("UPDATE imapsync SET is_running = 0")
+
+with open('/etc/sogo/sieve.creds') as f:
+  master_user, master_pass = [s.strip() for s in f.readline().split(':', 1)]
+
+# Worker pool: keep up to max_parallel workers busy, claim the next due job whenever a slot is free
+# (after every worker exit and at least every 10 seconds while workers run)
+try:
+  while True:
+    for proc, job in list(running.items()):
+      if proc.poll() is not None:
+        del running[proc]
+        reap_job(job)
+
+    max_parallel, global_bw = read_settings()
+    attempts = 0
+    while not stopping and len(running) < max_parallel and attempts < 100:
+      attempts += 1
+      j = query(next_sql).fetchone()
+      if not j:
+        break
+      if query(claim_sql, (j['id'],)).rowcount != 1:
+        continue
+
+      files = []
+      try:
+        job = prepare_job(j, global_bw, files)
+      except Exception as e:
+        print('imapsync job %d: %s' % (j['id'], e), file=sys.stderr)
+        finish_job(j['id'], 'Could not start or finish imapsync', 0, None)
+        job = None
+      if not job:
+        unlink(files)
+        continue
+      try:
+        running[start_job(job)] = job
+      except OSError as e:
+        # Release the claim, the job is picked up again later
+        print('imapsync job %d: %s' % (j['id'], e), file=sys.stderr)
+        unlink(files)
+        query("UPDATE imapsync SET is_running = 0 WHERE id = %s", (j['id'],))
+        break
+
+    if not running:
+      break
+    # A worker exit ends the wait early
+    for _ in range(50):
+      if any(proc.poll() is not None for proc in running):
+        break
+      time.sleep(0.2)
+finally:
+  # Never leave password/token files behind
+  for job in running.values():
+    unlink(job['files'])