Repository: kmaork/hypno (language: python)

Issue:
Use hard links instead of copying
This would save disk space and increase performance.

A possible edge case can occur if the library is installed on a mount, in which case hard links cannot be created directly to the library (this happened to me while working with wsl - the files were located on the mounted windows filesystem).

In that case we should copy the library to some location on the local disk and create hard links from there, but we have to consider possible collisions/races with other instances (and versions) of hypno.
Consider concatenating code to end of SO instead of dlfcn trick



Discussion / hints:

Fix patch:
diff --git a/hypno/__init__.py b/hypno/__init__.py
index 97e4c31..5f7cb0a 100644
--- a/hypno/__init__.py
+++ b/hypno/__init__.py
@@ -1,1 +1,1 @@
-from .hypno import inject_py
+from .hypno import inject_py, CodeTooLongException
diff --git a/hypno/hypno.py b/hypno/hypno.py
index 6d7feab..9297759 100644
--- a/hypno/hypno.py
+++ b/hypno/hypno.py
@@ -1,21 +1,33 @@
 from importlib.util import find_spec
 from typing import AnyStr
-
 from pyinjector import inject
-from os import urandom, unlink
-from shutil import copy
+from tempfile import NamedTemporaryFile
 from pathlib import Path
 
-DATA_DIR = Path('/tmp/hypno')
-DATA_DIR.mkdir(parents=True, exist_ok=True)
+INJECTION_LIB_PATH = Path(find_spec('.injection', __package__).origin)
+MAGIC = b'--- hypno code start ---'
+
 
-injection_lib_path = Path(find_spec('.injection', __package__).origin)
+class CodeTooLongException(Exception):
+    def __init__(self, code: bytes, max_size: int):
+        super().__init__(f'The given python code is too long ({len(code)}). The maximum length is {max_size}.\n'
+                         f'Consider writing the code to a file and executing it with runpy.run_path.\n'
+                         f'Also, please report this on https://github.com/kmaork/hypno/issues/new')
 
 
 def inject_py(pid: int, python_code: AnyStr) -> None:
     if isinstance(python_code, str):
         python_code = python_code.encode()
-    copied_lib = DATA_DIR / f'{urandom(3).hex()}-{python_code.hex()}'
-    copy(injection_lib_path, copied_lib)
-    inject(pid, str(copied_lib))
-    unlink(copied_lib)
+    lib = INJECTION_LIB_PATH.read_bytes()
+    magic_addr = lib.find(MAGIC)
+    code_addr = magic_addr - 1
+    max_size_addr = magic_addr + len(MAGIC)
+    max_size_end_addr = lib.find(b'\0', max_size_addr)
+    max_size = int(lib[max_size_addr:max_size_end_addr])
+    if len(python_code) > max_size:
+        raise CodeTooLongException(python_code, max_size)
+    with NamedTemporaryFile() as temp:
+        temp.write(lib[:code_addr])
+        temp.write(python_code)
+        temp.write(lib[code_addr + len(python_code):])
+        inject(pid, str(temp.name))
diff --git a/hypno/injection.c b/hypno/injection.c
index b874632..392d6e1 100644
--- a/hypno/injection.c
+++ b/hypno/injection.c
@@ -1,42 +1,21 @@
 #define _GNU_SOURCE
 
-#include <dlfcn.h>
 #include <Python.h>
 #include <errno.h>
-#include <string.h>
 
-static void init(void) __attribute__((constructor));
+#define MAX_PYTHON_CODE_SIZE 60500
+#define STR_EXPAND(tok) #tok
+#define STR(tok) STR_EXPAND(tok)
 
-const char *get_lib_name() {
-    Dl_info dl_info;
-    dladdr((void *) get_lib_name, &dl_info);
-    return dl_info.dli_fname;
-}
-
-void inject_py(const char *py_code) {
-    PyGILState_STATE gstate;
-    gstate = PyGILState_Ensure();
-    PyRun_SimpleString(py_code);
-    PyGILState_Release(gstate);
-}
+char PYTHON_CODE[MAX_PYTHON_CODE_SIZE + 1] = "\0--- hypno code start ---" STR(MAX_PYTHON_CODE_SIZE);
 
-char *get_py_code_from_lib_name() {
-    const char *lib_name = get_lib_name();
-    int lib_name_len = strlen(lib_name);
-    char *hex = strrchr(lib_name, '-') + 1;
-    int py_code_len = (lib_name + lib_name_len - hex) / 2;
-    char *py_code = malloc(py_code_len + 1);
-    for (int py_code_i = 0; py_code_i < py_code_len; ++py_code_i) {
-        hex += sscanf(hex, "%2hhx", &py_code[py_code_i]) * 2;
+__attribute__((constructor))
+static void init(void) {
+    if (PYTHON_CODE[0]) {
+        int saved_errno = errno;
+        PyGILState_STATE gstate = PyGILState_Ensure();
+        PyRun_SimpleString(PYTHON_CODE);
+        PyGILState_Release(gstate);
+        errno = saved_errno;
     }
-    py_code[py_code_len] = '\0';
-    return py_code;
 }
-
-static void init(void) {
-    int saved_errno = errno;
-    char *py_code = get_py_code_from_lib_name();
-    inject_py(py_code);
-    free(py_code);
-    errno = saved_errno;
-}
\ No newline at end of file


Test patch:
diff --git a/tests/test_hypno.py b/tests/test_hypno.py
index 90765c2..336e3cb 100644
--- a/tests/test_hypno.py
+++ b/tests/test_hypno.py
@@ -1,9 +1,11 @@
 import sys
 from subprocess import Popen, PIPE, TimeoutExpired
 from pathlib import Path
-from hypno import inject_py
+import pytest
+from hypno import inject_py, CodeTooLongException
 from time import sleep
 
+
 WHILE_TRUE_SCRIPT = Path(__file__).parent.resolve() / 'while_true.py'
 PROCESS_WAIT_TIMEOUT = 1
 WAIT_FOR_PYTHON_SECONDS = 0.5
@@ -22,3 +24,9 @@ def test_hypno():
             process.kill()
             raise
         assert process.stdout.read() == data
+
+
+def test_hypno_with_too_long_code():
+    code = b'^' * 100000
+    with pytest.raises(CodeTooLongException):
+        inject_py(-1, code)