Initial.

Author: mzhang28

Diff for: Dockerfile

@@ -0,0 +1,11 @@
+FROM python:3.5.2
+
+RUN apt-get update && apt-get install -y netcat
+
+RUN mkdir /judge
+COPY requirements.txt /judge/
+WORKDIR /judge
+RUN pip3 install -r requirements.txt
+COPY . /judge/
+
+CMD ["bash", "-c", "bash wait-for-db.sh && gunicorn --worker-class eventlet --bind 0.0.0.0:80 -w 1 main:app"]

Diff for: autoscale.py

@@ -0,0 +1,169 @@
+import os
+import logging
+import time
+from collections import deque
+
+import digitalocean
+from dotenv import load_dotenv, find_dotenv
+
+from main import app
+from models import APIKey, Job
+import util
+
+load_dotenv(find_dotenv())
+
+JUDGE_URL = os.getenv('JUDGE_URL', '')
+MAX_JURIES = 10
+
+DIGITALOCEAN_API_TOKEN = os.getenv('DIGITALOCEAN_API_TOKEN', '')
+
+# TODO: Add stop command for jury systemd service
+USER_DATA_TEMPLATE = '''#!/bin/bash
+
+cat > /etc/systemd/system/docker-jury.service <<EOF
+[Unit]
+Description=Jury container
+Requires=docker.service
+After=docker.service
+
+[Service]
+Restart=always
+ExecStart=/usr/bin/docker run --cap-add=SYS_PTRACE -e JUDGE_URL={judge_url} -e JUDGE_API_KEY={api_key} easyctf/openctf-jury:latest
+ExecStop=:
+
+[Install]
+WantedBy=default.target
+EOF
+
+systemctl daemon-reload
+systemctl enable docker-jury
+systemctl start docker-jury
+'''
+
+logging.getLogger().setLevel(logging.INFO)
+logger = logging.getLogger('autoscale')
+logger.setLevel(logging.DEBUG)
+logging.info('Starting up!')
+
+
+class Cloud:
+    def get_current_jury_count(self):
+        raise NotImplemented
+
+    def create_jury(self, n=1):
+        raise NotImplemented
+
+    def destroy_jury(self, n=1):
+        raise NotImplemented
+
+
+class DigitalOcean(Cloud):
+    def __init__(self, token):
+        self.token = token
+        self.manager = digitalocean.Manager(token=self.token)
+
+    def _get_juries(self):
+        return self.manager.get_all_droplets('jury')
+
+    def get_current_jury_count(self):
+        return len(self._get_juries())
+
+    def create_jury(self, n=1):
+        for _ in range(n):
+            name = 'jury-{}'.format(util.generate_hex_string(8))
+            with app.app_context():
+                api_key = APIKey.new(name=name, perm_jury=True).key
+
+            digitalocean.Droplet(
+                token=self.token,
+                name=name,
+                region='sfo2',
+                image='docker-16-04',
+                size_slug='2gb',
+                tags=['jury'],
+                user_data=USER_DATA_TEMPLATE.format(judge_url=JUDGE_URL, api_key=api_key)
+            ).create()
+
+    def destroy_jury(self, n=1):
+        juries = self._get_juries()
+        n = min(n, len(juries))
+        for _ in range(n):
+            juries.pop().destroy()
+        return n
+
+
+class LoadIndex:
+    def __init__(self, jury_count=1):
+        self.window_size = 10
+        self.last_n = deque()
+        self.jury_count = jury_count
+
+    def update(self, new_load):
+        self.last_n.append(new_load)
+        if len(self.last_n) > self.window_size:
+            self.last_n.popleft()
+
+    def update_jury_count(self, jury_count):
+        self.jury_count = jury_count
+
+    def optimal_change(self):
+        avg = sum(self.last_n) / len(self.last_n)
+        index = avg / self.jury_count
+        logger.info('Average enqueued is {} - {} per jury.'.format(avg, index))
+        if index >= 20:
+            return int(index) // 20
+        if index < 2:
+            return -1
+        return 0
+
+
+def get_enqueued_jobs():
+    with app.app_context():
+        return Job.query_can_claim().count()
+
+
+cloud = DigitalOcean(token=DIGITALOCEAN_API_TOKEN)
+
+load_index = LoadIndex()
+
+# TODO: better tracking of juries
+jury_count = cloud.get_current_jury_count()
+
+
+def tick():
+    global jury_count
+    enqueued_jobs = get_enqueued_jobs()
+    load_index.update(enqueued_jobs)
+    load_index.update_jury_count(jury_count)
+    optimal_change = load_index.optimal_change()
+
+    logger.info('{} juries currently exist and optimal change is {}.'.format(jury_count, optimal_change))
+    if optimal_change >= 2:
+        if jury_count < MAX_JURIES:
+            to_create = min(optimal_change, MAX_JURIES - jury_count)
+            logger.info('Spinning up {} juries.'.format(to_create))
+            cloud.create_jury(to_create)
+            jury_count += to_create
+        else:
+            logger.info('Maximum jury count reached.')
+    elif optimal_change <= -1:
+        # TODO: clean shutdown of juries or detection of jury's current job
+        if jury_count > 1:
+            to_destroy = min(-optimal_change, jury_count - 1)
+            logger.info('Destroying {} juries.'.format(to_destroy))
+            destroyed = cloud.destroy_jury(to_destroy)
+            jury_count -= destroyed
+            logger.info('Destroyed {} juries.'.format(destroyed))
+        else:
+            logger.info('Not enough juries to destroy!')
+
+
+if cloud.get_current_jury_count() == 0:
+    logger.info('Spinning up 1 jury because none previously existed.')
+    cloud.create_jury()
+    jury_count += 1
+
+while True:
+    tick()
+
+    time.sleep(5)

Diff for: config.py

@@ -0,0 +1,60 @@
+import os
+import pathlib
+
+from dotenv import load_dotenv, find_dotenv
+
+load_dotenv(find_dotenv())
+
+SUPPORTED_LANGUAGES = {
+    'cxx': 'C++',
+    'python2': 'Python 2',
+    'python3': 'Python 3',
+    'java': 'Java',
+}
+
+
+class JudgeConfig:
+    def __init__(self, app_root: str = None, testing: bool = False):
+        if app_root is None:
+            self.app_root = pathlib.Path(os.path.dirname(os.path.abspath(__file__)))
+        else:
+            self.app_root = pathlib.Path(app_root)
+
+        self.ENABLE_SOCKETIO = bool(int(os.getenv('ENABLE_SOCKETIO', 1)))
+
+        self.SECRET_KEY = None
+        self._load_secret_key()
+        self.SQLALCHEMY_DATABASE_URI = self._get_test_database_uri() if testing else self._get_database_uri()
+        self.SQLALCHEMY_TRACK_MODIFICATIONS = False
+        self.REDIS_URI = self._get_redis_uri()
+
+        if testing:
+            self.TESTING = True
+            self.WTF_CSRF_ENABLED = False
+
+    def _load_secret_key(self):
+        if 'SECRET_KEY' in os.environ:
+            self.SECRET_KEY = os.environ['SECRET_KEY']
+        else:
+            secret_path = self.app_root / '.secret_key'
+            with secret_path.open('a+b') as secret_file:
+                secret_file.seek(0)
+                contents = secret_file.read()
+                if not contents and len(contents) == 0:
+                    key = os.urandom(128)
+                    secret_file.write(key)
+                    secret_file.flush()
+                else:
+                    key = contents
+            self.SECRET_KEY = key
+
+        return self.SECRET_KEY
+
+    def _get_database_uri(self):
+        return os.getenv('DATABASE_URI', '')
+
+    def _get_test_database_uri(self):
+        return os.getenv('TEST_DATABASE_URI', '')
+
+    def _get_redis_uri(self):
+        return os.getenv('REDIS_URI', '')

Diff for: conftest.py

@@ -0,0 +1,71 @@
+import pytest
+
+import main
+from config import JudgeConfig
+from main import db as app_db
+
+
+@pytest.fixture(scope='session')
+def app(request):
+    app = main.app
+    app.config.from_object(JudgeConfig(testing=True))
+
+    ctx = app.app_context()
+    ctx.push()
+
+    def teardown():
+        ctx.pop()
+
+    request.addfinalizer(teardown)
+
+    return app
+
+
+@pytest.fixture(scope='function')
+def client(app, request_context):
+    return app.test_client()
+
+
+@pytest.fixture(scope='class')
+def db(request, app):
+    app_db.reflect()  # Weird hack
+    app_db.drop_all()
+
+    app_db.create_all()
+
+    def teardown():
+        app_db.session.close()
+        app_db.drop_all()
+
+    request.addfinalizer(teardown)
+    return app_db
+
+
+@pytest.fixture(scope='function')
+def request_context(request, app):
+    ctx = app.test_request_context()
+    ctx.push()
+
+    def teardown():
+        ctx.pop()
+
+    request.addfinalizer(teardown)
+
+
+@pytest.fixture(scope='class')
+def session(request, db):
+    connection = db.engine.connect()
+    transaction = connection.begin()
+
+    options = dict(bind=connection, binds={})
+    session = db.create_scoped_session(options=options)
+
+    db.session = session
+
+    def teardown():
+        transaction.rollback()
+        connection.close()
+        session.remove()
+
+    request.addfinalizer(teardown)
+    return session

Diff for: constants.py

@@ -0,0 +1,27 @@
+import enum
+
+
+class APIKeyType(enum.Enum):
+    jury = 'jury'
+    observer = 'observer'
+
+
+class JobStatus(enum.Enum):
+    queued = 'queued'
+    cancelled = 'cancelled'
+    started = 'started'
+    awaiting_verdict = 'awaiting_verdict'
+    finished = 'finished'
+
+
+class JobVerdict(enum.Enum):
+    accepted = 'AC'
+    ran = 'RAN'
+    invalid_source = 'IS'
+    wrong_answer = 'WA'
+    time_limit_exceeded = 'TLE'
+    memory_limit_exceeded = 'MLE'
+    runtime_error = 'RTE'
+    illegal_syscall = 'ISC'
+    compilation_error = 'CE'
+    judge_error = 'JE'

Diff for: docker-compose.yml

@@ -0,0 +1,35 @@
+version: '2'
+services:
+  judge:
+    image: app
+    env_file: .env
+    links:
+      - db
+      - redis
+    ports:
+      - "80:80"
+    depends_on:
+      - migrations
+      - redis
+
+  db:
+    image: mariadb:10.1.16
+    env_file: .env
+    expose:
+      - 3306
+    volumes:
+      - "./.data/db:/var/lib/mysql"
+
+  migrations:
+    build: .
+    image: app
+    env_file: .env
+    command: bash -c "bash wait-for-db.sh && python3 manage.py db upgrade"
+    links:
+      - db
+    depends_on:
+      - db
+
+  redis:
+    restart: "no"
+    image: redis