fine tuning; robust error handling; more debug messages

Signed-off-by: segfault <segfault@secmail.pro>

3 files changed, 42 insertions, 13 deletions

diff --git a/example.c b/example.c
index 4a7b3a7..48d02cc 100644
--- a/example.c
+++ b/example.c
@@ -21,7 +21,17 @@ static void crypted_fn2(void)
 {
     CRYPT_PROLOGUE(crypted_fn2);
     printf("Another decrypted fn..\n");
+    printf(
+#ifdef _DEBUG
+        "CRYPT_PROLOGUE crpyt_func returned: %s\n",
+        crypt_strs[ CRYPT_RETVAL() ]
+#else
+        "CRYPT_PROLOGUE crpyt_func returned: %d\n",
+        CRYPT_RETVAL()
+#endif
+    );
     CRYPT_EPILOGUE(crypted_fn2);
+    printf("This part stays unencrypted all the time.\n");
 }
 
 CRYPT_FNDEF(crypted_fn3, void *arg0, int arg1, const char *arg2)
diff --git a/funccrypt.c b/funccrypt.c
index 60c3edd..873aa40 100644
--- a/funccrypt.c
+++ b/funccrypt.c
@@ -2,15 +2,23 @@
 #include <unistd.h>
 #ifdef __linux__
 #include <sys/mman.h> /* PROT_* */
-#else
+#elif __WIN32__
 #include <windows.h>
+#else
+#error "Unsupported OS. Only __linux__ and __WIN32__ are supported."
 #endif
 
 #include "funccrypt.h"
 
 #ifdef _DEBUG
 const char *crypt_strs[] = {
-    "ERROR", "PROLOGUE", "EPILOGUE", "CHECK", "OK"
+    "ERROR",
+    "ERROR MEMORY",
+    "PROLOGUE",
+    "EPILOGUE",
+    "CHECK",
+    "OK ENCRYPTED",
+    "OK DECRYPTED"
 };
 
 
@@ -95,6 +103,7 @@ crypt_return crypt_func(void *fn_start)
 #endif
             {
                 if (hdr->crpyted == 0x00) {
+                    /* function not encrypted; encrypt it */
                     hdr->crpyted = 0xFF;
 #ifdef __linux__
                     hdr->key  = (uint64_t) rand() << 32;
@@ -105,17 +114,23 @@ crypt_return crypt_func(void *fn_start)
                     hdr->key |= (uint64_t) rand() << 16;
                     hdr->key |= (uint64_t) rand();
 #endif
+                    cret = CRET_OK_ENC;
+                } else {
+                    /* function encrypted, decrypt it */
+                    cret = CRET_OK_DEC;
                 }
+
+                /* (en|de)cryption */
                 for (i = 0; i < crypt_size / 0x8; ++i) {
                     hdr->func_body[i] ^= hdr->key;
                 }
 #ifdef __linux__
-                if (!mprotect(mbuf, sysconf(_SC_PAGESIZE), PROT_READ|PROT_EXEC))
+                if (mprotect(mbuf, sysconf(_SC_PAGESIZE), PROT_READ|PROT_EXEC))
 #else
-                if (VirtualProtect(mbuf, crypt_size, old_prot, &old_prot))
+                if (!VirtualProtect(mbuf, crypt_size, old_prot, &old_prot))
 #endif
-                    cret = CRET_OK;
-            }
+                    cret = CRET_ERROR_MEM;
+            } else cret = CRET_ERROR_MEM;
         }
     }
 
diff --git a/funccrypt.h b/funccrypt.h
index bdac0f3..56d1930 100644
--- a/funccrypt.h
+++ b/funccrypt.h
@@ -21,21 +21,24 @@ typedef struct crypt_header {
 } GCC_PACKED crypt_header;
 
 typedef enum crypt_return {
-    CRET_ERROR    /* Neither prologue marker nor epilogue marker found. */,
-    CRET_PROLOGUE /* prologue marker found */,
-    CRET_EPILOGUE /* epilogue marker found */,
-    CRET_CHECK    /* all pre (en|de)cryption checks successful */,
-    CRET_OK       /* (en|de)cryption succeeded */
+    CRET_ERROR,     /* Neither prologue marker nor epilogue marker found. */
+    CRET_ERROR_MEM, /* set memory page protection failed */
+    CRET_PROLOGUE,  /* prologue marker found */
+    CRET_EPILOGUE,  /* epilogue marker found */
+    CRET_CHECK,     /* all pre (en|de)cryption checks successful */
+    CRET_OK_ENC,    /* encryption succeeded */
+    CRET_OK_DEC,    /* decryption succeeded */
 } crypt_return;
 
 #define CRYPT_FUNC_MAXSIZ 0x100
 #define CRYPT_FUNC(fn) \
     crypt_func((void *)fn)
+#define CRYPT_RETVAL() __cr
 #define CRYPT_PROLOGUE(fn) \
     crypt_return __cr; \
     { \
         __cr = CRYPT_FUNC(fn); \
-        if (__cr != CRET_OK) \
+        if (__cr != CRET_OK_DEC) \
             asm volatile goto("jmp %l0          \n" \
                 : : : : cr_epilogue); \
         asm volatile goto("jmp %l0          \n" \
@@ -65,7 +68,8 @@ typedef enum crypt_return {
         ); \
     } \
     cr_epilogue: \
-    CRYPT_FUNC(fn);
+    if (CRYPT_RETVAL() == CRET_OK_DEC) \
+        CRYPT_FUNC(fn);
 
 #define CRYPT_FNDEF(name, ...) \
     void name( __VA_ARGS__ ) { \