u80211: sync with upstream and use new work system

kernel-src/io/net/u80211/include/u80211/kernel_interface.h

@@ -67,19 +67,30 @@ void u80211_kernel_release_rwlock_shared(void *rwlock);
 
 typedef void (*u80211_kernel_work_fn_t)(void *context);
 
+// allocates an opaque timer object
+// return NULL on allocation failure
+void *u80211_kernel_allocate_timer(void);
+
+// cancels and frees the timer object referenced by 'timer'
+// 'timer' will never be NULL
+void u80211_kernel_free_timer(void *timer);
+
 // allocates an opaque work object
 // return NULL on allocation failure
 void *u80211_kernel_allocate_work(void);
 
 // enqueue threaded work. this might be called from an interrupt context
-// 'ms' == 0 enqueues the work immediatelly. otherwise, 
-// work is enqueued 'ms' milliseconds in the future
 // if work is already pending, the new request must be ignored
-void u80211_kernel_enqueue_work(void *work, u80211_kernel_work_fn_t function, void *context, size_t ms);
+void u80211_kernel_enqueue_work(void *work, u80211_kernel_work_fn_t function, void *context);
+
+// enqueue threaded work 'ms' milliseconds in the future
+// 'timer' must not be used by another delayed work request
+// if work is already pending, the new request must be ignored
+void u80211_kernel_enqueue_delayed_work(void *work, void *timer, u80211_kernel_work_fn_t function, void *context, size_t ms);
 
 // frees the work object referenced by 'work'
 // 'work' will never be NULL
-// may be called by the work's callback
+// must not be called by the work's callback
 void u80211_kernel_free_work(void *work);
 
 // kernel packet receive callback

kernel-src/io/net/u80211/include/u80211/u80211.h

@@ -53,6 +53,7 @@ struct u80211_device {
 
 	void *scan_spinlock;
 	void *scan_context;
+	void *scan_cleanup_work;
 	u80211_list_t scan_waiters;
 
 	bss_cache_t bss_cache;
@@ -60,6 +61,8 @@ struct u80211_device {
 	void *association_spinlock;
 	void *association_context;
 	void *association_cleanup_work;
+	void *association_cleanup_context;
+	bool association_cleanup_pending;
 	u80211_list_t association_waiters;
 	unsigned int association_generation;
 	int association_result;

kernel-src/io/net/u80211/kernel_interface.c

@@ -1,26 +1,24 @@
 #include <u80211/kernel_interface.h>
 
+#include <errno.h>
 #include <kernel/alloc.h>
 #include <kernel/eth.h>
-#include <kernel/event.h>
 #include <kernel/init.h>
 #include <kernel/interrupt.h>
-#include <kernel/page.h>
+#include <kernel/itimer.h>
-#include <kernel/scheduler.h>
 #include <kernel/slab.h>
-#include <kernel/thread.h>
+#include <kernel/work.h>
-#include <kernel/timekeeper.h>
 #include <mutex.h>
 #include <pushlock.h>
 #include <semaphore.h>
 #include <spinlock.h>
 #include <string.h>
-#include <time.h>
 
 static scache_t *pushlock_cache;
 static scache_t *semaphore_cache;
 static scache_t *spinlock_cache;
 static scache_t *work_cache;
+static scache_t *timer_cache;
 
 void *u80211_kernel_allocate(size_t size) {
 	return alloc(size);
@@ -126,198 +124,162 @@ void u80211_kernel_release_rwlock_shared(void *rwlock) {
 	pushlock_release_shared(rwlock);
 }
 
-typedef struct u80211_work_t {
+typedef struct u80211_timer u80211_timer_t;
-	struct u80211_work_t *next;
+
-	timespec_t deadline;
+typedef struct u80211_work {
+	work_t work;
+	u80211_timer_t *timer;
 	u80211_kernel_work_fn_t function;
 	void *context;
-	eventheader_t finished_event;
 	bool pending;
-	bool running;
 	bool freeing;
-	bool destroy_on_return;
 } u80211_work_t;
 
-static spinlock_t work_queue_lock;
+struct u80211_timer {
-static eventheader_t work_queue_event;
+	itimer_t timer;
-static u80211_work_t *work_queue;
+	u80211_work_t *work;
-static u80211_work_t *current_work;
+};
-static thread_t *worker_thread;
 
-static long acquire_work_queue_lock(void) {
+static work_queue_t *u80211_work_queue;
-	return spinlock_acquire_raise_ipl(&work_queue_lock, IPL_NET);
+static spinlock_t work_state_lock;
+
+static long acquire_work_state_lock(void) {
+	return spinlock_acquire_raise_ipl(&work_state_lock, IPL_NET);
 }
 
-static void release_work_queue_lock(long old_ipl) {
+static void release_work_state_lock(long old_ipl) {
-	spinlock_release_lower_ipl(&work_queue_lock, old_ipl);
+	spinlock_release_lower_ipl(&work_state_lock, old_ipl);
 }
 
-static timespec_t deadline_after_ms(size_t ms) {
+static void dispatch_work(void *context, size_t) {
-	timespec_t delay = {
+	u80211_work_t *work = context;
-		.s = ms / 1000,
+	u80211_kernel_work_fn_t function;
-		.ns = (ms % 1000) * 1000000,
+	void *function_context;
-	};
 
-	return timespec_add(timekeeper_timefromboot(), delay);
+	long old_ipl = acquire_work_state_lock();
-}
-
-static void insert_work(u80211_work_t *work) {
-	// insert in the ordered list
-	u80211_work_t **entry = &work_queue;
-	while (*entry && !timespec_bigger((*entry)->deadline, work->deadline))
-		entry = &(*entry)->next;
-
-	work->next = *entry;
-	*entry = work;
-	work->pending = true;
-}
-
-static bool remove_work(u80211_work_t *work) {
-	// find previous in the ordered list
-	u80211_work_t **entry = &work_queue;
-	while (*entry && *entry != work)
-		entry = &(*entry)->next;
-
-	if (*entry == NULL)
-		return false;
-
-	*entry = work->next;
-	work->next = NULL;
 	work->pending = false;
-	return true;
+	if (work->freeing) {
-}
+		release_work_state_lock(old_ipl);
-
+		return;
-static time_t deadline_wait_us(timespec_t deadline, timespec_t now) {
-	return (deadline.s - now.s) * 1000000 + (deadline.ns - now.ns) / 1000;
-}
-
-static void wake_work_thread(void) {
-	EVENT_SIGNAL(&work_queue_event);
-}
-
-static void work_thread(void) {
-	for (;;) {
-		eventlistener_t listener;
-		EVENT_INITLISTENER(&listener);
-		EVENT_ATTACH(&listener, &work_queue_event);
-
-		long old_ipl = acquire_work_queue_lock();
-		u80211_work_t *work = work_queue;
-		if (work == NULL) {
-			release_work_queue_lock(old_ipl);
-			EVENT_WAIT(&listener, 0);
-			EVENT_DETACHALL(&listener);
-			continue;
-		}
-
-		timespec_t now = timekeeper_timefromboot();
-		if (timespec_bigger(work->deadline, now)) {
-			time_t wait_us = deadline_wait_us(work->deadline, now);
-			release_work_queue_lock(old_ipl);
-			EVENT_WAIT(&listener, wait_us);
-			EVENT_DETACHALL(&listener);
-			continue;
-		}
-
-		work_queue = work->next;
-
-		work->next = NULL;
-		work->pending = false;
-		work->running = true;
-
-		current_work = work;
-
-		u80211_kernel_work_fn_t function = work->function;
-		void *context = work->context;
-
-		release_work_queue_lock(old_ipl);
-		EVENT_DETACHALL(&listener);
-
-		function(context);
-
-		old_ipl = acquire_work_queue_lock();
-		current_work = NULL;
-		work->running = false;
-		bool destroy = work->freeing && work->destroy_on_return;
-		bool notify = work->freeing && !work->destroy_on_return;
-		release_work_queue_lock(old_ipl);
-
-		if (notify)
-			EVENT_SIGNAL(&work->finished_event);
-		else if (destroy)
-			slab_free(work_cache, work);
 	}
+
+	function = work->function;
+	function_context = work->context;
+	release_work_state_lock(old_ipl);
+
+	function(function_context);
+}
+
+static void delayed_work_timer(context_t *, dpcarg_t argument) {
+	u80211_timer_t *timer = argument;
+
+	long old_ipl = acquire_work_state_lock();
+
+	u80211_work_t *work = timer->work;
+	if (work) {
+		timer->work = NULL;
+		work->timer = NULL;
+		if (work->pending && !work->freeing)
+			work_enqueue(u80211_work_queue, &work->work);
+		else
+			work->pending = false;
+	}
+
+	release_work_state_lock(old_ipl);
+}
+
+void *u80211_kernel_allocate_timer(void) {
+	u80211_timer_t *timer = slab_allocate(timer_cache);
+	if (timer) {
+		timer->work = NULL;
+		itimer_init(&timer->timer, delayed_work_timer, timer);
+	}
+
+	return timer;
+}
+
+void u80211_kernel_free_timer(void *opaque_timer) {
+	u80211_timer_t *timer = opaque_timer;
+	itimer_pause(&timer->timer, NULL, NULL);
+
+	long old_ipl = acquire_work_state_lock();
+
+	u80211_work_t *work = timer->work;
+	if (work && work->timer == timer) {
+		work->timer = NULL;
+		work->pending = false;
+	}
+
+	timer->work = NULL;
+
+	release_work_state_lock(old_ipl);
+
+	slab_free(timer_cache, timer);
 }
 
 void *u80211_kernel_allocate_work(void) {
 	u80211_work_t *work = slab_allocate(work_cache);
 	if (work) {
 		memset(work, 0, sizeof(*work));
-		EVENT_INITHEADER(&work->finished_event);
+		WORK_INIT(&work->work, dispatch_work, work);
 	}
 
 	return work;
 }
 
-void u80211_kernel_enqueue_work(void *opaque_work, u80211_kernel_work_fn_t function, void *context, size_t ms) {
+void u80211_kernel_enqueue_work(void *opaque_work, u80211_kernel_work_fn_t function, void *context) {
 	u80211_work_t *work = opaque_work;
-	timespec_t deadline = deadline_after_ms(ms);
 
-	long old_ipl = acquire_work_queue_lock();
+	long old_ipl = acquire_work_state_lock();
+
 	if (!work->pending && !work->freeing) {
-		work->deadline = deadline;
 		work->function = function;
 		work->context = context;
+		work->pending = true;
+		work_enqueue(u80211_work_queue, &work->work);
+	}
 
-		insert_work(work);
+	release_work_state_lock(old_ipl);
-		release_work_queue_lock(old_ipl);
+}
 
-		wake_work_thread();
+void u80211_kernel_enqueue_delayed_work(void *opaque_work, void *opaque_timer, u80211_kernel_work_fn_t function, void *context, size_t ms) {
+	if (ms == 0) {
+		u80211_kernel_enqueue_work(opaque_work, function, context);
 		return;
 	}
 
-	release_work_queue_lock(old_ipl);
+	u80211_work_t *work = opaque_work;
+	u80211_timer_t *timer = opaque_timer;
+
+	long old_ipl = acquire_work_state_lock();
+
+	if (!work->pending && !work->freeing && timer->work == NULL) {
+		work->function = function;
+		work->context = context;
+		work->pending = true;
+		work->timer = timer;
+		timer->work = work;
+
+		itimer_set(&timer->timer, ms * 1000, 0);
+		itimer_resume(&timer->timer);
+	}
+
+	release_work_state_lock(old_ipl);
 }
 
 void u80211_kernel_free_work(void *opaque_work) {
 	u80211_work_t *work = opaque_work;
-	bool removed;
-	eventlistener_t listener;
-	EVENT_INITLISTENER(&listener);
-	EVENT_ATTACH(&listener, &work->finished_event);
 
-	long old_ipl = acquire_work_queue_lock();
+	long old_ipl = acquire_work_state_lock();
-	removed = remove_work(work);
-	if (!work->running) {
-		release_work_queue_lock(old_ipl);
-		EVENT_DETACHALL(&listener);
-
-		if (removed)
-			wake_work_thread();
-
-		slab_free(work_cache, work);
-		return;
-	}
 
 	work->freeing = true;
-	if (current_thread() == worker_thread && current_work == work) {
+	work->pending = false;
-		work->destroy_on_return = true;
-		release_work_queue_lock(old_ipl);
-		EVENT_DETACHALL(&listener);
 
-		if (removed)
+	release_work_state_lock(old_ipl);
-			wake_work_thread();
 
-		return;
+	while (work_dequeue(u80211_work_queue, &work->work) == EBUSY)
-	}
+		work_wait(u80211_work_queue, &work->work);
 
-	release_work_queue_lock(old_ipl);
-
-	if (removed)
-		wake_work_thread();
-
-	EVENT_WAIT(&listener, 0);
-	EVENT_DETACHALL(&listener);
 	slab_free(work_cache, work);
 }
 
@@ -326,17 +288,15 @@ static void u80211_kernel_interface_init(void) {
 	semaphore_cache = slab_newcache(sizeof(semaphore_t), 0, NULL, NULL);
 	spinlock_cache = slab_newcache(sizeof(u80211_spinlock_t), 0, NULL, NULL);
 	work_cache = slab_newcache(sizeof(u80211_work_t), 0, NULL, NULL);
-	__assert(pushlock_cache && semaphore_cache && spinlock_cache && work_cache);
+	timer_cache = slab_newcache(sizeof(u80211_timer_t), 0, NULL, NULL);
+	__assert(pushlock_cache && semaphore_cache && spinlock_cache && work_cache && timer_cache);
 
-	SPINLOCK_INIT(work_queue_lock);
+	SPINLOCK_INIT(work_state_lock);
-	EVENT_INITHEADER(&work_queue_event);
+	u80211_work_queue = work_queue_create("u80211", 1, IPL_NET);
-
+	__assert(u80211_work_queue);
-	worker_thread = sched_newthread(work_thread, PAGE_SIZE * 4, 0, NULL, NULL);
-	__assert(worker_thread);
-	sched_queue(worker_thread);
 }
 
-INIT_ROUTINE_DEFINE(u80211, INIT_ROUTINE_FLAGS_NONE, u80211_kernel_interface_init, scheduler);
+INIT_ROUTINE_DEFINE(u80211, INIT_ROUTINE_FLAGS_NONE, u80211_kernel_interface_init, work_queue);
 
 void u80211_kernel_receive_callback(u80211_device_t *device, void *buffer, size_t size) {
 	(void)size;

kernel-src/io/net/u80211/src/association.c

@@ -12,9 +12,10 @@
 typedef struct {
 	void *completion_work;
 	void *auth_timeout_work;
+	void *auth_timeout_timer;
 	void *assoc_timeout_work;
+	void *assoc_timeout_timer;
 	unsigned int generation;
-	unsigned int timeouts_in_flight;
 	u80211_device_t *device;
 	u80211_ap_t *ap; // this keeps a reference, the reference in the device is only incremented only when fully associated
 	void *information_elements;
@@ -37,6 +38,8 @@ static void wake_association_waiters(u80211_device_t *device, int result) {
 }
 
 static void destroy_association_context(u80211_association_context_t *association_context) {
+	u80211_kernel_free_timer(association_context->auth_timeout_timer);
+	u80211_kernel_free_timer(association_context->assoc_timeout_timer);
 	u80211_kernel_free_work(association_context->completion_work);
 	u80211_kernel_free_work(association_context->auth_timeout_work);
 	u80211_kernel_free_work(association_context->assoc_timeout_work);
@@ -74,6 +77,37 @@ static void release_disconnected_ap(void *ctx) {
 		u80211_ap_release(ap);
 }
 
+static void cleanup_association(void *ctx) {
+	u80211_device_t *device = ctx;
+
+	u80211_kernel_acquire_spinlock(device->association_spinlock);
+	u80211_association_context_t *association_context = device->association_cleanup_context;
+	u80211_ap_t *disconnected_ap = device->disconnected_ap;
+	device->disconnected_ap = NULL;
+	u80211_kernel_release_spinlock(device->association_spinlock);
+
+	if (association_context != NULL) {
+		u80211_ap_release(association_context->ap);
+		destroy_association_context(association_context);
+	}
+	if (disconnected_ap != NULL)
+		u80211_ap_release(disconnected_ap);
+
+	if (association_context != NULL) {
+		u80211_kernel_acquire_spinlock(device->association_spinlock);
+		device->association_cleanup_context = NULL;
+		device->association_cleanup_pending = false;
+		wake_association_waiters(device, device->association_result);
+		u80211_kernel_release_spinlock(device->association_spinlock);
+	}
+}
+
+static void schedule_association_cleanup(u80211_device_t *device, u80211_association_context_t *association_context) {
+	device->association_context = NULL;
+	device->association_cleanup_context = association_context;
+	device->association_cleanup_pending = true;
+}
+
 static void association_process_teardown(u80211_device_t *device, const u80211_mac_address_t *address, bool deauthentication) {
 	bool cleanup_needed = false;
 
@@ -91,9 +125,12 @@ static void association_process_teardown(u80211_device_t *device, const u80211_m
 		goto leave;
 
 	if (attempt_in_progress) {
-		device->association_context = NULL;
+		u80211_association_context_t *association_context = device->association_context;
+		if (association_context == NULL)
+			goto leave;
+		schedule_association_cleanup(device, association_context);
 		device->association_result = U80211_STATUS_REJECTED;
-		wake_association_waiters(device, device->association_result);
+		cleanup_needed = true;
 	} else {
 		device->disconnected_ap = device->ap;
 		cleanup_needed = true;
@@ -105,7 +142,7 @@ static void association_process_teardown(u80211_device_t *device, const u80211_m
 leave:
 	u80211_kernel_release_spinlock(device->association_spinlock);
 	if (cleanup_needed)
-		u80211_kernel_enqueue_work(device->association_cleanup_work, release_disconnected_ap, device, 0);
+		u80211_kernel_enqueue_work(device->association_cleanup_work, cleanup_association, device);
 }
 
 void u80211_association_process_deauthentication(u80211_device_t *device, u80211_deauthentication_data_t *deauthentication_data) {
@@ -117,6 +154,7 @@ void u80211_association_process_disassociation(u80211_device_t *device, u80211_d
 }
 
 void u80211_association_process_response(u80211_device_t *device, u80211_association_response_data_t *association_data) {
+	bool cleanup_needed = false;
 	u80211_kernel_acquire_spinlock(device->association_spinlock);
 
 	u80211_association_context_t *association_context = device->association_context;
@@ -135,43 +173,41 @@ void u80211_association_process_response(u80211_device_t *device, u80211_associa
 	else
 		u80211_ap_hold(device->ap);
 
-	device->association_context = NULL;
+	schedule_association_cleanup(device, association_context);
 	++device->association_generation;
 	device->association_result = association_data->status ? U80211_STATUS_REJECTED : U80211_STATUS_SUCCESS;
-	wake_association_waiters(device, device->association_result);
+	cleanup_needed = true;
 
 leave:
 	u80211_kernel_release_spinlock(device->association_spinlock);
+	if (cleanup_needed)
+		u80211_kernel_enqueue_work(device->association_cleanup_work, cleanup_association, device);
 }
 
 static void association_timeout(void *ctx, int expected_state) {
 	u80211_association_context_t *association_context = ctx;
 	u80211_device_t *device = association_context->device;
-	bool free = true;
+	bool cleanup_needed = false;
 
 	u80211_kernel_acquire_spinlock(device->association_spinlock);
 
 	// not associating/different association generation
-	if (!device->ap || device->association_generation != association_context->generation)
+	if (device->association_context != association_context || !device->ap ||
+			device->association_generation != association_context->generation)
 		goto leave;
 
 	if (u80211_set_device_state(device, expected_state, U80211_DEVICE_STATE_DOWN)) {
 		device->ap = NULL;
-		device->association_context = NULL;
+		schedule_association_cleanup(device, association_context);
 		++device->association_generation;
 		device->association_result = U80211_STATUS_TIMED_OUT;
-		wake_association_waiters(device, device->association_result);
+		cleanup_needed = true;
-	} else {
-		// already progressed
-		free = false;
 	}
 
 leave:
 	u80211_kernel_release_spinlock(device->association_spinlock);
-	if (__atomic_sub_fetch(&association_context->timeouts_in_flight, 1, __ATOMIC_ACQ_REL) == 0 && free) {
+	if (cleanup_needed)
-		u80211_ap_release(association_context->ap);
+		u80211_kernel_enqueue_work(device->association_cleanup_work, cleanup_association, device);
-		destroy_association_context(association_context);
-	}
 }
 
 static void assoc_timeout(void *ctx) {
@@ -183,11 +219,12 @@ static void auth_completion_work(void *ctx) {
 	u80211_device_t *device = association_context->device;
 
 	u80211_send_association_request(device, association_context->information_elements, association_context->information_elements_size);
-	__atomic_add_fetch(&association_context->timeouts_in_flight, 1, __ATOMIC_RELAXED);
+	u80211_kernel_enqueue_delayed_work(association_context->assoc_timeout_work, association_context->assoc_timeout_timer,
-	u80211_kernel_enqueue_work(association_context->assoc_timeout_work, assoc_timeout, association_context, ASSOC_TIMEOUT);
+		assoc_timeout, association_context, ASSOC_TIMEOUT);
 }
 
 void u80211_association_process_authentication(u80211_device_t *device, u80211_auth_data_t *auth_data) {
+	bool cleanup_needed = false;
 	u80211_kernel_acquire_spinlock(device->association_spinlock);
 
 	// not associating/different AP
@@ -203,21 +240,24 @@ void u80211_association_process_authentication(u80211_device_t *device, u80211_a
 		if (!u80211_set_device_state(device, U80211_DEVICE_STATE_AUTHENTICATING, U80211_DEVICE_STATE_DOWN))
 			goto leave;
 
+		u80211_association_context_t *association_context = device->association_context;
 		device->ap = NULL;
-		device->association_context = NULL;
+		schedule_association_cleanup(device, association_context);
 		++device->association_generation;
 		device->association_result = U80211_STATUS_REJECTED;
-		wake_association_waiters(device, device->association_result);
+		cleanup_needed = true;
 	} else {
 		if (!u80211_set_device_state(device, U80211_DEVICE_STATE_AUTHENTICATING, U80211_DEVICE_STATE_ASSOCIATING))
 			goto leave;
 
 		u80211_association_context_t *ctx = device->association_context;
-		u80211_kernel_enqueue_work(ctx->completion_work, auth_completion_work, ctx, 0);
+		u80211_kernel_enqueue_work(ctx->completion_work, auth_completion_work, ctx);
 	}
 
 leave:
 	u80211_kernel_release_spinlock(device->association_spinlock);
+	if (cleanup_needed)
+		u80211_kernel_enqueue_work(device->association_cleanup_work, cleanup_association, device);
 }
 
 static void auth_timeout(void *ctx) {
@@ -247,8 +287,27 @@ int u80211_associate(u80211_device_t *device, u80211_ap_t *ap, const void *infor
 		return U80211_STATUS_ENOMEM;
 	}
 
+	association_context->auth_timeout_timer = u80211_kernel_allocate_timer();
+	if (association_context->auth_timeout_timer == NULL) {
+		u80211_kernel_free_work(association_context->auth_timeout_work);
+		u80211_kernel_free_work(association_context->completion_work);
+		u80211_kernel_free(association_context);
+		return U80211_STATUS_ENOMEM;
+	}
+
 	association_context->assoc_timeout_work = u80211_kernel_allocate_work();
 	if (association_context->assoc_timeout_work == NULL) {
+		u80211_kernel_free_timer(association_context->auth_timeout_timer);
+		u80211_kernel_free_work(association_context->auth_timeout_work);
+		u80211_kernel_free_work(association_context->completion_work);
+		u80211_kernel_free(association_context);
+		return U80211_STATUS_ENOMEM;
+	}
+
+	association_context->assoc_timeout_timer = u80211_kernel_allocate_timer();
+	if (association_context->assoc_timeout_timer == NULL) {
+		u80211_kernel_free_work(association_context->assoc_timeout_work);
+		u80211_kernel_free_timer(association_context->auth_timeout_timer);
 		u80211_kernel_free_work(association_context->auth_timeout_work);
 		u80211_kernel_free_work(association_context->completion_work);
 		u80211_kernel_free(association_context);
@@ -269,10 +328,10 @@ int u80211_associate(u80211_device_t *device, u80211_ap_t *ap, const void *infor
 
 	association_context->device = device;
 	association_context->ap = ap;
-	__atomic_store_n(&association_context->timeouts_in_flight, 1, __ATOMIC_RELAXED);
 
 	u80211_kernel_acquire_spinlock(device->association_spinlock);
-	if (!u80211_set_device_state(device, U80211_DEVICE_STATE_DOWN, U80211_DEVICE_STATE_AUTHENTICATING)) {
+	if (device->association_cleanup_pending ||
+			!u80211_set_device_state(device, U80211_DEVICE_STATE_DOWN, U80211_DEVICE_STATE_AUTHENTICATING)) {
 		u80211_kernel_release_spinlock(device->association_spinlock);
 		destroy_association_context(association_context);
 		return U80211_STATUS_BUSY;
@@ -296,7 +355,8 @@ int u80211_associate(u80211_device_t *device, u80211_ap_t *ap, const void *infor
 	device->ops->set_channel(device, ap->channel);
 	u80211_send_authentication(device, &auth_data);
 
-	u80211_kernel_enqueue_work(association_context->auth_timeout_work, auth_timeout, association_context, AUTH_TIMEOUT);
+	u80211_kernel_enqueue_delayed_work(association_context->auth_timeout_work, association_context->auth_timeout_timer,
+		auth_timeout, association_context, AUTH_TIMEOUT);
 
 	return 0;
 }
@@ -356,7 +416,8 @@ int u80211_wait_for_association_completion(u80211_device_t *device) {
 
 	u80211_kernel_acquire_spinlock(device->association_spinlock);
 	int state = u80211_get_device_state(device);
-	if (state != U80211_DEVICE_STATE_AUTHENTICATING && state != U80211_DEVICE_STATE_ASSOCIATING) {
+	if (state != U80211_DEVICE_STATE_AUTHENTICATING && state != U80211_DEVICE_STATE_ASSOCIATING &&
+			!device->association_cleanup_pending) {
 		int result = device->association_result;
 		u80211_kernel_release_spinlock(device->association_spinlock);
 		u80211_kernel_free_semaphore(semaphore);

kernel-src/io/net/u80211/src/device.c

@@ -29,10 +29,18 @@ int u80211_register_device(const u80211_device_metadata_t *metadata, const u8021
 		return U80211_STATUS_ENOMEM;
 	}
 	device->scan_context = NULL;
+	device->scan_cleanup_work = u80211_kernel_allocate_work();
+	if (device->scan_cleanup_work == NULL) {
+		u80211_kernel_free_spinlock(device->scan_spinlock);
+		u80211_key_state_deinit(device);
+		u80211_kernel_free(device);
+		return U80211_STATUS_ENOMEM;
+	}
 	u80211_list_init(&device->scan_waiters);
 
 	device->association_spinlock = u80211_kernel_allocate_spinlock();
 	if (device->association_spinlock == NULL) {
+		u80211_kernel_free_work(device->scan_cleanup_work);
 		u80211_kernel_free_spinlock(device->scan_spinlock);
 		u80211_key_state_deinit(device);
 		u80211_kernel_free(device);
@@ -41,12 +49,15 @@ int u80211_register_device(const u80211_device_metadata_t *metadata, const u8021
 	device->association_cleanup_work = u80211_kernel_allocate_work();
 	if (device->association_cleanup_work == NULL) {
 		u80211_kernel_free_spinlock(device->association_spinlock);
+		u80211_kernel_free_work(device->scan_cleanup_work);
 		u80211_kernel_free_spinlock(device->scan_spinlock);
 		u80211_key_state_deinit(device);
 		u80211_kernel_free(device);
 		return U80211_STATUS_ENOMEM;
 	}
 	device->association_context = NULL;
+	device->association_cleanup_context = NULL;
+	device->association_cleanup_pending = false;
 	u80211_list_init(&device->association_waiters);
 	device->association_generation = 0;
 	device->association_result = U80211_STATUS_UNKNOWN_ERROR;
@@ -60,6 +71,7 @@ int u80211_register_device(const u80211_device_metadata_t *metadata, const u8021
 	if (status != U80211_STATUS_SUCCESS) {
 		u80211_kernel_free_work(device->association_cleanup_work);
 		u80211_kernel_free_spinlock(device->association_spinlock);
+		u80211_kernel_free_work(device->scan_cleanup_work);
 		u80211_kernel_free_spinlock(device->scan_spinlock);
 		u80211_key_state_deinit(device);
 		u80211_kernel_free(device);
@@ -72,6 +84,7 @@ int u80211_register_device(const u80211_device_metadata_t *metadata, const u8021
 
 void u80211_unregister_device(u80211_device_t *device) {
 	u80211_key_state_deinit(device);
+	u80211_kernel_free_work(device->scan_cleanup_work);
 	u80211_kernel_free_work(device->association_cleanup_work);
 	if (device->disconnected_ap != NULL)
 		u80211_ap_release(device->disconnected_ap);

kernel-src/io/net/u80211/src/scan.c

@@ -22,6 +22,7 @@ typedef struct {
 	bool receiving;
 	int current_channel;
 	void *work;
+	void *timer;
 } u80211_scan_state_t;
 
 typedef struct {
@@ -38,22 +39,34 @@ static void wake_scan_waiters(u80211_device_t *device) {
 }
 
 static void destroy_scan_state(u80211_scan_state_t *scan_state) {
+	u80211_kernel_free_timer(scan_state->timer);
 	u80211_kernel_free_work(scan_state->work);
 	u80211_ringbuffer_destroy(&scan_state->buffer);
 	u80211_kernel_free(scan_state);
 }
 
+static void cleanup_scan(void *ctx);
+
 static void complete_scan(u80211_scan_state_t *scan_state) {
 	u80211_device_t *device = scan_state->device;
 
 	u80211_kernel_acquire_spinlock(device->scan_spinlock);
-
+	device->scan_context = NULL;
-	u80211_set_device_state(device, U80211_DEVICE_STATE_SCANNING, U80211_DEVICE_STATE_DOWN);
-	wake_scan_waiters(device);
-
 	u80211_kernel_release_spinlock(device->scan_spinlock);
 
+	u80211_kernel_enqueue_work(device->scan_cleanup_work, cleanup_scan, scan_state);
+}
+
+static void cleanup_scan(void *ctx) {
+	u80211_scan_state_t *scan_state = ctx;
+	u80211_device_t *device = scan_state->device;
+
 	destroy_scan_state(scan_state);
+
+	u80211_kernel_acquire_spinlock(device->scan_spinlock);
+	u80211_set_device_state(device, U80211_DEVICE_STATE_SCANNING, U80211_DEVICE_STATE_DOWN);
+	wake_scan_waiters(device);
+	u80211_kernel_release_spinlock(device->scan_spinlock);
 }
 
 void u80211_scan_process_response(u80211_device_t *device, u80211_beacon_data_t *beacon_data) {
@@ -151,7 +164,7 @@ static void scan_work(void *ctx) {
 
 	u80211_send_probe_request(device);
 
-	u80211_kernel_enqueue_work(scan_state->work, scan_work, scan_state, WAIT_MS);
+	u80211_kernel_enqueue_delayed_work(scan_state->work, scan_state->timer, scan_work, scan_state, WAIT_MS);
 }
 
 int u80211_scan(u80211_device_t *device) {
@@ -172,6 +185,14 @@ int u80211_scan(u80211_device_t *device) {
 		return U80211_STATUS_ENOMEM;
 	}
 
+	scan_state->timer = u80211_kernel_allocate_timer();
+	if (scan_state->timer == NULL) {
+		u80211_kernel_free_work(scan_state->work);
+		u80211_ringbuffer_destroy(&scan_state->buffer);
+		u80211_kernel_free(scan_state);
+		return U80211_STATUS_ENOMEM;
+	}
+
 	scan_state->receiving = false;
 	scan_state->device = device;
 	scan_state->current_channel = next_channel(0);
@@ -199,7 +220,7 @@ int u80211_scan(u80211_device_t *device) {
 
 	u80211_send_probe_request(device);
 
-	u80211_kernel_enqueue_work(scan_state->work, scan_work, scan_state, WAIT_MS);
+	u80211_kernel_enqueue_delayed_work(scan_state->work, scan_state->timer, scan_work, scan_state, WAIT_MS);
 
 	return U80211_STATUS_SUCCESS;
 }