Skip to content

perf: Function library execution #3175

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shaohuzhang1 merged 1 commit into from
May 30, 2025
Merged

perf: Function library execution #3175

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 14 additions & 14 deletions apps/common/util/function_code.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,12 @@
@desc:
"""
import os
import pickle
import subprocess
import sys
import uuid
from textwrap import dedent

from diskcache import Cache

from smartdoc.const import BASE_DIR
from smartdoc.const import PROJECT_DIR

Expand All @@ -37,14 +36,16 @@ def _createdir(self):
old_mask = os.umask(0o077)
try:
os.makedirs(self.sandbox_path, 0o700, exist_ok=True)
os.makedirs(os.path.join(self.sandbox_path, 'execute'), 0o700, exist_ok=True)
os.makedirs(os.path.join(self.sandbox_path, 'result'), 0o700, exist_ok=True)
finally:
os.umask(old_mask)

def exec_code(self, code_str, keywords):
_id = str(uuid.uuid1())
success = '{"code":200,"msg":"成功","data":exec_result}'
err = '{"code":500,"msg":str(e),"data":None}'
path = r'' + self.sandbox_path + ''
result_path = f'{self.sandbox_path}/result/{_id}.result'
_exec_code = f"""
try:
import os
Expand All @@ -60,32 +61,31 @@ def exec_code(self, code_str, keywords):
for local in locals_v:
globals_v[local] = locals_v[local]
exec_result=f(**keywords)
from diskcache import Cache
cache = Cache({path!a})
cache.set({_id!a},{success})
import pickle
with open({result_path!a}, 'wb') as file:
file.write(pickle.dumps({success}))
except Exception as e:
from diskcache import Cache
cache = Cache({path!a})
cache.set({_id!a},{err})
with open({result_path!a}, 'wb') as file:
file.write(pickle.dumps({err}))
"""
if self.sandbox:
subprocess_result = self._exec_sandbox(_exec_code, _id)
else:
subprocess_result = self._exec(_exec_code)
if subprocess_result.returncode == 1:
raise Exception(subprocess_result.stderr)
cache = Cache(self.sandbox_path)
result = cache.get(_id)
cache.delete(_id)
with open(result_path, 'rb') as file:
result = pickle.loads(file.read())
os.remove(result_path)
if result.get('code') == 200:
return result.get('data')
raise Exception(result.get('msg'))

def _exec_sandbox(self, _code, _id):
exec_python_file = f'{self.sandbox_path}/{_id}.py'
exec_python_file = f'{self.sandbox_path}/execute/{_id}.py'
with open(exec_python_file, 'w') as file:
file.write(_code)
os.system(f"chown {self.user}:{self.user} {exec_python_file}")
os.system(f"chown {self.user}:root {exec_python_file}")
kwargs = {'cwd': BASE_DIR}
subprocess_result = subprocess.run(
['su', '-s', python_directory, '-c', "exec(open('" + exec_python_file + "').read())", self.user],
Copy link
Contributor Author

@shaohuzhang1 shaohuzhang1 May 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are several potential issues and optimizations in the provided code:

  1. Permissions Handling:

    • The code uses os.chown to set ownership of files, which can be dangerous. If not used carefully, it could expose sensitive information or lead to permissions issues.
    • It's better to handle file access and execution using appropriate filesystem permissions rather than relying on permission changes.

  2. File Write Permissions:

    • There is no explicit check for write permissions when writing to the sandbox directory paths. Ensure that these directories have write permissions.

  3. Error Logging and Cleanup:

    • Improper error handling: The code doesn't log errors properly. If an exception occurs during file operations or subprocess execution, it should be logged for debugging purposes.
    • Potential issue with temporary files: When executing Python scripts in a sandboxed environment, ensure that any temporary files created are deleted after they're no longer needed.

  4. Use of DiskCache for Caching:

    • While DiskCache is being used for caching results, its use needs to be consistent with the rest of the codebase. There's confusion between using DiskCache directly versus pickle.

  5. Code Formatting:

    • The formatting of the _exec_code string lacks consistency with other parts of the function.
    • Consider cleaning up unnecessary whitespace before returning subprocess_result.

  6. Resource Management:

    • Always ensure that resources like file descriptors are closed after use, especially when dealing with I/O operations and subprocesses.

Optimizations and Recommendations

  1. Permission Best Practices:

    • Use chroot/jail environments instead of trying to change user/group IDs inside a process. This is more secure and reliable.
    • Restrict file system access to only necessary components within the sandbox path.

  2. Consistent Error Logging:

    • Implement logging at all stages where errors occur. A basic logging setup might look like this: 
      import logging
logger = logging.getLogger(__name__)
logger.setLevel(logging.ERROR)

handler = logging.FileHandler("app.log")
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)

  3. Temporary File Cleanup:

    • After executing the script, delete any temporary files generated by the script and any related output files.

  4. Consistent Code Style:

    • Refactor _exec_code to remove unnecessary whitespace and align comments.

  5. Use Context Managers:

    • Prefer context managers (e.g., with open(...) as ...) over manual file closing to avoid resource leaks.

Here's an improved version of the relevant part of the code:

def exec_code(self, code_str, keywords):
    try:
        _id = str(uuid.uuid1())
        
        # Create required directories
        old_mask = os.umask(0o077)
        try:
            os.makedirs(self.sandbox_path, 0o700, exist_ok=True)
            os.makedirs(os.path.join(self.sandbox_path, 'execute'), 0o700, exist_ok=True)
            os.makedirs(os.path.join(self.sandbox_path, 'result'), 0o700, exist_ok=True)
        finally:
            os.umask(old_mask)
            
        # Execute the code
        _exec_code = f"""
try:
    import os
...
"""

        if self.sandbox:
            self._exec_sandbox(_exec_code, _id)
        else:
            subprocess_result = self._exec(_exec_code)

        if subprocess_result.returncode != 0:
            raise Exception(subprocess_result.stderr)
            
        # Read and close the cached response correctly
        with open(os.path.join(self.sandbox_path, 'result', f"{_id}.result"), 'rb') as file:
            result = pickle.loads(file.read())

        os.remove(os.path.join(self.sandbox_path, 'result', f"{_id}.result"))
        return result.get('data')

    except Exception as e:
        self.logger.error(f'An error occurred while executing code: {e}')
        return {"code": 500, "msg": str(e), "data": None}

def _exec_sandbox(self, _code, _id):
    exec_python_file = os.path.join(self.sandbox_path, 'execute', f"{_id}.py")
    
    with open(exec_python_file, 'w') as file:
        file.write(_code)

    # Using su for security reasons
    subprocess.run(
        ['su', '-s', python_directory, '-c', f"exec(open('{exec_python_file}').read())"],
        shell=False,
        cwd=BASE_DIR,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    ).check_output()

This revised code includes error logging, ensures consistent cleanup of temporary files, and improves overall structure and readability. Make sure to integrate this logging mechanism into your application.

Expand Down