1 files changed, 15 insertions, 6 deletions

diff --git a/drivers/iio/buffer_cb.c b/drivers/iio/buffer_cb.c
index 841fec1e78b2..2d9c6f8c06db 100644
--- a/drivers/iio/buffer_cb.c
+++ b/drivers/iio/buffer_cb.c
@@ -12,17 +12,27 @@ struct iio_cb_buffer {
         struct iio_channel *channels;
 };
 
-static int iio_buffer_cb_store_to(struct iio_buffer *buffer, const void *data)
+static struct iio_cb_buffer *buffer_to_cb_buffer(struct iio_buffer *buffer)
 {
-        struct iio_cb_buffer *cb_buff = container_of(buffer,
-                                                     struct iio_cb_buffer,
-                                                     buffer);
+        return container_of(buffer, struct iio_cb_buffer, buffer);
+}
 
+static int iio_buffer_cb_store_to(struct iio_buffer *buffer, const void *data)
+{
+        struct iio_cb_buffer *cb_buff = buffer_to_cb_buffer(buffer);
         return cb_buff->cb(data, cb_buff->private);
 }
 
+static void iio_buffer_cb_release(struct iio_buffer *buffer)
+{
+        struct iio_cb_buffer *cb_buff = buffer_to_cb_buffer(buffer);
+        kfree(cb_buff->buffer.scan_mask);
+        kfree(cb_buff);
+}
+
 static const struct iio_buffer_access_funcs iio_cb_access = {
         .store_to = &iio_buffer_cb_store_to,
+        .release = &iio_buffer_cb_release,
 };
 
 struct iio_cb_buffer *iio_channel_get_all_cb(struct device *dev,
@@ -104,9 +114,8 @@ EXPORT_SYMBOL_GPL(iio_channel_stop_all_cb);
 
 void iio_channel_release_all_cb(struct iio_cb_buffer *cb_buff)
 {
-        kfree(cb_buff->buffer.scan_mask);
         iio_channel_release_all(cb_buff->channels);
-        kfree(cb_buff);
+        iio_buffer_put(&cb_buff->buffer);
 }
 EXPORT_SYMBOL_GPL(iio_channel_release_all_cb);