Rust研究：Pythonスクリプトの呼び出し

Rust 初心者です。

スクリプトで実装した計算処理を Rust から呼び出せたら便利だと思いました。
PyO3 を使うと Python と連携できるようなので調べてみました。

今回もコンパイルできたところまでメモします。

お題

Rust コードに数値計算の関数をハードコーディングしてしまうと関数の処理内容を変更するたびに
再コンパイルが必要になるので、数値計算の関数は Rust コードから分離して Python スクリプトで
実装するようにし、それを Rust コードから呼び出すようにしたい。

スクリプトはこんな感じ↓で。

import math

def func(x):
    return math.cos(x)

def gaussian(x, a=1.0, b=0, c=1.0):
    y = (x-b)/c
    return a*math.exp(-y*y)

Rust コードからは "func" とか "gaussian" といった関数名で呼び出す。

スクリプトファイルの場所は利用者が指定できるものとする。

ソースコード

PyO3 を利用。

Cargo.toml

[package]
name = "study0203"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
pyo3 = { version = "0.18.0", features = ["auto-initialize"] }

// lib.rs
pub mod conf;
pub mod error;
pub mod python;

// python.rs

// [Referrence]
//
// PyO3 user guide / Calling Python in Rust code
//   https://pyo3.rs/v0.18.0/python_from_rust.html#executing-existing-python-code
//
// sys.path (モジュールを検索するパス)
//   https://docs.python.org/3/library/sys.html#sys.path

use pyo3::prelude::{Python,PyResult};
use pyo3::PyTryInto;
use pyo3::types::PyList;

// run_py は Python スクリプトで定義された関数に指定された引数を渡して計算させる。
// 引数：
//   script_path --- 処理するスクリプトファイルのあるフォルダのパス
//   script_name --- 処理するスクリプト名　（スクリプトファイルの拡張子 ".py" の手前まで）
//   func_name --- 処理する関数名
//   args --- スクリプトで定義された関数に渡す引数のリスト
// 返り値：
//   各引数で関数を計算した結果のリスト
// [WARNING]
// この関数は実行時のプロセスの権限の範囲内で任意の Python スクリプトを実行できてしまう
// ことから、スクリプトに想定外のことが実装されていると脆弱性となる恐れがあることに注意が必要である。
// 信頼できないスクリプトを実行させる場合には別途適切なサンドボックスなどを用意すること。

pub fn run_py(script_path:&str, script_name:&str, func_name:&str, args:Vec<f32>)
    -> PyResult<Vec<f32>> {

    Python::with_gil(|py| {

        // sys.path に処理するスクリプトのあるパスを追加する
        let syspath = py.import("sys")?.getattr("path")?;
        let syspath: &PyList = PyTryInto::try_into(syspath)?;
        syspath.insert(0, script_path)?;
    
        // 対象となる関数の取得
        let run = py.import(script_name)?.getattr(func_name)?;

        // 関数の計算結果を格納するリストを用意
        let mut results:Vec<f32> = vec![];

        // 各引数で関数を計算させ、その結果をリストに格納する
        for x in args {
            let result:f32 = run.call((x, ), None)?.extract()?;
            results.push(result)
        }

        // println!("results={:?}", results);

        Ok(results)
    })
}

// conf.rs
use std::env;
use std::path::Path;
use crate::error::ErrorCode;

// 設定
pub struct Conf {
    script_path: String, // 処理するスクリプトのあるフォルダのパス
    script_name: String, // 処理するスクリプト名（拡張子 .py なし）
    func_name: String, // 処理する関数名
}

impl Conf {
    pub fn new () -> Result<Self, ErrorCode> {

        // 引数をチェック
        let args:Vec<String> = env::args().collect();
        if args.len() < 3 {
            return Err(ErrorCode::ArgumentError)
        }

        // 処理する関数名
        let func_name = &args[2];

        // 処理するスクリプト名とスクリプトのあるフォルダのパス
        let (script_name, script_path) = get_script_name_and_path(&args[1])?;

        Ok(Conf{
            script_path: script_path,
            script_name: script_name, 
            func_name: func_name.to_string(),
        })
    }

    pub fn script_name(&self) -> &str {
        &self.script_name
    }

    pub fn script_path(&self) -> &str {
        &self.script_path
    }

    pub fn func_name(&self) -> &str {
        &self.func_name
    }
}

// 指定されたスクリプトのファイルパスからスクリプト名とスクリプトのあるフォルダのパスを返す。
// その際にいくつかのチェックを行う。
fn get_script_name_and_path(file_path: &str) -> Result<(String, String), ErrorCode> {
    let p = Path::new(file_path);

    // ファイルパスの有無の確認。ないときはエラー
    if !p.exists() {
        return Err(ErrorCode::FileNotFoundError(file_path.to_string()))
    }

    // スクリプトファイル名の取得
    let script_name:Option<&str> = match p.file_name() { // Option<&OsStr>
        None => return Err(ErrorCode::AbnormalScriptNameError(format!(
            "not a normal file. maybe a directory: \"{}\"", file_path))),
        Some(sn) => sn.to_str() // &OsStr から &str に変換 Option(&str)
    };

    // Option(&str) の unpack
    let script_name:&str = match script_name {
        Some(sn) => sn,
        None => return Err(ErrorCode::AbnormalScriptNameError(format!(
            "invalid utf8: \"{}\"", file_path))),
    };

    // スクリプトファイル名の拡張子が ".py" かを確認
    if ! script_name.ends_with(".py") {
        return Err(ErrorCode::AbnormalScriptNameError(format!(
            "not .py: \"{}\"", file_path)))
    }

    // スクリプトファイル名の末尾の ".py" をカット
    let script_name = &script_name[0..script_name.len()-3];

    // スクリプトファイルのあるフォルダのパスの取得
    let parent_path = match p.parent() {
        None => return Err(ErrorCode::AbnormalScriptNameError(format!(
            "the path terminates in a root or prefix, or if it’s the empty string: \"{}\"", 
            file_path))),
        Some(parent) => parent
    };


    // [TODO]
    // パスの有無以外にパスのチェックを行っていない。
    // 利用できるパスを制限する条件を追加すべきだが後で検討する。
    // 例：信頼できる特定のフォルダ配下のパスに限定する

    Ok((script_name.to_string(), format!("{}", parent_path.display())))
}

// error.rs
use std::process::{Termination,ExitCode};

pub enum ErrorCode {
    Success,
    ArgumentError,
    AbnormalScriptNameError(String),
    FileNotFoundError(String),
    PythonError(String),
}

impl ErrorCode {
    pub fn to_value(&self) -> u8 {
        match self {
            ErrorCode::Success => 0,
            ErrorCode::ArgumentError => 10,
            ErrorCode::AbnormalScriptNameError(_) => 20,
            ErrorCode::FileNotFoundError(_) => 21,
            ErrorCode::PythonError(_) => 30,
        }
    }

    pub fn to_string(&self) -> String {
        match self {
            ErrorCode::Success => "Success".to_string(),
            ErrorCode::ArgumentError => "[Usage] study0203.exe script.py func_name".to_string(),
            ErrorCode::AbnormalScriptNameError(file_path) => format!("[Abnormal Script Name] {}", file_path),
            ErrorCode::FileNotFoundError(file_path) => format!("[File Not Found] file:{}", file_path),
            ErrorCode::PythonError(msg) => format!("[Python Error] {}", msg),
        }
    }

}

// ErrorCode に Termination トレイトを実装する
impl Termination for ErrorCode {
    fn report(self) -> ExitCode {
        ExitCode::from(self.to_value())
    }
}

// main.rs

// [WARNING]
// このプログラムでは実行時の権限の範囲内で任意の Python スクリプトを実行できることから、
// スクリプトに想定外のことが実装されていると脆弱性となる恐れがあることに注意が必要である。
// 信頼できないスクリプトを実行させる場合には別途適切なサンドボックスなどを用意すること。

use study0203::conf::Conf;
use study0203::error::ErrorCode;
use study0203::python::run_py;

use std::f32::consts::PI;

fn main() -> ErrorCode {
    let conf = match Conf::new() {
        Ok(c) => c,
        Err(e) => return console(e)
    };

    console(run(conf))
}

fn run(conf:Conf) -> ErrorCode {

    // Pythonスクリプトの関数にあたえる引数のリスト(-π～π)
    let args:Vec<f32> = vec![-PI, -0.75*PI, -0.5*PI, -0.25*PI, 0.0, 0.25*PI, 0.5*PI, 0.75*PI, PI];

    // Pythonスクリプトの関数の実行
    let results = match run_py(conf.script_path(), conf.script_name(), conf.func_name(), args) {
        Ok(results) => results,
        Err(e) => return ErrorCode::PythonError(format!("{:?}", e))
    };
    println!("{:?}", results);

    ErrorCode::Success
}

fn console(e:ErrorCode) -> ErrorCode {
    match e {
        ErrorCode::Success => (),
        _ => eprintln!("[ERROR] {}", e.to_string())
    };
    e
}

実行例

C:\work\study0203>where python
C:\Users\hama2\AppData\Local\Programs\Python\Python311\python.exe

C:\work\study0203>python -V
Python 3.11.1

C:\work\study0203>target\debug\study0203.exe
[ERROR] [Usage] study0203.exe script.py func_name

C:\work\study0203>echo %errorlevel%
10

C:\work\study0203>target\debug\study0203.exe script.py func
[-1.0, -0.70710677, -4.371139e-8, 0.70710677, 1.0, 0.70710677, -4.371139e-8, -0.70710677, -1.0]

C:\work\study0203>target\debug\study0203.exe script.py gaussian
[5.1723157e-5, 0.0038810384, 0.08480496, 0.53964144, 1.0, 0.53964144, 0.08480496, 0.0038810384, 5.1723157e-5]

C:\work\study0203>echo %errorlevel%
0

ここまでの感想

• Rust の勉強をはじめて二週間目くらいです。Rust コードを読むのも書くのも大分慣れてきました。やはり Rust コンパイラが素晴らしいです。
• Rust から Python スクリプトを呼び出すところで少しはまりました。Stack overflow の古い記事やら PyO3 の API ドキュメントを探って試行錯誤しました。最後は Rust コンパイラのエラーメッセージで "PyTryInto::try_into" にたどり着きました。
• PyO3 の API で Python の dict オブジェクトも作れそうなので、Rust から Python スクリプトへもう少し複雑な構造データを受け渡すことにもいずれ挑戦してみたいです。
• ソースコードにもコメントいれましたが、Python スクリプトの中から色々なモジュールをインポートできることから、信頼できないスクリプトを動かさなければならない場合にはサンドボックスの仕組みを用意しないと場合によっては危険なことになるかもしれないと思いました。
• たとえば動作環境をコンテナ化して挙動を制限する、管理者権限のないユーザ権限で動作させる、libseccomp でシステムコールを制限する、capability で権限を制限する、chroot で使用するルートをシステムルートとは別にする、auditd で監視する、動作環境の hardening …。
• 今さらながらですが、ちょこちょこ計算内容を変更したいようなシチュエーションの場合には最初からすべて Python で実装すればいいだけの話なので、今回のようなシチュエーションは稀なのかもしれません。でも、Rust ソースに計算内容をハードコーディングしたくないケースもそれなりにありそうな気もします。