Abstract out BGP session stream type

src/bgp/mod.rs

@@ -27,7 +27,7 @@ use std::net::{IpAddr, SocketAddr};
 use std::rc::Rc;
 use strum::EnumDiscriminants;
 use thiserror::Error;
-use tokio::io::{AsyncWrite, BufReader};
+use tokio::io::{AsyncRead, AsyncWrite, BufReader};
 use tokio::net::TcpStream;
 use tokio::select;
 use tokio::time::{interval, Duration, Instant, Interval};
@@ -51,13 +51,13 @@ use State::*;
 /// - RFC 8955: Dissemination of Flow Specification Rules
 /// - RFC 8956: Dissemination of Flow Specification Rules for IPv6
 #[derive(Debug)]
-pub struct Session {
+pub struct Session<S: AsyncRead + AsyncWrite + Unpin> {
   config: RunArgs,
-  state: State,
+  state: State<S>,
   routes: Routes,
 }
 
-impl Session {
+impl<S: AsyncRead + AsyncWrite + Unpin> Session<S> {
   pub fn new(c: RunArgs) -> Result<Self> {
     let nft = (!c.dry_run)
       .then(|| Nft::new(c.table.clone(), c.chain.clone(), c.hooked, c.priority))
@@ -83,14 +83,14 @@ impl Session {
   pub fn config(&self) -> &RunArgs {
     &self.config
   }
-  pub fn state(&self) -> &State {
+  pub fn state(&self) -> &State<S> {
     &self.state
   }
   pub fn routes(&self) -> &Routes {
     &self.routes
   }
 
-  pub async fn accept(&mut self, mut stream: TcpStream, addr: SocketAddr) -> Result<()> {
+  pub async fn accept(&mut self, mut stream: S, addr: SocketAddr) -> Result<()> {
     let ip = addr.ip();
     if !self.config.allowed_ips.iter().any(|x| x.contains(ip)) {
       return Err(Error::UnacceptableAddr(ip));
@@ -103,7 +103,7 @@ impl Session {
       self.config.router_id.to_bits(),
     );
     open.send(&mut stream).await?;
-    replace_with_or_abort(&mut self.state, |_| OpenSent { stream: BufReader::new(stream) });
+    replace_with_or_abort(&mut self.state, |_| OpenSent { stream });
     info!("accepting BGP connection from {addr}");
     Ok(())
   }
@@ -212,30 +212,29 @@ impl Session {
 }
 
 #[derive(Debug)]
-pub enum State {
+pub enum State<S: AsyncRead + AsyncWrite + Unpin> {
   Idle,
   Connect, // never used in passive mode
   Active,
-  OpenSent {
-    stream: BufReader<TcpStream>,
-  },
-  OpenConfirm {
-    stream: BufReader<TcpStream>,
-    remote_open: OpenMessage<'static>,
-    timers: Option<Timers>,
-  },
-  Established {
-    stream: BufReader<TcpStream>,
-    remote_open: OpenMessage<'static>,
-    timers: Option<Timers>,
-  },
+  OpenSent { stream: S },
+  OpenConfirm { stream: S, remote_open: OpenMessage<'static>, timers: Option<Timers> },
+  Established { stream: S, remote_open: OpenMessage<'static>, timers: Option<Timers> },
 }
 
-impl State {
+impl<S: AsyncRead + AsyncWrite + Unpin> State<S> {
   pub fn kind(&self) -> StateKind {
-    self.view().into()
+    match self {
+      Idle => StateKind::Idle,
+      Connect => StateKind::Connect,
+      Active => StateKind::Active,
+      OpenSent { .. } => StateKind::OpenSent,
+      OpenConfirm { .. } => StateKind::OpenConfirm,
+      Established { .. } => StateKind::Established,
+    }
   }
+}
 
+impl State<BufReader<TcpStream>> {
   pub fn view(&self) -> StateView {
     match self {
       Idle => StateView::Idle,

src/ipc.rs

@@ -5,8 +5,8 @@ use std::ffi::CStr;
 use std::io;
 use std::mem::MaybeUninit;
 use std::path::{Path, PathBuf};
-use tokio::io::{AsyncReadExt, AsyncWrite, AsyncWriteExt};
-use tokio::net::{UnixListener, UnixStream};
+use tokio::io::{AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
+use tokio::net::{TcpStream, UnixListener, UnixStream};
 
 pub struct IpcServer {
   path: PathBuf,
@@ -31,7 +31,7 @@ impl Drop for IpcServer {
   }
 }
 
-impl Session {
+impl Session<BufReader<TcpStream>> {
   pub async fn write_states(&self, writer: &mut (impl AsyncWrite + Unpin)) -> anyhow::Result<()> {
     writer.write_all(&postcard::to_allocvec_cobs(self.config())?).await?;
     writer.write_all(&postcard::to_allocvec_cobs(&self.state().view())?).await?;

src/main.rs

@@ -19,10 +19,11 @@ use itertools::Itertools;
 use log::{error, info, warn, Level, LevelFilter, Record};
 use std::fs::{create_dir_all, File};
 use std::io::ErrorKind::UnexpectedEof;
-use std::io::{self, BufRead, BufReader, Write};
+use std::io::{self, BufRead, Write};
 use std::net::Ipv4Addr;
 use std::path::Path;
 use std::process::ExitCode;
+use tokio::io::BufReader;
 use tokio::net::TcpListener;
 use tokio::signal::unix::{signal, SignalKind};
 use tokio::{pin, select};
@@ -34,7 +35,7 @@ async fn run(mut args: RunArgs, sock_path: &str) -> anyhow::Result<ExitCode> {
     args = RunArgs::parse_from(
       Some(Ok(format!("{cmd} run")))
         .into_iter()
-        .chain(BufReader::new(File::open(file)?).lines())
+        .chain(std::io::BufReader::new(File::open(file)?).lines())
         .filter(|x| !x.as_ref().is_ok_and(|x| x.is_empty() || x.chars().next().unwrap() == '#'))
         .map_ok(|x| "--".to_string() + &x)
         .collect::<Result<Vec<_>, _>>()?,
@@ -72,7 +73,7 @@ async fn run(mut args: RunArgs, sock_path: &str) -> anyhow::Result<ExitCode> {
         result = listener.accept(), if matches!(bgp.state(), bgp::State::Active) => {
           let (stream, mut addr) = result.context("failed to accept TCP connection")?;
           addr.set_ip(addr.ip().to_canonical());
-          bgp.accept(stream, addr).await.context("failed to accept BGP connection")?;
+          bgp.accept(BufReader::new(stream), addr).await.context("failed to accept BGP connection")?;
         }
         result = bgp.process() => match result {
           Ok(()) => {}