From 274799872fdabff90b9ad865659742706da65e8d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Mon, 11 May 2020 14:13:14 +0200
Subject: [PATCH] aquatic_ws: clean up; remove connection on IO error; update
 TODO

---
 TODO.md                       |  2 ++
 aquatic_ws/src/lib/network.rs | 22 +++++++++++-----------
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/TODO.md b/TODO.md
index 3b86386..f50dead 100644
--- a/TODO.md
+++ b/TODO.md
@@ -7,6 +7,8 @@
 * test
 * torrent state cleaning
 * config
+* use src as field in addition to peer_id in peer map, since peers have
+  access to others' peer id's
 
 ## aquatic
 * mio: set oneshot for epoll and kqueue? otherwise, stop reregistering?
diff --git a/aquatic_ws/src/lib/network.rs b/aquatic_ws/src/lib/network.rs
index 5838653..b8b2869 100644
--- a/aquatic_ws/src/lib/network.rs
+++ b/aquatic_ws/src/lib/network.rs
@@ -307,17 +307,20 @@ pub fn read_and_forward_in_messages(
 
                             break;
                         },
-                        Err(err) => {
+                        Err(HandshakeError::Failure(err)) => {
                             dbg!(err);
                         },
                     }
                 },
                 _ => unreachable!(),
             }
-        } else if let Some(Connection{ stage: ConnectionStage::Established(connection), ..}) = connections.get_mut(&poll_token){
+        } else if let Some(Connection{
+            stage: ConnectionStage::Established(peer_connection),
+            ..
+        }) = connections.get_mut(&poll_token){
             println!("conn established");
 
-            match connection.ws.read_message(){
+            match peer_connection.ws.read_message(){
                 Ok(ws_message) => {
                     dbg!(ws_message.clone());
 
@@ -327,28 +330,25 @@ pub fn read_and_forward_in_messages(
                         let meta = ConnectionMeta {
                             socket_worker_index,
                             socket_worker_poll_token: poll_token,
-                            peer_socket_addr: connection.peer_socket_addr
+                            peer_socket_addr: peer_connection.peer_socket_addr
                         };
 
                         in_message_sender.send((meta, in_message));
                     }
 
-                    connection.valid_until = valid_until;
+                    peer_connection.valid_until = valid_until;
                 },
                 Err(tungstenite::Error::Io(err)) => {
                     if err.kind() == ErrorKind::WouldBlock {
                         break
                     }
 
+                    remove_connection_if_exists(poll, connections, poll_token);
+
                     eprint!("{}", err);
                 },
                 Err(tungstenite::Error::ConnectionClosed) => {
-                    // FIXME: necessary?
-                    poll.registry()
-                        .deregister(connection.ws.get_mut())
-                        .unwrap();
-
-                    connections.remove(&poll_token);
+                    remove_connection_if_exists(poll, connections, poll_token);
                 },
                 Err(err) => {
                     eprint!("{}", err);