From 2a73a27e682395d0805d34f58b4dd090edc438c1 Mon Sep 17 00:00:00 2001
From: Ludvig Strigeus <strigeus@gmail.com>
Date: Sat, 17 Nov 2018 17:24:48 +0100
Subject: [PATCH] Forget endpoint for incoming connections after a little while

---
 wireguard.cpp       | 11 ++++++++++-
 wireguard_proto.cpp | 12 +++++++++---
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/wireguard.cpp b/wireguard.cpp
index 1947dd1..4be0024 100644
--- a/wireguard.cpp
+++ b/wireguard.cpp
@@ -632,7 +632,7 @@ void WireguardProcessor::RunAllMainThreadScheduled() {
 void WireguardProcessor::SendHandshakeInitiation(WgPeer *peer) {
   assert(dev_.IsMainThread());
 
-  if (!peer->CheckHandshakeRateLimit())
+  if (!peer->CheckHandshakeRateLimit() || peer->endpoint_.sin.sin_family == 0)
     return;
   stats_.handshakes_out++;
   Packet *packet = AllocPacket();
@@ -646,6 +646,15 @@ void WireguardProcessor::SendHandshakeInitiation(WgPeer *peer) {
     packet->addr = peer->endpoint_;
     packet->protocol = peer->endpoint_protocol_;
     peer->tx_bytes_ += packet->size;
+
+    // If this is an incoming oneway connection (such as tcp), forget the
+    // endpoint after a number of attempts.
+    if (attempts >= 3 && peer->allow_endpoint_change_ &&
+        (peer->endpoint_protocol_ & kPacketProtocolIncomingConnection)) {
+      peer->endpoint_protocol_ = 0;
+      peer->endpoint_.sin.sin_family = 0;
+    }
+
     WG_RELEASE_LOCK(peer->mutex_);
     DoWriteUdpPacket(packet);
     if (attempts > 1 && attempts <= 20)
diff --git a/wireguard_proto.cpp b/wireguard_proto.cpp
index fc5efb3..214ac4b 100644
--- a/wireguard_proto.cpp
+++ b/wireguard_proto.cpp
@@ -747,6 +747,10 @@ WgPeer *WgPeer::ParseMessageHandshakeResponse(WgDevice *dev, const Packet *packe
   peer_and_keypair->second = keypair;
 
   WG_ACQUIRE_LOCK(peer->mutex_);
+  if (peer->allow_endpoint_change_) {
+    peer->endpoint_ = packet->addr;
+    peer->endpoint_protocol_ = packet->protocol;
+  }
   peer->rx_bytes_ += packet->size;
   peer->InsertKeypairInPeer_Locked(keypair);
   WG_RELEASE_LOCK(peer->mutex_);
@@ -1182,7 +1186,7 @@ uint32 WgPeer::CheckTimeouts_Locked(uint64 now) {
   if (t & 0x1F) {
     if ((t & (1 << TIMER_RETRANSMIT_HANDSHAKE)) && (now32 - timer_value_[TIMER_RETRANSMIT_HANDSHAKE]) >= REKEY_TIMEOUT_MS) {
       t ^= (1 << TIMER_RETRANSMIT_HANDSHAKE);
-      if (handshake_attempts_ > MAX_HANDSHAKE_ATTEMPTS) {
+      if (handshake_attempts_ > MAX_HANDSHAKE_ATTEMPTS || endpoint_.sin.sin_family == 0) {
         t &= ~(1 << TIMER_SEND_KEEPALIVE);
         ClearPacketQueue_Locked();
       } else {
@@ -1208,8 +1212,10 @@ uint32 WgPeer::CheckTimeouts_Locked(uint64 now) {
     }
     if ((t & (1 << TIMER_NEW_HANDSHAKE)) && (now32 - timer_value_[TIMER_NEW_HANDSHAKE]) >= KEEPALIVE_TIMEOUT_MS + REKEY_TIMEOUT_MS) {
       t &= ~(1 << TIMER_NEW_HANDSHAKE);
-      handshake_attempts_ = 0;
-      rv |= ACTION_SEND_HANDSHAKE;
+      if (endpoint_.sin.sin_family != 0) {
+        handshake_attempts_ = 0;
+        rv |= ACTION_SEND_HANDSHAKE;
+      }
     }
     if ((t & (1 << TIMER_ZERO_KEYS)) && (now32 - timer_value_[TIMER_ZERO_KEYS]) >= REJECT_AFTER_TIME_MS * 3) {
       RINFO("Expiring all keys for peer");