From 9a9e55be72ed8261a5bd5f0f30a768c2a89f170c Mon Sep 17 00:00:00 2001
From: Wirtos_new <Wirtos.new@gmail.com>
Date: Wed, 7 Apr 2021 15:16:06 +0300
Subject: [PATCH] Various fixes and improvements (#21)

* fix potential write to read only memory

* fix infinite loop

* move macros to the source file, make PUT a statement rather than a block
---
 tgcrypto/aes256.c | 6 ++++++
 tgcrypto/aes256.h | 6 ------
 tgcrypto/cbc256.h | 2 +-
 tgcrypto/ctr256.c | 3 ++-
 tgcrypto/ctr256.h | 2 +-
 5 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/tgcrypto/aes256.c b/tgcrypto/aes256.c
index dbaf8a2..6f515cd 100644
--- a/tgcrypto/aes256.c
+++ b/tgcrypto/aes256.c
@@ -20,6 +20,12 @@
 
 #include "aes256.h"
 
+#define LROTL(x) (((x) << 8) | ((x) >> 24))
+#define LROTR(x) (((x) >> 8) | ((x) << 24))
+#define SWAP(x) ((LROTL((x)) & 0x00ff00ff) | (LROTR((x)) & 0xff00ff00))
+#define GET(p) SWAP(*((uint32_t *)(p)))
+#define PUT(ct, st) (*((uint32_t *)(ct)) = SWAP((st)))
+
 static const uint32_t Te0[256] = {
     0xc66363a5, 0xf87c7c84, 0xee777799, 0xf67b7b8d, 0xfff2f20d, 0xd66b6bbd, 0xde6f6fb1, 0x91c5c554,
     0x60303050, 0x02010103, 0xce6767a9, 0x562b2b7d, 0xe7fefe19, 0xb5d7d762, 0x4dababe6, 0xec76769a,
diff --git a/tgcrypto/aes256.h b/tgcrypto/aes256.h
index 2ab7edd..603a745 100644
--- a/tgcrypto/aes256.h
+++ b/tgcrypto/aes256.h
@@ -28,12 +28,6 @@
 #define AES_BLOCK_SIZE 16
 #define EXPANDED_KEY_SIZE 60
 
-#define LROTL(x) (((x) << 8) | ((x) >> 24))
-#define LROTR(x) (((x) >> 8) | ((x) << 24))
-#define SWAP(x) ((LROTL((x)) & 0x00ff00ff) | (LROTR((x)) & 0xff00ff00))
-#define GET(p) SWAP(*((uint32_t *)(p)))
-#define PUT(ct, st) {*((uint32_t *)(ct)) = SWAP((st));}
-
 void aes256_set_encryption_key(const uint8_t key[32], uint32_t expandedKey[60]);
 
 void aes256_set_decryption_key(const uint8_t key[32], uint32_t expandedKey[60]);
diff --git a/tgcrypto/cbc256.h b/tgcrypto/cbc256.h
index 201514e..891cd6f 100644
--- a/tgcrypto/cbc256.h
+++ b/tgcrypto/cbc256.h
@@ -21,6 +21,6 @@
 #ifndef CBC256_H
 #define CBC256_H
 
-uint8_t *cbc256(const uint8_t in[], uint32_t length, const uint8_t key[32], const uint8_t iv[16], uint8_t encrypt);
+uint8_t *cbc256(const uint8_t in[], uint32_t length, const uint8_t key[32], uint8_t iv[16], uint8_t encrypt);
 
 #endif  // CBC256_H
diff --git a/tgcrypto/ctr256.c b/tgcrypto/ctr256.c
index 64d3ca6..d9e4843 100644
--- a/tgcrypto/ctr256.c
+++ b/tgcrypto/ctr256.c
@@ -41,7 +41,8 @@ uint8_t *ctr256(const uint8_t in[], uint32_t length, const uint8_t key[32], uint
                 *state = 0;
 
             if (*state == 0) {
-                for (k = AES_BLOCK_SIZE - 1; k >= 0; --k)
+                k = AES_BLOCK_SIZE;
+                while(k--)
                     if (++iv[k])
                         break;
 
diff --git a/tgcrypto/ctr256.h b/tgcrypto/ctr256.h
index 87d9a3c..ec6f69a 100644
--- a/tgcrypto/ctr256.h
+++ b/tgcrypto/ctr256.h
@@ -21,6 +21,6 @@
 #ifndef CTR256_H
 #define CTR256_H
 
-uint8_t *ctr256(const uint8_t in[], uint32_t length, const uint8_t key[32], const uint8_t iv[16], uint8_t *state);
+uint8_t *ctr256(const uint8_t in[], uint32_t length, const uint8_t key[32], uint8_t iv[16], uint8_t *state);
 
 #endif  // CTR256_H