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