Fix marshal and unmarshal

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

1 files changed, 21 insertions, 10 deletions

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
 }