Crypt the text in the undo file if the file itself is crypted.

5 files changed, 251 insertions, 44 deletions

diff --git a/src/fileio.c b/src/fileio.c
index 101c804de..ff21783ff 100644
--- a/src/fileio.c
+++ b/src/fileio.c
@@ -2827,7 +2827,7 @@ check_marks_read()
 }
 #endif
 
-#ifdef FEAT_CRYPT
+#if defined(FEAT_CRYPT) || defined(PROTO)
 /*
  * Get the crypt method used for a file from "ptr[len]", the magic text at the
  * start of the file.
@@ -2926,7 +2926,129 @@ check_for_cryptkey(cryptkey, ptr, sizep, filesizep, newfile, did_ask)
 
     return cryptkey;
 }
-#endif
+
+/*
+ * Check for magic number used for encryption.  Applies to the current buffer.
+ * If found and decryption is possible returns OK;
+ */
+    int
+prepare_crypt_read(fp)
+    FILE	*fp;
+{
+    int		method;
+    char_u	buffer[CRYPT_MAGIC_LEN + CRYPT_SEED_LEN_MAX + 2];
+
+    if (fread(buffer, CRYPT_MAGIC_LEN, 1, fp) != 1)
+	return FAIL;
+    method = get_crypt_method((char *)buffer,
+					CRYPT_MAGIC_LEN + CRYPT_SEED_LEN_MAX);
+    if (method < 0 || method != curbuf->b_p_cm)
+	return FAIL;
+
+    if (method == 0)
+	crypt_init_keys(curbuf->b_p_key);
+    else
+    {
+	int seed_len = crypt_seed_len[method];
+
+	if (fread(buffer, seed_len, 1, fp) != 1)
+	    return FAIL;
+	bf_key_init(curbuf->b_p_key);
+	bf_ofb_init(buffer, seed_len);
+    }
+    return OK;
+}
+
+/*
+ * Prepare for writing encrypted bytes for buffer "buf".
+ * Returns a pointer to an allocated header of length "*lenp".
+ */
+    char_u *
+prepare_crypt_write(buf, lenp)
+    buf_T *buf;
+    int   *lenp;
+{
+    char_u  *header;
+    int	    seed_len = crypt_seed_len[buf->b_p_cm];
+
+    header = alloc_clear(CRYPT_MAGIC_LEN + CRYPT_SEED_LEN_MAX + 2);
+    if (header != NULL)
+    {
+	use_crypt_method = buf->b_p_cm;  /* select pkzip or blowfish */
+	vim_strncpy(header, (char_u *)crypt_magic[use_crypt_method],
+							     CRYPT_MAGIC_LEN);
+	if (buf->b_p_cm == 0)
+	    crypt_init_keys(buf->b_p_key);
+	else
+	{
+	    /* Using blowfish, add seed. */
+	    sha2_seed(header + CRYPT_MAGIC_LEN, seed_len); /* create iv */
+	    bf_ofb_init(header + CRYPT_MAGIC_LEN, seed_len);
+	    bf_key_init(buf->b_p_key);
+	}
+    }
+    *lenp = CRYPT_MAGIC_LEN + seed_len;
+    return header;
+}
+
+/*
+ * Like fwrite() but crypt the bytes when 'key' is set.
+ * Returns 1 if successful.
+ */
+    size_t
+fwrite_crypt(buf, ptr, len, fp)
+    buf_T	*buf;
+    char_u	*ptr;
+    size_t	len;
+    FILE	*fp;
+{
+    char_u  *copy;
+    char_u  small_buf[100];
+    int	    ztemp, t;
+    size_t  i;
+
+    if (*buf->b_p_key == NUL)
+	return fwrite(ptr, len, (size_t)1, fp);
+    if (len < 100)
+	copy = small_buf;  /* no malloc()/free() for short strings */
+    else
+    {
+	copy = lalloc(len, FALSE);
+	if (copy == NULL)
+	    return 0;
+    }
+    for (i = 0; i < len; ++i)
+    {
+	ztemp = ptr[i];
+	copy[i] = ZENCODE(ztemp, t);
+    }
+    i = fwrite(copy, len, (size_t)1, fp);
+    if (copy != small_buf)
+	vim_free(copy);
+    return i;
+}
+
+/*
+ * Read a string of length "len" from "fd".
+ * When 'key' is set decrypt the bytes.
+ */
+    char_u *
+read_string_decrypt(buf, fd, len)
+    buf_T   *buf;
+    FILE    *fd;
+    int	    len;
+{
+    char_u  *ptr;
+    char_u  *p;
+
+    ptr = read_string(fd, len);
+    if (ptr != NULL || *buf->b_p_key != NUL)
+	for (p = ptr; p < ptr + len; ++p)
+	    ZDECODE(*p);
+    return ptr;
+}
+
+#endif  /* FEAT_CRYPT */
 
 #ifdef UNIX
     static void
@@ -4323,34 +4445,25 @@ restore_backup:
 #ifdef FEAT_CRYPT
     if (*buf->b_p_key && !filtering)
     {
-	char_u header[CRYPT_MAGIC_LEN + CRYPT_SEED_LEN_MAX + 2];
-	int seed_len = crypt_seed_len[buf->b_p_cm];
-
-	use_crypt_method = buf->b_p_cm;  /* select pkzip or blowfish */
+	char_u *header;
+	int    header_len;
 
-	vim_memset(header, 0, sizeof(header));
-	vim_strncpy(header, (char_u *)crypt_magic[use_crypt_method],
-							     CRYPT_MAGIC_LEN);
-
-	if (buf->b_p_cm == 0)
-	    crypt_init_keys(buf->b_p_key);
+	header = prepare_crypt_write(buf, &header_len);
+	if (header == NULL)
+	    end = 0;
 	else
 	{
-	    /* Using blowfish, add seed. */
-	    sha2_seed(header + CRYPT_MAGIC_LEN, seed_len); /* create iv */
-	    bf_ofb_init(header + CRYPT_MAGIC_LEN, seed_len);
-	    bf_key_init(buf->b_p_key);
+	    /* Write magic number, so that Vim knows that this file is
+	     * encrypted when reading it again.  This also undergoes utf-8 to
+	     * ucs-2/4 conversion when needed. */
+	    write_info.bw_buf = header;
+	    write_info.bw_len = header_len;
+	    write_info.bw_flags = FIO_NOCONVERT;
+	    if (buf_write_bytes(&write_info) == FAIL)
+		end = 0;
+	    wb_flags |= FIO_ENCRYPTED;
+	    vim_free(header);
 	}
-
-	/* Write magic number, so that Vim knows that this file is
-	 * encrypted when reading it again.  This also undergoes utf-8 to
-	 * ucs-2/4 conversion when needed. */
-	write_info.bw_buf = (char_u *)header;
-	write_info.bw_len = CRYPT_MAGIC_LEN + seed_len;
-	write_info.bw_flags = FIO_NOCONVERT;
-	if (buf_write_bytes(&write_info) == FAIL)
-	    end = 0;
-	wb_flags |= FIO_ENCRYPTED;
     }
 #endif
 
@@ -5558,7 +5671,7 @@ buf_write_bytes(ip)
 
 	for (i = 0; i < len; i++)
 	{
-	    ztemp  = buf[i];
+	    ztemp = buf[i];
 	    buf[i] = ZENCODE(ztemp, t);
 	}
     }
diff --git a/src/proto/fileio.pro b/src/proto/fileio.pro
index e5a1b284a..32e304552 100644
--- a/src/proto/fileio.pro
+++ b/src/proto/fileio.pro
@@ -2,6 +2,10 @@
 void filemess __ARGS((buf_T *buf, char_u *name, char_u *s, int attr));
 int readfile __ARGS((char_u *fname, char_u *sfname, linenr_T from, linenr_T lines_to_skip, linenr_T lines_to_read, exarg_T *eap, int flags));
 int prep_exarg __ARGS((exarg_T *eap, buf_T *buf));
+int prepare_crypt_read __ARGS((FILE *fp));
+char_u *prepare_crypt_write __ARGS((buf_T *buf, int *lenp));
+size_t fwrite_crypt __ARGS((buf_T *buf, char_u *ptr, size_t len, FILE *fp));
+char_u *read_string_decrypt __ARGS((buf_T *buf, FILE *fd, int len));
 int check_file_readonly __ARGS((char_u *fname, int perm));
 int buf_write __ARGS((buf_T *buf, char_u *fname, char_u *sfname, linenr_T start, linenr_T end, exarg_T *eap, int append, int forceit, int reset_changed, int filtering));
 void msg_add_fname __ARGS((buf_T *buf, char_u *fname));
diff --git a/src/testdir/test72.in b/src/testdir/test72.in
index a5674b527..b6c6bfab3 100644
--- a/src/testdir/test72.in
+++ b/src/testdir/test72.in
@@ -49,6 +49,55 @@ dd:set ul=100
 :e Xtestfile
 uuu:w >>test.out
 :"
+:" And now with encryption, cryptmethod=0
+:e! Xtestfile
+:set undofile cm=0
+ggdG
+imonday
+tuesday
+wednesday
+thursday
+friday:set ul=100
+kkkdd:set ul=100
+dd:set ul=100
+dd:set ul=100
+:X
+foobar
+foobar
+:w!
+:bwipe!
+:e Xtestfile
+foobar
+:set key=
+uu:w >>test.out
+:"
+:"
+:" With encryption, cryptmethod=1
+:e! Xtestfile
+:set undofile cm=1
+ggdG
+ijan
+feb
+mar
+apr
+jun:set ul=100
+kk0ifoo :set ul=100
+dd:set ul=100
+ibar :set ul=100
+:X
+foobar
+foobar
+:w!
+:bwipe!
+:e Xtestfile
+foobar
+:set key=
+/bar
+:.w >>test.out
+u:.w >>test.out
+u:.w >>test.out
+u:.w >>test.out
+:"
 :" Rename the undo file so that it gets cleaned up.
 :call rename(".Xtestfile.un~", "Xtestundo")
 :qa!
diff --git a/src/testdir/test72.ok b/src/testdir/test72.ok
index 66e7250df..44210c942 100644
--- a/src/testdir/test72.ok
+++ b/src/testdir/test72.ok
@@ -7,3 +7,11 @@ seven
 eight
 nine
 ten
+monday
+wednesday
+thursday
+friday
+bar apr
+apr
+foo mar
+mar
diff --git a/src/undo.c b/src/undo.c
index 35b66c3db..86a4d1753 100644
--- a/src/undo.c
+++ b/src/undo.c
@@ -103,9 +103,9 @@ static void u_freeentry __ARGS((u_entry_T *, long));
 static void corruption_error __ARGS((char *msg, char_u *file_name));
 static void u_free_uhp __ARGS((u_header_T *uhp));
 static int serialize_header __ARGS((FILE *fp, buf_T *buf, char_u *hash));
-static int serialize_uhp __ARGS((FILE *fp, u_header_T *uhp));
+static int serialize_uhp __ARGS((FILE *fp, buf_T *buf, u_header_T *uhp));
 static u_header_T *unserialize_uhp __ARGS((FILE *fp, char_u *file_name));
-static int serialize_uep __ARGS((u_entry_T *uep, FILE *fp));
+static int serialize_uep __ARGS((FILE *fp, buf_T *buf, u_entry_T *uep));
 static u_entry_T *unserialize_uep __ARGS((FILE *fp, int *error, char_u *file_name));
 static void serialize_pos __ARGS((pos_T pos, FILE *fp));
 static void unserialize_pos __ARGS((pos_T *pos, FILE *fp));
@@ -670,6 +670,7 @@ nomem:
 # define UF_ENTRY_MAGIC		0xf518	/* magic at start of entry */
 # define UF_ENTRY_END_MAGIC	0x3581	/* magic after last entry */
 # define UF_VERSION		1	/* 2-byte undofile version number */
+# define UF_VERSION_CRYPT	0x8001	/* idem, encrypted */
 
 static char_u e_not_open[] = N_("E828: Cannot open undo file for writing: %s");
 
@@ -811,7 +812,28 @@ serialize_header(fp, buf, hash)
     /* Start writing, first the magic marker and undo info version. */
     if (fwrite(UF_START_MAGIC, (size_t)UF_START_MAGIC_LEN, (size_t)1, fp) != 1)
 	return FAIL;
-    put_bytes(fp, (long_u)UF_VERSION, 2);
+
+    /* If the buffer is encrypted then all text bytes following will be
+     * encrypted.  Numbers and other info is not crypted. */
+#ifdef FEAT_CRYPT
+    if (*buf->b_p_key)
+    {
+	char_u *header;
+	int    header_len;
+
+	put_bytes(fp, (long_u)UF_VERSION_CRYPT, 2);
+	header = prepare_crypt_write(buf, &header_len);
+	if (header == NULL)
+	    return FAIL;
+	len = fwrite(header, (size_t)header_len, (size_t)1, fp);
+	vim_free(header);
+	if (len != 1)
+	    return FAIL;
+    }
+    else
+#endif
+	put_bytes(fp, (long_u)UF_VERSION, 2);
+
 
     /* Write a hash of the buffer text, so that we can verify it is still the
      * same when reading the buffer text. */
@@ -822,7 +844,7 @@ serialize_header(fp, buf, hash)
     put_bytes(fp, (long_u)buf->b_ml.ml_line_count, 4);
     len = buf->b_u_line_ptr != NULL ? (int)STRLEN(buf->b_u_line_ptr) : 0;
     put_bytes(fp, (long_u)len, 4);
-    if (len > 0 && fwrite(buf->b_u_line_ptr, (size_t)len, (size_t)1, fp) != 1)
+    if (len > 0 && fwrite_crypt(buf, buf->b_u_line_ptr, (size_t)len, fp) != 1)
 	return FAIL;
     put_bytes(fp, (long_u)buf->b_u_line_lnum, 4);
     put_bytes(fp, (long_u)buf->b_u_line_colnr, 4);
@@ -841,8 +863,9 @@ serialize_header(fp, buf, hash)
 }
 
     static int
-serialize_uhp(fp, uhp)
+serialize_uhp(fp, buf, uhp)
     FILE	*fp;
+    buf_T	*buf;
     u_header_T	*uhp;
 {
     int		i;
@@ -882,7 +905,7 @@ serialize_uhp(fp, uhp)
     for (uep = uhp->uh_entry; uep != NULL; uep = uep->ue_next)
     {
 	put_bytes(fp, (long_u)UF_ENTRY_MAGIC, 2);
-	if (serialize_uep(uep, fp) == FAIL)
+	if (serialize_uep(fp, buf, uep) == FAIL)
 	    return FAIL;
     }
     put_bytes(fp, (long_u)UF_ENTRY_END_MAGIC, 2);
@@ -971,9 +994,10 @@ unserialize_uhp(fp, file_name)
  * Serialize "uep" to "fp".
  */
     static int
-serialize_uep(uep, fp)
-    u_entry_T	*uep;
+serialize_uep(fp, buf, uep)
     FILE	*fp;
+    buf_T	*buf;
+    u_entry_T	*uep;
 {
     int		i;
     size_t	len;
@@ -987,7 +1011,7 @@ serialize_uep(uep, fp)
 	len = STRLEN(uep->ue_array[i]);
         if (put_bytes(fp, (long_u)len, 4) == FAIL)
 	    return FAIL;
-        if (len > 0 && fwrite(uep->ue_array[i], len, (size_t)1, fp) != 1)
+        if (len > 0 && fwrite_crypt(buf, uep->ue_array[i], len, fp) != 1)
 	    return FAIL;
     }
     return OK;
@@ -1034,7 +1058,7 @@ unserialize_uep(fp, error, file_name)
     {
 	line_len = get4c(fp);
 	if (line_len >= 0)
-	    line = read_string(fp, line_len);
+	    line = read_string_decrypt(curbuf, fp, line_len);
 	else
 	{
 	    line = NULL;
@@ -1333,7 +1357,7 @@ u_write_undo(name, forceit, buf, hash)
 #ifdef U_DEBUG
 	    ++headers_written;
 #endif
-	    if (serialize_uhp(fp, uhp) == FAIL)
+	    if (serialize_uhp(fp, buf, uhp) == FAIL)
 		goto write_error;
         }
 
@@ -1478,7 +1502,20 @@ u_read_undo(name, hash, orig_name)
         goto error;
     }
     version = get2c(fp);
-    if (version != UF_VERSION)
+    if (version == UF_VERSION_CRYPT)
+    {
+#ifdef FEAT_CRYPT
+	if (prepare_crypt_read(fp) == FAIL)
+	{
+	    EMSG2(_("E826: Undo file decryption failed: %s"), file_name);
+	    goto error;
+	}
+#else
+        EMSG2(_("E826: Undo file is encrypted: %s"), file_name);
+        goto error;
+#endif
+    }
+    else if (version != UF_VERSION)
     {
         EMSG2(_("E824: Incompatible undo file: %s"), file_name);
         goto error;
@@ -1510,7 +1547,7 @@ u_read_undo(name, hash, orig_name)
     if (str_len < 0)
         goto error;
     if (str_len > 0)
-	line_ptr = read_string(fp, str_len);
+	line_ptr = read_string_decrypt(curbuf, fp, str_len);
     line_lnum = (linenr_T)get4c(fp);
     line_colnr = (colnr_T)get4c(fp);
     if (line_lnum < 0 || line_colnr < 0)
@@ -1634,10 +1671,6 @@ u_read_undo(name, hash, orig_name)
     curbuf->b_u_oldhead = old_idx < 0 ? NULL : uhp_table[old_idx];
     curbuf->b_u_newhead = new_idx < 0 ? NULL : uhp_table[new_idx];
     curbuf->b_u_curhead = cur_idx < 0 ? NULL : uhp_table[cur_idx];
-#ifdef U_DEBUG
-    if (curbuf->b_u_curhead != NULL)
-	corruption_error("curhead not NULL", file_name);
-#endif
     curbuf->b_u_line_ptr = line_ptr;
     curbuf->b_u_line_lnum = line_lnum;
     curbuf->b_u_line_colnr = line_colnr;