mojo

以安全的视角浅谈新生代专为AI设计的语言Mojo

0x01 介绍

在刚刚过去的defcon quals 2024上出现了Mojo[1]写的应用，看见了小伙伴对它的吐槽，我也很好奇它到底是怎样的一个语言，决定深入探索一下。Mojo的主推者Chris Lattner同时是LLVM和Swift的创始人，我想这样优秀的编程语言领域工程师，品控一定不会太差。很早之前就听闻过Mojo，但是一直没有尝试过去了解它，对它的印象仅仅是来源于它本身的一个宣传 "专为AI设计的语言，兼容Python，并且要比Python快xxx倍"。那么我觉得它的定位，或者说试用人群，应当是那些以Python为主，并且想要写出高质量的代码的AI工程师。

接下来文章，首先我会介绍Mojo的一些基础知识和包括调试过程，然后去理解defcon中出现的Mojo应用中存在的问题以及利用的方式，最后给出我个人对Mojo的一些看法。

0x02 Mojo基础

2.1 def和fn函数

Mojo支持以两种函数定义，但是实际上def是fn一个语法糖。前者对标了Python中的def，相对来说更加的灵活，主要体现函数参数和返回值不需要显示地指明类型，函数体中的变量定义不需要显示使用var, 下面是一个简单的def函数例子:

1
2
3
def greet(name):
    greeting = "Hello, " + name + "!"
    return greeting

看起来和Python似乎并没有太大的区别，可以说一模一样。而fn则是需要让def可选的东西全部变成必须，对标上面的例子的fn函数:

1
2
3
fn greet(name: String) -> String:
    var greeting = "Hello, " + name + "!"
    return greeting

使用了强制的类型检查，可以看做是fn是def的strict版本。

2.2 变量，类型与结构

Mojo允许定义一个指定类型的变量 (i.e.,var id : Int)。Mojo拥有所有的基本类型，对于number类型分的比较细(i.e.,Int8, Int16 Int32, Int64, Float32, Float64等等)，并且用于拥有一些特殊类型:

  1. SIMD类型，支持各种SIMD操作 (single instruction, multiple data)。你可以通过它来定一个固定长度的向量，并且高效地现实各种向量操作。
  2. Register-passable和Memory-only类型，根据变量所在的位置来区分它们。

更多的类型可以查看它们的官方文档。Mojo没有class的概念，与之对应是struct，这一点和Go比较类似，但它拥有和Python类似的magic methods.

2.4 值语义和引用语义

Mojo同时支持值语义和引用语义。我觉得在了解一个语言的时候，是非常有必要去弄清楚它当中的值是如何传递的。值语义另外一个通俗的说法是值传递，而引用语义即为引用传递。值传递意味接收方并不会对原值产生影响，常见的值传递包括

  1. 传递一个copied value,
  2. 传递一个immutable reference,

常见的引用传递对应mutable reference. Mojo想要做到以值传递为主，并且安全地进行引用传递。

2.5 所有权机制及变量生命周期

所有权机制是当下避免GC的一个相当火热内存管理实现方案，Mojo也采用了这种方案。Mojo为此提出了三个agrument specifier:

  1. borrowed : 接受一个immutable reference作为参数。
  2. inout : 接受一个mutable reference作为参数。
  3. owned : 接受一个value, 并且当前参数是其唯一的owner。

重点理解一下owned, 这里会出现两种情况:

  1. caller把某个值的所有权交给calle。Mojo使用^来作为transfer opertor, 例如f(v^)就将v对应的value所有权传递给了f。
  2. callee获得是一个copied value. 当caller并没有使用^的时候，就会以copied value来传递。

最终配合以值传递为主，Mojo就可以构建一个所有权机制。def所有参数默认为owned，而fn所有参数默认为borrowed。因此下面两个函数是等价的

1
2
3
4
5
def example(borrowed a: Int, inout b: Int, c):
    pass

fn example(a: Int, inout b: Int, owned c: object):
    pass

另外Mojo中显式的引用是以Reference类型出现的，它对应的dereference operator为[]，这一点对于我来说比较奇怪，因为它经常和数组操作绑定在一起。比较奇特是Mojo对于变量生命周期的规定: 当一个变量不活跃之后，Mojo会马上释放它，而不是在等到某个destruction阶段，比如function desctruction。所以值得注意是的，你如果想在调试的时候一直hold某个值到函数结束，必须在函数结尾添加一个特殊的use。

本文需要的Mojo基础知识就这些了，后问会再提及一些Mojo基本类型的相关知识。 更多的Mojo features比如调用python代码，函数参数化和Traits等等可以查看文档，这里就不在累述了。

0x03 调试Mojo

这一节我们将介绍Mojo的编译和调试，主要对象为defcon中的应用，我后面称其为Star。

Mojo是一个比较新的语言，目前该有的基础设施都比较匮乏，所以第一个让我比较头疼的事是如何调试它。官方提供了Debugging指导[2], 使用是LLDB，区别与我常用的GDB，试了一下感觉有些别扭还有一些限制。但是配合VSCode插件还是可以用的，于是我基本路线是:

  1. VSCode+LLDB: 弄清楚Mojo基本类型的内存布局和猜一下Mojo的内存管理。
  2. GDB+Pwntools+Print大法: 检查相关应用运行时的内存。

首先给出对Mojo的第一个吐槽，没有完善的异常处理。这表现出问题直接就是segmentfault，也不知道是为啥。比如需要在编译时通过环境变量MOJO_PYTHON_LIBRARY指定libpython来引入Python runtime，如果你Mojo里面调用了Python代码。当你没有指定时，也会编译通过，然后运行就segfault，通过GBD看backtrace和Star的编译脚本最终才确定原因。我们可以通过下面的方式来实现addressof

1
Reference(target_var).get_unsafe_pointer()

UnsafePointer支持直接打印。第二个需要吐槽是Mojo是编译速度极慢，甚至是一个极小的程序。

3.1 Patches

Star的目标Mojo版本为mojo-24.2.0，但是Star所呈现的问题，在latest版本依然存在，所以我直接使用了它。因此在编译Star之前需要做一些简单的patches:

  • 将所有的Reference.get_unsafe_pointer()换成Referecen.get_legacy_pointer().

  • 使用Python代码中的对象之前，都需要对其解引用，例如

    1
    2
    3
    4
    5
    6
    7
    py_app = Python.import_module('app').App
    py_app.value()

    # 换成

    py_app = Python.import_module('app').App
    py_app.value()[]

3.2 基本类型List

基本类型的定义都可以在官方stblib[3]中找到。这里重点介绍一下List，它对标Python中的List。它的基本定义为

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
#https://github.com/modularml/mojo/blob/main/stdlib/src/collections/list.mojo#L90
struct List[T: CollectionElement](CollectionElement, Sized, Boolable):
	"""The `List` type is a dynamically-allocated list.

    It supports pushing and popping from the back resizing the underlying
    storage as needed.  When it is deallocated, it frees its memory.

    Parameters:
        T: The type of the elements.
    """
	var data: UnsafePointer[T]
    """The underlying storage for the list."""
    var size: Int
    """The number of elements in the list."""
    var capacity: Int
    ...

可以看到Mojo的List实现是unsafe的。它在内存的分布为

即struct的字段分布是顺序的。值得注意是List的操作中并没有越界检测以及释放后将data pointer置NULL，已经用户反应了的此类问题[4]. 官方打算支持但是目前还没有时间。这一点是非常震惊我的，这意味OOB和UAF在Mojo是很容易做的。

3.3 基本类型String

它本质是一个List。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
#https://github.com/modularml/mojo/blob/main/stdlib/src/builtin/string.mojo#L328
struct String(
    Sized,
    Stringable,
    IntableRaising,
    KeyElement,
    Boolable,
    Formattable,
    ToFormatter,
):
    """Represents a mutable string."""

    alias _buffer_type = List[Int8]
    var _buffer: Self._buffer_type
    ...

0x03 Star中的问题

Star本身涵盖了defcon中的两个题目，第一个是一个越权问题，第二问题发生在privileged功能中，我不会过多的阐述如何利用第一个越权问题，假设我们已经拥有了想要的权限。Star是一个类似Web应用，注册了相关的路由，用来处理用户请求，其中某些路由是privileged。而Star本身的功能是实现了一个Database, 具有collections和fragments的概念，用户可以在database中新增和修改它们。Star不仅包括Mojo代码，也包括Python相关代码，前者会调用后者。

第二个问题和前面提到的owned标识符有关。因为原应用相对比较复杂，我用伪代码来阐述这里的问题, 如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
struct App:
    var db: Database

    async fn add_collections (inout self: Self, name: String):
        self.db.collections.add(Collection(name))
        
    async fn add_fragment (inout self: Self, col_index: Int, var data: String):
        if col_index < len(self.db.collections):
            self.db.gef_ref(col_index)[].fragments.add(Fragment(data))
            
    async fn do_modify(inout self: Self, col_index: Int, fra_index, filter : String, value: String):
        if re.match(self.db.collection[col_index][fra_index]):
            self.db.collection.gef_ref(col_index)[].fragments.history.get_ref(fra_index)[].add(value)
        
    async fn modify_fragments (inout self: Self, owned db: Databse, owned filter: String, owned value: string, dict Dict[int,int]):
        for e in dict.items():
            var col_index = e.key()
            var fra_index = e.value()
            if col_index < len(self.db.collections) and fra_index < len(db.collections[col_index].fragments)
                await do_modify(col_index, fra_index, filter, value)
        ...

这里App含一个Database实例， 它里面有4个async函数 (async函数运行实例在Mojo中视为一个coroutine):

  1. 添加一个collection.
  2. 在指定的collection中添加一个fragment.
  3. 更新指定的fragments，用一个dict来给定它们的位置。此外还有一个正则表达式filter，用来匹配指定的fragments.
  4. 根据filter匹配fragment，对匹配到的fragment，在其history (类型List) 中新增给定的value.

这里的问题出现在一个非常细微的地方，在modfiy_fragments中传进来了一个db，它的标识符为owned，并且在后面的if第二个条件使用了它。当调用形式为self.modify_fragments(self.db, ...)，db就会是一个copied value，它和self.db是两个独立的东西。这就造成了if的两个条件使用的database实例是不一样， 可能潜在的导致上面的bound check if失败。 这个地方看起来似乎是比较刻意的，但是我觉得这样的隐式拷贝是非常有可能出现在实际中的，如果用户没有正确地理解Mojo中的值传递概念。

0x04 利用

[5]给出的solution需要通过race来绕过bound check实现OOB。具体的方案是构造两个coroutine:

  1. coroutine1: 调用modify_fragments，指定fragment dict形如{0:0, 0:0, ...., N:M}，其中col:0, fra:0是存在的fragment，而col:N, fra:M是不存在的fragment，我们假设只有一个collection。 通过精心地构造存在的fragment内容和传入filter，我们可以在do_modify中实现RE-Dos，用来延迟最后关于col:1, fra:20的访问。
  2. coroutine2: 调用add_collections，往app.db中新增collections，使得当coroutine1访问col:N, fra:M通过if的第一个条件，而在第二个条件上发生OOB（因为copied db不会发生变化）。

但是即使发生了OOB，也需要让绕过第二条件，由此达到do_motify中的OOB write。我们可以通过调整N来找到一个合适的faked collection, 使得它当中用于记录fragments的List对应的len字段大于M。这一点很容易做到，因为堆上往往有很大的随机数字，因此这里大概率不需要使用堆喷。并且Mojo runtime没有对List的额外效验，所以伪造List比较简单。

注意modfiy_fragments具体操作是对fragment中的history list新增一个string。[5]中使用了堆喷来布置大量的faked fragments, 并且使得其中list的data_pointer指向Star的路由表对应list，然后调整M，使得正好访问这些faked fragments中的一个, 实现对路由表的修改。此时路由表所在位置实际存储是一个string，但是前面我们说过，string本身还是一个list, 因此我们可以这个string上未知任意的路由表。最终将其劫持到一个Star本身给定的exploitable的函数上。其完整的利用过程如下:

0x05 总结

回到最初我们提到的Mojo的使命，它是否可以让以Python为主的AI工程师进行高效地工作呢？ 目前看来它还达不到。

引入的所有权机制无疑会增加工程师们心智上的负担。比如在写某个函数的时候，就会陷入我这个参数应该用哪个标识符呢？在Python里面你似乎没有这样的顾虑，而你忧虑的性能也许Python已经帮你顾及到了，例如reference counting + copy on write可以在一定程度上弥补使用immutable reference带来的性能。其次，如果某个地方涉及到了ownership的传递，而你不小心少写了一个transfer operator，可能性能也因此不会变好。但是，我觉得Mojo上ownership model的实现上要比Rust看起来更加的清爽。

相比于一次性写出高效的算法实现，当下我更加同意另外一个观点，有人负责算法的实现，有人负责实现上的优化，彼此分工合作，那这是不是也意味Python只是缺少一些更加高效的JIT技术呢？另外一个，在stdlib中各种乱飞的unsafe和直接使用MLIR的操作，比较令人担忧。但是无论如何它是在MLIR上的一个native language, 相比于那些喊着口号为AI设计的语言，要更贴切主题一些。

总之，Mojo目前依然还不是一个成熟的语言，但是它在往前走，是否哪一个语言不需要个把年来沉淀？对此我还是比较期待它的未来的，愿意接受新事物的同学，可以大胆一探究竟。

0x06 引用

  1. Mojo官方文档, https://docs.modular.com/mojo
  2. Mojo Debugging, https://docs.modular.com/mojo/manual/values/ownership
  3. Mojo stdlib https://github.com/modularml/mojo/tree/main/stdlib
  4. Use-after-free / Out-of-bound on Mojo Pointer, https://github.com/modularml/mojo/issues/142
  5. Star solver, https://github.com/Nautilus-Institute/quals-2024/blob/main/%F0%9F%8C%8C/solver/solver.py