Break when socket is unexpectedly shut down

hickory-dns · Mar 16, 2024 · d2e64d8 · d2e64d8
1 parent da41e4a
commit d2e64d8
Showing 1 changed file with 42 additions and 4 deletions.
diff --git a/crates/server/src/server/server_future.rs b/crates/server/src/server/server_future.rs
@@ -95,6 +95,9 @@ impl<T: RequestHandler> ServerFuture<T> {
                     let message = match message {
                         Err(e) => {
                             warn!("error receiving message on udp_socket: {}", e);
+                            if is_unrecoverable_socket_error(&e) {
+                                break;
+                            }
                             continue;
                         }
                         Ok(message) => message,
@@ -169,6 +172,9 @@ impl<T: RequestHandler> ServerFuture<T> {
                         Ok((t, s)) => (t, s),
                         Err(e) => {
                             debug!("error receiving TCP tcp_stream error: {}", e);
+                            if is_unrecoverable_socket_error(&e) {
+                                break;
+                            }
                             continue;
                         },
                     },
@@ -227,7 +233,11 @@ impl<T: RequestHandler> ServerFuture<T> {
                 reap_tasks(&mut inner_join_set);
             }
 
-            Ok(())
+            if shutdown.is_cancelled() {
+                Ok(())
+            } else {
+                Err(ProtoError::from("unexpected close of socket"))
+            }
         });
     }
 
@@ -298,6 +308,9 @@ impl<T: RequestHandler> ServerFuture<T> {
                         Ok((t, s)) => (t, s),
                         Err(e) => {
                             debug!("error receiving TLS tcp_stream error: {}", e);
+                            if is_unrecoverable_socket_error(&e) {
+                                break;
+                            }
                             continue;
                         },
                     },
@@ -373,7 +386,11 @@ impl<T: RequestHandler> ServerFuture<T> {
                 reap_tasks(&mut inner_join_set);
             }
 
-            Ok(())
+            if shutdown.is_cancelled() {
+                Ok(())
+            } else {
+                Err(ProtoError::from("unexpected close of socket"))
+            }
         });
 
         Ok(())
@@ -451,6 +468,9 @@ impl<T: RequestHandler> ServerFuture<T> {
                         Ok((t, s)) => (t, s),
                         Err(e) => {
                             debug!("error receiving TLS tcp_stream error: {}", e);
+                            if is_unrecoverable_socket_error(&e) {
+                                break;
+                            }
                             continue;
                         },
                     },
@@ -519,7 +539,11 @@ impl<T: RequestHandler> ServerFuture<T> {
                 reap_tasks(&mut inner_join_set);
             }
 
-            Ok(())
+            if shutdown.is_cancelled() {
+                Ok(())
+            } else {
+                Err(ProtoError::from("unexpected close of socket"))
+            }
         });
 
         Ok(())
@@ -646,6 +670,9 @@ impl<T: RequestHandler> ServerFuture<T> {
                         Ok((t, s)) => (t, s),
                         Err(e) => {
                             debug!("error receiving HTTPS tcp_stream error: {}", e);
+                            if is_unrecoverable_socket_error(&e) {
+                                break;
+                            }
                             continue;
                         },
                     },
@@ -696,7 +723,11 @@ impl<T: RequestHandler> ServerFuture<T> {
                 reap_tasks(&mut inner_join_set);
             }
 
-            Ok(())
+            if shutdown.is_cancelled() {
+                Ok(())
+            } else {
+                Err(ProtoError::from("unexpected close of socket"))
+            }
         });
 
         Ok(())
@@ -1214,6 +1245,13 @@ fn sanitize_src_address(src: SocketAddr) -> Result<(), String> {
     }
 }
 
+fn is_unrecoverable_socket_error(err: &io::Error) -> bool {
+    matches!(
+        err.kind(),
+        io::ErrorKind::NotConnected | io::ErrorKind::ConnectionAborted
+    )
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;