handshake: Add cleanup function for handshake_state

To allow _auto_(handshake_state_free) variables to be used.

2 files changed, 9 insertions, 1 deletions

diff --git a/src/handshake.c b/src/handshake.c
index 6b93774ab..1c5ed2c9b 100644
--- a/src/handshake.c
+++ b/src/handshake.c
@@ -105,7 +105,12 @@ void __handshake_set_install_ext_tk_func(handshake_install_ext_tk_func_t func)
 
 void handshake_state_free(struct handshake_state *s)
 {
-	__typeof__(s->free) destroy = s->free;
+	__typeof__(s->free) destroy;
+
+	if (!s)
+		return;
+
+	destroy = s->free;
 
 	if (s->in_event) {
 		s->in_event = false;
diff --git a/src/handshake.h b/src/handshake.h
index 7200c3617..815eb44ff 100644
--- a/src/handshake.h
+++ b/src/handshake.h
@@ -24,6 +24,7 @@
 #include <stdbool.h>
 #include <asm/byteorder.h>
 #include <linux/types.h>
+#include <ell/cleanup.h>
 
 struct handshake_state;
 enum crypto_cipher;
@@ -298,3 +299,5 @@ const uint8_t *handshake_util_find_pmkid_kde(const uint8_t *data,
 					size_t data_len);
 void handshake_util_build_gtk_kde(enum crypto_cipher cipher, const uint8_t *key,
 					unsigned int key_index, uint8_t *to);
+
+DEFINE_CLEANUP_FUNC(handshake_state_free);