diff --git a/Makefile.am b/Makefile.am
index 2c5ae6f5..8f579624 100644
--- a/Makefile.am
+++ b/Makefile.am
@@ -449,7 +449,8 @@ unit_test_sae_SOURCES = unit/test-sae.c \
 				src/crypto.h src/crypto.c \
 				src/ie.h src/ie.c \
 				src/handshake.h src/handshake.c \
-				src/util.h src/util.c
+				src/util.h src/util.c \
+				src/mpdu.h src/mpdu.c
 unit_test_sae_LDADD = $(ell_ldadd)
 
 TESTS = $(unit_tests)
diff --git a/unit/test-sae.c b/unit/test-sae.c
index 33357573..88a2b8c4 100644
--- a/unit/test-sae.c
+++ b/unit/test-sae.c
@@ -33,6 +33,7 @@
 #include "src/handshake.h"
 #include "src/mpdu.h"
 #include "src/sae.h"
+#include "src/auth-proto.h"
 
 struct test_handshake_state {
 	struct handshake_state super;
@@ -48,7 +49,8 @@ struct test_data {
 	/* true if a rejection packet was sent */
 	bool tx_reject_occurred;
 	/* true if the tx function was called */
-	bool tx_called;
+	bool tx_auth_called;
+	bool tx_assoc_called;
 	/* copy of clogging token (if present) */
 	uint8_t test_clogging_token[32];
 	/* copy of last packet sent */
@@ -60,6 +62,16 @@ struct test_data {
 	struct handshake_state *handshake;
 };
 
+struct authenticate_frame {
+	struct mmpdu_header hdr;
+	struct mmpdu_authentication auth;
+} __attribute__ ((packed));
+
+struct associate_frame {
+	struct mmpdu_header hdr;
+	struct mmpdu_association_response assoc;
+} __attribute__ ((packed));
+
 static uint8_t spa[] = {2, 0, 0, 0, 0, 0};
 static uint8_t aa[] = {2, 0, 0, 0, 0, 1};
 static char *passphrase = "secret123";
@@ -84,30 +96,20 @@ static struct handshake_state *test_handshake_state_new(uint32_t ifindex)
 	return &ths->super;
 }
 
-static void test_complete_func(uint16_t status, void *user_data)
-{
-	struct test_data *td = user_data;
-
-	td->status = status;
-}
-
-static int test_tx_func(const uint8_t *dest, const uint8_t *frame, size_t len,
-					void *user_data)
+static void test_tx_auth_func(const uint8_t *frame, size_t len, void *user_data)
 {
 	struct test_data *td = user_data;
 	uint16_t trans;
 
-	td->tx_called = true;
+	td->tx_auth_called = true;
 
 	memset(td->tx_packet, 0, sizeof(td->tx_packet));
 	memcpy(td->tx_packet, frame, len);
 	td->tx_packet_len = len;
 
-	assert(!memcmp(dest, aa, 6));
-
 	if (len <= 6 && l_get_le16(frame + 2) != 0) {
 		td->tx_reject_occurred = true;
-		return 0;
+		return;
 	}
 
 	trans = l_get_le16(frame);	/* transaction */
@@ -127,24 +129,29 @@ static int test_tx_func(const uint8_t *dest, const uint8_t *frame, size_t len,
 
 		td->commit_success = true;
 
-		return 0;
+		return;
 	case 2:
 		assert(l_get_le16(frame + 2) == 0);
 		assert(len == 38);
 
 		td->confirm_success = true;
 
-		return 0;
+		return;
 	}
 
 	assert(false);
-
-	return 0;
 }
 
-static struct sae_sm *test_initialize(struct test_data *td)
+static void test_tx_assoc_func(void *user_data)
 {
-	struct sae_sm *sm;
+	struct test_data *td = user_data;
+
+	td->tx_assoc_called = true;
+}
+
+static struct auth_proto *test_initialize(struct test_data *td)
+{
+	struct auth_proto *ap;
 	struct handshake_state *hs = test_handshake_state_new(1);
 
 	td->handshake = hs;
@@ -155,14 +162,14 @@ static struct sae_sm *test_initialize(struct test_data *td)
 
 	memset(td->test_clogging_token, 0xde, 32);
 
-	sm = sae_sm_new(hs, test_tx_func, test_complete_func, td);
+	ap = sae_sm_new(hs, test_tx_auth_func, test_tx_assoc_func, td);
 
 	td->commit_success = false;
-	sae_start(sm);
+	auth_proto_start(ap);
 
 	assert(td->commit_success == true);
 
-	return sm;
+	return ap;
 }
 
 static void test_destruct(struct test_data *td)
@@ -172,7 +179,7 @@ static void test_destruct(struct test_data *td)
 }
 
 static uint8_t aa_commit[] = {
-	0x01, 0x00, 0x00, 0x00, 0x13, 0x00, 0x50, 0x5b, 0xb2, 0x1f, 0xaf, 0x7d,
+	0x13, 0x00, 0x50, 0x5b, 0xb2, 0x1f, 0xaf, 0x7d,
 	0xaf, 0x14, 0x7c, 0x7b, 0x19, 0xc9, 0x72, 0x82, 0xbc, 0x1a, 0xdb, 0xa1,
 	0xbd, 0x6e, 0x5a, 0xc7, 0x58, 0x0a, 0x65, 0x1f, 0xd2, 0xde, 0xb0, 0x66,
 	0xa5, 0xf9, 0x3e, 0x95, 0x4a, 0xe1, 0x83, 0xdb, 0x8a, 0xf5, 0x47, 0x8a,
@@ -184,160 +191,218 @@ static uint8_t aa_commit[] = {
 };
 
 static uint8_t aa_confirm[] = {
-	0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0x0e, 0xf7, 0x5c, 0x1c, 0xab,
+	0x00, 0x00, 0x03, 0x0e, 0xf7, 0x5c, 0x1c, 0xab,
 	0x7c, 0x29, 0xa1, 0x79, 0x22, 0xe4, 0x3b, 0x64, 0xb8, 0xf0, 0x70, 0x25,
 	0x40, 0xcc, 0x78, 0x81, 0x27, 0x12, 0xca, 0xa9, 0xf5, 0xe5, 0x0f, 0xa7,
 	0x73, 0x6d
 };
 
+static size_t setup_auth_frame(struct authenticate_frame *frame,
+				const uint8_t *addr,
+				uint16_t trans, uint16_t status,
+				const uint8_t *data, size_t len)
+{
+	memcpy(frame->hdr.address_2, addr, 6);
+
+	frame->hdr.fc.type = MPDU_TYPE_MANAGEMENT;
+	frame->hdr.fc.subtype = MPDU_MANAGEMENT_SUBTYPE_AUTHENTICATION;
+	frame->hdr.fc.order = 1;
+
+	l_put_le16(MMPDU_AUTH_ALGO_SAE, &frame->auth.algorithm);
+	l_put_le16(trans, &frame->auth.transaction_sequence);
+	l_put_le16(status, &frame->auth.status);
+
+	if (data)
+		memcpy(frame->auth.ies, data, len);
+
+	return sizeof(frame->hdr) + sizeof(frame->auth) + len;
+}
+
+static size_t setup_assoc_frame(struct associate_frame *frame, uint16_t status)
+{
+	/*
+	 * Only need the frame to verify with mpdu_validate and have status
+	 * code set.
+	 */
+	memset(frame, 0, sizeof(struct associate_frame));
+
+	frame->hdr.fc.type = MPDU_TYPE_MANAGEMENT;
+	frame->hdr.fc.subtype = MPDU_MANAGEMENT_SUBTYPE_ASSOCIATION_REQUEST;
+	frame->hdr.fc.order = 1;
+
+	l_put_le16(status, &frame->assoc.status_code);
+
+	return sizeof(frame->hdr) + sizeof(frame->assoc);
+}
+
 static void test_confirm_timeout(const void *arg)
 {
 	struct test_data *td = l_new(struct test_data, 1);
-	struct sae_sm *sm = test_initialize(td);
+	struct auth_proto *ap = test_initialize(td);
+	struct authenticate_frame *frame = alloca(
+					sizeof(struct authenticate_frame) +
+					sizeof(aa_commit));
+	size_t len;
 	int i;
 
-	sae_rx_packet(sm, aa, aa_commit, sizeof(aa_commit));
+	len = setup_auth_frame(frame, aa, 1, 0, aa_commit, sizeof(aa_commit));
+
+	assert(auth_proto_rx_authenticate(ap, (uint8_t *)frame, len) == 0);
 
 	assert(td->confirm_success);
 
 	assert(l_get_le16(td->tx_packet + 4) == 1);
 
 	for (i = 1; i < 5; i++) {
-		sae_timeout(sm);
+		assert(auth_proto_auth_timeout(ap));
 		assert(l_get_le16(td->tx_packet + 4) == i + 1);
 	}
 
-	sae_timeout(sm);
-
-	assert(td->status != 0);
+	assert(!auth_proto_auth_timeout(ap));
 
 	test_destruct(td);
-	sae_sm_free(sm);
+
+	auth_proto_free(ap);
 }
 
 static void test_commit_timeout(const void *arg)
 {
 	struct test_data *td = l_new(struct test_data, 1);
-	struct sae_sm *sm = test_initialize(td);
+	struct auth_proto *ap = test_initialize(td);
 	uint8_t last_packet[512];
 	int i;
 
 	memcpy(last_packet, td->tx_packet, td->tx_packet_len);
 
 	for (i = 0; i < 4; i++) {
-		sae_timeout(sm);
+		assert(auth_proto_auth_timeout(ap));
 
 		assert(!memcmp(last_packet, td->tx_packet, td->tx_packet_len));
 
 		memcpy(last_packet, td->tx_packet, td->tx_packet_len);
 	}
 
-	sae_timeout(sm);
-
-	assert(td->status != 0);
+	assert(!auth_proto_auth_timeout(ap));
 
 	test_destruct(td);
-	sae_sm_free(sm);
+	auth_proto_free(ap);
 }
 
 static void test_clogging(const void *arg)
 {
-	uint8_t frame[38];
 	struct test_data *td = l_new(struct test_data, 1);
-	struct sae_sm *sm = test_initialize(td);
+	struct auth_proto *ap = test_initialize(td);
+	struct authenticate_frame *frame = alloca(
+					sizeof(struct authenticate_frame) + 34);
+	uint8_t extra[34];
+	size_t len;
 
-	l_put_le16(1, frame);
-	l_put_le16(MMPDU_STATUS_CODE_ANTI_CLOGGING_TOKEN_REQ, frame + 2);
-	l_put_le16(19, frame + 4);
-	memcpy(frame + 6, td->test_clogging_token, 32);
+	l_put_le16(19, extra);
+	memcpy(extra + 2, td->test_clogging_token, 32);
+
+	len = setup_auth_frame(frame, aa, 1,
+				MMPDU_STATUS_CODE_ANTI_CLOGGING_TOKEN_REQ,
+				extra, sizeof(extra));
 
 	td->test_anti_clogging = true;
 	td->commit_success = false;
 
-	sae_rx_packet(sm, aa, frame, 38);
+	assert(auth_proto_rx_authenticate(ap, (uint8_t *)frame, len) ==
+						-EAGAIN);
 
 	assert(td->commit_success == true);
 
 	test_destruct(td);
-	sae_sm_free(sm);
+	auth_proto_free(ap);
 }
 
 static void test_early_confirm(const void *arg)
 {
 	struct test_data *td = l_new(struct test_data, 1);
-	struct sae_sm *sm = test_initialize(td);
-
-	uint8_t frame[38];
+	struct auth_proto *ap = test_initialize(td);
 	uint8_t first_commit[102];
+	struct authenticate_frame *frame = alloca(
+					sizeof(struct authenticate_frame) + 32);
+	size_t len;
 
 	/* save the initial commit */
 	memcpy(first_commit, td->tx_packet, td->tx_packet_len);
 
-	l_put_u16(2, frame);
-	l_put_u16(0, frame + 2);
-
-	memset(frame + 4, 0xfe, 32);
+	len = setup_auth_frame(frame, aa, 2, 0, NULL, 32);
+	memset(frame->auth.ies, 0xfe, 32);
 
 	td->test_anti_clogging = false;
 
-	sae_rx_packet(sm, aa, frame, 36);
+	assert(auth_proto_rx_authenticate(ap, (uint8_t *)frame, len) ==
+						-EAGAIN);
 
 	/* verify earlier commit matched most recent */
 	assert(!memcmp(td->tx_packet, first_commit, td->tx_packet_len));
 
 	test_destruct(td);
-	sae_sm_free(sm);
+	auth_proto_free(ap);
 }
 
 static void test_reflection(const void *arg)
 {
 	struct test_data *td = l_new(struct test_data, 1);
-	struct sae_sm *sm = test_initialize(td);
+	struct auth_proto *ap = test_initialize(td);
 
-	td->tx_called = false;
+	td->tx_auth_called = false;
 	/* send reflect same commit */
-	sae_rx_packet(sm, aa, td->tx_packet, td->tx_packet_len);
+	ap->rx_authenticate(ap, td->tx_packet, td->tx_packet_len);
 
-	assert(td->tx_called == false);
+	assert(td->tx_auth_called == false);
 
 	test_destruct(td);
-	sae_sm_free(sm);
+	auth_proto_free(ap);
 }
 
 static void test_malformed_commit(const void *arg)
 {
 	struct test_data *td = l_new(struct test_data, 1);
-	struct sae_sm *sm = test_initialize(td);
+	struct auth_proto *ap = test_initialize(td);
+	struct authenticate_frame *frame = alloca(
+					sizeof(struct authenticate_frame) +
+					sizeof(aa_commit));
+	size_t len;
+
+	len = setup_auth_frame(frame, aa, 1, 0, aa_commit, sizeof(aa_commit));
 
 	/* dont send entire commit */
-	sae_rx_packet(sm, aa, aa_commit, sizeof(aa_commit) - 20);
-
-	assert(td->status != 0);
+	assert(auth_proto_rx_authenticate(ap, (uint8_t *)frame, len - 20) != 0);
 
 	test_destruct(td);
-	sae_sm_free(sm);
+	auth_proto_free(ap);
 }
 
 static void test_malformed_confirm(const void *arg)
 {
 	struct test_data *td = l_new(struct test_data, 1);
-	struct sae_sm *sm = test_initialize(td);
+	struct auth_proto *ap = test_initialize(td);
+	struct authenticate_frame *frame = alloca(
+					sizeof(struct authenticate_frame) +
+					sizeof(aa_commit));
+	size_t len;
 
-	sae_rx_packet(sm, aa, aa_commit, sizeof(aa_commit));
+	len = setup_auth_frame(frame, aa, 1, 0, aa_commit, sizeof(aa_commit));
+
+	assert(auth_proto_rx_authenticate(ap, (uint8_t *)frame, len) == 0);
 
 	assert(td->commit_success);
 
-	/* dont send entire confirm */
-	sae_rx_packet(sm, aa, aa_confirm, sizeof(aa_confirm) - 10);
+	frame = alloca(sizeof(struct authenticate_frame) + sizeof(aa_confirm));
+	len = setup_auth_frame(frame, aa, 2, 0, aa_confirm, sizeof(aa_confirm));
 
-	assert(td->status != 0);
+	/* dont send entire confirm */
+	assert(auth_proto_rx_authenticate(ap, (uint8_t *)frame, len - 10) != 0);
 
 	test_destruct(td);
-	sae_sm_free(sm);
+	auth_proto_free(ap);
 }
 
 static uint8_t aa_commit_bad_group[] = {
-	0x01, 0x00, 0x00, 0x00, 0xff, 0x00, 0x50, 0x5b, 0xb2, 0x1f, 0xaf, 0x7d,
+	0xff, 0x00, 0x50, 0x5b, 0xb2, 0x1f, 0xaf, 0x7d,
 	0xaf, 0x14, 0x7c, 0x7b, 0x19, 0xc9, 0x72, 0x82, 0xbc, 0x1a, 0xdb, 0xa1,
 	0xbd, 0x6e, 0x5a, 0xc7, 0x58, 0x0a, 0x65, 0x1f, 0xd2, 0xde, 0xb0, 0x66,
 	0xa5, 0xf9, 0x3e, 0x95, 0x4a, 0xe1, 0x83, 0xdb, 0x8a, 0xf5, 0x47, 0x8a,
@@ -351,36 +416,43 @@ static uint8_t aa_commit_bad_group[] = {
 static void test_bad_group(const void *arg)
 {
 	struct test_data *td = l_new(struct test_data, 1);
-	struct sae_sm *sm = test_initialize(td);
+	struct auth_proto *ap = test_initialize(td);
+	struct authenticate_frame *frame = alloca(
+					sizeof(struct authenticate_frame) +
+					sizeof(aa_commit_bad_group));
+	size_t len;
 
-	sae_rx_packet(sm, aa, aa_commit_bad_group, sizeof(aa_commit_bad_group));
+	len = setup_auth_frame(frame, aa, 1, 0, aa_commit_bad_group,
+				sizeof(aa_commit_bad_group));
+
+	assert(auth_proto_rx_authenticate(ap, (uint8_t *)frame, len) ==
+			MMPDU_STATUS_CODE_UNSUPP_FINITE_CYCLIC_GROUP);
 
 	assert(td->tx_reject_occurred);
-	assert(td->status == MMPDU_STATUS_CODE_UNSUPP_FINITE_CYCLIC_GROUP);
 
 	test_destruct(td);
-	sae_sm_free(sm);
+	auth_proto_free(ap);
 }
 
-static int end_to_end_tx_func(const uint8_t *dest, const uint8_t *frame,
-						size_t len, void *user_data)
+static void end_to_end_tx_func(const uint8_t *frame, size_t len, void *user_data)
 {
 	struct test_data *td = user_data;
 
 	memcpy(td->tx_packet, frame, len);
 	td->tx_packet_len = len;
-
-	return 0;
 }
 
 static void test_bad_confirm(const void *arg)
 {
-	struct sae_sm *sm1;
-	struct sae_sm *sm2;
+	struct auth_proto *ap1;
+	struct auth_proto *ap2;
 	struct test_data *td1 = l_new(struct test_data, 1);
 	struct test_data *td2 = l_new(struct test_data, 1);
 	struct handshake_state *hs1 = test_handshake_state_new(1);
 	struct handshake_state *hs2 = test_handshake_state_new(2);
+	struct authenticate_frame *frame = alloca(
+				sizeof(struct authenticate_frame) + 512);
+	size_t frame_len;
 	uint8_t tmp_commit[512];
 	size_t tmp_commit_len;
 
@@ -396,35 +468,48 @@ static void test_bad_confirm(const void *arg)
 	handshake_state_set_passphrase(hs2, passphrase);
 	handshake_state_set_authenticator(hs2, true);
 
-	sm1 = sae_sm_new(hs1, end_to_end_tx_func, test_complete_func, td1);
-	sm2 = sae_sm_new(hs2, end_to_end_tx_func, test_complete_func, td2);
+	ap1 = sae_sm_new(hs1, end_to_end_tx_func, test_tx_assoc_func, td1);
+	ap2 = sae_sm_new(hs2, end_to_end_tx_func, test_tx_assoc_func, td2);
 
 	/* both peers send out commit */
-	sae_start(sm1);
-	sae_start(sm2);
+	ap1->start(ap1);
+	ap2->start(ap2);
 
 	/* save sm1 commit, tx_packet will get overwritten with confirm */
 	memcpy(tmp_commit, td1->tx_packet, td1->tx_packet_len);
 	tmp_commit_len = td1->tx_packet_len;
 
+	/* Setup whole frame */
+	frame_len = setup_auth_frame(frame, aa, 1, 0, td2->tx_packet + 4,
+					td2->tx_packet_len - 4);
+
 	/* rx commit for both peers */
-	sae_rx_packet(sm1, aa, td2->tx_packet, td2->tx_packet_len);
-	sae_rx_packet(sm2, spa, tmp_commit, tmp_commit_len);
+	ap1->rx_authenticate(ap1, (uint8_t *) frame, frame_len);
+
+	frame_len = setup_auth_frame(frame, spa, 1, 0, tmp_commit + 4,
+					tmp_commit_len - 4);
+	ap2->rx_authenticate(ap2, (uint8_t *)frame, frame_len);
 	/* both peers should now have sent confirm */
 
 	/* rx confirm for both peers */
-	sae_rx_packet(sm1, aa, td2->tx_packet, td2->tx_packet_len);
+	frame_len = setup_auth_frame(frame, aa, 2, 0, td2->tx_packet + 4,
+					td2->tx_packet_len - 4);
+	ap1->rx_authenticate(ap1, (uint8_t *)frame, frame_len);
+
 	/* muck with a byte in the confirm */
 	td1->tx_packet[10] = ~td1->tx_packet[10];
-	sae_rx_packet(sm2, spa, td1->tx_packet, td1->tx_packet_len);
+	frame_len = setup_auth_frame(frame, spa, 2, 0, td1->tx_packet + 4,
+					td1->tx_packet_len - 4);
+	ap2->rx_authenticate(ap2, (uint8_t *)frame, frame_len);
 
-	assert(td1->status == 0);
+	assert(td1->tx_assoc_called);
 	assert(td2->status != 0);
 
 	handshake_state_free(hs1);
 	handshake_state_free(hs2);
-	sae_sm_free(sm1);
-	sae_sm_free(sm2);
+
+	ap1->free(ap1);
+	ap2->free(ap2);
 
 	/* sm2 gets freed by sae since it failed */
 	l_free(td1);
@@ -433,12 +518,16 @@ static void test_bad_confirm(const void *arg)
 
 static void test_confirm_after_accept(const void *arg)
 {
-	struct sae_sm *sm1;
-	struct sae_sm *sm2;
+	struct auth_proto *ap1;
+	struct auth_proto *ap2;
 	struct test_data *td1 = l_new(struct test_data, 1);
 	struct test_data *td2 = l_new(struct test_data, 1);
 	struct handshake_state *hs1 = test_handshake_state_new(1);
 	struct handshake_state *hs2 = test_handshake_state_new(2);
+	struct authenticate_frame *frame = alloca(
+				sizeof(struct authenticate_frame) + 512);
+	struct associate_frame *assoc = alloca(sizeof(struct associate_frame));
+	size_t frame_len;
 	uint8_t tmp_commit[512];
 	size_t tmp_commit_len;
 
@@ -454,54 +543,85 @@ static void test_confirm_after_accept(const void *arg)
 	handshake_state_set_passphrase(hs2, passphrase);
 	handshake_state_set_authenticator(hs2, true);
 
-	sm1 = sae_sm_new(hs1, end_to_end_tx_func, test_complete_func, td1);
-	sm2 = sae_sm_new(hs2, end_to_end_tx_func, test_complete_func, td2);
+	ap1 = sae_sm_new(hs1, end_to_end_tx_func, test_tx_assoc_func, td1);
+	ap2 = sae_sm_new(hs2, end_to_end_tx_func, test_tx_assoc_func, td2);
 
 	/* both peers send out commit */
-	sae_start(sm1);
-	sae_start(sm2);
+	auth_proto_start(ap1);
+	auth_proto_start(ap2);
 
 	/* save sm1 commit, tx_packet will get overwritten with confirm */
 	memcpy(tmp_commit, td1->tx_packet, td1->tx_packet_len);
 	tmp_commit_len = td1->tx_packet_len;
 
 	/* rx commit for both peers */
-	sae_rx_packet(sm1, aa, td2->tx_packet, td2->tx_packet_len);
-	sae_rx_packet(sm2, spa, tmp_commit, tmp_commit_len);
+	frame_len = setup_auth_frame(frame, aa, 1, 0, td2->tx_packet + 4,
+					td2->tx_packet_len - 4);
+	assert(auth_proto_rx_authenticate(ap1, (uint8_t *)frame,
+						frame_len) == 0);
+
+	frame_len = setup_auth_frame(frame, spa, 1, 0, tmp_commit + 4,
+					tmp_commit_len - 4);
+	assert(auth_proto_rx_authenticate(ap2, (uint8_t *)frame,
+						frame_len) == 0);
 	/* both peers should now have sent confirm */
 
 	/* rx confirm for one peer, sm1 should accept confirm */
-	sae_rx_packet(sm1, aa, td2->tx_packet, td2->tx_packet_len);
-	assert(td1->status == 0);
+	frame_len = setup_auth_frame(frame, aa, 2, 0, td2->tx_packet + 4,
+					td2->tx_packet_len - 4);
+	assert(auth_proto_rx_authenticate(ap1, (uint8_t *)frame,
+						frame_len) == 0);
+
+	assert(td1->tx_assoc_called);
 
 	/* simulate sm2 not receiving confirm and resending its confirm */
-	sae_timeout(sm2);
-	sae_rx_packet(sm1, aa, td2->tx_packet, td2->tx_packet_len);
+	ap2->auth_timeout(ap2);
+	frame_len = setup_auth_frame(frame, aa, 2, 0, td2->tx_packet + 4,
+					td2->tx_packet_len - 4);
+	assert(auth_proto_rx_authenticate(ap1, (uint8_t *)frame,
+						frame_len) == 0);
 
 	/* sc should be set to 0xffff */
 	assert(l_get_u16(td1->tx_packet + 4) == 0xffff);
 	/* sm1 should respond with a new confirm, and accept */
-	sae_rx_packet(sm2, spa, td1->tx_packet, td1->tx_packet_len);
+	frame_len = setup_auth_frame(frame, spa, 2, 0, td1->tx_packet + 4,
+					td1->tx_packet_len - 4);
+	assert(auth_proto_rx_authenticate(ap2, (uint8_t *)frame,
+						frame_len) == 0);
 
-	assert(td1->status == 0);
-	assert(td2->status == 0);
+	assert(td1->tx_assoc_called);
+
+	frame_len = setup_assoc_frame(assoc, 0);
+
+	/*
+	 * This is just to complete the connection, SAE just verifies status
+	 * so the same frame can be used for both SMs
+	 */
+	assert(auth_proto_rx_associate(ap1, (uint8_t *)assoc, frame_len) == 0);
+	assert(auth_proto_rx_associate(ap2, (uint8_t *)assoc, frame_len) == 0);
 
 	handshake_state_free(hs1);
 	handshake_state_free(hs2);
-	sae_sm_free(sm1);
-	sae_sm_free(sm2);
+
+	auth_proto_free(ap1);
+	auth_proto_free(ap2);
+
 	l_free(td1);
 	l_free(td2);
 }
 
 static void test_end_to_end(const void *arg)
 {
-	struct sae_sm *sm1;
-	struct sae_sm *sm2;
+	struct auth_proto *ap1;
+	struct auth_proto *ap2;
 	struct test_data *td1 = l_new(struct test_data, 1);
 	struct test_data *td2 = l_new(struct test_data, 1);
 	struct handshake_state *hs1 = test_handshake_state_new(1);
 	struct handshake_state *hs2 = test_handshake_state_new(2);
+	struct authenticate_frame *frame = alloca(
+				sizeof(struct authenticate_frame) + 512);
+	struct associate_frame *assoc = alloca(sizeof(struct associate_frame));
+	size_t frame_len;
 	uint8_t tmp_commit[512];
 	size_t tmp_commit_len;
 
@@ -517,33 +637,53 @@ static void test_end_to_end(const void *arg)
 	handshake_state_set_passphrase(hs2, passphrase);
 	handshake_state_set_authenticator(hs2, true);
 
-	sm1 = sae_sm_new(hs1, end_to_end_tx_func, test_complete_func, td1);
-	sm2 = sae_sm_new(hs2, end_to_end_tx_func, test_complete_func, td2);
+	ap1 = sae_sm_new(hs1, end_to_end_tx_func, test_tx_assoc_func, td1);
+	ap2 = sae_sm_new(hs2, end_to_end_tx_func, test_tx_assoc_func, td2);
 
 	/* both peers send out commit */
-	sae_start(sm1);
-	sae_start(sm2);
+	auth_proto_start(ap1);
+	auth_proto_start(ap2);
 
 	/* save sm1 commit, tx_packet will get overwritten with confirm */
 	memcpy(tmp_commit, td1->tx_packet, td1->tx_packet_len);
 	tmp_commit_len = td1->tx_packet_len;
 
 	/* rx commit for both peers */
-	sae_rx_packet(sm1, aa, td2->tx_packet, td2->tx_packet_len);
-	sae_rx_packet(sm2, spa, tmp_commit, tmp_commit_len);
+	frame_len = setup_auth_frame(frame, aa, 1, 0, td2->tx_packet + 4,
+					td2->tx_packet_len - 4);
+	assert(auth_proto_rx_authenticate(ap1, (uint8_t *)frame,
+						frame_len) == 0);
+
 	/* both peers should now have sent confirm */
+	frame_len = setup_auth_frame(frame, spa, 1, 0, tmp_commit + 4,
+					tmp_commit_len - 4);
+	assert(auth_proto_rx_authenticate(ap2, (uint8_t *)frame,
+						frame_len) == 0);
 
 	/* rx confirm for both peers */
-	sae_rx_packet(sm1, aa, td2->tx_packet, td2->tx_packet_len);
-	sae_rx_packet(sm2, spa, td1->tx_packet, td1->tx_packet_len);
+	frame_len = setup_auth_frame(frame, aa, 2, 0, td2->tx_packet + 4,
+					td2->tx_packet_len - 4);
+	assert(auth_proto_rx_authenticate(ap1, (uint8_t *)frame,
+						frame_len) == 0);
 
-	assert(td1->status == 0);
-	assert(td2->status == 0);
+	frame_len = setup_auth_frame(frame, spa, 2, 0, td1->tx_packet + 4,
+					td1->tx_packet_len - 4);
+	assert(auth_proto_rx_authenticate(ap2, (uint8_t *)frame,
+						frame_len) == 0);
+
+	assert(td1->tx_assoc_called);
+	assert(td2->tx_assoc_called);
+
+	frame_len = setup_assoc_frame(assoc, 0);
+	assert(auth_proto_rx_associate(ap1, (uint8_t *)assoc, frame_len) == 0);
+	assert(auth_proto_rx_associate(ap2, (uint8_t *)assoc, frame_len) == 0);
 
 	handshake_state_free(hs1);
 	handshake_state_free(hs2);
-	sae_sm_free(sm1);
-	sae_sm_free(sm2);
+
+	auth_proto_free(ap1);
+	auth_proto_free(ap2);
+
 	l_free(td1);
 	l_free(td2);
 }