В прошлый раз я писал про потерянное пробуждение и получил в комментариях справедливое замечание, что примитив там классический. Так и есть. Сегодня про место, где классики меньше: как в моём ядре устроено переключение контекста, и почему ассемблер там знает смещение поля в структуре Rust.

Это неприятная связь. Ассемблеру нужен байтовый доступ к сохранённым регистрам, а раскладку структуры определяет компилятор. Если эти двое разойдутся во мнениях, вы не получите ошибку компиляции. Вы получите ядро, которое пишет регистры мимо и падает где-то далеко от места ошибки.

Сам переключатель

asm

switch_context:
    .byte 0xf3, 0x0f, 0x1e, 0xfa
    mov qword ptr [rdi + {context_off} + 0],  r15
    mov qword ptr [rdi + {context_off} + 8],  r14
    mov qword ptr [rdi + {context_off} + 16], r13
    mov qword ptr [rdi + {context_off} + 24], r12
    mov qword ptr [rdi + {context_off} + 32], rbx
    mov qword ptr [rdi + {context_off} + 40], rbp
    lea rax, [rip + .Lswitch_return]
    mov qword ptr [rdi + {context_off} + 48], rax
    mov qword ptr [rdi + {context_off} + 56], rsp

    mov r15, qword ptr [rsi + {context_off} + 0]
    mov r14, qword ptr [rsi + {context_off} + 8]
    mov r13, qword ptr [rsi + {context_off} + 16]
    mov r12, qword ptr [rsi + {context_off} + 24]
    mov rbx, qword ptr [rsi + {context_off} + 32]
    mov rbp, qword ptr [rsi + {context_off} + 40]
    mov rsp, qword ptr [rsi + {context_off} + 56]
    jmp qword ptr [rsi + {context_off} + 48]
.Lswitch_return:
    ret

Двадцать инструкций, и в них помещается вся смена исполняемого потока. rdi это уходящий поток, rsi приходящий, так требует System V ABI.

Первая строка это endbr64, метка допустимой цели непрямого перехода для Intel CET. Записана байтами.

Почему регистров всего шесть

Первый вопрос, который обычно задают: где остальные регистры? Где rax, rcx, rdx, где вся арифметика?

Их сохранять не нужно, и это не оптимизация, а следствие соглашения о вызовах. По System V регистры делятся на два лагеря. Callee-saved (rbx, rbp, r12–r15) обязана сохранить вызываемая функция. Caller-saved всё остальное, и о них заботится вызывающая сторона.

switch_context вызывается из Rust как обычная функция. Значит компилятор уже сам расставил вокруг вызова сохранение всего, что ему было дорого из caller-saved. Если бы я сохранял их ещё раз, я бы просто дублировал работу, которую уже сделал компилятор.

Остаются шесть callee-saved плюс указатель стека и адрес возврата. Восемь по восемь байт, ровно 64:

rust

#[repr(C)]
#[derive(Clone, Copy)]
pub struct Context {
    pub r15: u64,
    pub r14: u64,
    pub r13: u64,
    pub r12: u64,
    pub rbx: u64,
    pub rbp: u64,
    pub rip: u64,
    pub rsp: u64,
}

Почему адрес возврата сохраняется руками

Обратите внимание на две строки в середине:

asm

    lea rax, [rip + .Lswitch_return]
    mov qword ptr [rdi + {context_off} + 48], rax

Функция берёт адрес метки внутри себя и кладёт его в поле rip уходящего потока. А заканчивается она не через ret, а прыжком по сохранённому rip приходящего:

asm

    jmp qword ptr [rsi + {context_off} + 48]

Причина в том, что обычный ret вернул бы управление туда, откуда пришли, а нам нужно ровно противоположное: уйти туда, где в прошлый раз остановился другой поток. Стек уже переключен строкой выше, так что ret вернулся бы по чужому адресу возврата, и это работало бы почти правильно, что хуже, чем не работало бы вовсе.

Метка .Lswitch_return это точка, в которую поток вернётся, когда его когда-нибудь возобновят. Там стоит ret, и он выполнится уже в контексте того потока, который сюда прыгнет. Получается симметрично: каждый поток входит через jmp и выходит через ret, только в разные моменты времени.

Отдельно есть switch_first для первого запуска потока: сохранять нечего, поэтому он только загружает next.

Та самая константа

{context_off} подставляется вот отсюда:

rust

pub const CONTEXT_OFF: usize = core::mem::offset_of!(Thread, context);

и передаётся в global_asm! как context_off = const CONTEXT_OFF.

Ассемблер не умеет обращаться к полю по имени, ему нужно число. Это число вычисляет компилятор из раскладки структуры Thread. То есть ассемблерный код напрямую зависит от того, в каком порядке в структуре лежат поля.

Thread объявлен как #[repr(C, align(64))], и repr(C) здесь принципиален: он фиксирует порядок полей. У repr(Rust) компилятор вправе переставлять поля как ему удобно, и тогда смещение поехало бы от версии к версии компилятора.

Что мешает всё сломать

Даже с repr(C) остаётся сценарий, в котором всё разваливается молча: кто-то добавляет поле в середину структуры. Смещение context уезжает, константа пересчитывается, ассемблер продолжает работать, только пишет теперь по другим адресам.

Поэтому над вставкой стоят проверки времени компиляции:

rust

const _: () = assert!(core::mem::size_of::<Context>() == 64);
const _: () = assert!(core::mem::offset_of!(Context, r15) == 0, "...");
const _: () = assert!(core::mem::offset_of!(Context, rip) == 48, "...");
const _: () = assert!(core::mem::offset_of!(Context, rsp) == 56, "...");

Смещения 0, 48 и 56 это ровно те числа, которые вписаны в ассемблер руками. Если раскладка изменится, сборка упадёт до того, как получится загрузиться.

Это, кстати, единственная разновидность защиты, которая тут вообще возможна. Система типов Rust ничего не знает про строку mov qword ptr [rdi + 48], для неё это просто текст. Компилятор проверяет то, что можно проверить, а именно числа, а связь между числами и текстом остаётся на моей совести.

Почему область FPU лежит последним полем

Вот сама структура:

rust

#[repr(C, align(64))]
pub struct Thread {
    pub id: ThreadId,
    pub state: ThreadState,
    pub cpu_id: u32,
    pub priority: u8,
    pub ticks_used: u32,
    pub kstack_top: u64,
    pub kstack_phys: u64,
    pub context: Context,
    pub fpu: FpuArea,
    _pad: [u8; 0],
}

FpuArea это килобайт под состояние x87 и SSE, и он стоит после context сознательно. Положи я его выше, смещение context сдвинулось бы на 1024 байта. Ассемблер это переживёт (константа же пересчитается), а вот проверки выше поймают несоответствие и сборка упадёт.

То есть порядок полей здесь не косметика, а часть контракта с ассемблером. В исходнике над полем fpu стоит комментарий ровно об этом, чтобы следующий человек (скорее всего, я через полгода) не переставил его машинально.

Сама область:

rust

#[repr(C, align(64))]
pub struct FpuArea {
    bytes: [u8; FPU_AREA_SIZE],
}

Выравнивание на 64 байта не эстетика, а требование инструкции XSAVE: она работает только с 64-байтно выровненным абсолютным адресом. Поэтому выровнен и сам Thread, и размещается он тоже по 64-выровненному адресу.

Инициализация нулями, кроме двух байт:

rust

bytes[24] = 0x80; // MXCSR low byte
bytes[25] = 0x1f; // MXCSR high byte -> 0x1F80

0x1F80 это значение MXCSR по умолчанию: все исключения SSE замаскированы. Ноль там означал бы включённые исключения, и первое же деление в пользовательском коде улетело бы в обработчик.

А XSTATE_BV в заголовке остаётся нулём намеренно. По правилам XRSTOR это значит «инициализируй компоненты в состояние по умолчанию», то есть свежий поток получает чистый x87 и SSE без того, чтобы я руками собирал корректный образ.

Размер тоже проверяется на этапе компиляции:

rust

const _: () = assert!(FPU_AREA_SIZE >= kernel_arch_x86_64::xsave::XSAVE_X87_SSE_SIZE, "...");
const _: () = assert!(FPU_AREA_SIZE.is_multiple_of(64), "...");

Сохранение сделано жадным: XSAVE уходящему и XRSTOR приходящему на каждом переключении, без ленивых схем с ловлей исключения при первом обращении к FPU. Ленивое сохранение экономит работу на потоках, которые FPU не трогают, но приносит собственный класс гонок при многоядерном исполнении. Мне сейчас важнее предсказуемость.

Что происходит до переключения регистров

Полная последовательность в хвосте schedule() выглядит так:

rust

    // Lock released - switch_context must follow immediately.
    run_context_switch_hook(next, cpu_id);

    #[cfg(target_os = "none")]
    {
        const FPU_MASK: u64 = kernel_arch_x86_64::xsave::XCR0_X87_SSE;
        if !cur.is_null() {
            unsafe { kernel_arch_x86_64::xsave::xsave64((*cur).fpu.as_mut_ptr(), FPU_MASK) };
        }
        unsafe { kernel_arch_x86_64::xsave::xrstor64((*next).fpu.as_ptr(), FPU_MASK) };
    }

    if cur.is_null() {
        unsafe { switch::switch_first(next) };
    } else {
        unsafe { switch::switch_context(cur, next) };
    }

Сначала адресное пространство, потом FPU, потом регистры. Всё это при выключенных прерываниях.

Хук нужен потому, что планировщик ничего не знает про процессы. Он оперирует потоками, а адресные пространства живут этажом выше, и связывает их зарегистрированный колбэк.

Самое интересное место

Внутри хука вызывается вот это:

rust

pub(crate) unsafe fn activate_address_space(new_as: *mut AddressSpace, cpu_id: u32) {
    let slot = &CPU_CURRENT_AS[cpu_id as usize % MAX_CPUS];
    let old_as = slot.load(Ordering::Acquire) as *mut AddressSpace;
    if old_as == new_as {
        return;
    }
    unsafe { (*new_as).active_cpus.mark_active(cpu_id) };
    slot.store(new_as as usize, Ordering::Release);
    unsafe { (*new_as).activate() };   // здесь пишется CR3
    if !old_as.is_null() {
        unsafe { (*old_as).active_cpus.mark_inactive(cpu_id) };
    }
}

Обратите внимание на порядок: новое адресное пространство помечается активным до записи CR3, а старое помечается неактивным после.

Почему не наоборот. Множество active_cpus используется при инвалидации TLB: когда где-то меняется отображение страниц, надо разослать межпроцессорные прерывания тем ядрам, которые это адресное пространство используют. Если пометить старое неактивным раньше записи CR3, появляется окно, в котором ядро всё ещё исполняется со старым CR3, но в множестве его уже нет. Пришедшая в это окно инвалидация обойдёт нас стороной, и в TLB останется устаревшая запись.

Текущий порядок даёт временный период, когда ядро числится в обоих множествах сразу. Это лишнее прерывание, то есть чуть-чуть лишней работы. Потерянная инвалидация это тихо неверная трансляция адреса. Между «иногда лишний IPI» и «иногда неправильная страница» выбор очевиден.

И заодно ответ на вопрос, который мне уже задавали: CR3 при переключении между процессами загружается, просто не внутри switch_context, а до него.

Чисто ядерные потоки (idle и подобные) этот путь пропускают: своего адресного пространства у них нет, а верхняя половина ядра отображена во всех PML4, так что предыдущий CR3 остаётся загруженным и всё продолжает работать.

Чего я тут не проверяю

Раз уж говорим про безопасность, скажу и про её пределы.

Всё, что описано выше, держится на трёх вещах: repr(C), проверках смещений на этапе компиляции и на том, что я не ошибся, вписывая числа в ассемблер. Первые две проверяет компилятор. Третью не проверяет никто.

Соответствие между mov qword ptr [rdi + {context_off} + 32], rbx и полем rbx по смещению 32 существует только потому, что я так написал. Перепутай я местами две строки, всё соберётся, а поведение станет удивительным. Здесь помогают только внимательность и то, что этот код почти не меняется.

Ещё одно ограничение: FPU-ветка идёт под #[cfg(target_os = "none")], то есть на хостовой сборке, где гоняются тесты логики, её нет вовсе, там switch_context заглушка. Значит, проверить эту часть можно только на реальной загрузке в эмуляторе. У меня для неё есть отдельный тест, и главное, что я про него знаю: если выключить XSAVE, он падает. Тест, который не падает при удалении починки, ничего не доказывает.

Код открыт, писал один, стадия ранняя: github.com/ScioFuturum/rumicos

Если увидите в порядке операций дыру, которую я не заметил, напишите. Это ровно тот код, где ошибка не проявляется там, где сделана.