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)
try:
shutil.rmtree('container/standard/cache')
except FileNotFoundError:
pass
os.mkdir('container/standard/cache', 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)
StdChal.build_cache = {}
StdChal.build_cache_refcount = {}
def comp_cxx(self, callback):
'''GCC, Clang compile.
Args:
callback (function): Callback of return_future.
Returns:
None
'''
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):
compile_path = self.chal_path + '/compile'
os.mkdir(compile_path, mode=0o770)
shutil.copyfile(self.code_path, compile_path + '/test.cpp', \
follow_symlinks=False)
with StackContext(Privilege.fullaccess):
os.chown(compile_path, self.compile_uid, self.compile_gid)
if self.comp_typ == 'g++':
compiler = '/usr/bin/g++'
elif self.comp_typ == 'clang++':
compiler = '/usr/bin/clang++'
task_id = PyExt.create_task(compiler, \
[
'-O2',
'-std=c++14',
'-o', './a.out',
'./test.cpp',
], \
[
'PATH=/usr/bin',
'TMPDIR=/home/%d/compile'%self.uniqid,
], \
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 comp_make(self, callback=None):
'''Makefile compile.
Args:
callback (function): Callback of return_future.
Returns:
None
'''
def _done_cb(task_id, stat):
'''Done callback.
Args:
task_id (int): Task ID.
stat (dict): Task result.
Returns:
None
'''
callback((stat['detect_error'], ''))
make_path = self.chal_path + '/compile'
FileUtils.copydir(self.res_path + '/make', make_path)
with StackContext(Privilege.fileaccess):
shutil.copyfile(self.code_path, make_path + '/main.cpp', \
follow_symlinks=False)
FileUtils.setperm(make_path, self.compile_uid, self.compile_gid)
with StackContext(Privilege.fullaccess):
os.chmod(make_path, mode=0o770)
task_id = PyExt.create_task('/usr/bin/make', \
[], \
[
'PATH=/usr/bin:/bin',
'TMPDIR=/home/%d/compile'%self.uniqid,
'OUT=./a.out',
], \
{
0: StdChal.null_fd,
1: StdChal.null_fd,
2: StdChal.null_fd,
}, \
'/home/%d/compile'%self.uniqid, 'container/standard', \
self.compile_uid, self.compile_gid, 60000, 1024 * 1024 * 1024, \
PyExt.RESTRICT_LEVEL_LOW)
if task_id is None:
callback((PyExt.DETECT_INTERNALERR, ''))
else:
PyExt.start_task(task_id, _done_cb)
def make_wrapped_function():
"""Wraps a function in three stack contexts, and returns
the function along with the deactivation functions.
"""
# Remove the test's stack context to make sure we can cover
# the case where the last context is deactivated.
with NullContext():
partial = functools.partial
with StackContext(partial(self.context, 'c0')) as c0:
with StackContext(partial(self.context, 'c1')) as c1:
with StackContext(partial(self.context, 'c2')) as c2:
return (wrap(check_contexts), [c0, c1, c2])
def comp_python(self, callback):
'''Python3.4 compile.
Args:
callback (function): Callback of return_future.
Returns:
None
'''
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):
compile_path = self.chal_path + '/compile'
os.mkdir(compile_path, mode=0o770)
shutil.copyfile(self.code_path, compile_path + '/test.py', \
follow_symlinks=False)
with StackContext(Privilege.fullaccess):
os.chown(compile_path, self.compile_uid, self.compile_gid)
task_id = PyExt.create_task('/usr/bin/python3.4', \
[
'-m',
'py_compile',
'./test.py'
], \
[
'HOME=/home/%d/compile'%self.uniqid,
'LANG=en_US.UTF-8'
], \
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 testNestedContexts(self):
"""Test the nesting of a single ContextLocal subclass."""
with util.Barrier(self._OnSuccess,
on_exception=self._OnException) as b:
with StackContext(ExampleContext(1)):
self.io_loop.add_callback(
partial(self._VerifyExampleContext, 1, b.Callback()))
with StackContext(ExampleContext(2)):
self._VerifyExampleContext(2, util.NoCallback)
self.io_loop.add_callback(
partial(self._VerifyExampleContext, 2, b.Callback()))
self._VerifyExampleContext(1, util.NoCallback)
self.wait()
def testMultipleContextTypes(self):
"""Test the usage of multiple ContextLocal subclasses in tandem."""
with util.Barrier(self._OnSuccess,
on_exception=self._OnException) as b:
with StackContext(ExampleContext(1)):
with StackContext(ExampleContextTwoParams(2, 3)):
self._VerifyExampleContext(1, util.NoCallback)
self._VerifyExampleContextTwoParams(2, 3, util.NoCallback)
self.io_loop.add_callback(
partial(self._VerifyExampleContext, 1, b.Callback()))
self.io_loop.add_callback(
partial(self._VerifyExampleContextTwoParams, 2, 3,
b.Callback()))
self.wait()
def test_pre_wrap(self):
# A pre-wrapped callback is run in the context in which it was
# wrapped, not when it was added to the IOLoop.
def f1():
self.assertIn('c1', self.active_contexts)
self.assertNotIn('c2', self.active_contexts)
self.stop()
with StackContext(functools.partial(self.context, 'c1')):
wrapped = wrap(f1)
with StackContext(functools.partial(self.context, 'c2')):
self.add_callback(wrapped)
self.wait()
def __call__(self, request):
request_id = request.headers.get('X-Request-Id')
if request_id is None:
request_id = FrontikApplication.next_request_id()
with StackContext(partial(RequestContext, {'request_id': request_id})):
return super(FrontikApplication, self).__call__(request)
def wrapper(self, transforms, *args, **kwargs):
"""
the wraper function
"""
request_id = uuid4().hex
with StackContext(lambda: RequestContext(request_id)):
return _execute(self, transforms, *args, **kwargs)
def f():
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)
def run(self):
with StackContext(die_on_error):
if self.enable_prom_exporter:
start_http_server(8000)
self.io_loop = ioloop.IOLoop.current()
self.io_loop.run_sync(self._run)
sys.exit(errno.EINTR)
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 wrapper(*args, **kwargs):
request_context = None
if hasattr(method.__self__, "_request_context"):
request_context = method.__self__._request_context
with StackContext(
RequestContextManager(request_context).context_manager):
return method(*args, **kwargs)
def run(self):
with StackContext(self._stack_context):
self.pre_setup()
self.setup()
for test in self._tests:
yield test()
self.teardown()
self.post_teardown()
def wrapper(*args, **kwargs):
token = request_context.initialize(request, request_id)
try:
with StackContext(context):
return func(*args, **kwargs)
finally:
request_context.reset(token)
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 _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 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 _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 _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):
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 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))