def test_pre_wrap_with_args(self):
# Same as test_pre_wrap, but the function takes arguments.
# Implementation note: The function must not be wrapped in a
# functools.partial until after it has been passed through
# stack_context.wrap
def f1(foo, bar):
self.assertIn('c1', self.active_contexts)
self.assertNotIn('c2', self.active_contexts)
self.stop((foo, bar))
with StackContext(functools.partial(self.context, 'c1')):
wrapped = wrap(f1)
with StackContext(functools.partial(self.context, 'c2')):
self.add_callback(wrapped, 1, bar=2)
result = self.wait()
self.assertEqual(result, (1, 2))
def init():
'''Initialize the module.'''
with StackContext(Privilege.fileaccess):
try:
shutil.rmtree('container/standard/home')
except FileNotFoundError:
pass
os.mkdir('container/standard/home', mode=0o771)
ffi = FFI()
ffi.cdef('''int mount(const char source[], const char target[],
const char filesystemtype[], unsigned long mountflags,
const void *data);''')
ffi.cdef('''int umount(const char *target);''')
libc = ffi.dlopen('libc.so.6')
with StackContext(Privilege.fullaccess):
libc.umount(b'container/standard/dev')
libc.mount(b'/dev', b'container/standard/dev', b'', MS_BIND, \
ffi.NULL)
StdChal.null_fd = os.open('/dev/null', os.O_RDWR | os.O_CLOEXEC)
def _done_cb():
'''Done callback.'''
nonlocal result_stat
nonlocal result_pass
nonlocal verdict_path
if result_pass is not None and result_stat is not None:
with StackContext(Privilege.fileaccess):
verfile = open(verdict_path, 'r')
verdict = verfile.read(140)
verfile.close()
callback((result_pass, result_stat, verdict))
return
def request_context(self, parent_tracing):
"""
Factory method meant to be used as:
.. code-block:: python
with tchannel.context_provider.request_context(parent_tracing):
handler_fn()
:param parent_tracing:
:return:
"""
# TODO should this be using a thread-safe version of StackContext?
return StackContext(lambda: RequestContext(parent_tracing))
def wrapper(self, *args, **kwargs):
with StackContext(functools.partial(ctx_man, self)) as cm:
w = fn #wrap(fn)
result = w(*args, **kwargs)
if isinstance(result, TemplateProxy):
if self._template_engine == 'tornado':
self.render(*result.args, **result.kwargs)
else:
template = self._template_env.get_template(
result.args[0])
self.finish(
template.render(handler=self, **result.kwargs))
else:
self.finish(result)
def _execute(self, transforms, *args, **kwargs):
if options.enable_appstats:
start_recording(tornado.wsgi.WSGIContainer.environ(self.request))
recorder = save()
@contextlib.contextmanager
def transfer_recorder():
restore(recorder)
yield
with StackContext(transfer_recorder):
super(RecordingRequestHandler,
self)._execute(transforms, *args, **kwargs)
else:
super(RecordingRequestHandler,
self)._execute(transforms, *args, **kwargs)
def f():
try:
with StackContext(functools.partial(self.context, 'c1')):
# This yield is a problem: the generator will be suspended
# and the StackContext's __exit__ is not called yet, so
# the context will be left on _state.contexts for anything
# that runs before the yield resolves.
yield gen.Task(self.io_loop.add_callback)
except StackContextInconsistentError:
# In python <= 3.3, this suspended generator is never garbage
# collected, so it remains suspended in the 'yield' forever.
# Starting in 3.4, it is made collectable by raising
# a GeneratorExit exception from the yield, which gets
# converted into a StackContextInconsistentError by the
# exit of the 'with' block.
pass
def prepare(self):
if options.enable_appstats:
recording.start_recording(
tornado.wsgi.WSGIContainer.environ(self.request))
recorder = save()
@contextlib.contextmanager
def transfer_recorder():
restore(recorder)
yield
with StackContext(transfer_recorder):
super(RecordingFallbackHandler, self).prepare()
recording.end_recording(self._status_code)
else:
super(RecordingFallbackHandler, self).prepare()
def _copy_fn(src, dst, follow_symlinks=True):
'''Copytree helper function.
Args:
src (string): Source path.
dst (string): Destination path.
follow_symlinks: Follow symbolic link or not.
Returns:
None
'''
shutil.copy(src, dst, follow_symlinks=False)
with StackContext(Privilege.fullaccess):
os.chown(dst, self.compile_uid, self.compile_gid)
def build_cache_decref(cache_hash):
'''Decrement the refcount of the build cache.
Delete the build cache if the refcount = 0.
Args:
cache_hash (int): Cache hash.
Returns:
None
'''
StdChal.build_cache_refcount[cache_hash] -= 1
if StdChal.build_cache_refcount[cache_hash] == 0:
with StackContext(Privilege.fileaccess):
shutil.rmtree('container/standard/cache/%x'%cache_hash)
def test_run_with_stack_context(self):
@gen.coroutine
def f1():
self.assertEqual(self.active_contexts, ['c1'])
yield run_with_stack_context(
StackContext(functools.partial(self.context, 'c2')), f2)
self.assertEqual(self.active_contexts, ['c1'])
@gen.coroutine
def f2():
self.assertEqual(self.active_contexts, ['c1', 'c2'])
yield gen.Task(self.io_loop.add_callback)
self.assertEqual(self.active_contexts, ['c1', 'c2'])
self.assertEqual(self.active_contexts, [])
yield run_with_stack_context(
StackContext(functools.partial(self.context, 'c1')), f1)
self.assertEqual(self.active_contexts, [])
def prefetch(self):
'''Prefetch files.'''
path_set = set([self.code_path])
for root, _, files in os.walk(self.res_path):
for filename in files:
path_set.add(os.path.abspath(os.path.join(root, filename)))
path_list = list(path_set)
proc_list = []
with StackContext(Privilege.fileaccess):
for idx in range(0, len(path_list), 16):
proc_list.append(process.Subprocess(
['./Prefetch.py'] + path_list[idx:idx + 16],
stdout=process.Subprocess.STREAM))
for proc in proc_list:
yield proc.stdout.read_bytes(2)
def _done_cb(task_id, stat):
'''Done callback.
Args:
task_id (int): Task ID.
stat (dict): Task result.
Returns:
None
'''
nonlocal compile_path
with StackContext(Privilege.fileaccess):
verfile = open(compile_path + '/verdict.txt', 'r')
verdict = verfile.read(140)
verfile.close()
callback((stat['detect_error'], verdict))
def test_exit_library_context(self):
def library_function(callback):
# capture the caller's context before introducing our own
callback = wrap(callback)
with StackContext(functools.partial(self.context, 'library')):
self.io_loop.add_callback(
functools.partial(library_inner_callback, callback))
def library_inner_callback(callback):
self.assertEqual(self.active_contexts[-2:],
['application', 'library'])
callback()
def final_callback():
# implementation detail: the full context stack at this point
# is ['application', 'library', 'application']. The 'library'
# context was not removed, but is no longer innermost so
# the application context takes precedence.
self.assertEqual(self.active_contexts[-1], 'application')
self.stop()
with StackContext(functools.partial(self.context, 'application')):
library_function(final_callback)
self.wait()
def test_exit_library_context(self):
def library_function(callback):
# capture the caller's context before introducing our own
callback = wrap(callback)
with StackContext(functools.partial(self.context, 'library')):
self.io_loop.add_callback(
functools.partial(library_inner_callback, callback))
def library_inner_callback(callback):
assert 'application' in self.active_contexts
assert 'library' in self.active_contexts
# pass the callback out to the IOLoop to get out of the library
# context (could also use a NullContext here, but that would result
# in multiple instantiations of the application context)
self.io_loop.add_callback(callback)
def final_callback():
assert 'application' in self.active_contexts
assert 'library' not in self.active_contexts
self.stop()
with StackContext(functools.partial(self.context, 'application')):
library_function(final_callback)
self.wait()
def judge_diff(self, src_path, exe_path, argv, envp, in_path, ans_path, \
timelimit, memlimit, callback):
'''Diff judge.
Args:
src_path (string): Executable source path.
exe_path (string): Executable or interpreter path in the sandbox.
argv ([string]): List of arguments.
envp ([string]): List of environment variables.
in_path (string): Input file path.
ans_path (string): Answer file path.
timelimit (int): Timelimit.
memlimit (int): Memlimit.
callback (function): Callback of return_future.
Returns:
None
'''
with StackContext(Privilege.fileaccess):
infile_fd = os.open(in_path, os.O_RDONLY | os.O_CLOEXEC)
ansfile = open(ans_path, 'rb')
outpipe_fd = os.pipe2(os.O_CLOEXEC)
fcntl.fcntl(outpipe_fd[0], fcntl.F_SETFL, os.O_NONBLOCK)
result_stat = None
result_pass = None
def _started_cb(task_id):
'''Started callback.
Close unused file descriptor after the task is started.
Args:
task_id (int): Task ID.
Returns:
None
'''
nonlocal infile_fd
nonlocal outpipe_fd
os.close(infile_fd)
os.close(outpipe_fd[1])
def _done_cb(task_id, stat):
'''Done callback.
Args:
task_id (int): Task ID.
stat (dict): Task result.
Returns:
None
'''
nonlocal result_stat
nonlocal result_pass
result_stat = (stat['utime'], stat['peakmem'],
stat['detect_error'])
if result_pass is not None:
callback((result_pass, result_stat))
def _diff_out(evfd, events):
'''Diff the output of the task.
Args:
evfd (int): Event file descriptor.
events (int): Event flags.
Returns:
None
'''
nonlocal outpipe_fd
nonlocal ansfile
nonlocal result_stat
nonlocal result_pass
end_flag = False
if events & IOLoop.READ:
while True:
try:
data = os.read(outpipe_fd[0], 65536)
except BlockingIOError:
break
ansdata = ansfile.read(len(data))
if data != ansdata:
result_pass = False
end_flag = True
break
if len(ansdata) == 0:
if len(ansfile.read(1)) == 0:
result_pass = True
else:
result_pass = False
end_flag = True
break
if (events & IOLoop.ERROR) or end_flag:
if result_pass is None:
if len(ansfile.read(1)) == 0:
result_pass = True
else:
result_pass = False
IOLoop.instance().remove_handler(evfd)
os.close(outpipe_fd[0])
ansfile.close()
if result_stat is not None:
callback((result_pass, result_stat))
StdChal.last_judge_uid += 1
judge_uid = StdChal.last_judge_uid
judge_gid = judge_uid
with StackContext(Privilege.fileaccess):
judge_path = self.chal_path + '/run_%d' % judge_uid
os.mkdir(judge_path, mode=0o771)
shutil.copyfile(src_path, judge_path + '/a.out', \
follow_symlinks=False)
with StackContext(Privilege.fullaccess):
os.chown(judge_path + '/a.out', judge_uid, judge_gid)
os.chmod(judge_path + '/a.out', 0o500)
IOLoop.instance().add_handler(outpipe_fd[0], _diff_out, \
IOLoop.READ | IOLoop.ERROR)
task_id = PyExt.create_task(exe_path, argv, envp, \
infile_fd, outpipe_fd[1], outpipe_fd[1], \
'/home/%d/run_%d'%(self.uniqid, judge_uid), 'container/standard', \
judge_uid, judge_gid, timelimit, memlimit, \
PyExt.RESTRICT_LEVEL_HIGH)
if task_id is None:
callback((False, (0, 0, PyExt.DETECT_INTERNALERR)))
PyExt.start_task(task_id, _done_cb, _started_cb)
def test_yield_outside_with(self):
# This pattern avoids the problem in the previous test.
cb = yield gen.Callback('k1')
with StackContext(functools.partial(self.context, 'c1')):
self.io_loop.add_callback(cb)
yield gen.Wait('k1')
def f1():
with NullContext():
wrapped = wrap(f2)
with StackContext(functools.partial(self.context, 'c2')):
wrapped()
def f3():
with StackContext(functools.partial(self.context, 'c3')) as c3:
deactivate_callbacks.append(c3)
self.io_loop.add_callback(f4)
def f2():
with StackContext(functools.partial(self.context, 'c2')) as c2:
deactivate_callbacks.append(c2)
self.io_loop.add_callback(f3)
def f1():
with StackContext(functools.partial(self.context, 'c1')) as c1:
deactivate_callbacks.append(c1)
self.io_loop.add_callback(f2)
def run(self, result=None):
with StackContext(self._stack_context):
super(AsyncTestCase, self).run(result)
def run(self, result=None):
with StackContext(self._stack_context):
super(AsyncTestCase, self).run(result)
# In case an exception escaped super.run or the StackContext caught
# an exception when there wasn't a wait() to re-raise it, do so here.
self.__rethrow()
def comp_make(self, callback):
'''Makefile compile.
Args:
callback (function): Callback of return_future.
Returns:
None
'''
def _copy_fn(src, dst, follow_symlinks=True):
'''Copytree helper function.
Args:
src (string): Source path.
dst (string): Destination path.
follow_symlinks: Follow symbolic link or not.
Returns:
None
'''
shutil.copy(src, dst, follow_symlinks=False)
with StackContext(Privilege.fullaccess):
os.chown(dst, self.compile_uid, self.compile_gid)
def _done_cb(task_id, stat):
'''Done callback.
Args:
task_id (int): Task ID.
stat (dict): Task result.
Returns:
None
'''
callback(stat['detect_error'])
with StackContext(Privilege.fileaccess):
make_path = self.chal_path + '/compile'
shutil.copytree(self.res_path + '/make', make_path, symlinks=True, \
copy_function=_copy_fn)
shutil.copyfile(self.code_path, make_path + '/main.cpp', \
follow_symlinks=False)
with StackContext(Privilege.fullaccess):
os.chown(make_path, self.compile_uid, self.compile_gid)
os.chmod(make_path, mode=0o770)
task_id = PyExt.create_task('/usr/bin/make', \
[], \
[
'PATH=/usr/bin',
'TMPDIR=/home/%d/compile'%self.uniqid,
'OUT=./a.out',
], \
StdChal.null_fd, StdChal.null_fd, StdChal.null_fd, \
'/home/%d/compile'%self.uniqid, 'container/standard', \
self.compile_uid, self.compile_gid, 60000, 256 * 1024 * 1024, \
PyExt.RESTRICT_LEVEL_LOW)
if task_id is None:
callback(-1)
else:
PyExt.start_task(task_id, _done_cb)
def start(self):
'''Start the challenge.
Returns:
dict: Challenge result.
'''
print('StdChal %d started' % self.chal_id)
self.chal_path = 'container/standard/home/%d' % self.uniqid
with StackContext(Privilege.fileaccess):
os.mkdir(self.chal_path, mode=0o771)
yield self.prefetch()
print('StdChal %d prefetched' % self.chal_id)
if self.comp_typ in ['g++', 'clang++']:
ret = yield self.comp_cxx()
elif self.comp_typ == 'makefile':
ret = yield self.comp_make()
elif self.comp_typ == 'python3':
ret = yield self.comp_python()
if ret != PyExt.DETECT_NONE:
ret_result = [(0, 0, STATUS_CE)] * len(self.test_list)
else:
print('StdChal %d compiled' % self.chal_id)
if self.comp_typ == 'python3':
exefile_path = self.chal_path \
+ '/compile/__pycache__/test.cpython-34.pyc'
exe_path = '/usr/bin/python3.4'
argv = ['./a.out']
envp = ['HOME=/', 'LANG=en_US.UTF-8']
else:
exefile_path = self.chal_path + '/compile/a.out'
exe_path = './a.out'
argv = []
envp = []
test_future = []
for test in self.test_list:
test_future.append(
self.judge_diff(exefile_path, exe_path, argv, envp,
test['in'], test['ans'], test['timelimit'],
test['memlimit']))
test_result = yield gen.multi(test_future)
ret_result = list()
for result in test_result:
test_pass, data = result
runtime, peakmem, error = data
status = STATUS_ERR
if error == PyExt.DETECT_NONE:
if test_pass is True:
status = STATUS_AC
else:
status = STATUS_WA
elif error == PyExt.DETECT_OOM:
status = STATUS_MLE
elif error == PyExt.DETECT_TIMEOUT \
or error == PyExt.DETECT_FORCETIMEOUT:
status = STATUS_TLE
elif error == PyExt.DETECT_EXITERR:
status = STATUS_RE
else:
status = STATUS_ERR
ret_result.append((runtime, peakmem, status))
with StackContext(Privilege.fileaccess):
shutil.rmtree(self.chal_path)
print('StdChal %d done' % self.chal_id)
return ret_result
def run(self):
with StackContext(self._stack_context):
super(AsyncTestCase, self).run()
def execution_wrapper():
# enable RequestContext on current thread
with StackContext(RequestContextManager(request_ctx).context_manager):
return fn(*args, **kwargs)
def f1():
self.assertEqual(self.active_contexts, ['c1'])
yield run_with_stack_context(
StackContext(functools.partial(self.context, 'c2')), f2)
self.assertEqual(self.active_contexts, ['c1'])
def library_function(callback):
# capture the caller's context before introducing our own
callback = wrap(callback)
with StackContext(functools.partial(self.context, 'library')):
self.io_loop.add_callback(
functools.partial(library_inner_callback, callback))
def request_context(parent_tracing):
return StackContext(lambda: RequestContext(parent_tracing))
