From 2dc1d597e486844e94987a15763f8bda8f1daf06 Mon Sep 17 00:00:00 2001
From: Christian Brabandt <cb@256bit.org>
Date: Tue, 27 Jun 2023 08:45:31 +0200
Subject: [PATCH] Fix a few more problems with sodium

1) Make sure to sodium_mlock the key, even when using :set key=
Prevents a segfault in Windows, when the key hasn't been mlocked before

2) Prevent segfault in ml_preserve on Windows
---
 src/crypt.c                | 10 ++++++++++
 src/memline.c              | 13 +++++++++++--
 src/optionstr.c            |  4 ++++
 src/proto/crypt.pro        |  2 ++
 src/testdir/test_crypt.vim |  2 --
 5 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/crypt.c b/src/crypt.c
index daa608ae68099..a1d9fa5345b36 100644
--- a/src/crypt.c
+++ b/src/crypt.c
@@ -907,6 +907,16 @@ crypt_append_msg(
     }
 }
 
+    void
+crypt_sodium_lock_key(
+       char_u  *key)
+{
+    if (sodium_init() < 0)
+       return;
+    sodium_mlock(key, STRLEN(key));
+    return;
+}
+
     static int
 crypt_sodium_init_(
     cryptstate_T	*state UNUSED,
diff --git a/src/memline.c b/src/memline.c
index d19b8cdba66ae..fed34f681c8d3 100644
--- a/src/memline.c
+++ b/src/memline.c
@@ -2504,8 +2504,6 @@ ml_sync_all(int check_file, int check_char)
 	    // close the swapfile
 	    mf_close_file(buf, TRUE);
 	    buf->b_p_swf = FALSE;
-	    vim_free(buf->b_p_key);
-	    buf->b_p_key = empty_option;
 	    continue;
 	}
 #endif
@@ -2567,6 +2565,17 @@ ml_preserve(buf_T *buf, int message)
 	    emsg(_(e_cannot_preserve_there_is_no_swap_file));
 	return;
     }
+#ifdef FEAT_CRYPT
+       // Safety Check
+       if (crypt_method_is_sodium(crypt_get_method_nr(buf))
+                                       && *buf->b_p_key != NUL)
+       {
+           // close the swapfile
+           mf_close_file(buf, TRUE);
+           buf->b_p_swf = FALSE;
+           return;
+       }
+#endif
 
     // We only want to stop when interrupted here, not when interrupted
     // before.
diff --git a/src/optionstr.c b/src/optionstr.c
index fa7084c65cf6e..f5cadaeb1e2a0 100644
--- a/src/optionstr.c
+++ b/src/optionstr.c
@@ -1174,6 +1174,10 @@ did_set_cryptkey(optset_T *args)
 		*curbuf->b_p_cm == NUL ? p_cm : curbuf->b_p_cm);
 	changed_internal();
     }
+#ifdef FEAT_SODIUM
+    if (crypt_method_is_sodium(crypt_get_method_nr(curbuf)))
+       crypt_sodium_lock_key(args->os_newval.string);
+#endif
 
     return NULL;
 }
diff --git a/src/proto/crypt.pro b/src/proto/crypt.pro
index d599238593e57..39e8397a75d78 100644
--- a/src/proto/crypt.pro
+++ b/src/proto/crypt.pro
@@ -30,4 +30,6 @@ int crypt_sodium_munlock(void *const addr, const size_t len);
 void crypt_sodium_randombytes_buf(void *const buf, const size_t size);
 int crypt_sodium_init(void);
 uint32_t crypt_sodium_randombytes_random(void);
+void crypt_sodium_lock_key(char_u *key);
+
 /* vim: set ft=c : */
diff --git a/src/testdir/test_crypt.vim b/src/testdir/test_crypt.vim
index d837f2760615d..4af148513ae1d 100644
--- a/src/testdir/test_crypt.vim
+++ b/src/testdir/test_crypt.vim
@@ -394,8 +394,6 @@ endfunc
 
 func Test_crypt_set_key_segfault()
   CheckFeature sodium
-  " This test fails on Windows, why?
-  CheckNotMSWindows
 
   defer delete('Xtest2.txt')
   new Xtest2.txt