From 6a9e1f5e72a00176ee15c24cde7d46a83d65b66e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Sat, 1 Aug 2020 04:27:31 +0200
Subject: [PATCH] aquatic_ws: keep valid_until up to date in established
 connections

---
 TODO.md                           | 1 -
 aquatic_ws/src/lib/network/mod.rs | 5 +++++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/TODO.md b/TODO.md
index a927d27..052bb02 100644
--- a/TODO.md
+++ b/TODO.md
@@ -33,7 +33,6 @@
   * tungstenite 0.11
 * is 'key' sent in announce request? if so, maybe handle it like in
   aquatic_http (including ip uniqueness part of peer map key)
-* established connections do not get valid_until updated, I think?
 * tests
 * use enum as return type for handshake machine
 
diff --git a/aquatic_ws/src/lib/network/mod.rs b/aquatic_ws/src/lib/network/mod.rs
index 824871d..c4706e2 100644
--- a/aquatic_ws/src/lib/network/mod.rs
+++ b/aquatic_ws/src/lib/network/mod.rs
@@ -182,6 +182,11 @@ pub fn run_handshakes_and_read_messages(
 ){
     loop {
         if let Some(established_ws) = connections.get_mut(&poll_token)
+            .map(|c| { // Ugly but works
+                c.valid_until = valid_until;
+
+                c
+            })
             .and_then(Connection::get_established_ws)
         {
             use ::tungstenite::Error::Io;