RustPython/crates/vm/src/stdlib/_imp.rs

use crate::builtins::{PyCode, PyStrInterned};
use crate::frozen::FrozenModule;
use crate::{Py, VirtualMachine, builtins::PyBaseExceptionRef};
use core::borrow::Borrow;

pub(crate) use _imp::module_def;

pub(super) use crate::vm::resolve_frozen_alias;

mod resources;

#[cfg(feature = "threading")]
#[pymodule(sub, name = "_imp")]
mod lock {
    use crate::{PyResult, VirtualMachine, stdlib::_thread::RawRMutex};
    use core::cell::Cell;

    static IMP_LOCK: RawRMutex = RawRMutex::INIT;

    thread_local! {
        static IMP_LOCK_DEPTH: Cell<usize> = const { Cell::new(0) };
    }

    fn bump_depth() {
        IMP_LOCK_DEPTH.with(|c| c.set(c.get() + 1));
    }

    fn drop_depth() {
        IMP_LOCK_DEPTH.with(|c| c.set(c.get().saturating_sub(1)));
    }

    #[pyfunction]
    fn acquire_lock(vm: &VirtualMachine) {
        // Detach while blocking on IMP_LOCK. The import lock is held across
        // bytecode by the importlib bootstrap, so its holder can be parked at a
        // safepoint mid-hold. Blocking here while attached would keep this
        // thread from honoring a stop-the-world request, so a requester could
        // wait for this thread while this thread waits for the parked holder.
        // Detaching makes the wait park-friendly.
        vm.allow_threads(acquire_lock_for_fork);
    }

    #[pyfunction]
    fn release_lock(vm: &VirtualMachine) -> PyResult<()> {
        if !IMP_LOCK.is_locked() || !IMP_LOCK.is_owned_by_current_thread() {
            Err(vm.new_runtime_error("Global import lock not held"))
        } else {
            unsafe { IMP_LOCK.unlock() };
            drop_depth();
            Ok(())
        }
    }

    #[pyfunction]
    fn lock_held(_vm: &VirtualMachine) -> bool {
        IMP_LOCK.is_locked()
    }

    pub(super) fn acquire_lock_for_fork() {
        IMP_LOCK.lock();
        bump_depth();
    }

    #[cfg(all(unix, feature = "host_env"))]
    pub(super) fn release_lock_after_fork_parent() {
        if IMP_LOCK.is_locked() && IMP_LOCK.is_owned_by_current_thread() {
            unsafe { IMP_LOCK.unlock() };
            drop_depth();
        }
    }

    /// Reset the import lock after `fork()`.
    ///
    /// Zero the lock so waiter queues from dead threads are discarded, then
    /// re-acquire it as many times as this thread held it. `unlock()` on the
    /// inherited lock would unpark those waiters.
    ///
    /// # Safety
    ///
    /// Must only be called from single-threaded child after fork().
    #[cfg(all(unix, feature = "host_env"))]
    pub(crate) unsafe fn reinit_after_fork() {
        let depth = IMP_LOCK_DEPTH.with(Cell::get);
        unsafe { rustpython_common::lock::zero_reinit_after_fork(&IMP_LOCK) };
        for _ in 0..depth {
            IMP_LOCK.lock();
        }
    }

    /// Restore the child's import lock, then drop the extra hold taken
    /// before `fork()`. Nested `imp.acquire_lock()` holds stay so the
    /// child's `release_lock()` can unwind them.
    #[cfg(all(unix, feature = "host_env"))]
    pub(super) unsafe fn after_fork_child_reinit_and_release() {
        unsafe { reinit_after_fork() };
        if IMP_LOCK.is_locked() && IMP_LOCK.is_owned_by_current_thread() {
            unsafe { IMP_LOCK.unlock() };
            drop_depth();
        }
    }
}

/// Re-export for fork safety code in posix.rs
///
/// Runs pre-fork on a normal attached VM thread. Detach while blocking so the
/// wait honors a concurrent stop-the-world request instead of pinning this
/// thread attached on IMP_LOCK; re-attach completes before `stop_the_world`, so
/// the fork requester protocol is unaffected.
#[cfg(all(unix, feature = "threading", feature = "host_env"))]
pub(crate) fn acquire_imp_lock_for_fork(vm: &VirtualMachine) {
    vm.allow_threads(lock::acquire_lock_for_fork);
}

#[cfg(all(unix, feature = "threading", feature = "host_env"))]
pub(crate) fn release_imp_lock_after_fork_parent() {
    lock::release_lock_after_fork_parent();
}

#[cfg(all(unix, feature = "threading", feature = "host_env"))]
pub(crate) unsafe fn reinit_imp_lock_after_fork() {
    unsafe { lock::reinit_after_fork() }
}

#[cfg(all(unix, feature = "threading", feature = "host_env"))]
pub(crate) unsafe fn after_fork_child_imp_lock_release() {
    unsafe { lock::after_fork_child_reinit_and_release() }
}

#[cfg(not(feature = "threading"))]
#[pymodule(sub, name = "_imp")]
mod lock {
    use crate::vm::VirtualMachine;
    #[pyfunction]
    pub(super) const fn acquire_lock(_vm: &VirtualMachine) {}
    #[pyfunction]
    pub(super) const fn release_lock(_vm: &VirtualMachine) {}
    #[pyfunction]
    pub(super) const fn lock_held(_vm: &VirtualMachine) -> bool {
        false
    }
}

#[allow(dead_code)]
enum FrozenError {
    BadName,  // The given module name wasn't valid.
    NotFound, // It wasn't in PyImport_FrozenModules.
    Disabled, // -X frozen_modules=off (and not essential)
    Excluded, // The PyImport_FrozenModules entry has NULL "code"
    //        (module is present but marked as unimportable, stops search).
    Invalid, // The PyImport_FrozenModules entry is bogus
             //          (eg. does not contain executable code).
}

impl FrozenError {
    fn to_pyexception(&self, mod_name: &str, vm: &VirtualMachine) -> PyBaseExceptionRef {
        let msg = match self {
            Self::BadName | Self::NotFound => format!("No such frozen object named {mod_name}"),
            Self::Disabled => format!(
                "Frozen modules are disabled and the frozen object named {mod_name} is not essential"
            ),
            Self::Excluded => format!("Excluded frozen object named {mod_name}"),
            Self::Invalid => format!("Frozen object named {mod_name} is invalid"),
        };
        vm.new_import_error(msg, vm.ctx.new_utf8_str(mod_name))
    }
}

// look_up_frozen + use_frozen in import.c
fn find_frozen(name: &str, vm: &VirtualMachine) -> Result<FrozenModule, FrozenError> {
    let frozen = vm
        .state
        .frozen
        .get(name)
        .copied()
        .ok_or(FrozenError::NotFound)?;

    // Bootstrap modules are always available regardless of override flag
    if matches!(
        name,
        "_frozen_importlib" | "_frozen_importlib_external" | "zipimport"
    ) {
        return Ok(frozen);
    }

    // use_frozen(): override > 0 → true, override < 0 → false, 0 → startup setting
    // When disabled, non-bootstrap modules are simply not found (same as look_up_frozen)
    let override_val = vm.state.override_frozen_modules.load();
    if override_val < 0
        || (override_val == 0
            && vm.import_bootstrap_complete
            && !vm.state.config.settings.use_frozen_modules)
    {
        return Err(FrozenError::NotFound);
    }

    Ok(frozen)
}

#[pymodule(with(lock, resources::resources))]
mod _imp {
    use crate::{
        AsObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult, VirtualMachine,
        builtins::{
            PyBytesRef, PyCode, PyDict, PyImportError, PyMemoryView, PyModule, PyStrRef,
            PyUtf8StrRef,
        },
        import, version,
    };

    use super::FrozenError;

    // Private exact-dict relocation primitive for future import-cache ordering.
    #[pyfunction]
    fn _dict_move_to_end(
        modules: PyRefExact<PyDict>,
        key: PyObjectRef,
        vm: &VirtualMachine,
    ) -> PyResult<PyObjectRef> {
        modules.move_to_end(key, vm)
    }

    #[pyattr]
    fn check_hash_based_pycs(vm: &VirtualMachine) -> PyStrRef {
        vm.ctx
            .new_str(vm.state.config.settings.check_hash_pycs_mode.to_string())
    }

    #[pyattr(name = "pyc_magic_number_token")]
    use version::PYC_MAGIC_NUMBER_TOKEN;

    #[pyfunction]
    const fn extension_suffixes() -> Vec<PyObjectRef> {
        Vec::new()
    }

    // This hook is also used by eager imports that satisfy pending lazy targets.
    #[pyfunction]
    fn _set_lazy_attributes(
        _modobj: PyObjectRef,
        name: PyStrRef,
        vm: &VirtualMachine,
    ) -> PyResult<()> {
        crate::lazy_import::clear_submodule(&name, false, vm)?;
        vm.state.lazy_imports.modules.discard(name.into(), vm)
    }

    #[pyfunction]
    fn is_builtin(name: PyStrRef, vm: &VirtualMachine) -> i32 {
        match name.to_str() {
            Some("sys" | "builtins") => -1,
            Some(name) => i32::from(vm.state.module_defs.contains_key(name)),
            None => 0,
        }
    }

    #[pyfunction]
    fn is_frozen(name: PyStrRef, vm: &VirtualMachine) -> bool {
        name.to_str()
            .is_some_and(|name| super::find_frozen(name, vm).is_ok())
    }

    #[pyfunction]
    fn create_builtin(spec: PyObjectRef, vm: &VirtualMachine) -> PyResult {
        let name_obj = spec.get_attr("name", vm)?;
        let name: PyStrRef = name_obj.downcast().map_err(|obj| {
            vm.new_type_error(format!(
                "name must be string, not {:.200}",
                obj.class().name()
            ))
        })?;
        if name.is_empty() {
            return Err(vm.new_value_error("name must not be empty"));
        }
        // Builtin module names use ASCII even when an unknown name is supplied.
        if !name.as_bytes().is_ascii() {
            vm.state
                .codec_registry
                .encode_text(name.clone(), "ascii", None, vm)?;
        }
        let name_str = name.as_utf8().expect("ASCII module name").as_str();
        if name_str.contains('\0') {
            return Err(vm.new_value_error("embedded null character"));
        }
        if !matches!(name_str, "sys" | "builtins") && !vm.state.module_defs.contains_key(name_str) {
            let exc = vm.new_payload_exception::<PyImportError>(
                vm.ctx.exceptions.module_not_found_error.to_owned(),
                vec![
                    vm.ctx
                        .new_str(format!("{name_str} module not found"))
                        .into(),
                ]
                .into(),
            )?;
            exc.as_object().set_attr("name", name.clone(), vm)?;
            return Err(exc.upcast());
        }

        let sys_modules = vm.sys_module.get_attr("modules", vm)?;
        // Only registered builtins (or the two core modules) may use this cache.
        if let Ok(module) = sys_modules.get_item(&*name, vm) {
            return Ok(module);
        }

        if let Some(&def) = vm.state.module_defs.get(name_str) {
            // Phase 1: Create module (use create slot if provided, else default creation)
            let module = if let Some(create) = def.slots.create {
                // Custom module creation
                create(vm, &spec, def)?
            } else {
                // Default module creation
                PyModule::from_def(def).into_ref(&vm.ctx)
            };

            // Initialize module dict and methods
            // Corresponds to PyModule_FromDefAndSpec: md_def, _add_methods_to_object, PyModule_SetDocString
            PyModule::__init_dict_from_def(vm, &module);
            module.__init_methods(vm)?;

            // Add to sys.modules BEFORE exec (critical for circular import handling)
            sys_modules.set_item(&*name, module.clone().into(), vm)?;

            // Phase 2: Call exec slot (can safely import other modules now)
            if let Some(exec) = def.slots.exec {
                exec(vm, &module)?;
            }

            return Ok(module.into());
        }

        Ok(vm.ctx.none())
    }

    #[derive(FromArgs)]
    struct CreateDynamicArgs {
        #[pyarg(positional)]
        spec: PyObjectRef,
        #[pyarg(positional, optional)]
        _file: crate::function::OptionalArg<PyObjectRef>,
    }

    #[pyfunction]
    fn create_dynamic(args: CreateDynamicArgs, vm: &VirtualMachine) -> PyResult {
        let name_obj = args.spec.get_attr("name", vm)?;
        let name: PyUtf8StrRef = name_obj.try_into_value(vm)?;
        if name.as_str().contains('\0') {
            return Err(vm.new_value_error("embedded null character".to_owned()));
        }

        let origin_obj = args.spec.get_attr("origin", vm)?;
        if vm.is_none(&origin_obj) {
            return Err(vm.new_value_error("origin must be set".to_owned()));
        }
        let origin: PyUtf8StrRef = origin_obj.try_into_value(vm)?;
        if origin.as_str().contains('\0') {
            return Err(vm.new_value_error("embedded null character".to_owned()));
        }

        let sys_modules = vm.sys_module.get_attr("modules", vm)?;
        if let Ok(module) = sys_modules.get_item(&*name, vm) {
            return Ok(module);
        }

        #[cfg(all(feature = "host_env", any(unix, windows)))]
        {
            let origin_str = origin.as_str();
            let short_name = name.as_str().rsplit('.').next().unwrap_or(name.as_str());
            let export_func_name = format!("PyModExport_{short_name}");

            #[cfg(unix)]
            let handle_res = {
                let mode = rustpython_host_env::ctypes::dlopen_mode(None);
                rustpython_host_env::ctypes::open_library_with_mode(origin_str, mode)
            };
            #[cfg(windows)]
            let handle_res = rustpython_host_env::ctypes::open_library(origin_str);

            if let Ok(handle) = handle_res
                && let Ok(export_fn_addr) = rustpython_host_env::ctypes::lookup_function_symbol_addr(
                    handle,
                    export_func_name.as_bytes(),
                )
                && export_fn_addr != 0
            {
                // abi3t PyModExport entry point
                type ModExportFn = unsafe extern "C" fn() -> *mut crate::PyObject;
                let export_fn: ModExportFn =
                    unsafe { core::mem::transmute(export_fn_addr as *const ()) };
                let mod_ptr = unsafe { export_fn() };
                if let Some(mod_nonnull) = core::ptr::NonNull::new(mod_ptr) {
                    let py_obj = unsafe { crate::PyObjectRef::from_raw(mod_nonnull) };
                    return Ok(py_obj);
                }
            }
        }

        Err(vm.new_import_error(
            format!(
                "dynamic module does not define module export function (PyModExport_{})",
                name.as_str()
            ),
            name.into_wtf8(),
        ))
    }

    #[pyfunction]
    fn exec_dynamic(_module: PyRef<PyModule>) -> i32 {
        0
    }

    #[pyfunction]
    fn exec_builtin(_mod: PyRef<PyModule>) -> i32 {
        // For multi-phase init modules, exec is already called in create_builtin
        0
    }

    #[derive(FromArgs)]
    struct FrozenObjectArgs {
        #[pyarg(positional)]
        name: PyUtf8StrRef,
        #[pyarg(positional, optional)]
        data: Option<PyObjectRef>,
    }

    #[pyfunction]
    fn get_frozen_object(args: FrozenObjectArgs, vm: &VirtualMachine) -> PyResult<PyRef<PyCode>> {
        let FrozenObjectArgs { name, data } = args;
        if let Some(data) = data
            && !vm.is_none(&data)
        {
            let invalid_err = || {
                vm.new_import_error(
                    format!("Frozen object named '{}' is invalid", name.as_str()),
                    name.clone().into_wtf8(),
                )
            };
            // A non-buffer is a TypeError, not invalid frozen data. The request
            // is the one marshal.loads() makes, so that what passes here is
            // exactly what it accepts.
            crate::protocol::PyBuffer::from_object(
                vm,
                &data,
                crate::protocol::BufferFlags::SIMPLE,
            )?;
            // The data is a marshalled code object: a whole marshal value, which
            // deserialize_code() does not read — it takes the code body alone,
            // without the type byte the writer puts in front of it.
            let loads = vm.import("marshal", 0)?.get_attr("loads", vm)?;
            let code = loads.call((data,), vm).map_err(|_| invalid_err())?;
            return code.downcast::<PyCode>().map_err(|_| invalid_err());
        }
        import::make_frozen(vm, name.as_str())
    }

    #[pyfunction]
    fn init_frozen(name: PyUtf8StrRef, vm: &VirtualMachine) -> PyResult {
        import::import_frozen(vm, name.as_str())
    }

    #[pyfunction]
    fn is_frozen_package(name: PyUtf8StrRef, vm: &VirtualMachine) -> PyResult<bool> {
        let name_str = name.as_str();
        super::find_frozen(name_str, vm)
            .map(|frozen| frozen.package)
            .map_err(|e| e.to_pyexception(name_str, vm))
    }

    #[pyfunction]
    fn _override_frozen_modules_for_tests(r#override: isize, vm: &VirtualMachine) {
        vm.state.override_frozen_modules.store(r#override);
    }

    #[pyfunction]
    fn _fix_co_filename(code: PyRef<PyCode>, path: PyStrRef, vm: &VirtualMachine) {
        let old_name = code.source_path();
        let new_name = vm.ctx.intern_str(path.as_wtf8());
        super::update_code_filenames(&code, old_name, new_name);
    }

    #[pyfunction]
    fn _frozen_module_names(vm: &VirtualMachine) -> Vec<PyObjectRef> {
        vm.state
            .frozen
            .keys()
            .map(|&name| vm.ctx.new_utf8_str(name).into())
            .collect()
    }

    #[derive(FromArgs)]
    struct FindFrozenArgs {
        #[pyarg(positional)]
        name: PyStrRef,
        #[pyarg(named, default)]
        withdata: bool,
    }

    #[allow(clippy::type_complexity)]
    #[pyfunction]
    fn find_frozen(
        args: FindFrozenArgs,
        vm: &VirtualMachine,
    ) -> PyResult<Option<(Option<PyRef<PyMemoryView>>, bool, Option<PyStrRef>)>> {
        let FindFrozenArgs { name, withdata } = args;
        let Some(name_str) = name.to_str() else {
            return Ok(None);
        };
        let info = match super::find_frozen(name_str, vm) {
            Ok(info) => info,
            Err(FrozenError::NotFound | FrozenError::Disabled | FrozenError::BadName) => {
                return Ok(None);
            }
            Err(e) => return Err(e.to_pyexception(name_str, vm)),
        };

        // The data is what get_frozen_object() takes back, i.e. marshalled code.
        // Frozen modules are stored in their own encoding, so it has to be
        // re-serialized rather than handed out as a view of the stored bytes.
        let data = if withdata {
            let code = PyCode::new_ref_from_frozen(vm, info.code);
            let dumps = vm.import("marshal", 0)?.get_attr("dumps", vm)?;
            let bytes = dumps.call((code,), vm)?;
            Some(PyMemoryView::from_object(&bytes, vm)?.into_ref(&vm.ctx))
        } else {
            None
        };

        // When origname is empty (e.g. __hello_only__), return None.
        // Otherwise return the resolved alias name.
        let origname_str = super::resolve_frozen_alias(name_str);
        let origname = if origname_str.is_empty() {
            None
        } else {
            Some(vm.ctx.new_utf8_str(origname_str).into())
        };
        Ok(Some((data, info.package, origname)))
    }

    #[derive(FromArgs)]
    struct SourceHashArgs {
        #[pyarg(any)]
        key: u64,
        #[pyarg(any)]
        source: PyBytesRef,
    }

    #[pyfunction]
    fn source_hash(SourceHashArgs { key, source }: SourceHashArgs) -> Vec<u8> {
        let hash: u64 = crate::common::hash::keyed_hash(key, source.as_bytes());
        hash.to_le_bytes().to_vec()
    }
}

fn update_code_filenames(
    code: &Py<PyCode>,
    old_name: &'static PyStrInterned,
    new_name: &'static PyStrInterned,
) {
    let current = code.source_path();
    if !core::ptr::eq(current, old_name) && current.as_str() != old_name.as_str() {
        return;
    }
    code.set_source_path(new_name);
    for constant in code.code.constants.iter() {
        let obj: &crate::PyObject = constant.borrow();
        if let Some(inner_code) = obj.downcast_ref::<PyCode>() {
            update_code_filenames(inner_code, old_name, new_name);
        }
    }
}