From: Josef Bacik <josef@toxicpanda.com>

The fscrypt_extent_info will be tied to the extent_map lifetime, so it
will be created when we create the IO em, or it'll already exist in the
NOCOW case.  Use this fscrypt_info when creating the ordered extent to
make sure everything is passed through properly.

Signed-off-by: Josef Bacik <josef@toxicpanda.com>
Signed-off-by: Daniel Vacek <neelx@suse.com>
---

v5: https://lore.kernel.org/linux-btrfs/33e650f3e91ed0318211301beb27fa613382f28e.1706116485.git.josef@toxicpanda.com/
 * Splitted the dio-related hunks from inode.c to direct-io.c as upstream
   refactored in the meantime.
 * Pass fscrypt_info using the file_extent structure to follow the
   upstream cleanup.
---
 fs/btrfs/direct-io.c |  4 +++-
 fs/btrfs/inode.c     | 36 ++++++++++++++++++++++++++++++------
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/fs/btrfs/direct-io.c b/fs/btrfs/direct-io.c
index f3efc451d9a5..c95b4e768043 100644
--- a/fs/btrfs/direct-io.c
+++ b/fs/btrfs/direct-io.c
@@ -140,7 +140,7 @@ static int lock_extent_direct(struct inode *inode, u64 lockstart, u64 lockend,
 static struct extent_map *btrfs_create_dio_extent(struct btrfs_inode *inode,
 						  struct btrfs_dio_data *dio_data,
 						  const u64 start,
-						  const struct btrfs_file_extent *file_extent,
+						  struct btrfs_file_extent *file_extent,
 						  const int type)
 {
 	struct extent_map *em = NULL;
@@ -150,6 +150,7 @@ static struct extent_map *btrfs_create_dio_extent(struct btrfs_inode *inode,
 		em = btrfs_create_io_em(inode, start, file_extent, type);
 		if (IS_ERR(em))
 			goto out;
+		file_extent->fscrypt_info = em->fscrypt_info;
 	}
 
 	ordered = btrfs_alloc_ordered_extent(inode, start, file_extent,
@@ -276,6 +277,7 @@ static int btrfs_get_blocks_direct_write(struct extent_map **map,
 		}
 		space_reserved = true;
 
+		file_extent.fscrypt_info = em->fscrypt_info;
 		em2 = btrfs_create_dio_extent(BTRFS_I(inode), dio_data, start,
 					      &file_extent, type);
 		btrfs_dec_nocow_writers(bg);
diff --git a/fs/btrfs/inode.c b/fs/btrfs/inode.c
index aef95d6e02bf..f81bdb97b212 100644
--- a/fs/btrfs/inode.c
+++ b/fs/btrfs/inode.c
@@ -1190,10 +1190,11 @@ static void submit_one_async_extent(struct async_chunk *async_chunk,
 		ret = PTR_ERR(em);
 		goto out_free_reserve;
 	}
-	btrfs_free_extent_map(em);
 
+	file_extent.fscrypt_info = em->fscrypt_info;
 	ordered = btrfs_alloc_ordered_extent(inode, start, &file_extent,
 					     1U << BTRFS_ORDERED_COMPRESSED);
+	btrfs_free_extent_map(em);
 	if (IS_ERR(ordered)) {
 		btrfs_drop_extent_map_range(inode, start, end, false);
 		ret = PTR_ERR(ordered);
@@ -1333,10 +1334,11 @@ static int cow_one_range(struct btrfs_inode *inode, struct folio *locked_folio,
 		ret = PTR_ERR(em);
 		goto free_reserved;
 	}
-	btrfs_free_extent_map(em);
 
+	file_extent.fscrypt_info = em->fscrypt_info;
 	ordered = btrfs_alloc_ordered_extent(inode, file_offset, &file_extent,
 					     1U << BTRFS_ORDERED_REGULAR);
+	btrfs_free_extent_map(em);
 	if (IS_ERR(ordered)) {
 		btrfs_drop_extent_map_range(inode, file_offset, cur_end, false);
 		ret = PTR_ERR(ordered);
@@ -2006,18 +2008,32 @@ static int nocow_one_range(struct btrfs_inode *inode, struct folio *locked_folio
 			   u64 file_pos, bool is_prealloc)
 {
 	struct btrfs_ordered_extent *ordered;
+	struct extent_map *em;
 	const u64 len = nocow_args->file_extent.num_bytes;
 	const u64 end = file_pos + len - 1;
 	int ret = 0;
 
 	btrfs_lock_extent(&inode->io_tree, file_pos, end, cached);
 
-	if (is_prealloc) {
-		struct extent_map *em;
+	/*
+	 * We only want to do this lookup if we're encrypted, otherwise
+	 * fsrypt_info will be null and we can avoid this lookup.
+	 */
+	if (IS_ENCRYPTED(&inode->vfs_inode)) {
+		em = btrfs_get_extent(inode, NULL, file_pos, len);
+		if (IS_ERR(em)) {
+			btrfs_unlock_extent(&inode->io_tree, file_pos, end, cached);
+			return PTR_ERR(em);
+		}
+		nocow_args->file_extent.fscrypt_info = fscrypt_get_extent_info(em->fscrypt_info);
+		btrfs_free_extent_map(em);
+	}
 
+	if (is_prealloc) {
 		em = btrfs_create_io_em(inode, file_pos, &nocow_args->file_extent,
 					BTRFS_ORDERED_PREALLOC);
 		if (IS_ERR(em)) {
+			fscrypt_put_extent_info(nocow_args->file_extent.fscrypt_info);
 			ret = PTR_ERR(em);
 			goto error;
 		}
@@ -2028,6 +2044,7 @@ static int nocow_one_range(struct btrfs_inode *inode, struct folio *locked_folio
 					     is_prealloc
 					     ? (1U << BTRFS_ORDERED_PREALLOC)
 					     : (1U << BTRFS_ORDERED_NOCOW));
+	fscrypt_put_extent_info(nocow_args->file_extent.fscrypt_info);
 	if (IS_ERR(ordered)) {
 		if (is_prealloc)
 			btrfs_drop_extent_map_range(inode, file_pos, end, false);
@@ -7582,7 +7599,13 @@ struct extent_map *btrfs_create_io_em(struct btrfs_inode *inode, u64 start,
 	em->flags |= EXTENT_FLAG_PINNED;
 	if (type == BTRFS_ORDERED_COMPRESSED)
 		btrfs_extent_map_set_compression(em, file_extent->compression);
-	btrfs_extent_map_set_encryption(em, BTRFS_ENCRYPTION_NONE);
+
+	if (file_extent->fscrypt_info) {
+		btrfs_extent_map_set_encryption(em, BTRFS_ENCRYPTION_FSCRYPT);
+		em->fscrypt_info = fscrypt_get_extent_info(file_extent->fscrypt_info);
+	} else {
+		btrfs_extent_map_set_encryption(em, BTRFS_ENCRYPTION_NONE);
+	}
 
 	ret = btrfs_replace_extent_map_range(inode, em, true);
 	if (ret) {
@@ -10125,11 +10148,12 @@ ssize_t btrfs_do_encoded_write(struct kiocb *iocb, struct iov_iter *from,
 		ret = PTR_ERR(em);
 		goto out_free_reserved;
 	}
-	btrfs_free_extent_map(em);
 
+	file_extent.fscrypt_info = em->fscrypt_info;
 	ordered = btrfs_alloc_ordered_extent(inode, start, &file_extent,
 				       (1U << BTRFS_ORDERED_ENCODED) |
 				       (1U << BTRFS_ORDERED_COMPRESSED));
+	btrfs_free_extent_map(em);
 	if (IS_ERR(ordered)) {
 		btrfs_drop_extent_map_range(inode, start, end, false);
 		ret = PTR_ERR(ordered);
-- 
2.51.0