unit: fix test-sae after auth_proto changes

The SAE unit test needed to be updated to use the handshake_driver,
but in addition all the packet building needed a major overhaul. SAE
was changed to behave more like OWE/FILS, in that netdev passes the
raw mpdu frame into the RX callbacks. Before, only the authentication
data was passed. This requires the unit tests to now build up the
entire authentication frame, and in some cases append the header
to the data coming from the TX functions.

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)

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)
+static void test_tx_auth_func(const uint8_t *frame, size_t len, 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)
 {
 	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(!auth_proto_auth_timeout(ap));
-
-	assert(td->status != 0);
 
 	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(!auth_proto_auth_timeout(ap));
-
-	assert(td->status != 0);
 
 	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(19, extra);
-	l_put_le16(MMPDU_STATUS_CODE_ANTI_CLOGGING_TOKEN_REQ, frame + 2);
+	memcpy(extra + 2, td->test_clogging_token, 32);
-	l_put_le16(19, frame + 4);
+
-	memcpy(frame + 6, 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);
+	struct auth_proto *ap = test_initialize(td);
-
-	uint8_t frame[38];
 	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);
+	len = setup_auth_frame(frame, aa, 2, 0, NULL, 32);
-	l_put_u16(0, frame + 2);
+	memset(frame->auth.ies, 0xfe, 32);
-
-	memset(frame + 4, 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(auth_proto_rx_authenticate(ap, (uint8_t *)frame, len - 20) != 0);
-
-	assert(td->status != 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 */
+	frame = alloca(sizeof(struct authenticate_frame) + sizeof(aa_confirm));
-	sae_rx_packet(sm, aa, aa_confirm, sizeof(aa_confirm) - 10);
+	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,
+static void end_to_end_tx_func(const uint8_t *frame, size_t len, void *user_data)
-						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 auth_proto *ap1;
-	struct sae_sm *sm2;
+	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);
+	ap1 = sae_sm_new(hs1, end_to_end_tx_func, test_tx_assoc_func, td1);
-	sm2 = sae_sm_new(hs2, end_to_end_tx_func, test_complete_func, td2);
+	ap2 = sae_sm_new(hs2, end_to_end_tx_func, test_tx_assoc_func, td2);
 
 	/* both peers send out commit */
-	sae_start(sm1);
+	ap1->start(ap1);
-	sae_start(sm2);
+	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);
+	ap1->rx_authenticate(ap1, (uint8_t *) frame, frame_len);
-	sae_rx_packet(sm2, spa, tmp_commit, tmp_commit_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 auth_proto *ap1;
-	struct sae_sm *sm2;
+	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);
+	ap1 = sae_sm_new(hs1, end_to_end_tx_func, test_tx_assoc_func, td1);
-	sm2 = sae_sm_new(hs2, end_to_end_tx_func, test_complete_func, td2);
+	ap2 = sae_sm_new(hs2, end_to_end_tx_func, test_tx_assoc_func, td2);
 
 	/* both peers send out commit */
-	sae_start(sm1);
+	auth_proto_start(ap1);
-	sae_start(sm2);
+	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);
+	frame_len = setup_auth_frame(frame, aa, 1, 0, td2->tx_packet + 4,
-	sae_rx_packet(sm2, spa, tmp_commit, tmp_commit_len);
+					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);
+	frame_len = setup_auth_frame(frame, aa, 2, 0, td2->tx_packet + 4,
-	assert(td1->status == 0);
+					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);
+	ap2->auth_timeout(ap2);
-	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);
+	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(td1->tx_assoc_called);
-	assert(td2->status == 0);
+
+	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 auth_proto *ap1;
-	struct sae_sm *sm2;
+	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);
+	ap1 = sae_sm_new(hs1, end_to_end_tx_func, test_tx_assoc_func, td1);
-	sm2 = sae_sm_new(hs2, end_to_end_tx_func, test_complete_func, td2);
+	ap2 = sae_sm_new(hs2, end_to_end_tx_func, test_tx_assoc_func, td2);
 
 	/* both peers send out commit */
-	sae_start(sm1);
+	auth_proto_start(ap1);
-	sae_start(sm2);
+	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);
+	frame_len = setup_auth_frame(frame, aa, 1, 0, td2->tx_packet + 4,
-	sae_rx_packet(sm2, spa, tmp_commit, tmp_commit_len);
+					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);
+	frame_len = setup_auth_frame(frame, aa, 2, 0, td2->tx_packet + 4,
-	sae_rx_packet(sm2, spa, td1->tx_packet, td1->tx_packet_len);
+					td2->tx_packet_len - 4);
+	assert(auth_proto_rx_authenticate(ap1, (uint8_t *)frame,
+						frame_len) == 0);
 
-	assert(td1->status == 0);
+	frame_len = setup_auth_frame(frame, spa, 2, 0, td1->tx_packet + 4,
-	assert(td2->status == 0);
+					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);
 }