4 files changed, 23 insertions, 21 deletions

diff --git a/subhook.c b/subhook.c
index 7d9ea6b..d69bc0d 100644
--- a/subhook.c
+++ b/subhook.c
@@ -28,14 +28,6 @@
 #include "subhook.h"
 #include "subhook_private.h"
 
-SUBHOOK_EXPORT void *SUBHOOK_API subhook_get_src(subhook_t hook) {
-	return hook->src;
-}
-
-SUBHOOK_EXPORT void *SUBHOOK_API subhook_get_dst(subhook_t hook) {
-	return hook->dst;
-}
-
 SUBHOOK_EXPORT void SUBHOOK_API subhook_set_src(subhook_t hook, void *src) {
 	hook->src = src;
 }
@@ -44,6 +36,14 @@ SUBHOOK_EXPORT void SUBHOOK_API subhook_set_dst(subhook_t hook, void *dst) {
 	hook->dst = dst;
 }
 
+SUBHOOK_EXPORT void *SUBHOOK_API subhook_get_src(subhook_t hook) {
+	return hook->src;
+}
+
+SUBHOOK_EXPORT void *SUBHOOK_API subhook_get_dst(subhook_t hook) {
+	return hook->dst;
+}
+
 SUBHOOK_EXPORT int SUBHOOK_API subhook_is_installed(subhook_t hook) {
 	return hook->installed;
 }
diff --git a/subhook.h b/subhook.h
index 3797a78..757218c 100644
--- a/subhook.h
+++ b/subhook.h
@@ -91,13 +91,15 @@ typedef struct subhook *subhook_t;
 SUBHOOK_EXPORT subhook_t SUBHOOK_API subhook_new();
 SUBHOOK_EXPORT void SUBHOOK_API subhook_free(subhook_t hook);
 
+/* Set hook source and destination.
+ * Do NOT call these functions after subhook_install() !!
+ */
 SUBHOOK_EXPORT void SUBHOOK_API subhook_set_src(subhook_t hook, void *src);
-SUBHOOK_EXPORT void *SUBHOOK_API subhook_get_src(subhook_t hook);
-
 SUBHOOK_EXPORT void SUBHOOK_API subhook_set_dst(subhook_t hook, void *dst);
+
+SUBHOOK_EXPORT void *SUBHOOK_API subhook_get_src(subhook_t hook);
 SUBHOOK_EXPORT void *SUBHOOK_API subhook_get_dst(subhook_t hook);
 
-/* These return 0 on failure and 1 on success. */
 SUBHOOK_EXPORT int SUBHOOK_API subhook_install(subhook_t hook);
 SUBHOOK_EXPORT int SUBHOOK_API subhook_remove(subhook_t hook);
 
diff --git a/subhook_private.h b/subhook_private.h
index c3de967..3c6c461 100644
--- a/subhook_private.h
+++ b/subhook_private.h
@@ -27,7 +27,8 @@
 #define SUBHOOK_PRIVATE_H
 
 struct subhook {
-	int installed: 1;
+	int unlocked  : 1;
+	int installed : 1;
 	void *src;
 	void *dst;
 };
diff --git a/subhook_x86.c b/subhook_x86.c
index 4f03c7b..18ba6c5 100644
--- a/subhook_x86.c
+++ b/subhook_x86.c
@@ -66,22 +66,21 @@ SUBHOOK_EXPORT void SUBHOOK_API subhook_free(subhook_t hook) {
 }
 
 SUBHOOK_EXPORT int SUBHOOK_API subhook_install(subhook_t hook) {
-	void *src;
-	void *dst;
 	intptr_t offset;
 
 	if (hook->installed)
 		return -EINVAL;
 
-	src = hook->src;
-	dst = hook->dst;
+	if (!hook->unlocked) {
+		subhook_unprotect(hook->src, sizeof(jmp_instr));
+		hook->unlocked = 1;
+	}
 
-	subhook_unprotect(src, sizeof(jmp_instr));
-	memcpy(((struct subhook_x86 *)hook)->code, src, sizeof(jmp_instr));
-	memcpy(src, &jmp_instr, sizeof(jmp_instr));
+	memcpy(((struct subhook_x86 *)hook)->code, hook->src, sizeof(jmp_instr));
+	memcpy(hook->src, &jmp_instr, sizeof(jmp_instr));
 
-	offset = (intptr_t)dst - ((intptr_t)src + sizeof(jmp_instr));
-	memcpy((void *)((intptr_t)src + sizeof(jmp_opcode)), &offset,
+	offset = (intptr_t)hook->dst - ((intptr_t)hook->src + sizeof(jmp_instr));
+	memcpy((void *)((intptr_t)hook->src + sizeof(jmp_opcode)), &offset,
 	       sizeof(jmp_instr) - sizeof(jmp_opcode));
 
 	hook->installed = 1;