From 7d8958e0e75d84b9a3711d2756ddeaff3b90ebd5 Mon Sep 17 00:00:00 2001
From: Jett Chen <jettchen12345@gmail.com>
Date: Sun, 21 Jan 2024 03:44:58 +0800
Subject: [PATCH] Add handling for connection expiration

---
 burrow/src/wireguard/pcb.rs | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/burrow/src/wireguard/pcb.rs b/burrow/src/wireguard/pcb.rs
index c6ebaa6..dfa6d2a 100755
--- a/burrow/src/wireguard/pcb.rs
+++ b/burrow/src/wireguard/pcb.rs
@@ -5,7 +5,9 @@ use fehler::throws;
 use ip_network::IpNetwork;
 use rand::random;
 use tokio::{net::UdpSocket, sync::RwLock, task::JoinHandle};
+use tokio::io::AsyncWrite;
 use tun::tokio::TunInterface;
+use crate::wireguard::noise::errors::WireGuardError;
 
 use super::{
     noise::{TunnResult, Tunnel},
@@ -62,6 +64,7 @@ impl PeerPcb {
             tracing::debug!("{}: waiting for packet", rid);
             let guard = self.socket.read().await;
             let Some(socket) = guard.as_ref() else {
+                self.open_if_closed().await?;
                 continue
             };
             let mut res_buf = [0; 1500];
@@ -136,6 +139,10 @@ impl PeerPcb {
     pub async fn update_timers(&self, dst: &mut [u8]) -> Result<(), Error> {
         match self.tunnel.write().await.update_timers(dst) {
             TunnResult::Done => {}
+            TunnResult::Err(WireGuardError::ConnectionExpired) => {
+                tracing::debug!("Connection expired, closing socket");
+                self.socket.write().await.take();
+            }
             TunnResult::Err(e) => {
                 tracing::error!(message = "Update timers error", error = ?e)
             }