Fix KDH decryption.

Signed-off-by: yzrh <yzrh@noema.org>

1 files changed, 19 insertions, 17 deletions

diff --git a/src/cnki_kdh.c b/src/cnki_kdh.c
index b13434d..af453a7 100644
--- a/src/cnki_kdh.c
+++ b/src/cnki_kdh.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2021, yzrh <yzrh@noema.org>
+ * Copyright (c) 2020-2023, yzrh <yzrh@noema.org>
  *
  * SPDX-License-Identifier: Apache-2.0
  */
@@ -15,16 +15,18 @@ cnki_kdh(cnki_t **param)
 	if ((*param)->stat > 0)
 		printf("Begin 'KDH' decryption\n");
 
-	fseek((*param)->fp_i, 0, SEEK_END);
-
-	long size = ftell((*param)->fp_i);
+	long cur = ADDRESS_KDH_BODY;
+	long end;
 
-	fseek((*param)->fp_i, ADDRESS_KDH_BODY, SEEK_SET);
+	fseek((*param)->fp_i, 0, SEEK_END);
+	end = ftell((*param)->fp_i);
+	fseek((*param)->fp_i, cur, SEEK_SET);
 
 	const char key[] = KEY_KDH;
 	const int key_len = KEY_KDH_LENGTH;
 	long key_cur = 0;
 
+	int buf_size;
 	char buf[(*param)->size_buf];
 
 	FILE *tmp = tmpfile();
@@ -33,32 +35,32 @@ cnki_kdh(cnki_t **param)
 		return 1;
 
 	for (;;) {
-		fread(buf, (*param)->size_buf, 1, (*param)->fp_i);
+		if (cur + (*param)->size_buf < end)
+			buf_size = (*param)->size_buf;
+		else
+			buf_size = end - cur;
+
+		fread(buf, buf_size, 1, (*param)->fp_i);
 
-		for (int i = 0; i < (*param)->size_buf; i++) {
-			buf[i] ^= key[key_cur % key_len];
-			key_cur++;
-		}
+		for (int i = 0; i < buf_size; i++)
+			buf[i] ^= key[key_cur++ % key_len];
 
-		fwrite(buf, (*param)->size_buf, 1, tmp);
+		fwrite(buf, buf_size, 1, tmp);
 
-		if (ftell((*param)->fp_i) == size)
+		if ((cur = ftell((*param)->fp_i)) >= end)
 			break;
 	}
 
 	if ((*param)->stat > 0)
 		printf("Decrypted %ld byte(s)\n", ftell(tmp));
 
-	fseek(tmp, 0, SEEK_SET);
+	fclose((*param)->fp_i);
 
-	FILE *orig = (*param)->fp_i;
+	fseek(tmp, 0, SEEK_SET);
 	(*param)->fp_i = tmp;
 
 	cnki_pdf(param);
 
-	(*param)->fp_i = orig;
-	fclose(tmp);
-
 	if ((*param)->stat > 0)
 		printf("Conversion ended\n");