swaylock: Handle possible fail on password realloc

1 files changed, 15 insertions, 2 deletions

diff --git a/swaylock/main.c b/swaylock/main.c
index 57fe0321..1885b511 100644
--- a/swaylock/main.c
+++ b/swaylock/main.c
@@ -113,6 +113,8 @@ bool verify_password() {
 
 void notify_key(enum wl_keyboard_key_state state, xkb_keysym_t sym, uint32_t code, uint32_t codepoint) {
         int redraw_screen = 0;
+        char *password_realloc;
+
         if (state == WL_KEYBOARD_KEY_STATE_PRESSED) {
                 switch (sym) {
                 case XKB_KEY_Return:
@@ -129,6 +131,7 @@ void notify_key(enum wl_keyboard_key_state state, xkb_keysym_t sym, uint32_t cod
                         redraw_screen = 1;
 
                         password_size = 1024;
+                        free(password);
                         password = malloc(password_size);
                         password[0] = '\0';
                         break;
@@ -149,7 +152,17 @@ void notify_key(enum wl_keyboard_key_state state, xkb_keysym_t sym, uint32_t cod
                                 int i = strlen(password);
                                 if (i + 1 == password_size) {
                                         password_size += 1024;
-                                        password = realloc(password, password_size);
+                                        password_realloc = realloc(password, password_size);
+                                        // reset password if realloc fails.
+                                        if (password_realloc == NULL) {
+                                                password_size = 1024;
+                                                free(password);
+                                                password = malloc(password_size);
+                                                password[0] = '\0';
+                                                break;
+                                        } else {
+                                                password = password_realloc;
+                                        }
                                 }
                                 password[i] = (char)codepoint;
                                 password[i + 1] = '\0';
@@ -305,7 +318,7 @@ int main(int argc, char **argv) {
                         break;
                 }
                 switch (c) {
-                case 'c': 
+                case 'c':
                 {
                         int colorlen = strlen(optarg);
                         if (colorlen < 6 || colorlen == 7 || colorlen > 8) {