Directly use peer when receiving on associated sockets

2 files changed, 15 insertions, 7 deletions

diff --git a/src/fastd.c b/src/fastd.c
index 1d74cb3..f121e49 100644
--- a/src/fastd.c
+++ b/src/fastd.c
@@ -256,7 +256,7 @@ static void bind_sockets(fastd_context *ctx) {
 	}
 }
 
-fastd_socket* fastd_socket_open(fastd_context *ctx, const fastd_peer *peer, int af) {
+fastd_socket* fastd_socket_open(fastd_context *ctx, fastd_peer *peer, int af) {
 	const fastd_bind_address any_address = { .addr.sa.sa_family = af };
 
 	int fd = bind_socket(ctx, &any_address, true);
@@ -697,10 +697,18 @@ static void handle_socket(fastd_context *ctx, fastd_socket *sock) {
 
 	fastd_buffer_push_head(&buffer, 1);
 
-	fastd_peer *peer;
-	for (peer = ctx->peers; peer; peer = peer->next) {
-		if (fastd_peer_address_equal(&peer->address, &recvaddr))
-			break;
+	fastd_peer *peer = NULL;
+
+	if (sock->peer) {
+		if (fastd_peer_address_equal(&sock->peer->address, &recvaddr)) {
+			peer = sock->peer;
+		}
+	}
+	else {
+		for (peer = ctx->peers; peer; peer = peer->next) {
+			if (fastd_peer_address_equal(&peer->address, &recvaddr))
+				break;
+		}
 	}
 
 	if (peer) {
diff --git a/src/fastd.h b/src/fastd.h
index 1395df8..9779de8 100644
--- a/src/fastd.h
+++ b/src/fastd.h
@@ -141,7 +141,7 @@ struct _fastd_bind_address {
 struct _fastd_socket {
 	int fd;
 	const fastd_bind_address *addr;
-	const fastd_peer *peer;
+	fastd_peer *peer;
 };
 
 struct _fastd_config {
@@ -269,7 +269,7 @@ void fastd_send(fastd_context *ctx, const fastd_socket *sock, const fastd_peer_a
 void fastd_send_handshake(fastd_context *ctx, const fastd_socket *sock, const fastd_peer_address *address, fastd_buffer buffer);
 void fastd_handle_receive(fastd_context *ctx, fastd_peer *peer, fastd_buffer buffer);
 
-fastd_socket* fastd_socket_open(fastd_context *ctx, const fastd_peer *peer, int af);
+fastd_socket* fastd_socket_open(fastd_context *ctx, fastd_peer *peer, int af);
 
 void fastd_resolve_peer(fastd_context *ctx, fastd_peer *peer);