diff --git a/lib/swiotlb.c b/lib/swiotlb.c
index 240a67c..f8eebd48 100644
--- a/lib/swiotlb.c
+++ b/lib/swiotlb.c
@@ -276,11 +276,7 @@
 static int
 address_needs_mapping(struct device *hwdev, dma_addr_t addr, size_t size)
 {
-	dma_addr_t mask = 0xffffffff;
-	/* If the device has a mask, use it, otherwise default to 32 bits */
-	if (hwdev && hwdev->dma_mask)
-		mask = *hwdev->dma_mask;
-	return !is_buffer_dma_capable(mask, addr, size);
+	return !is_buffer_dma_capable(dma_get_mask(hwdev), addr, size);
 }
 
 static int is_swiotlb_buffer(char *addr)