3 files changed, 24 insertions, 5 deletions

diff --git a/src/methods/composed_gmac/composed_gmac.c b/src/methods/composed_gmac/composed_gmac.c
index cb0d5d6..b8ee2f6 100644
--- a/src/methods/composed_gmac/composed_gmac.c
+++ b/src/methods/composed_gmac/composed_gmac.c
@@ -147,7 +147,13 @@ static fastd_method_session_state_t* method_session_init(fastd_context_t *ctx, c
 	uint8_t zeroiv[session->gmac_cipher_info->iv_length];
 	memset(zeroiv, 0, session->gmac_cipher_info->iv_length);
 
-	session->gmac_cipher->crypt(ctx, session->gmac_cipher_state, &H, &ZERO_BLOCK, sizeof(fastd_block128_t), zeroiv);
+	if (!session->gmac_cipher->crypt(ctx, session->gmac_cipher_state, &H, &ZERO_BLOCK, sizeof(fastd_block128_t), zeroiv)) {
+		session->cipher->free_state(ctx, session->cipher_state);
+		session->gmac_cipher->free_state(ctx, session->gmac_cipher_state);
+		free(session);
+
+		return NULL;
+	}
 
 	session->ghash = fastd_mac_get_by_name(ctx, "ghash", &session->ghash_info, &session->ghash_ctx);
 	if (!session->ghash)
diff --git a/src/methods/generic_gcm/generic_gcm.c b/src/methods/generic_gcm/generic_gcm.c
index 9ce8522..9481100 100644
--- a/src/methods/generic_gcm/generic_gcm.c
+++ b/src/methods/generic_gcm/generic_gcm.c
@@ -113,7 +113,11 @@ static fastd_method_session_state_t* method_session_init(fastd_context_t *ctx, c
 	uint8_t zeroiv[session->cipher_info->iv_length];
 	memset(zeroiv, 0, session->cipher_info->iv_length);
 
-	session->cipher->crypt(ctx, session->cipher_state, &H, &zeroblock, sizeof(fastd_block128_t), zeroiv);
+	if (!session->cipher->crypt(ctx, session->cipher_state, &H, &zeroblock, sizeof(fastd_block128_t), zeroiv)) {
+		session->cipher->free_state(ctx, session->cipher_state);
+		free(session);
+		return NULL;
+	}
 
 	session->ghash = fastd_mac_get_by_name(ctx, "ghash", &session->ghash_info, &session->ghash_ctx);
 	if (!session->ghash)
diff --git a/src/protocols/ec25519_fhmqvc/handshake.c b/src/protocols/ec25519_fhmqvc/handshake.c
index 5606463..c359fb1 100644
--- a/src/protocols/ec25519_fhmqvc/handshake.c
+++ b/src/protocols/ec25519_fhmqvc/handshake.c
@@ -91,7 +91,7 @@ static inline void supersede_session(fastd_context_t *ctx, fastd_peer_t *peer, c
 	}
 }
 
-static inline void new_session(fastd_context_t *ctx, fastd_peer_t *peer, const char *method_name, const fastd_method_t *method, bool initiator,
+static inline bool new_session(fastd_context_t *ctx, fastd_peer_t *peer, const char *method_name, const fastd_method_t *method, bool initiator,
 			       const aligned_int256_t *A, const aligned_int256_t *B, const aligned_int256_t *X, const aligned_int256_t *Y,
 			       const aligned_int256_t *sigma, const uint32_t *salt, uint64_t serial) {
 
@@ -110,11 +110,16 @@ static inline void new_session(fastd_context_t *ctx, fastd_peer_t *peer, const c
 		peer->protocol_state->session.method_state = method->session_init_compat(ctx, method_name, hash.b, HASHBYTES, initiator);
 	}
 
+	if (!peer->protocol_state->session.method_state)
+		return false;
+
 	peer->protocol_state->session.established = ctx->now;
 	peer->protocol_state->session.handshakes_cleaned = false;
 	peer->protocol_state->session.refreshing = false;
 	peer->protocol_state->session.method = method;
 	peer->protocol_state->last_serial = serial;
+
+	return true;
 }
 
 static bool establish(fastd_context_t *ctx, fastd_peer_t *peer, const char *method_name, fastd_socket_t *sock,
@@ -128,7 +133,7 @@ static bool establish(fastd_context_t *ctx, fastd_peer_t *peer, const char *meth
 
 	const fastd_method_t *method = fastd_method_get_by_name(method_name);
 	if (!salt && !method->session_init_compat) {
-		pr_warn(ctx, "can't establish session with %P[%I] (method without compat support)");
+		pr_warn(ctx, "can't establish session with %P[%I] (method without compat support)", peer, remote_addr);
 		return false;
 	}
 
@@ -140,7 +145,11 @@ static bool establish(fastd_context_t *ctx, fastd_peer_t *peer, const char *meth
 		return false;
 	}
 
-	new_session(ctx, peer, method_name, method, initiator, A, B, X, Y, sigma, salt, serial);
+	if (!new_session(ctx, peer, method_name, method, initiator, A, B, X, Y, sigma, salt, serial)) {
+		pr_error(ctx, "failed to initialize method session for %P (method `%s'%s)", peer, method_name, salt ? "" : " (compat mode)");
+		fastd_peer_reset(ctx, peer);
+		return false;
+	}
 
 	fastd_peer_seen(ctx, peer);
 	fastd_peer_set_established(ctx, peer);