Vulnerabilidad TOCTOU en Diffusers permite bypass de trust_remote_code y ejecución remota de código

Contexto

Esta vulnerabilidad se encuentra en el paquete diffusers, la biblioteca equivalente a transformers para modelos de difusión.

Se localiza en el flujo DiffusionPipeline.from_pretrained, que se utiliza para cargar un pipeline desde el Hub de HuggingFace.

Esta función tiene una protección trust_remote_code: si el archivo model_index.json del repositorio hace referencia a una clase de pipeline personalizada definida en un archivo .py en el repositorio, la carga se bloquea a menos que se pase explícitamente trust_remote_code=True:

ValueError: The repository for attacker/repo contains custom code in pipeline.py
which must be executed to correctly load the model. You can inspect the repository
content at https://hf.co/attacker/repo/blob/main/pipeline.py.
Please pass the argument `trust_remote_code=True` to allow custom code to be run.

La vulnerabilidad permite la ejecución arbitraria de código a través del flujo de pipeline personalizado desde un repositorio del Hub, sin pasar los argumentos custom_pipeline o trust_remote_code. La llamada from_pretrained tiene éxito y devuelve un pipeline funcional.

Flujo Normal

DiffusionPipeline.from_pretrained comienza extrayendo todos los argumentos relevantes de kwargs en variables locales, luego llama a DiffusionPipeline.download() para obtener los archivos del repositorio:

# pipeline_utils.py:853
cached_folder = cls.download(
    pretrained_model_name_or_path,
    ...
    custom_pipeline=custom_pipeline,
    trust_remote_code=trust_remote_code,
    ...
)

Dentro de download(), model_index.json se obtiene primero como un archivo independiente vía hf_hub_download:

# pipeline_utils.py:1636
config_file = hf_hub_download(
    pretrained_model_name,
    cls.config_name,
    ...
)
config_dict = cls._dict_from_json_file(config_file)

Esta configuración se utiliza para detectar código de pipeline personalizado y aplicar la verificación de confianza:

# pipeline_utils.py:1672
if custom_pipeline is None and isinstance(config_dict["_class_name"], (list, tuple)):
    custom_pipeline = config_dict["_class_name"][0]

load_pipe_from_hub = custom_pipeline is not None and f"{custom_pipeline}.py" in filenames

if load_pipe_from_hub and not trust_remote_code:
    raise ValueError(...)

Después de que la verificación pasa, snapshot_download obtiene todos los archivos y los guarda en disco:

# pipeline_utils.py:1778
cached_folder = snapshot_download(
    pretrained_model_name,
    ...
    revision=revision,
    allow_patterns=allow_patterns,
    ...
)

De vuelta en from_pretrained, la configuración se lee una segunda vez desde la instantánea descargada, y _resolve_custom_pipeline_and_cls lee la configuración para verificar nuevamente si se necesita cargar código personalizado:

# pipeline_loading_utils.py:974
def _resolve_custom_pipeline_and_cls(folder, config, custom_pipeline):
    custom_class_name = None
    if os.path.isfile(os.path.join(folder, f"{custom_pipeline}.py")):
        custom_pipeline = os.path.join(folder, f"{custom_pipeline}.py")
    elif isinstance(config["_class_name"], (list, tuple)) and os.path.isfile(
        os.path.join(folder, f"{config['_class_name'][0]}.py")
    ):
        custom_pipeline = os.path.join(folder, f"{config['_class_name'][0]}.py")
        custom_class_name = config["_class_name"][1]

    return custom_pipeline, custom_class_name

Si la configuración apunta a un archivo .py, se importa.

La Vulnerabilidad

hf_hub_download y snapshot_download son dos llamadas HTTP independientes al Hub, ambas resolviendo la rama predeterminada del repositorio (si revision=None) a su HEAD actual en el momento de la llamada. No hay garantía de atomicidad entre ellas: si el repositorio se actualiza entre las dos llamadas, resolverán a commits diferentes y descargarán contenido diferente, sin mostrar advertencia al usuario.

La verificación de confianza en download() opera sobre el contenido obtenido por hf_hub_download (commit A). La llamada snapshot_download que sigue inmediatamente puede obtener silenciosamente un commit más nuevo (commit B). La configuración en el commit más nuevo será la que analice _resolve_custom_pipeline_and_cls.

Por lo tanto, es posible introducir código remoto en el repositorio entre las dos llamadas, evadiendo la verificación de confianza.

La ventana de carrera es todo lo que está entre las dos llamadas al Hub dentro de download():

# pipeline_utils.py:1636
config_file = hf_hub_download(...)   # ← ve commit A, verificación de confianza pasa

# ... procesamiento de nombres de archivo, construcción de patrones, verificación pipeline_is_cached ...
# ~~~ EL ATACANTE EMPUJA COMMIT B AQUÍ ~~~

# pipeline_utils.py:1778
cached_folder = snapshot_download(...)  # ← ve commit B, descarga pipeline.py

Para el exploit, el commit A lleva una configuración limpia con _class_name como una cadena simple, lo que causa que load_pipe_from_hub sea False y la verificación de confianza pase. El commit B cambia _class_name a una lista y agrega pipeline.py:

Commit A - model_index.json:

{
  "_class_name": "FluxPipeline",
  "_diffusers_version": "0.31.0"
}

Commit B - model_index.json:

{
  "_class_name": ["pipeline", "FluxPipeline"],
  "_diffusers_version": "0.31.0"
}

Cuando from_pretrained lee la instantánea después de que download() retorna, config["_class_name"] es ahora una lista, pipeline.py existe en disco (obtenido por snapshot_download), y _resolve_custom_pipeline_and_cls resuelve custom_pipeline a la ruta local de ese archivo. _get_pipeline_class luego lo importa, sin verificación de confianza en este punto del código.

Prueba de Concepto

1. Crear un repositorio del Hub con el model_index.json del commit A (cadena simple _class_name).
2. Ejecutar DiffusionPipeline.from_pretrained("attacker/repo") con un breakpoint establecido en pipeline_utils.py:1778 (la llamada snapshot_download). Esto es para que la ventana sea lo suficientemente grande para responder manualmente.
3. Cuando la ejecución se pause en el breakpoint, empujar el commit B: actualizar model_index.json para usar una lista _class_name y agregar pipeline.py.
4. Reanudar la ejecución.
5. snapshot_download obtiene el commit B; /tmp/pwned se escribe durante la subsecuente llamada _get_pipeline_class.

Limitaciones

• No aplica cuando revision está fijado a un hash de commit específico: ambas llamadas al Hub resuelven al mismo contenido.
• No aplica cuando se carga desde un directorio local.
• Si todos los archivos esperados ya están presentes en el caché local de HF, download() retorna temprano antes de alcanzar snapshot_download (retorno temprano línea 1767), cerrando la ventana de carrera. El exploit por tanto requiere una primera descarga (o forzada).

Explotabilidad

La ventana entre las dos llamadas es muy corta. Las pruebas locales resultaron en una ventana de aproximadamente ~0.5 segundos para que el atacante empuje el cambio. Esto es, por supuesto, inviable de lograr para cada nueva descarga. Sin embargo, dado un repositorio popular con muchas descargas por día, uno puede lograr éxito estadístico cambiando el estado del repositorio de vez en cuando o cada pocos segundos, con algún porcentaje de descargadores cayendo en la ventana exacta.

Impacto

La vulnerabilidad es un RCE silencioso: permite que se cargue código arbitrario a través del flujo de pipeline personalizado desde un repositorio del Hub, sin argumentos custom_pipeline o trust_remote_code. La llamada from_pretrained tiene éxito y devuelve un pipeline completamente funcional.