ta: pkcs11: helper for serial arguments with allocation

Helper functions for serial arguments that expect memory allocation.

Reviewed-by: Ricardo Salveti <ricardo@foundries.io>
Co-developed-by: Etienne Carriere <etienne.carriere@linaro.org>
Signed-off-by: Etienne Carriere <etienne.carriere@linaro.org>
Signed-off-by: Jens Wiklander <jens.wiklander@linaro.org>

2 files changed, 109 insertions, 23 deletions

diff --git a/ta/pkcs11/src/serializer.c b/ta/pkcs11/src/serializer.c
index 2cc807b7..dac930a9 100644
--- a/ta/pkcs11/src/serializer.c
+++ b/ta/pkcs11/src/serializer.c
@@ -6,8 +6,10 @@
 #include <pkcs11_ta.h>
 #include <stdbool.h>
 #include <stdint.h>
+#include <tee_internal_api_extensions.h>
 #include <tee_internal_api.h>
 #include <trace.h>
+#include <util.h>
 
 #include "pkcs11_token.h"
 #include "serializer.h"
@@ -24,58 +26,72 @@ void serialargs_init(struct serialargs *args, void *in, size_t size)
 
 enum pkcs11_rc serialargs_get(struct serialargs *args, void *out, size_t size)
 {
-	if (args->next + size > args->start + args->size) {
-		EMSG("arg too short: full %zd, remain %zd, expect %zd",
-		     args->size, args->size - (args->next - args->start), size);
-		return PKCS11_CKR_ARGUMENTS_BAD;
-	}
-
-	TEE_MemMove(out, args->next, size);
+	enum pkcs11_rc rc = PKCS11_CKR_OK;
+	void *src = NULL;
 
-	args->next += size;
+	rc = serialargs_get_ptr(args, &src, size);
+	if (!rc)
+		TEE_MemMove(out, src, size);
 
-	return PKCS11_CKR_OK;
+	return rc;
 }
 
-enum pkcs11_rc serialargs_alloc_and_get(struct serialargs *args,
-					void **out, size_t size)
+static enum pkcs11_rc alloc_and_get(struct serialargs *args, char *orig_next,
+				    const void *buf0, size_t buf0_sz,
+				    void **out, size_t size)
 {
-	void *ptr = NULL;
+	enum pkcs11_rc rc = PKCS11_CKR_OK;
+	uint8_t *ptr = NULL;
+	void *src = NULL;
+	size_t sz = 0;
 
-	if (!size) {
+	if (ADD_OVERFLOW(buf0_sz, size, &sz))
+		return PKCS11_CKR_ARGUMENTS_BAD;
+
+	if (!sz) {
 		*out = NULL;
 		return PKCS11_CKR_OK;
 	}
 
-	if (args->next + size > args->start + args->size) {
-		EMSG("arg too short: full %zd, remain %zd, expect %zd",
-		     args->size, args->size - (args->next - args->start), size);
-		return PKCS11_CKR_ARGUMENTS_BAD;
-	}
+	rc = serialargs_get_ptr(args, &src, size);
+	if (rc)
+		return rc;
 
-	ptr = TEE_Malloc(size, TEE_MALLOC_FILL_ZERO);
-	if (!ptr)
+	ptr = TEE_Malloc(sz, TEE_MALLOC_FILL_ZERO);
+	if (!ptr) {
+		args->next = orig_next;
 		return PKCS11_CKR_DEVICE_MEMORY;
+	}
 
-	TEE_MemMove(ptr, args->next, size);
+	TEE_MemMove(ptr, buf0, buf0_sz);
+	TEE_MemMove(ptr + buf0_sz, src, size);
 
-	args->next += size;
 	*out = ptr;
 
 	return PKCS11_CKR_OK;
 }
 
+enum pkcs11_rc serialargs_alloc_and_get(struct serialargs *args,
+					void **out, size_t size)
+{
+	return alloc_and_get(args, args->next, NULL, 0, out, size);
+}
+
 enum pkcs11_rc serialargs_get_ptr(struct serialargs *args, void **out,
 				  size_t size)
 {
 	void *ptr = args->next;
+	vaddr_t next_end = 0;
+
+	if (ADD_OVERFLOW((vaddr_t)args->next, size, &next_end))
+		return PKCS11_CKR_ARGUMENTS_BAD;
 
 	if (!size) {
 		*out = NULL;
 		return PKCS11_CKR_OK;
 	}
 
-	if (args->next + size > args->start + args->size) {
+	if ((char *)next_end > args->start + args->size) {
 		EMSG("arg too short: full %zd, remain %zd, expect %zd",
 		     args->size, args->size - (args->next - args->start), size);
 		return PKCS11_CKR_ARGUMENTS_BAD;
@@ -87,6 +103,50 @@ enum pkcs11_rc serialargs_get_ptr(struct serialargs *args, void **out,
 	return PKCS11_CKR_OK;
 }
 
+enum pkcs11_rc
+serialargs_alloc_get_one_attribute(struct serialargs *args,
+				   struct pkcs11_attribute_head **out)
+{
+	struct pkcs11_attribute_head head = { };
+	enum pkcs11_rc rc = PKCS11_CKR_OK;
+	char *orig_next = args->next;
+	void *p = NULL;
+
+	rc = serialargs_get(args, &head, sizeof(head));
+	if (rc)
+		return rc;
+
+	rc = alloc_and_get(args, orig_next, &head, sizeof(head), &p, head.size);
+	if (rc)
+		return rc;
+
+	*out = p;
+
+	return PKCS11_CKR_OK;
+}
+
+enum pkcs11_rc serialargs_alloc_get_attributes(struct serialargs *args,
+					       struct pkcs11_object_head **out)
+{
+	struct pkcs11_object_head attr = { };
+	enum pkcs11_rc rc = PKCS11_CKR_OK;
+	char *orig_next = args->next;
+	void *p = NULL;
+
+	rc = serialargs_get(args, &attr, sizeof(attr));
+	if (rc)
+		return rc;
+
+	rc = alloc_and_get(args, orig_next, &attr, sizeof(attr), &p,
+			   attr.attrs_size);
+	if (rc)
+		return rc;
+
+	*out = p;
+
+	return PKCS11_CKR_OK;
+}
+
 bool serialargs_remaining_bytes(struct serialargs *args)
 {
 	return args->next < args->start + args->size;
diff --git a/ta/pkcs11/src/serializer.h b/ta/pkcs11/src/serializer.h
index a0a7c351..3e6fe09e 100644
--- a/ta/pkcs11/src/serializer.h
+++ b/ta/pkcs11/src/serializer.h
@@ -22,6 +22,9 @@ struct serialargs {
 	size_t size;
 };
 
+struct pkcs11_client;
+struct pkcs11_session;
+
 /*
  * serialargs_init() - Initialize with a new input buffer
  * @args:	serializing state
@@ -65,6 +68,29 @@ enum pkcs11_rc serialargs_get_ptr(struct serialargs *args, void **out,
 				  size_t size);
 
 /*
+ * serialargs_alloc_get_one_attribute() - allocate and extract one attribute
+ * @args:	serializing state
+ * @out:	Pointer to the allocated and extracted attribute in *@out
+ *
+ * Returns PKCS11_CKR_OK on success or an error code from enum pkcs11_rc on
+ * failure.
+ */
+enum pkcs11_rc
+serialargs_alloc_get_one_attribute(struct serialargs *args,
+				   struct pkcs11_attribute_head **out);
+
+/*
+ * serialargs_alloc_get_attributes() - allocate and extract an object
+ * @args:	serializing state
+ * @out:	Pointer to the allocated and extracted object in *@out
+ *
+ * Returns PKCS11_CKR_OK on success or an error code from enum pkcs11_rc on
+ * failure.
+ */
+enum pkcs11_rc serialargs_alloc_get_attributes(struct serialargs *args,
+					       struct pkcs11_object_head **out);
+
+/*
  * serialargs_alloc_and_get() - allocate and extract data
  * @args:	serializing state
  * @out:	Pointer to the allocated and extracted data in *@out