HACK: dma-direct: Add some sort of inline encryption

Change-Id: I463fcb773340167013163ba61cdac843ddc8e35a
Signed-off-by: Mostafa Saleh <smostafa@google.com>

kernel/dma/direct.h

diff --git a/kernel/dma/direct.h b/kernel/dma/direct.h
index f3fc28f..c61e4f2 100644
--- a/kernel/dma/direct.h
+++ b/kernel/dma/direct.h
@@ -82,6 +82,8 @@ static inline void dma_direct_sync_single_for_cpu(struct device *dev,
 	swiotlb_sync_single_for_cpu(dev, paddr, size, dir);
 }
 
+#define N_SIZE		(10 * PAGE_SIZE)
+
 static inline dma_addr_t dma_direct_map_phys(struct device *dev,
 		phys_addr_t phys, size_t size, enum dma_data_direction dir,
 		unsigned long attrs, bool flush)
@@ -115,6 +117,25 @@ static inline dma_addr_t dma_direct_map_phys(struct device *dev,
 
 	if (unlikely(!dma_capable(dev, dma_addr, size, true, attrs)) ||
 	    dma_kmalloc_needs_bounce(dev, size, dir)) {
+
+		/*
+		 * TBD:
+		 * - CONFIG
+		 * - dma_capable() returns special error for missing encryption
+		 * - check it is safe to call phys_to_virt
+		 */
+		unsigned long vaddr = (unsigned long)phys_to_virt(phys);
+
+		if (PAGE_ALIGNED(size) && PAGE_ALIGNED(phys) && (size >= N_SIZE) && 
+			!dma_kmalloc_needs_bounce(dev, size, dir)) {
+			int ret = set_memory_decrypted(vaddr, size >> PAGE_SHIFT);
+			if (!ret) {
+				//pr_err("PKVM_DEBUG: in-place map phys=%llx size=%lx\n", phys, size);
+				goto dma_mapped;
+			}
+			pr_err("PKVM_DEBUG: in-place map FAILED phys=%llx size=%lx ret=%d\n", phys, size, ret);
+		}
+		//pr_err("have to bounce %llx %lx\n", phys, size);
 		if (is_swiotlb_active(dev) &&
 		    !(attrs & DMA_ATTR_REQUIRE_COHERENT))
 			return swiotlb_map(dev, phys, size, dir, attrs);
@@ -154,5 +175,14 @@ static inline void dma_direct_unmap_phys(struct device *dev, dma_addr_t addr,
 
 	swiotlb_tbl_unmap_single(dev, phys, size, dir,
 					 attrs | DMA_ATTR_SKIP_CPU_SYNC);
+	if (unlikely(!dma_capable(dev, addr, size, true, attrs))) {
+		unsigned long vaddr = (unsigned long)phys_to_virt(phys);
+
+		if (PAGE_ALIGNED(size) && PAGE_ALIGNED(phys) && (size >= N_SIZE) && !swiotlb_find_pool(dev, phys)) {
+			set_memory_encrypted(vaddr, size >> PAGE_SHIFT);
+			//pr_err("PKVM_DEBUG: in-place unmap phys=%llx size=%lx ret=%d\n", phys, size, ret);
+		}
+	}
+	/* TBD check if the memory was encrypted some how??? and reverse */
 }
 #endif /* _KERNEL_DMA_DIRECT_H */