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

AI 1个月前 admin
19 0 0

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 deffn函数

Mojo支持以两种函数定义,但是实际上deffn一个语法糖。前者对标了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

使用了强制的类型检查,可以看做是fndef的strict版本。

2.2 变量,类型与结构

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

  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的。它在内存的分布为

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

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的函数上。其完整的利用过程如下:

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

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

原文始发于maplgebra:以安全的视角浅谈新生代专为AI设计的语言Mojo

版权声明:admin 发表于 2024年5月21日 上午8:49。
转载请注明:以安全的视角浅谈新生代专为AI设计的语言Mojo | CTF导航

相关文章