make adding kmods easier

examples/c_kmod/README.md

@@ -1,15 +1,16 @@
 # README
 
-This directory houses example kernel modules that can automatically be compiled into a kernel by
-tweaking your `LIKE_DBG` config:
+This directory houses example kernel modules that can automatically be compiled into a kernel by tweaking your `LIKE_DBG` config:
 
 ```ini
 [kernel_builder]
 # Provide a path to a parent directory that houses custom kernel modules (see the example)
 custom_modules = examples/c_kmod/
 ```
 
-Every module located in `examples/c_kmod` is to be placed in a proper subdirectory, each with a sound `Makefile` and `Kconfig`.
+You can either provide a path to a folder housing a single kernel module such as `examples/c_kmod/ioctl_test_drv/`.
+Alternatively, you can provide a path to a folder housing multiple kernel modules as well, e.g.: `examples/c_kmod/`.
+Then, every module located in `examples/c_kmod` is to be placed in a proper subdirectory, each with a sound `Makefile` and `Kconfig`.
 For example like so:
 
 ```bash

examples/c_kmod/echo_service/echoservice.c

@@ -1,23 +1,26 @@
 #include <asm-generic/errno-base.h>
-#include <linux/kernel.h>
-#include <linux/module.h>
-#include <linux/device.h>
-#include <linux/init.h>
 #include <linux/cdev.h>
+#include <linux/device.h>
 #include <linux/fs.h>
-#include <linux/uaccess.h>
+#include <linux/init.h>
+#include <linux/kernel.h>
+#include <linux/module.h>
 #include <linux/slab.h>
+#include <linux/uaccess.h>
 #include <stddef.h>
 
 MODULE_LICENSE("GPL");
 MODULE_AUTHOR("0x434b");
-MODULE_DESCRIPTION("Dummy kernel module that highlights how to incorporate a kernel module into LIKE-DBG");
+MODULE_DESCRIPTION("Dummy kernel module that highlights how to incorporate a "
+                   "kernel module into LIKE-DBG");
+MODULE_VERSION("0.1");
 
 /* Prototyes section*/
 static int device_open(struct inode *, struct file *);
 static int device_release(struct inode *, struct file *);
 static ssize_t device_read(struct file *, char __user *, size_t, loff_t *);
-static ssize_t device_write(struct file *, const char __user *, size_t, loff_t *);
+static ssize_t device_write(struct file *, const char __user *, size_t,
+                            loff_t *);
 
 /* Device name as populated in /dev/ */
 #define DEV_NAME "likedbg"
@@ -29,57 +32,57 @@ static ssize_t device_write(struct file *, const char __user *, size_t, loff_t *
 
 char *gbuf = NULL;
 
-
 static int device_open(struct inode *inode, struct file *file) {
-    pr_info("DEVICE_OPEN CALLED\n");
-    gbuf = kmalloc(BUF_SZ, GFP_KERNEL);
-    if (!gbuf) {
-        pr_warn("KMALLOC FAILED\n");
-        return -ENOMEM;
-    }
-    memcpy((void *) gbuf, (void *) BUF_CONTENT, sizeof(BUF_CONTENT));
-
-    return EXIT_SUCCESS;
+  pr_info("DEVICE_OPEN CALLED\n");
+  gbuf = kmalloc(BUF_SZ, GFP_KERNEL);
+  if (!gbuf) {
+    pr_warn("KMALLOC FAILED\n");
+    return -ENOMEM;
+  }
+  memcpy((void *)gbuf, (void *)BUF_CONTENT, sizeof(BUF_CONTENT));
+
+  return EXIT_SUCCESS;
 }
 
-static ssize_t device_read(struct file *file, char __user *buf, size_t count, loff_t *f_pos) {
-    size_t len = count < (BUF_SZ - (*f_pos)) ? count : (BUF_SZ - (*f_pos));
-    pr_info("DEVICE_READ CALLED\n\tREADING %lu bytes (Requested: %lu)\n", len, count);
-    if (copy_to_user(buf, gbuf, len)) {
-        pr_warn("COPY_TO_USER FAILED\n");
-        return -EINVAL;
-    }
-    (*f_pos) += len;
-    return len;
+static ssize_t device_read(struct file *file, char __user *buf, size_t count,
+                           loff_t *f_pos) {
+  size_t len = count < (BUF_SZ - (*f_pos)) ? count : (BUF_SZ - (*f_pos));
+  pr_info("DEVICE_READ CALLED\n\tREADING %lu bytes (Requested: %lu)\n", len,
+          count);
+  if (copy_to_user(buf, gbuf, len)) {
+    pr_warn("COPY_TO_USER FAILED\n");
+    return -EINVAL;
+  }
+  (*f_pos) += len;
+  return len;
 }
 
-static ssize_t device_write(struct file *file, const char __user *buf, size_t count, loff_t *f_pos) {
-    size_t len = count < BUF_SZ ? count : BUF_SZ;
-    pr_info("DEVICE_WRITE CALLED\n");
-    if (copy_from_user(gbuf, buf, len)) {
-        pr_warn("COPY_FROM_USER FAILED\n");
-        return -EINVAL;
-    }
-    return len;
+static ssize_t device_write(struct file *file, const char __user *buf,
+                            size_t count, loff_t *f_pos) {
+  size_t len = count < BUF_SZ ? count : BUF_SZ;
+  pr_info("DEVICE_WRITE CALLED\n");
+  if (copy_from_user(gbuf, buf, len)) {
+    pr_warn("COPY_FROM_USER FAILED\n");
+    return -EINVAL;
+  }
+  return len;
 }
 
 static int device_release(struct inode *inode, struct file *file) {
-    pr_info("DEVICE_RELEASE CALLED\n");
-    kfree(gbuf);
-    return EXIT_SUCCESS;
+  pr_info("DEVICE_RELEASE CALLED\n");
+  kfree(gbuf);
+  return EXIT_SUCCESS;
 }
 
-struct file_operations fops = {
-    .owner = THIS_MODULE,
-    .read = device_read,
-    .write = device_write,
-    .open = device_open,
-    .release = device_release
-};
+struct file_operations echo_fops = {.owner = THIS_MODULE,
+                                    .read = device_read,
+                                    .write = device_write,
+                                    .open = device_open,
+                                    .release = device_release};
 
 static int likedbgdev_uevent(struct device *dev, struct kobj_uevent_env *env) {
-    add_uevent_var(env, "DEVMODE=%#o", 0666);
-    return EXIT_SUCCESS;
+  add_uevent_var(env, "DEVMODE=%#o", 0666);
+  return EXIT_SUCCESS;
 }
 
 static dev_t dev_id;
@@ -88,42 +91,42 @@ static int dev_major = 0;
 static struct class *likedbgdev_class = NULL;
 
 static int __init echo_init(void) {
-    pr_info("HELLO");
-    if (alloc_chrdev_region(&dev_id, 0, 1, DEV_NAME)) {
-        pr_warn("FAILED TO REGISTER CHAR DEVICE: '%s'\n", DEV_NAME);
-        return -EBUSY;
-    }
-    dev_major = MAJOR(dev_id);
-    likedbgdev_class = class_create(THIS_MODULE, DEV_NAME);
-    if (IS_ERR(likedbgdev_class)) {
-        pr_warn("FAILED TO CREATE CLASS\n");
-        return -EBUSY;
-    }
-    likedbgdev_class->dev_uevent = likedbgdev_uevent;
-
-    cdev_init(&c_dev, &fops);
-    c_dev.owner = THIS_MODULE;
-
-    if (cdev_add(&c_dev, MKDEV(dev_major, 0), 1)) {
-        pr_warn("FAILED TO ADD CDEV\n");
-        unregister_chrdev_region(MKDEV(dev_major, 0), MINORMASK);
-        return -EBUSY;
-    }
-    device_create(likedbgdev_class, NULL, MKDEV(dev_major, 0), NULL, DEV_NAME);
-    if (IS_ERR(likedbgdev_class)) {
-        pr_warn("FAILED TO CREATE DEVICE\n");
-        return -EBUSY;
-    }
-    return 0;
+  pr_info("HELLO");
+  if (alloc_chrdev_region(&dev_id, 0, 1, DEV_NAME)) {
+    pr_warn("FAILED TO REGISTER CHAR DEVICE: '%s'\n", DEV_NAME);
+    return -EBUSY;
+  }
+  dev_major = MAJOR(dev_id);
+  likedbgdev_class = class_create(THIS_MODULE, DEV_NAME);
+  if (IS_ERR(likedbgdev_class)) {
+    pr_warn("FAILED TO CREATE CLASS\n");
+    return -EBUSY;
+  }
+  likedbgdev_class->dev_uevent = likedbgdev_uevent;
+
+  cdev_init(&c_dev, &echo_fops);
+  c_dev.owner = THIS_MODULE;
+
+  if (cdev_add(&c_dev, MKDEV(dev_major, 0), 1)) {
+    pr_warn("FAILED TO ADD CDEV\n");
+    unregister_chrdev_region(MKDEV(dev_major, 0), MINORMASK);
+    return -EBUSY;
+  }
+  device_create(likedbgdev_class, NULL, MKDEV(dev_major, 0), NULL, DEV_NAME);
+  if (IS_ERR(likedbgdev_class)) {
+    pr_warn("FAILED TO CREATE DEVICE\n");
+    return -EBUSY;
+  }
+  return 0;
 }
 
 static void __exit echo_exit(void) {
-    device_destroy(likedbgdev_class, MKDEV(dev_major, 0));
-    class_destroy(likedbgdev_class);
+  device_destroy(likedbgdev_class, MKDEV(dev_major, 0));
+  class_destroy(likedbgdev_class);
 
-    cdev_del(&c_dev);
-    unregister_chrdev_region(MKDEV(dev_major, 0), MINORMASK);
-    pr_info("GOODBYE\n");
+  cdev_del(&c_dev);
+  unregister_chrdev_region(MKDEV(dev_major, 0), MINORMASK);
+  pr_info("GOODBYE\n");
 }
 
 module_init(echo_init);

src/kernel_builder.py

@@ -5,7 +5,6 @@
 import os
 from pathlib import Path
 import subprocess as sp
-from invoke.exceptions import UnexpectedExit
 
 from loguru import logger
 
@@ -85,12 +84,15 @@ def _configure_kernel(self) -> int:
 
     def _get_params(self) -> str:
         params = ""
+        if self.llvm_flag:
+            # TODO: Allow LTO_CLANG_FULL & LTO_CLANG_THIN options once they're not experiment anymore
+            params += "-e LTO_NONE -d LTO_CLANG_FULL -d LTO_CLANG_THIN "
         if self.mode == "syzkaller":
-            params = self.syzkaller_args
+            params += self.syzkaller_args
         elif self.mode == "generic":
-            params = self.generic_args
+            params += self.generic_args
         elif self.mode == "custom":
-            params = self._custom_args()
+            params += self._custom_args()
         if self.extra_args:
             params = self._extra_args(params)
         if params:
@@ -138,27 +140,39 @@ def _wait_for_container(self) -> None:
             else:
                 break
 
-    def _add_modules(self) -> None:
-        for d in Path(self.custom_modules).iterdir():
+    def _add_multiple_mods(self, modules: list[Path]) -> None:
+        for d in modules:
             if not d.is_dir():
                 continue
-            dst = f"{Path(self.kernel_root) / MISC_DRVS_PATH}"
-            sp.run(f"cp -fr {d} {dst}", shell=True)
-            kcfg_mod_path = Path(dst) / d.name / "Kconfig"
-            mod_kcfg_content = kcfg_mod_path.read_text()
-            tmp = "_".join(re.search(r"config .*", mod_kcfg_content)[0].upper().split())
-            ins = f"obj-$({tmp}) += {d.name}/\n"
-            if ins.strip() not in Path(f"{dst}/Makefile").read_text():
-                with open(f"{dst}/Makefile", "a") as g:
-                    g.write(ins)
-            with open(f"{dst}/Kconfig", "r") as f:
-                contents = f.readlines()
-            ins = f"""source "{MISC_DRVS_PATH / d.name / 'Kconfig'}"\n"""
-            if ins not in contents:
-                contents.insert(len(contents) - 1, ins)
-                with open(f"{dst}/Kconfig", "w") as kc:
-                    kc.writelines(contents)
-            logger.debug(f"Added module {d} to the kernel")
+            logger.debug(f"Adding module: {d}")
+            self._add_single_mod(Path(d))
+
+    def _add_single_mod(self, mod: Path) -> None:
+        dst = f"{Path(self.kernel_root) / MISC_DRVS_PATH}"
+        sp.run(f"cp -fr {mod} {dst}", shell=True)
+        kcfg_mod_path = Path(dst) / mod.name / "Kconfig"
+        mod_kcfg_content = kcfg_mod_path.read_text()
+        tmp = "_".join(re.search(r"config .*", mod_kcfg_content)[0].upper().split())
+        ins = f"obj-$({tmp}) += {mod.name}/\n"
+        if ins.strip() not in Path(f"{dst}/Makefile").read_text():
+            with open(f"{dst}/Makefile", "a") as g:
+                g.write(ins)
+        with open(f"{dst}/Kconfig", "r") as f:
+            contents = f.readlines()
+        ins = f"""source "{MISC_DRVS_PATH / mod.name / 'Kconfig'}"\n"""
+        if ins not in contents:
+            contents.insert(len(contents) - 1, ins)
+            with open(f"{dst}/Kconfig", "w") as kc:
+                kc.writelines(contents)
+        logger.debug(f"Added module {mod} to the kernel")
+
+    def _add_modules(self) -> None:
+        mods = list(Path(self.custom_modules).iterdir())
+        logger.error(mods)
+        if all(ele in [x.name for x in mods] for ele in ["Kconfig", "Makefile"]):
+            self._add_single_mod(Path(self.custom_modules))
+        else:
+            self._add_multiple_mods(mods)
 
     def run_container(self) -> None:
         logger.info("Building kernel. This may take a while...")
@@ -192,16 +206,19 @@ def run_container(self) -> None:
         except FileNotFoundError as e:
             logger.error(f"Failed to find file: {e}")
             exit(-1)
-        except UnexpectedExit as e:
+        except Exception as e:
             logger.error(f"A command caused an unexpected exit: {e}")
-            exit(-1)
+            exit(-2)
         else:
             logger.info("Successfully build the kernel")
             if self.arch == "x86_64":
                 cmd = self.make_sudo("ln -s bzImage Image")
                 self.ssh_conn.run(f"cd {self.docker_mnt}/{self.kernel_root}/arch/{self.arch}/boot && {cmd}", echo=True)
         finally:
-            self.stop_container()
+            try:
+                self.stop_container()
+            except AttributeError:
+                pass
 
     def run(self) -> None:
         super().run(check_existing=True)

src/tests/test_kernel_builder.py

@@ -9,7 +9,7 @@
 
 
 USER_INI = Path("configs/user.ini")
-CUSTOM_MODULE = Path("examples/like_dbg_confs/echo_module.ini")
+CUSTOM_MODULE = Path("examples/like_dbg_confs/echo_module_x86.ini")
 
 
 def fetch_cfg_value_from_section_and_key(c: Path, sect: str, key: str) -> str:
@@ -113,6 +113,30 @@ def test_add_modules() -> None:
     assert data[-2] != fst
 
 
+def test_add_module() -> None:
+    p = Path(f"/tmp/{uuid.uuid1().hex}")
+    kb = KernelBuilder(**{"kroot": p})
+    kb.custom_modules = fetch_cfg_value_from_section_and_key(CUSTOM_MODULE, "kernel_builder", "custom_modules")
+    kb.custom_modules += "echo_service"
+    Path(p / MISC_DRVS_PATH).mkdir(parents=True)
+    fst = "This is the 1st line.\n"
+    lst = "This is the last line.\n"
+    q = Path(p / MISC_DRVS_PATH / "Makefile")
+    q.touch()
+    q.write_text(fst)
+    r = Path(p / MISC_DRVS_PATH / "Kconfig")
+    r.touch()
+    r.write_text(f"{fst}\n{lst}")
+    kb._add_modules()
+    with open(q, "r") as f:
+        data = f.readlines()
+    assert data[-1] != fst
+    with open(r, "r") as f:
+        data = f.readlines()
+    assert data[-1] == lst
+    assert data[-2] != fst
+
+
 @patch.object(KernelBuilder, "_run_ssh")
 def test_build_arch_no_args(mock_m) -> None:
     kb = KernelBuilder(**{"kroot": "foo"})
@@ -175,6 +199,18 @@ def test_make_sucess(mock_m) -> None:
     mock_m.assert_called_with("CC=gcc ARCH=x86_64  make -j$(nproc) modules")
 
 
+def throw_with_cause():
+    raise Exception("Failed") from ValueError("That was unexpected.")
+
+
+def test_general_exception() -> None:
+    kb = KernelBuilder(**{"kroot": "foo"})
+    with pytest.raises(SystemExit) as ext:
+        kb.run_container()
+    assert ext.type == SystemExit
+    assert ext.value.code == -2
+
+
 @patch.object(KernelBuilder, "_run_ssh", return_value=1)
 @patch.object(KernelBuilder, "stop_container", return_value=0)
 def test_make_fail(mock_m, mock_k) -> None: