From d03865d92e84db35b4608779740b3ff981760e63 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Tue, 11 Jul 2023 00:05:35 +0100
Subject: [PATCH] fix(relay): reuse `delete_allocation` function (#1743)

Previously, we would access the state around allocations from different
places. This actually led to a minor memory leak where we wouldn't clean
up the `allocations_by_port` table. We refactor the code slightly to
avoid this.

---------

Co-authored-by: Jamil <jamilbk@users.noreply.github.com>
---
 rust/relay/src/server.rs | 24 +++++++++++-------------
 1 file changed, 11 insertions(+), 13 deletions(-)

diff --git a/rust/relay/src/server.rs b/rust/relay/src/server.rs
index 044892014..1597509f1 100644
--- a/rust/relay/src/server.rs
+++ b/rust/relay/src/server.rs
@@ -47,7 +47,6 @@ pub struct Server<R> {
 
     /// All client allocations, indexed by client's socket address.
     allocations: HashMap<SocketAddr, Allocation>,
-
     clients_by_allocation: HashMap<AllocationId, SocketAddr>,
     allocations_by_port: HashMap<u16, AllocationId>,
 
@@ -450,23 +449,14 @@ where
             .ok_or(error_response(AllocationMismatch, &request))?;
 
         if effective_lifetime.lifetime().is_zero() {
-            let port = allocation.port;
+            let id = allocation.id;
 
-            self.pending_commands
-                .push_back(Command::FreeAddresses { id: allocation.id });
-            self.allocations.remove(&sender);
-            self.allocations_by_port.remove(&port);
+            self.delete_allocation(id);
             self.send_message(
                 refresh_success_response(effective_lifetime, request.transaction_id()),
                 sender,
             );
 
-            tracing::info!(
-                target: "relay",
-                %port,
-                "Deleted allocation"
-            );
-
             return Ok(());
         }
 
@@ -743,9 +733,17 @@ where
             return;
         };
 
-        self.allocations.remove(&client);
+        let allocation = self
+            .allocations
+            .remove(&client)
+            .expect("internal state mismatch");
+        let port = allocation.port;
+
+        self.allocations_by_port.remove(&port);
         self.pending_commands
             .push_back(Command::FreeAddresses { id });
+
+        tracing::info!(target: "relay", %port, "Deleted allocation");
     }
 
     fn delete_channel_binding(&mut self, chan: u16) {