From c188d75cd2e6bc1b461da9165a38296d2223092c Mon Sep 17 00:00:00 2001
From: Paul Cercueil <paul@crapouillou.net>
Date: Tue, 13 Apr 2021 16:50:00 +0100
Subject: [PATCH] iiod-client: Add support for ZSTD-compressed metadata

If compiled WITH_ZSTD, use the ZPRINT command to receive a
ZSTD-compressed payload that can be decompressed into the original XML
string of the remote context. If the remote IIOD does not support
ZPRINT, it will simply revert to using PRINT instead.

This can dramatically speed up context creation when using a slow
transport (e.g. serial), as the XML strings can be really big.

For instance, the XML string of the ADALM-Pluto is close to 25 KiB.
On a standard 57600 bps UART connection, this would take about ~3.4
seconds to transfer. The compressed payload amounts to only about
3.7 KiB, which would take about ~0.5 seconds to transfer.

Signed-off-by: Paul Cercueil <paul@crapouillou.net>
---
 iiod-client.c | 63 +++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 59 insertions(+), 4 deletions(-)

diff --git a/iiod-client.c b/iiod-client.c
index dbf3b2bea..abd5f71c7 100644
--- a/iiod-client.c
+++ b/iiod-client.c
@@ -8,6 +8,7 @@
 
 #include "debug.h"
 #include "iiod-client.h"
+#include "iio-config.h"
 #include "iio-lock.h"
 #include "iio-private.h"
 
@@ -15,6 +16,9 @@
 #include <inttypes.h>
 #include <string.h>
 #include <stdio.h>
+#if WITH_ZSTD
+#include <zstd.h>
+#endif
 
 struct iiod_client {
 	struct iio_context_pdata *pdata;
@@ -504,18 +508,29 @@ ssize_t iiod_client_write_attr(struct iiod_client *client, void *desc,
 	return ret;
 }
 
-struct iio_context * iiod_client_create_context(
-		struct iiod_client *client, void *desc)
+static struct iio_context *
+iiod_client_create_context_private(struct iiod_client *client,
+				   void *desc, bool zstd)
 {
+	const char *cmd = zstd ? "ZPRINT\r\n" : "PRINT\r\n";
 	struct iio_context *ctx = NULL;
 	size_t xml_len;
 	char *xml;
 	int ret;
 
 	iio_mutex_lock(client->lock);
-	ret = iiod_client_exec_command(client, desc, "PRINT\r\n");
-	if (ret < 0)
+	ret = iiod_client_exec_command(client, desc, cmd);
+	if (ret < 0) {
+		if (ret == -EINVAL && zstd) {
+			/* If the ZPRINT command does not exist, try again
+			 * with the regular PRINT command. */
+			iio_mutex_unlock(client->lock);
+
+			return iiod_client_create_context_private(client, desc, false);
+		}
+
 		goto out_unlock;
+	}
 
 	xml_len = (size_t) ret;
 	xml = malloc(xml_len + 1);
@@ -529,6 +544,41 @@ struct iio_context * iiod_client_create_context(
 	if (ret < 0)
 		goto out_free_xml;
 
+#if WITH_ZSTD
+	if (zstd) {
+		unsigned long long len;
+		char *xml_zstd;
+
+		IIO_DEBUG("Received ZSTD-compressed XML string.\n");
+
+		len = ZSTD_getFrameContentSize(xml, xml_len);
+		if (len == ZSTD_CONTENTSIZE_UNKNOWN ||
+		    len == ZSTD_CONTENTSIZE_ERROR) {
+			ret = -EIO;
+			goto out_free_xml;
+		}
+
+		xml_zstd = malloc(len);
+		if (!xml_zstd) {
+			ret = -ENOMEM;
+			goto out_free_xml;
+		}
+
+		xml_len = ZSTD_decompress(xml_zstd, len, xml, xml_len);
+		if (ZSTD_isError(xml_len)) {
+			IIO_ERROR("Unable to decompress ZSTD data: %s\n",
+				  ZSTD_getErrorName(xml_len));
+			ret = -EIO;
+			free(xml_zstd);
+			goto out_free_xml;
+		}
+
+		/* Free compressed data, make "xml" point to uncompressed data */
+		free(xml);
+		xml = xml_zstd;
+	}
+#endif
+
 	ctx = iio_create_xml_context_mem(xml, xml_len);
 	if (!ctx)
 		ret = -errno;
@@ -542,6 +592,11 @@ struct iio_context * iiod_client_create_context(
 	return ctx;
 }
 
+struct iio_context * iiod_client_create_context(struct iiod_client *client, void *desc)
+{
+	return iiod_client_create_context_private(client, desc, WITH_ZSTD);
+}
+
 int iiod_client_open_unlocked(struct iiod_client *client, void *desc,
 		const struct iio_device *dev, size_t samples_count, bool cyclic)
 {