Make sure to send a proper reset code when resetting a connection

so the other side can know if reconnecting is OK

Author: rklaehn

src/provider.rs

@@ -20,13 +20,12 @@ use tracing::{debug, debug_span, warn, Instrument};
 
 use crate::{
     api::{
-        self,
         blobs::{Bitfield, WriteProgress},
-        Store,
+        ExportBaoResult, Store,
     },
     hashseq::HashSeq,
     protocol::{GetManyRequest, GetRequest, ObserveItem, ObserveRequest, PushRequest, Request},
-    provider::events::{ClientConnected, ClientError, ConnectionClosed, RequestTracker},
+    provider::events::{ClientConnected, ConnectionClosed, RequestTracker},
     Hash,
 };
 pub mod events;
@@ -94,7 +93,7 @@ impl StreamPair {
     }
 
     /// We are done with reading. Return a ProgressWriter that contains the read stats and connection id
-    pub async fn into_writer(
+    async fn into_writer(
         mut self,
         tracker: RequestTracker,
     ) -> Result<ProgressWriter, ReadToEndError> {
@@ -118,7 +117,7 @@ impl StreamPair {
         ))
     }
 
-    pub async fn into_reader(
+    async fn into_reader(
         mut self,
         tracker: RequestTracker,
     ) -> Result<ProgressReader, ClosedStream> {
@@ -141,39 +140,71 @@ impl StreamPair {
     }
 
     pub async fn get_request(
-        &self,
+        mut self,
         f: impl FnOnce() -> GetRequest,
-    ) -> Result<RequestTracker, ClientError> {
-        self.events
+    ) -> anyhow::Result<ProgressWriter> {
+        let res = self
+            .events
             .request(f, self.connection_id, self.request_id)
-            .await
+            .await;
+        match res {
+            Err(e) => {
+                self.writer.reset(e.code()).ok();
+                Err(e.into())
+            }
+            Ok(tracker) => Ok(self.into_writer(tracker).await?),
+        }
     }
 
     pub async fn get_many_request(
-        &self,
+        mut self,
         f: impl FnOnce() -> GetManyRequest,
-    ) -> Result<RequestTracker, ClientError> {
-        self.events
+    ) -> anyhow::Result<ProgressWriter> {
+        let res = self
+            .events
             .request(f, self.connection_id, self.request_id)
-            .await
+            .await;
+        match res {
+            Err(e) => {
+                self.writer.reset(e.code()).ok();
+                Err(e.into())
+            }
+            Ok(tracker) => Ok(self.into_writer(tracker).await?),
+        }
     }
 
     pub async fn push_request(
-        &self,
+        mut self,
         f: impl FnOnce() -> PushRequest,
-    ) -> Result<RequestTracker, ClientError> {
-        self.events
+    ) -> anyhow::Result<ProgressReader> {
+        let res = self
+            .events
             .request(f, self.connection_id, self.request_id)
-            .await
+            .await;
+        match res {
+            Err(e) => {
+                self.writer.reset(e.code()).ok();
+                Err(e.into())
+            }
+            Ok(tracker) => Ok(self.into_reader(tracker).await?),
+        }
     }
 
     pub async fn observe_request(
-        &self,
+        mut self,
         f: impl FnOnce() -> ObserveRequest,
-    ) -> Result<RequestTracker, ClientError> {
-        self.events
+    ) -> anyhow::Result<ProgressWriter> {
+        let res = self
+            .events
             .request(f, self.connection_id, self.request_id)
-            .await
+            .await;
+        match res {
+            Err(e) => {
+                self.writer.reset(e.code()).ok();
+                Err(e.into())
+            }
+            Ok(tracker) => Ok(self.into_writer(tracker).await?),
+        }
     }
 
     fn stats(&self) -> TransferStats {
@@ -299,7 +330,8 @@ pub async fn handle_connection(
             })
             .await
         {
-            debug!("client not authorized to connect: {cause}");
+            connection.close(cause.code(), cause.reason());
+            debug!("closing connection: {cause}");
             return;
         }
         while let Ok(context) = StreamPair::accept(&connection, &progress).await {
@@ -323,35 +355,32 @@ async fn handle_stream(store: Store, mut context: StreamPair) -> anyhow::Result<
 
     match request {
         Request::Get(request) => {
-            let tracker = context.get_request(|| request.clone()).await?;
-            let mut writer = context.into_writer(tracker).await?;
-            if handle_get(store, request, &mut writer).await.is_ok() {
+            let mut writer = context.get_request(|| request.clone()).await?;
+            let res = handle_get(store, request, &mut writer).await;
+            if res.is_ok() {
                 writer.transfer_completed().await;
             } else {
                 writer.transfer_aborted().await;
             }
         }
         Request::GetMany(request) => {
-            let tracker = context.get_many_request(|| request.clone()).await?;
-            let mut writer = context.into_writer(tracker).await?;
+            let mut writer = context.get_many_request(|| request.clone()).await?;
             if handle_get_many(store, request, &mut writer).await.is_ok() {
                 writer.transfer_completed().await;
             } else {
                 writer.transfer_aborted().await;
             }
         }
         Request::Observe(request) => {
-            let tracker = context.observe_request(|| request.clone()).await?;
-            let mut writer = context.into_writer(tracker).await?;
+            let mut writer = context.observe_request(|| request.clone()).await?;
             if handle_observe(store, request, &mut writer).await.is_ok() {
                 writer.transfer_completed().await;
             } else {
                 writer.transfer_aborted().await;
             }
         }
         Request::Push(request) => {
-            let tracker = context.push_request(|| request.clone()).await?;
-            let mut reader = context.into_reader(tracker).await?;
+            let mut reader = context.push_request(|| request.clone()).await?;
             if handle_push(store, request, &mut reader).await.is_ok() {
                 reader.transfer_completed().await;
             } else {
@@ -464,11 +493,11 @@ pub(crate) async fn send_blob(
     hash: Hash,
     ranges: ChunkRanges,
     writer: &mut ProgressWriter,
-) -> api::Result<()> {
-    Ok(store
+) -> ExportBaoResult<()> {
+    store
         .export_bao(hash, ranges)
         .write_quinn_with_progress(&mut writer.inner, &mut writer.context, &hash, index)
-        .await?)
+        .await
 }
 
 /// Handle a single push request.

src/provider/events.rs

@@ -8,7 +8,10 @@ use serde::{Deserialize, Serialize};
 use snafu::Snafu;
 
 use crate::{
-    protocol::{GetManyRequest, GetRequest, ObserveRequest, PushRequest},
+    protocol::{
+        GetManyRequest, GetRequest, ObserveRequest, PushRequest, ERR_INTERNAL, ERR_LIMIT,
+        ERR_PERMISSION,
+    },
     provider::{events::irpc_ext::IrpcClientExt, TransferStats},
     Hash,
 };
@@ -82,6 +85,24 @@ pub enum ClientError {
     },
 }
 
+impl ClientError {
+    pub fn code(&self) -> quinn::VarInt {
+        match self {
+            ClientError::RateLimited => ERR_LIMIT.into(),
+            ClientError::Permission => ERR_PERMISSION.into(),
+            ClientError::Irpc { .. } => ERR_INTERNAL.into(),
+        }
+    }
+
+    pub fn reason(&self) -> &'static [u8] {
+        match self {
+            ClientError::RateLimited => b"limit",
+            ClientError::Permission => b"permission",
+            ClientError::Irpc { .. } => b"internal",
+        }
+    }
+}
+
 impl From<AbortReason> for ClientError {
     fn from(value: AbortReason) -> Self {
         match value {
@@ -211,11 +232,14 @@ impl RequestTracker {
     /// Transfer for index `index` started, size `size`
     pub async fn transfer_started(&self, index: u64, hash: &Hash, size: u64) -> irpc::Result<()> {
         if let RequestUpdates::Active(tx) = &self.updates {
-            tx.send(RequestUpdate::Started(TransferStarted {
-                index,
-                hash: *hash,
-                size,
-            }))
+            tx.send(
+                TransferStarted {
+                    index,
+                    hash: *hash,
+                    size,
+                }
+                .into(),
+            )
             .await?;
         }
         Ok(())
@@ -224,8 +248,7 @@ impl RequestTracker {
     /// Transfer progress for the previously reported blob, end_offset is the new end offset in bytes.
     pub async fn transfer_progress(&mut self, len: u64, end_offset: u64) -> ClientResult {
         if let RequestUpdates::Active(tx) = &mut self.updates {
-            tx.try_send(RequestUpdate::Progress(TransferProgress { end_offset }))
-                .await?;
+            tx.try_send(TransferProgress { end_offset }.into()).await?;
         }
         if let Some((throttle, connection_id, request_id)) = &self.throttle {
             throttle
@@ -242,17 +265,15 @@ impl RequestTracker {
     /// Transfer completed for the previously reported blob.
     pub async fn transfer_completed(&self, f: impl Fn() -> Box<TransferStats>) -> irpc::Result<()> {
         if let RequestUpdates::Active(tx) = &self.updates {
-            tx.send(RequestUpdate::Completed(TransferCompleted { stats: f() }))
-                .await?;
+            tx.send(TransferCompleted { stats: f() }.into()).await?;
         }
         Ok(())
     }
 
     /// Transfer aborted for the previously reported blob.
     pub async fn transfer_aborted(&self, f: impl Fn() -> Box<TransferStats>) -> irpc::Result<()> {
         if let RequestUpdates::Active(tx) = &self.updates {
-            tx.send(RequestUpdate::Aborted(TransferAborted { stats: f() }))
-                .await?;
+            tx.send(TransferAborted { stats: f() }.into()).await?;
         }
         Ok(())
     }
@@ -583,7 +604,7 @@ mod proto {
     }
 
     /// Stream of updates for a single request
-    #[derive(Debug, Serialize, Deserialize)]
+    #[derive(Debug, Serialize, Deserialize, derive_more::From)]
     pub enum RequestUpdate {
         /// Start of transfer for a blob, mandatory event
         Started(TransferStarted),