Fix marshal and unmarshal

Signed-off-by: Mattias Andrée <maandree@kth.se>

2 files changed, 45 insertions, 28 deletions

diff --git a/libkeccak_state_marshal.c b/libkeccak_state_marshal.c
index e4d2efa..6563038 100644
--- a/libkeccak_state_marshal.c
+++ b/libkeccak_state_marshal.c
@@ -17,24 +17,30 @@
 size_t
 libkeccak_state_marshal(const struct libkeccak_state *restrict state, void *restrict data_)
 {
-#define set(type, var) *((type *)data) = state->var, data += sizeof(type)
-	unsigned char *restrict data = data_;
-	if (data) {
-		set(long int, r);
-		set(long int, c);
-		set(long int, n);
-		set(long int, b);
-		set(long int, w);
-		set(int64_t, wmod);
-		set(long int, l);
-		set(long int, nr);
-		memcpy(data, state->S, sizeof(state->S));
-		data += sizeof(state->S);
-		set(size_t, mptr);
-		set(size_t, mlen);
-		memcpy(data, state->M, state->mptr * sizeof(char));
-		data += state->mptr;
+#define set(type, var) *(type *)data = state->var, data += sizeof(type)
+	unsigned char *restrict start = data_;
+	unsigned char *restrict data = start;
+	if (!data) {
+		return 7 * sizeof(long int) +
+		       1 * sizeof(int64_t) +
+		       sizeof(state->S) +
+		       2 * sizeof(size_t) +
+		       state->mptr;
 	}
-	return sizeof(struct libkeccak_state) - sizeof(char *) + state->mptr * sizeof(char);
+	set(long int, r);
+	set(long int, c);
+	set(long int, n);
+	set(long int, b);
+	set(long int, w);
+	set(int64_t, wmod);
+	set(long int, l);
+	set(long int, nr);
+	memcpy(data, state->S, sizeof(state->S));
+	data += sizeof(state->S);
+	set(size_t, mptr);
+	set(size_t, mlen);
+	memcpy(data, state->M, state->mptr * sizeof(char));
+	data += state->mptr;
+	return (size_t)(data - start);
 #undef set
 }
diff --git a/libkeccak_state_unmarshal.c b/libkeccak_state_unmarshal.c
index 4e0df28..f221fad 100644
--- a/libkeccak_state_unmarshal.c
+++ b/libkeccak_state_unmarshal.c
@@ -18,11 +18,18 @@
 size_t
 libkeccak_state_unmarshal(struct libkeccak_state *restrict state, const void *restrict data_)
 {
-#define get(type, var) state->var = *((const type *)data), data += sizeof(type)
-	const unsigned char *restrict data = data_;
+#define get(type, var) state->var = *(const type *)data, data += sizeof(type)
+	const unsigned char *restrict start = data_;
+	const unsigned char *restrict data = start;
+	size_t mptr;
 	if (!state) {
-		data += (7 * sizeof(long int) + 26 * sizeof(int64_t));
-		return sizeof(struct libkeccak_state) - sizeof(char *) + *(const size_t *)data * sizeof(char);
+		data += 7 * sizeof(long int);
+		data += 1 * sizeof(int64_t);
+		data += sizeof(state->S);
+		mptr = *(const size_t *)data;
+		data += 2 * sizeof(size_t);
+		data += mptr;
+		return (size_t)(data - start);
 	}
 	get(long int, r);
 	get(long int, c);
@@ -36,11 +43,15 @@ libkeccak_state_unmarshal(struct libkeccak_state *restrict state, const void *re
 	data += sizeof(state->S);
 	get(size_t, mptr);
 	get(size_t, mlen);
-	state->M = malloc(state->mptr * sizeof(char));
-	if (!state->M)
-		return 0;
-	memcpy(state->M, data, state->mptr * sizeof(char));
-	data += state->mptr;
-	return sizeof(struct libkeccak_state) - sizeof(char *) + state->mptr * sizeof(char);
+	if (state->mptr) {
+		state->M = malloc(state->mptr * sizeof(char));
+		if (!state->M)
+			return 0;
+		memcpy(state->M, data, state->mptr * sizeof(char));
+		data += state->mptr;
+	} else {
+		state->M = NULL;
+	}
+	return (size_t)(data - start);
 #undef get
 }