socket: keep track of src address in sending packets

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

11 files changed, 135 insertions, 122 deletions

diff --git a/src/config.c b/src/config.c
index a98ecd9..829c0aa 100644
--- a/src/config.c
+++ b/src/config.c
@@ -76,7 +76,7 @@ static int set_peer(struct wireguard_device *wg, void __user *user_peer, size_t
 		if (netdev_pub(wg)->flags & IFF_UP)
 			timers_init_peer(peer);
 	} else
-		pr_debug("Peer %Lu (%pISpfsc) modified\n", peer->internal_id, &peer->endpoint_addr);
+		pr_debug("Peer %Lu (%pISpfsc) modified\n", peer->internal_id, &peer->endpoint.addr_storage);
 
 	if (in_peer.remove_me) {
 		peer_put(peer);
@@ -84,8 +84,10 @@ static int set_peer(struct wireguard_device *wg, void __user *user_peer, size_t
 		return 0;
 	}
 
-	if (in_peer.endpoint.ss_family == AF_INET || in_peer.endpoint.ss_family == AF_INET6)
-		socket_set_peer_addr(peer, &in_peer.endpoint);
+	if (in_peer.endpoint.ss_family == AF_INET || in_peer.endpoint.ss_family == AF_INET6) {
+		struct endpoint endpoint = { .addr_storage = in_peer.endpoint };
+		socket_set_peer_endpoint(peer, &endpoint);
+	}
 
 	if (in_peer.replace_ipmasks)
 		routing_table_remove_by_peer(&wg->peer_routing_table, peer);
@@ -235,7 +237,7 @@ static int populate_peer(struct wireguard_peer *peer, void *ctx)
 
 	memcpy(out_peer.public_key, peer->handshake.remote_static, NOISE_PUBLIC_KEY_LEN);
 	read_lock_bh(&peer->endpoint_lock);
-	out_peer.endpoint = peer->endpoint_addr;
+	out_peer.endpoint = peer->endpoint.addr_storage;
 	read_unlock_bh(&peer->endpoint_lock);
 	out_peer.last_handshake_time = peer->walltime_last_handshake;
 	out_peer.tx_bytes = peer->tx_bytes;
diff --git a/src/data.c b/src/data.c
index f53bbc6..b6d20b8 100644
--- a/src/data.c
+++ b/src/data.c
@@ -34,9 +34,9 @@ struct encryption_ctx {
 struct decryption_ctx {
 	struct padata_priv padata;
 	struct sk_buff *skb;
-	void (*callback)(struct sk_buff *skb, struct wireguard_peer *, struct sockaddr_storage *, bool used_new_key, int err);
+	void (*callback)(struct sk_buff *skb, struct wireguard_peer *, struct endpoint *, bool used_new_key, int err);
 	struct noise_keypair *keypair;
-	struct sockaddr_storage addr;
+	struct endpoint endpoint;
 	uint64_t nonce;
 	uint8_t num_frags;
 	int ret;
@@ -372,7 +372,7 @@ static void finish_decrypt_packet(struct decryption_ctx *ctx)
 	}
 
 	noise_keypair_put(ctx->keypair);
-	ctx->callback(ctx->skb, ctx->keypair->entry.peer, &ctx->addr, used_new_key, 0);
+	ctx->callback(ctx->skb, ctx->keypair->entry.peer, &ctx->endpoint, used_new_key, 0);
 	return;
 
 err:
@@ -403,10 +403,10 @@ static inline int start_decryption(struct padata_instance *padata, struct padata
 }
 #endif
 
-void packet_consume_data(struct sk_buff *skb, size_t offset, struct wireguard_device *wg, void(*callback)(struct sk_buff *skb, struct wireguard_peer *, struct sockaddr_storage *, bool used_new_key, int err))
+void packet_consume_data(struct sk_buff *skb, size_t offset, struct wireguard_device *wg, void(*callback)(struct sk_buff *skb, struct wireguard_peer *, struct endpoint *, bool used_new_key, int err))
 {
 	int ret;
-	struct sockaddr_storage addr = { 0 };
+	struct endpoint endpoint;
 	unsigned int num_frags;
 	struct sk_buff *trailer;
 	struct message_data *header;
@@ -414,7 +414,7 @@ void packet_consume_data(struct sk_buff *skb, size_t offset, struct wireguard_de
 	uint64_t nonce;
 	__le32 idx;
 
-	ret = socket_addr_from_skb(&addr, skb);
+	ret = socket_endpoint_from_skb(&endpoint, skb);
 	if (unlikely(ret < 0))
 		goto err;
 
@@ -459,7 +459,7 @@ void packet_consume_data(struct sk_buff *skb, size_t offset, struct wireguard_de
 		ctx->callback = callback;
 		ctx->nonce = nonce;
 		ctx->num_frags = num_frags;
-		ctx->addr = addr;
+		ctx->endpoint = endpoint;
 		ret = start_decryption(wg->parallel_receive, &ctx->padata, cpu);
 		if (unlikely(ret)) {
 			kmem_cache_free(decryption_ctx_cache, ctx);
@@ -474,7 +474,7 @@ void packet_consume_data(struct sk_buff *skb, size_t offset, struct wireguard_de
 			.callback = callback,
 			.nonce = nonce,
 			.num_frags = num_frags,
-			.addr = addr
+			.endpoint = endpoint
 		};
 		begin_decrypt_packet(&ctx);
 		finish_decrypt_packet(&ctx);
diff --git a/src/device.c b/src/device.c
index 46d6704..8774d80 100644
--- a/src/device.c
+++ b/src/device.c
@@ -112,16 +112,16 @@ static netdev_tx_t xmit(struct sk_buff *skb, struct net_device *dev)
 	peer = routing_table_lookup_dst(&wg->peer_routing_table, skb);
 	if (unlikely(!peer)) {
 #if defined(CONFIG_DYNAMIC_DEBUG) || defined(DEBUG)
-		struct sockaddr_storage addr;
-		socket_addr_from_skb(&addr, skb);
-		net_dbg_ratelimited("No peer is configured for %pISc\n", &addr);
+		struct endpoint endpoint;
+		socket_endpoint_from_skb(&endpoint, skb);
+		net_dbg_ratelimited("No peer is configured for %pISc\n", &endpoint.addr_storage);
 #endif
 		skb_unsendable(skb, dev);
 		return -ENOKEY;
 	}
 
 	read_lock_bh(&peer->endpoint_lock);
-	ret = peer->endpoint_addr.ss_family != AF_INET && peer->endpoint_addr.ss_family != AF_INET6;
+	ret = peer->endpoint.addr_storage.ss_family != AF_INET && peer->endpoint.addr_storage.ss_family != AF_INET6;
 	read_unlock_bh(&peer->endpoint_lock);
 	if (unlikely(ret)) {
 		net_dbg_ratelimited("No valid endpoint has been configured or discovered for device\n");
diff --git a/src/packets.h b/src/packets.h
index c491640..22ce8fa 100644
--- a/src/packets.h
+++ b/src/packets.h
@@ -37,7 +37,7 @@ void packet_send_handshake_cookie(struct wireguard_device *wg, struct sk_buff *i
 
 /* data.c */
 int packet_create_data(struct sk_buff_head *queue, struct wireguard_peer *peer, void(*callback)(struct sk_buff_head *, struct wireguard_peer *));
-void packet_consume_data(struct sk_buff *skb, size_t offset, struct wireguard_device *wg, void(*callback)(struct sk_buff *, struct wireguard_peer *, struct sockaddr_storage *, bool used_new_key, int err));
+void packet_consume_data(struct sk_buff *skb, size_t offset, struct wireguard_device *wg, void(*callback)(struct sk_buff *, struct wireguard_peer *, struct endpoint *, bool, int));
 
 #ifdef CONFIG_WIREGUARD_PARALLEL
 int packet_init_data_caches(void);
diff --git a/src/peer.c b/src/peer.c
index e9162a5..f50a9e8 100644
--- a/src/peer.c
+++ b/src/peer.c
@@ -85,7 +85,7 @@ void peer_remove(struct wireguard_peer *peer)
 static void rcu_release(struct rcu_head *rcu)
 {
 	struct wireguard_peer *peer = container_of(rcu, struct wireguard_peer, rcu);
-	pr_debug("Peer %Lu (%pISpfsc) destroyed\n", peer->internal_id, &peer->endpoint_addr);
+	pr_debug("Peer %Lu (%pISpfsc) destroyed\n", peer->internal_id, &peer->endpoint.addr_storage);
 	skb_queue_purge(&peer->tx_packet_queue);
 	dst_cache_destroy(&peer->endpoint_cache);
 	kzfree(peer);
diff --git a/src/peer.h b/src/peer.h
index 895eb77..07c4d16 100644
--- a/src/peer.h
+++ b/src/peer.h
@@ -16,9 +16,21 @@
 
 struct wireguard_device;
 
+struct endpoint {
+	union {
+		struct sockaddr_storage addr_storage;
+		struct sockaddr_in addr4;
+		struct sockaddr_in6 addr6;
+	};
+	union {
+		struct in_addr src4;
+		struct in6_addr src6;
+	};
+};
+
 struct wireguard_peer {
 	struct wireguard_device *device;
-	struct sockaddr_storage endpoint_addr;
+	struct endpoint endpoint;
 	struct dst_cache endpoint_cache;
 	rwlock_t endpoint_lock;
 	struct noise_handshake handshake;
diff --git a/src/receive.c b/src/receive.c
index 9f5c3d1..33c8ebe 100644
--- a/src/receive.c
+++ b/src/receive.c
@@ -25,19 +25,17 @@ static inline void rx_stats(struct wireguard_peer *peer, size_t len)
 
 static inline void update_latest_addr(struct wireguard_peer *peer, struct sk_buff *skb)
 {
-	struct sockaddr_storage addr = { 0 };
-	if (!socket_addr_from_skb(&addr, skb))
-		socket_set_peer_addr(peer, &addr);
+	struct endpoint endpoint;
+	if (!socket_endpoint_from_skb(&endpoint, skb))
+		socket_set_peer_endpoint(peer, &endpoint);
 }
 
 static inline int skb_data_offset(struct sk_buff *skb, size_t *data_offset, size_t *data_len)
 {
 	struct udphdr *udp;
+	struct endpoint endpoint;
 #if defined(CONFIG_DYNAMIC_DEBUG) || defined(DEBUG)
-	struct sockaddr_storage addr = { 0 };
-	socket_addr_from_skb(&addr, skb);
-#else
-	static const u8 addr;
+	socket_endpoint_from_skb(&endpoint, skb);
 #endif
 
 	if (unlikely(skb->len < sizeof(struct iphdr)))
@@ -50,26 +48,26 @@ static inline int skb_data_offset(struct sk_buff *skb, size_t *data_offset, size
 	udp = udp_hdr(skb);
 	*data_offset = (u8 *)udp - skb->data;
 	if (unlikely(*data_offset > U16_MAX)) {
-		net_dbg_ratelimited("Packet has offset at impossible location from %pISpfsc\n", &addr);
+		net_dbg_ratelimited("Packet has offset at impossible location from %pISpfsc\n", &endpoint.addr_storage);
 		return -EINVAL;
 	}
 	if (unlikely(*data_offset + sizeof(struct udphdr) > skb->len)) {
-		net_dbg_ratelimited("Packet isn't big enough to have UDP fields from %pISpfsc\n", &addr);
+		net_dbg_ratelimited("Packet isn't big enough to have UDP fields from %pISpfsc\n", &endpoint.addr_storage);
 		return -EINVAL;
 	}
 	*data_len = ntohs(udp->len);
 	if (unlikely(*data_len < sizeof(struct udphdr))) {
-		net_dbg_ratelimited("UDP packet is reporting too small of a size from %pISpfsc\n", &addr);
+		net_dbg_ratelimited("UDP packet is reporting too small of a size from %pISpfsc\n", &endpoint.addr_storage);
 		return -EINVAL;
 	}
 	if (unlikely(*data_len > skb->len - *data_offset)) {
-		net_dbg_ratelimited("UDP packet is lying about its size from %pISpfsc\n", &addr);
+		net_dbg_ratelimited("UDP packet is lying about its size from %pISpfsc\n", &endpoint.addr_storage);
 		return -EINVAL;
 	}
 	*data_len -= sizeof(struct udphdr);
 	*data_offset = (u8 *)udp + sizeof(struct udphdr) - skb->data;
 	if (!pskb_may_pull(skb, *data_offset + sizeof(struct message_header))) {
-		net_dbg_ratelimited("Could not pull header into data section from %pISpfsc\n", &addr);
+		net_dbg_ratelimited("Could not pull header into data section from %pISpfsc\n", &endpoint.addr_storage);
 		return -EINVAL;
 	}
 
@@ -83,18 +81,16 @@ static void receive_handshake_packet(struct wireguard_device *wg, void *data, si
 	bool under_load;
 	enum cookie_mac_state mac_state;
 	bool packet_needs_cookie;
+	struct endpoint endpoint;
 
 #if defined(CONFIG_DYNAMIC_DEBUG) || defined(DEBUG)
-	struct sockaddr_storage addr = { 0 };
-	socket_addr_from_skb(&addr, skb);
-#else
-	static const u8 addr;
+	socket_endpoint_from_skb(&endpoint, skb);
 #endif
 
 	message_type = message_determine_type(data, len);
 
 	if (message_type == MESSAGE_HANDSHAKE_COOKIE) {
-		net_dbg_ratelimited("Receiving cookie response from %pISpfsc\n", &addr);
+		net_dbg_ratelimited("Receiving cookie response from %pISpfsc\n", &endpoint.addr_storage);
 		cookie_message_consume(data, wg);
 		return;
 	}
@@ -106,7 +102,7 @@ static void receive_handshake_packet(struct wireguard_device *wg, void *data, si
 	else if (under_load && mac_state == VALID_MAC_BUT_NO_COOKIE)
 		packet_needs_cookie = true;
 	else {
-		net_dbg_ratelimited("Invalid MAC of handshake, dropping packet from %pISpfsc\n", &addr);
+		net_dbg_ratelimited("Invalid MAC of handshake, dropping packet from %pISpfsc\n", &endpoint.addr_storage);
 		return;
 	}
 
@@ -119,10 +115,10 @@ static void receive_handshake_packet(struct wireguard_device *wg, void *data, si
 		}
 		peer = noise_handshake_consume_initiation(data, wg);
 		if (unlikely(!peer)) {
-			net_dbg_ratelimited("Invalid handshake initiation from %pISpfsc\n", &addr);
+			net_dbg_ratelimited("Invalid handshake initiation from %pISpfsc\n", &endpoint.addr_storage);
 			return;
 		}
-		net_dbg_ratelimited("Receiving handshake initiation from peer %Lu (%pISpfsc)\n", peer->internal_id, &addr);
+		net_dbg_ratelimited("Receiving handshake initiation from peer %Lu (%pISpfsc)\n", peer->internal_id, &endpoint.addr_storage);
 		update_latest_addr(peer, skb);
 		packet_send_handshake_response(peer);
 		break;
@@ -134,10 +130,10 @@ static void receive_handshake_packet(struct wireguard_device *wg, void *data, si
 		}
 		peer = noise_handshake_consume_response(data, wg);
 		if (unlikely(!peer)) {
-			net_dbg_ratelimited("Invalid handshake response from %pISpfsc\n", &addr);
+			net_dbg_ratelimited("Invalid handshake response from %pISpfsc\n", &endpoint.addr_storage);
 			return;
 		}
-		net_dbg_ratelimited("Receiving handshake response from peer %Lu (%pISpfsc)\n", peer->internal_id, &addr);
+		net_dbg_ratelimited("Receiving handshake response from peer %Lu (%pISpfsc)\n", peer->internal_id, &endpoint.addr_storage);
 		if (noise_handshake_begin_session(&peer->handshake, &peer->keypairs, true)) {
 			timers_ephemeral_key_created(peer);
 			timers_handshake_complete(peer);
@@ -150,7 +146,7 @@ static void receive_handshake_packet(struct wireguard_device *wg, void *data, si
 		}
 		break;
 	default:
-		net_err_ratelimited("Somehow a wrong type of packet wound up in the handshake queue from %pISpfsc!\n", &addr);
+		net_err_ratelimited("Somehow a wrong type of packet wound up in the handshake queue from %pISpfsc!\n", &endpoint.addr_storage);
 		BUG();
 		return;
 	}
@@ -207,13 +203,13 @@ struct packet_cb {
 };
 #define PACKET_CB(skb) ((struct packet_cb *)skb->cb)
 
-static void receive_data_packet(struct sk_buff *skb, struct wireguard_peer *peer, struct sockaddr_storage *addr, bool used_new_key, int err)
+static void receive_data_packet(struct sk_buff *skb, struct wireguard_peer *peer, struct endpoint *endpoint, bool used_new_key, int err)
 {
 	struct net_device *dev;
 	struct wireguard_peer *routed_peer;
 	struct wireguard_device *wg;
 
-	if (unlikely(err < 0 || !peer || !addr)) {
+	if (unlikely(err < 0 || !peer || !endpoint)) {
 		dev_kfree_skb(skb);
 		return;
 	}
@@ -230,21 +226,21 @@ static void receive_data_packet(struct sk_buff *skb, struct wireguard_peer *peer
 
 	/* A packet with length 0 is a keepalive packet */
 	if (unlikely(!skb->len)) {
-		net_dbg_ratelimited("Receiving keepalive packet from peer %Lu (%pISpfsc)\n", peer->internal_id, addr);
+		net_dbg_ratelimited("Receiving keepalive packet from peer %Lu (%pISpfsc)\n", peer->internal_id, &endpoint->addr_storage);
 		goto packet_processed;
 	}
 
 	if (unlikely(skb->len < sizeof(struct iphdr))) {
 		++dev->stats.rx_errors;
 		++dev->stats.rx_length_errors;
-		net_dbg_ratelimited("Packet missing ip header from peer %Lu (%pISpfsc)\n", peer->internal_id, addr);
+		net_dbg_ratelimited("Packet missing ip header from peer %Lu (%pISpfsc)\n", peer->internal_id, &endpoint->addr_storage);
 		goto packet_processed;
 	}
 
 	if (!pskb_may_pull(skb, 1 /* For checking the ip version below */)) {
 		++dev->stats.rx_errors;
 		++dev->stats.rx_length_errors;
-		net_dbg_ratelimited("Packet missing IP version from peer %Lu (%pISpfsc)\n", peer->internal_id, addr);
+		net_dbg_ratelimited("Packet missing IP version from peer %Lu (%pISpfsc)\n", peer->internal_id, &endpoint->addr_storage);
 		goto packet_processed;
 	}
 
@@ -258,7 +254,7 @@ static void receive_data_packet(struct sk_buff *skb, struct wireguard_peer *peer
 		if (unlikely(skb->len < sizeof(struct ipv6hdr))) {
 			++dev->stats.rx_errors;
 			++dev->stats.rx_length_errors;
-			net_dbg_ratelimited("Packet missing ipv6 header from peer %Lu (%pISpfsc)\n", peer->internal_id, addr);
+			net_dbg_ratelimited("Packet missing ipv6 header from peer %Lu (%pISpfsc)\n", peer->internal_id, &endpoint->addr_storage);
 			goto packet_processed;
 		}
 		skb->protocol = htons(ETH_P_IPV6);
@@ -267,7 +263,7 @@ static void receive_data_packet(struct sk_buff *skb, struct wireguard_peer *peer
 	} else {
 		++dev->stats.rx_errors;
 		++dev->stats.rx_length_errors;
-		net_dbg_ratelimited("Packet neither ipv4 nor ipv6 from peer %Lu (%pISpfsc)\n", peer->internal_id, addr);
+		net_dbg_ratelimited("Packet neither ipv4 nor ipv6 from peer %Lu (%pISpfsc)\n", peer->internal_id, &endpoint->addr_storage);
 		goto packet_processed;
 	}
 
@@ -278,9 +274,9 @@ static void receive_data_packet(struct sk_buff *skb, struct wireguard_peer *peer
 
 	if (unlikely(routed_peer != peer)) {
 #if defined(CONFIG_DYNAMIC_DEBUG) || defined(DEBUG)
-		struct sockaddr_storage unencrypted_addr = { 0 };
-		socket_addr_from_skb(&unencrypted_addr, skb);
-		net_dbg_ratelimited("Packet has unallowed src IP (%pISc) from peer %Lu (%pISpfsc)\n", &unencrypted_addr, peer->internal_id, addr);
+		struct endpoint unencrypted_endpoint;
+		socket_endpoint_from_skb(&unencrypted_endpoint, skb);
+		net_dbg_ratelimited("Packet has unallowed src IP (%pISc) from peer %Lu (%pISpfsc)\n", &unencrypted_endpoint.addr_storage, peer->internal_id, &endpoint->addr_storage);
 #endif
 		++dev->stats.rx_errors;
 		++dev->stats.rx_frame_errors;
@@ -292,7 +288,7 @@ static void receive_data_packet(struct sk_buff *skb, struct wireguard_peer *peer
 		rx_stats(peer, skb->len);
 	else {
 		++dev->stats.rx_dropped;
-		net_dbg_ratelimited("Failed to give packet to userspace from peer %Lu (%pISpfsc)\n", peer->internal_id, addr);
+		net_dbg_ratelimited("Failed to give packet to userspace from peer %Lu (%pISpfsc)\n", peer->internal_id, &endpoint->addr_storage);
 	}
 	goto continue_processing;
 
@@ -301,18 +297,16 @@ packet_processed:
 continue_processing:
 	timers_any_authenticated_packet_received(peer);
 	timers_any_authenticated_packet_traversal(peer);
-	socket_set_peer_addr(peer, addr);
+	socket_set_peer_endpoint(peer, endpoint);
 	peer_put(peer);
 }
 
 void packet_receive(struct wireguard_device *wg, struct sk_buff *skb)
 {
 	size_t len, offset;
+	struct endpoint endpoint;
 #if defined(CONFIG_DYNAMIC_DEBUG) || defined(DEBUG)
-	struct sockaddr_storage addr = { 0 };
-	socket_addr_from_skb(&addr, skb);
-#else
-	static const u8 addr;
+	socket_endpoint_from_skb(&endpoint, skb);
 #endif
 
 	if (unlikely(skb_data_offset(skb, &offset, &len) < 0))
@@ -322,11 +316,11 @@ void packet_receive(struct wireguard_device *wg, struct sk_buff *skb)
 	case MESSAGE_HANDSHAKE_RESPONSE:
 	case MESSAGE_HANDSHAKE_COOKIE:
 		if (skb_queue_len(&wg->incoming_handshakes) > MAX_QUEUED_HANDSHAKES) {
-			net_dbg_ratelimited("Too many handshakes queued, dropping packet from %pISpfsc\n", &addr);
+			net_dbg_ratelimited("Too many handshakes queued, dropping packet from %pISpfsc\n", &endpoint.addr_storage);
 			goto err;
 		}
 		if (skb_linearize(skb) < 0) {
-			net_dbg_ratelimited("Unable to linearize handshake skb from %pISpfsc\n", &addr);
+			net_dbg_ratelimited("Unable to linearize handshake skb from %pISpfsc\n", &endpoint.addr_storage);
 			goto err;
 		}
 		skb_queue_tail(&wg->incoming_handshakes, skb);
@@ -338,7 +332,7 @@ void packet_receive(struct wireguard_device *wg, struct sk_buff *skb)
 		packet_consume_data(skb, offset, wg, receive_data_packet);
 		break;
 	default:
-		net_dbg_ratelimited("Invalid packet from %pISpfsc\n", &addr);
+		net_dbg_ratelimited("Invalid packet from %pISpfsc\n", &endpoint.addr_storage);
 		goto err;
 	}
 	return;
diff --git a/src/send.c b/src/send.c
index cd67068..78119f0 100644
--- a/src/send.c
+++ b/src/send.c
@@ -27,7 +27,7 @@ static void packet_send_handshake_initiation(struct wireguard_peer *peer)
 	peer->last_sent_handshake = get_jiffies_64();
 	up_write(&peer->handshake.lock);
 
-	net_dbg_ratelimited("Sending handshake initiation to peer %Lu (%pISpfsc)\n", peer->internal_id, &peer->endpoint_addr);
+	net_dbg_ratelimited("Sending handshake initiation to peer %Lu (%pISpfsc)\n", peer->internal_id, &peer->endpoint.addr_storage);
 
 	if (noise_handshake_create_initiation(&packet, &peer->handshake)) {
 		cookie_add_mac_to_packet(&packet, sizeof(packet), peer);
@@ -64,7 +64,7 @@ void packet_send_handshake_response(struct wireguard_peer *peer)
 {
 	struct message_handshake_response packet;
 
-	net_dbg_ratelimited("Sending handshake response to peer %Lu (%pISpfsc)\n", peer->internal_id, &peer->endpoint_addr);
+	net_dbg_ratelimited("Sending handshake response to peer %Lu (%pISpfsc)\n", peer->internal_id, &peer->endpoint.addr_storage);
 	peer->last_sent_handshake = get_jiffies_64();
 
 	if (noise_handshake_create_response(&packet, &peer->handshake)) {
@@ -82,9 +82,9 @@ void packet_send_handshake_cookie(struct wireguard_device *wg, struct sk_buff *i
 	struct message_handshake_cookie packet;
 
 #if defined(CONFIG_DYNAMIC_DEBUG) || defined(DEBUG)
-	struct sockaddr_storage addr = { 0 };
-	socket_addr_from_skb(&addr, initiating_skb);
-	net_dbg_ratelimited("Sending cookie response for denied handshake message for %pISpfsc\n", &addr);
+	struct endpoint endpoint;
+	socket_endpoint_from_skb(&endpoint, initiating_skb);
+	net_dbg_ratelimited("Sending cookie response for denied handshake message for %pISpfsc\n", &endpoint.addr_storage);
 #endif
 	cookie_message_create(&packet, initiating_skb, data, data_len, sender_index, &wg->cookie_checker);
 	socket_send_buffer_as_reply_to_skb(wg, initiating_skb, &packet, sizeof(packet));
@@ -117,7 +117,7 @@ void packet_send_keepalive(struct wireguard_peer *peer)
 		skb_reserve(skb, DATA_PACKET_HEAD_ROOM);
 		skb->dev = netdev_pub(peer->device);
 		skb_queue_tail(&peer->tx_packet_queue, skb);
-		net_dbg_ratelimited("Sending keepalive packet to peer %Lu (%pISpfsc)\n", peer->internal_id, &peer->endpoint_addr);
+		net_dbg_ratelimited("Sending keepalive packet to peer %Lu (%pISpfsc)\n", peer->internal_id, &peer->endpoint.addr_storage);
 	}
 	packet_send_queue(peer);
 }
diff --git a/src/socket.c b/src/socket.c
index 069d45d..20e82fd 100644
--- a/src/socket.c
+++ b/src/socket.c
@@ -13,11 +13,12 @@
 #include <net/udp_tunnel.h>
 #include <net/ipv6.h>
 
-static inline int send4(struct wireguard_device *wg, struct sk_buff *skb, struct sockaddr_in *addr, uint8_t ds, struct dst_cache *cache)
+static inline int send4(struct wireguard_device *wg, struct sk_buff *skb, struct endpoint *endpoint, uint8_t ds, struct dst_cache *cache)
 {
 	struct flowi4 fl = {
-		.daddr = addr->sin_addr.s_addr,
-		.fl4_dport = addr->sin_port,
+		.saddr = endpoint->src4.s_addr,
+		.daddr = endpoint->addr4.sin_addr.s_addr,
+		.fl4_dport = endpoint->addr4.sin_port,
 		.fl4_sport = htons(wg->incoming_port),
 		.flowi4_proto = IPPROTO_UDP
 	};
@@ -44,12 +45,12 @@ static inline int send4(struct wireguard_device *wg, struct sk_buff *skb, struct
 		rt = ip_route_output_flow(sock_net(sock), &fl, sock);
 		if (unlikely(IS_ERR(rt))) {
 			ret = PTR_ERR(rt);
-			net_dbg_ratelimited("No route to %pISpfsc, error %d\n", addr, ret);
+			net_dbg_ratelimited("No route to %pISpfsc, error %d\n", &endpoint->addr_storage, ret);
 			goto err;
 		} else if (unlikely(rt->dst.dev == skb->dev)) {
 			dst_release(&rt->dst);
 			ret = -ELOOP;
-			net_dbg_ratelimited("Avoiding routing loop to %pISpfsc\n", addr);
+			net_dbg_ratelimited("Avoiding routing loop to %pISpfsc\n", &endpoint->addr_storage);
 			goto err;
 		}
 		if (cache)
@@ -70,14 +71,15 @@ out:
 	return ret;
 }
 
-static inline int send6(struct wireguard_device *wg, struct sk_buff *skb, struct sockaddr_in6 *addr, uint8_t ds, struct dst_cache *cache)
+static inline int send6(struct wireguard_device *wg, struct sk_buff *skb, struct endpoint *endpoint, uint8_t ds, struct dst_cache *cache)
 {
 #if IS_ENABLED(CONFIG_IPV6)
 	struct flowi6 fl = {
-		.daddr = addr->sin6_addr,
-		.fl6_dport = addr->sin6_port,
+		.saddr = endpoint->src6,
+		.daddr = endpoint->addr6.sin6_addr,
+		.fl6_dport = endpoint->addr6.sin6_port,
 		.fl6_sport = htons(wg->incoming_port),
-		.flowi6_oif = addr->sin6_scope_id,
+		.flowi6_oif = endpoint->addr6.sin6_scope_id,
 		.flowi6_proto = IPPROTO_UDP
 		/* TODO: addr->sin6_flowinfo */
 	};
@@ -103,12 +105,12 @@ static inline int send6(struct wireguard_device *wg, struct sk_buff *skb, struct
 		security_sk_classify_flow(sock, flowi6_to_flowi(&fl));
 		ret = ipv6_stub->ipv6_dst_lookup(sock_net(sock), sock, &dst, &fl);
 		if (unlikely(ret)) {
-			net_dbg_ratelimited("No route to %pISpfsc, error %d\n", addr, ret);
+			net_dbg_ratelimited("No route to %pISpfsc, error %d\n", &endpoint->addr_storage, ret);
 			goto err;
 		} else if (unlikely(dst->dev == skb->dev)) {
 			dst_release(dst);
 			ret = -ELOOP;
-			net_dbg_ratelimited("Avoiding routing loop to %pISpfsc\n", addr);
+			net_dbg_ratelimited("Avoiding routing loop to %pISpfsc\n", &endpoint->addr_storage);
 			goto err;
 		}
 		if (cache)
@@ -138,10 +140,10 @@ int socket_send_skb_to_peer(struct wireguard_peer *peer, struct sk_buff *skb, ui
 	int ret = -EAFNOSUPPORT;
 
 	read_lock_bh(&peer->endpoint_lock);
-	if (peer->endpoint_addr.ss_family == AF_INET)
-		ret = send4(peer->device, skb, (struct sockaddr_in *)&peer->endpoint_addr, ds, &peer->endpoint_cache);
-	else if (peer->endpoint_addr.ss_family == AF_INET6)
-		ret = send6(peer->device, skb, (struct sockaddr_in6 *)&peer->endpoint_addr, ds, &peer->endpoint_cache);
+	if (peer->endpoint.addr_storage.ss_family == AF_INET)
+		ret = send4(peer->device, skb, &peer->endpoint, ds, &peer->endpoint_cache);
+	else if (peer->endpoint.addr_storage.ss_family == AF_INET6)
+		ret = send6(peer->device, skb, &peer->endpoint, ds, &peer->endpoint_cache);
 	if (likely(!ret))
 		peer->tx_bytes += skb_len;
 	read_unlock_bh(&peer->endpoint_lock);
@@ -163,11 +165,11 @@ int socket_send_buffer_as_reply_to_skb(struct wireguard_device *wg, struct sk_bu
 {
 	int ret = 0;
 	struct sk_buff *skb;
-	struct sockaddr_storage addr = { 0 };
+	struct endpoint endpoint;
 
 	if (unlikely(!in_skb))
 		return -EINVAL;
-	ret = socket_addr_from_skb(&addr, in_skb);
+	ret = socket_endpoint_from_skb(&endpoint, in_skb);
 	if (unlikely(ret < 0))
 		return ret;
 
@@ -177,60 +179,62 @@ int socket_send_buffer_as_reply_to_skb(struct wireguard_device *wg, struct sk_bu
 	skb_reserve(skb, SKB_HEADER_LEN);
 	memcpy(skb_put(skb, len), out_buffer, len);
 
-	if (addr.ss_family == AF_INET)
-		ret = send4(wg, skb, (struct sockaddr_in *)&addr, 0, NULL);
-	else if (addr.ss_family == AF_INET6)
-		ret = send6(wg, skb, (struct sockaddr_in6 *)&addr, 0, NULL);
+	if (endpoint.addr_storage.ss_family == AF_INET)
+		ret = send4(wg, skb, &endpoint, 0, NULL);
+	else if (endpoint.addr_storage.ss_family == AF_INET6)
+		ret = send6(wg, skb, &endpoint, 0, NULL);
 	else
 		ret = -EAFNOSUPPORT;
 
 	return ret;
 }
 
-int socket_addr_from_skb(struct sockaddr_storage *sockaddr, struct sk_buff *skb)
+int socket_endpoint_from_skb(struct endpoint *endpoint, struct sk_buff *skb)
 {
-	struct iphdr *ip4;
-	struct ipv6hdr *ip6;
-	struct udphdr *udp;
-	struct sockaddr_in *addr4;
-	struct sockaddr_in6 *addr6;
-
-	addr4 = (struct sockaddr_in *)sockaddr;
-	addr6 = (struct sockaddr_in6 *)sockaddr;
-	ip4 = ip_hdr(skb);
-	ip6 = ipv6_hdr(skb);
-	udp = udp_hdr(skb);
-	if (ip4->version == 4) {
-		addr4->sin_family = AF_INET;
-		addr4->sin_port = udp->source;
-		addr4->sin_addr.s_addr = ip4->saddr;
-	} else if (ip4->version == 6) {
-		addr6->sin6_family = AF_INET6;
-		addr6->sin6_port = udp->source;
-		addr6->sin6_addr = ip6->saddr;
-		addr6->sin6_scope_id = ipv6_iface_scope_id(&ip6->saddr, skb->skb_iif);
-		/* TODO: addr6->sin6_flowinfo */
+	memset(endpoint, 0, sizeof(struct endpoint));
+	if (ip_hdr(skb)->version == 4) {
+		endpoint->addr4.sin_family = AF_INET;
+		endpoint->addr4.sin_port = udp_hdr(skb)->source;
+		endpoint->addr4.sin_addr.s_addr = ip_hdr(skb)->saddr;
+		endpoint->src4.s_addr = ip_hdr(skb)->daddr;
+	} else if (ip_hdr(skb)->version == 6) {
+		endpoint->addr6.sin6_family = AF_INET6;
+		endpoint->addr6.sin6_port = udp_hdr(skb)->source;
+		endpoint->addr6.sin6_addr = ipv6_hdr(skb)->saddr;
+		endpoint->addr6.sin6_scope_id = ipv6_iface_scope_id(&ipv6_hdr(skb)->saddr, skb->skb_iif);
+		/* TODO: endpoint->addr6.sin6_flowinfo */
+		endpoint->src6 = ipv6_hdr(skb)->daddr;
 	} else
 		return -EINVAL;
 	return 0;
 }
 
-void socket_set_peer_addr(struct wireguard_peer *peer, struct sockaddr_storage *sockaddr)
+void socket_set_peer_endpoint(struct wireguard_peer *peer, struct endpoint *endpoint)
 {
-	if (sockaddr->ss_family == AF_INET) {
+	if (endpoint->addr_storage.ss_family == AF_INET) {
 		read_lock_bh(&peer->endpoint_lock);
-		if (!memcmp(sockaddr, &peer->endpoint_addr, sizeof(struct sockaddr_in)))
+		if (likely(peer->endpoint.addr4.sin_family == AF_INET &&
+			   peer->endpoint.addr4.sin_port == endpoint->addr4.sin_port &&
+			   peer->endpoint.addr4.sin_addr.s_addr == endpoint->addr4.sin_addr.s_addr &&
+			   peer->endpoint.src4.s_addr == endpoint->src4.s_addr))
 			goto out;
 		read_unlock_bh(&peer->endpoint_lock);
 		write_lock_bh(&peer->endpoint_lock);
-		memcpy(&peer->endpoint_addr, sockaddr, sizeof(struct sockaddr_in));
-	} else if (sockaddr->ss_family == AF_INET6) {
+		peer->endpoint.addr4 = endpoint->addr4;
+		peer->endpoint.src4 = endpoint->src4;
+	} else if (endpoint->addr_storage.ss_family == AF_INET6) {
 		read_lock_bh(&peer->endpoint_lock);
-		if (!memcmp(sockaddr, &peer->endpoint_addr, sizeof(struct sockaddr_in6)))
+		if (likely(peer->endpoint.addr6.sin6_family == AF_INET6 &&
+			   peer->endpoint.addr6.sin6_port == endpoint->addr6.sin6_port &&
+			   /* TODO: peer->endpoint.addr6.sin6_flowinfo == endpoint->addr6.sin6_flowinfo && */
+			   ipv6_addr_equal(&peer->endpoint.addr6.sin6_addr, &endpoint->addr6.sin6_addr) &&
+			   peer->endpoint.addr6.sin6_scope_id == endpoint->addr6.sin6_scope_id &&
+			   ipv6_addr_equal(&peer->endpoint.src6, &endpoint->src6)))
 			goto out;
 		read_unlock_bh(&peer->endpoint_lock);
 		write_lock_bh(&peer->endpoint_lock);
-		memcpy(&peer->endpoint_addr, sockaddr, sizeof(struct sockaddr_in6));
+		peer->endpoint.addr6 = endpoint->addr6;
+		peer->endpoint.src6 = endpoint->src6;
 	} else
 		return;
 	dst_cache_reset(&peer->endpoint_cache);
diff --git a/src/socket.h b/src/socket.h
index 36cb4ed..ce668ad 100644
--- a/src/socket.h
+++ b/src/socket.h
@@ -9,6 +9,7 @@
 #include <linux/if_ether.h>
 
 struct wireguard_device;
+struct endpoint;
 
 #define SKB_HEADER_LEN (max(sizeof(struct iphdr), sizeof(struct ipv6hdr)) + sizeof(struct udphdr) + NET_SKB_PAD)
 
@@ -18,7 +19,7 @@ int socket_send_buffer_to_peer(struct wireguard_peer *peer, void *data, size_t l
 int socket_send_skb_to_peer(struct wireguard_peer *peer, struct sk_buff *skb, uint8_t ds);
 int socket_send_buffer_as_reply_to_skb(struct wireguard_device *wg, struct sk_buff *in_skb, void *out_buffer, size_t len);
 
-int socket_addr_from_skb(struct sockaddr_storage *sockaddr, struct sk_buff *skb);
-void socket_set_peer_addr(struct wireguard_peer *peer, struct sockaddr_storage *sockaddr);
+int socket_endpoint_from_skb(struct endpoint *endpoint, struct sk_buff *skb);
+void socket_set_peer_endpoint(struct wireguard_peer *peer, struct endpoint *endpoint);
 
 #endif
diff --git a/src/timers.c b/src/timers.c
index 8cc2104..f002a65 100644
--- a/src/timers.c
+++ b/src/timers.c
@@ -27,7 +27,7 @@ static inline unsigned long slack_time(unsigned long time)
 static void expired_retransmit_handshake(unsigned long ptr)
 {
 	peer_get_from_ptr(ptr);
-	pr_debug("Handshake for peer %Lu (%pISpfsc) did not complete after %d seconds, retrying\n", peer->internal_id, &peer->endpoint_addr, REKEY_TIMEOUT / HZ);
+	pr_debug("Handshake for peer %Lu (%pISpfsc) did not complete after %d seconds, retrying\n", peer->internal_id, &peer->endpoint.addr_storage, REKEY_TIMEOUT / HZ);
 	if (peer->timer_handshake_attempts > MAX_TIMER_HANDSHAKES) {
 		del_timer(&peer->timer_send_keepalive);
 		/* We remove all existing packets and don't try again,
@@ -59,7 +59,7 @@ static void expired_send_keepalive(unsigned long ptr)
 static void expired_new_handshake(unsigned long ptr)
 {
 	peer_get_from_ptr(ptr);
-	pr_debug("Retrying handshake with peer %Lu (%pISpfsc) because we stopped hearing back after %d seconds\n", peer->internal_id, &peer->endpoint_addr, (KEEPALIVE_TIMEOUT + REKEY_TIMEOUT) / HZ);
+	pr_debug("Retrying handshake with peer %Lu (%pISpfsc) because we stopped hearing back after %d seconds\n", peer->internal_id, &peer->endpoint.addr_storage, (KEEPALIVE_TIMEOUT + REKEY_TIMEOUT) / HZ);
 	packet_queue_handshake_initiation(peer);
 	peer_put(peer);
 }
@@ -73,7 +73,7 @@ static void expired_kill_ephemerals(unsigned long ptr)
 static void queued_expired_kill_ephemerals(struct work_struct *work)
 {
 	struct wireguard_peer *peer = container_of(work, struct wireguard_peer, clear_peer_work);
-	pr_debug("Zeroing out all keys for peer %Lu (%pISpfsc), since we haven't received a new one in %d seconds\n", peer->internal_id, &peer->endpoint_addr, (REJECT_AFTER_TIME * 3) / HZ);
+	pr_debug("Zeroing out all keys for peer %Lu (%pISpfsc), since we haven't received a new one in %d seconds\n", peer->internal_id, &peer->endpoint.addr_storage, (REJECT_AFTER_TIME * 3) / HZ);
 	noise_handshake_clear(&peer->handshake);
 	noise_keypairs_clear(&peer->keypairs);
 	peer_put(peer);