diff --git a/src/hardware/configuration.rs b/src/hardware/configuration.rs
index 513e851..e84723d 100644
--- a/src/hardware/configuration.rs
+++ b/src/hardware/configuration.rs
@@ -9,7 +9,6 @@ use stm32h7xx_hal::{
 
 const NUM_SOCKETS: usize = 5;
 
-use heapless::{consts, Vec};
 use smoltcp_nal::smoltcp;
 
 use embedded_hal::digital::v2::{InputPin, OutputPin};
@@ -41,14 +40,14 @@ pub struct NetStorage {
 #[derive(Copy, Clone)]
 pub struct SocketStorage {
     rx_storage: [u8; 1024],
-    tx_storage: [u8; 1024],
+    tx_storage: [u8; 1024 * 3],
 }
 
 impl SocketStorage {
     const fn new() -> Self {
         Self {
             rx_storage: [0; 1024],
-            tx_storage: [0; 1024],
+            tx_storage: [0; 1024 * 3],
         }
     }
 }
@@ -181,6 +180,8 @@ pub fn setup(
     let gpiof = device.GPIOF.split(ccdr.peripheral.GPIOF);
     let mut gpiog = device.GPIOG.split(ccdr.peripheral.GPIOG);
 
+    let _uart_tx = gpiod.pd8.into_push_pull_output().set_speed(hal::gpio::Speed::VeryHigh);
+
     let dma_streams =
         hal::dma::dma::StreamsTuple::new(device.DMA1, ccdr.peripheral.DMA1);
 
@@ -589,12 +590,10 @@ pub fn setup(
             .routes(routes)
             .finalize();
 
-        let (mut sockets, handles) = {
+        let mut sockets = {
             let mut sockets =
                 smoltcp::socket::SocketSet::new(&mut store.sockets[..]);
 
-            let mut handles: Vec<smoltcp::socket::SocketHandle, consts::U64> =
-                Vec::new();
             for storage in store.socket_storage.iter_mut() {
                 let tcp_socket = {
                     let rx_buffer = smoltcp::socket::TcpSocketBuffer::new(
@@ -606,12 +605,10 @@ pub fn setup(
 
                     smoltcp::socket::TcpSocket::new(rx_buffer, tx_buffer)
                 };
-                let handle = sockets.add(tcp_socket);
-
-                handles.push(handle).unwrap();
+                sockets.add(tcp_socket);
             }
 
-            (sockets, handles)
+            sockets
         };
 
         let dhcp_client = {
@@ -647,7 +644,6 @@ pub fn setup(
         let mut stack = smoltcp_nal::NetworkStack::new(
             interface,
             sockets,
-            &handles,
             Some(dhcp_client),
         );
 
diff --git a/src/net/data_stream.rs b/src/net/data_stream.rs
index 2b025b0..3796440 100644
--- a/src/net/data_stream.rs
+++ b/src/net/data_stream.rs
@@ -3,7 +3,7 @@ use heapless::{
     spsc::{Consumer, Producer, Queue},
     Vec,
 };
-use smoltcp_nal::{smoltcp, embedded_nal::{Mode, SocketAddr, TcpStack}};
+use smoltcp_nal::embedded_nal::{SocketAddr, TcpClientStack};
 
 use super::NetworkReference;
 use crate::hardware::design_parameters::SAMPLE_BUFFER_SIZE;
@@ -66,7 +66,7 @@ impl BlockGenerator {
 
 pub struct DataStream {
     stack: NetworkReference,
-    socket: Option<<NetworkReference as TcpStack>::TcpSocket>,
+    socket: Option<<NetworkReference as TcpClientStack>::TcpSocket>,
     queue: Consumer<'static, AdcDacData, BlockBufferSize>,
     remote: Option<SocketAddr>,
 }
@@ -122,18 +122,18 @@ impl DataStream {
             self.close();
         }
 
-        let socket =
+        let mut socket =
             self.stack
-                .open(Mode::NonBlocking)
+                .socket()
                 .map_err(|err| match err {
-                    <NetworkReference as TcpStack>::Error::NoIpAddress => (),
+                    <NetworkReference as TcpClientStack>::Error::NoIpAddress => (),
                     _ => ()
                 })?;
 
         // TODO: How should we handle a connection failure?
-        let socket = self.stack.connect(socket, remote).unwrap();
+        self.stack.connect(&mut socket, remote).unwrap();
 
-        //log::info!("Stream connecting to {:?}", remote);
+        log::info!("Stream connecting to {:?}", remote);
 
         // Note(unwrap): The socket will be empty before we replace it.
         self.socket.replace(socket);
@@ -166,7 +166,7 @@ impl DataStream {
         };
 
         if !connected {
-            self.socket.replace(self.stack.connect(socket, self.remote.unwrap()).unwrap());
+            self.stack.connect(&mut socket, self.remote.unwrap()).unwrap();
         }
     }
 
@@ -185,12 +185,7 @@ impl DataStream {
         }
 
         let mut handle = self.socket.borrow_mut().unwrap();
-
-        let capacity = self.stack.lock(|stack| {
-            let mut all_sockets = stack.sockets.borrow_mut();
-            let socket: &mut smoltcp::socket::TcpSocket = &mut *all_sockets.get(handle);
-            socket.send_capacity() - socket.send_queue()
-        });
+        let capacity = self.stack.lock(|stack| stack.get_remaining_send_buffer(handle)).unwrap();
 
         // TODO: Clean up magic numbers.
         if capacity < 72 {
@@ -218,7 +213,7 @@ impl DataStream {
 
             // Transmit the data block.
             // TODO: Should we measure how many packets get dropped as telemetry?
-            match self.stack.write(&mut handle, &data) {
+            match self.stack.send(&mut handle, &data) {
                 Ok(len) => {
                     if len != data.len() {
                         log::error!("Short message: {} {}", len, data.len());
diff --git a/src/net/shared.rs b/src/net/shared.rs
index 4f86793..4761da0 100644
--- a/src/net/shared.rs
+++ b/src/net/shared.rs
@@ -69,6 +69,21 @@ where
     forward! {close(socket: S::TcpSocket) -> Result<(), S::Error>}
 }
 
+impl<'a, S> embedded_nal::UdpClientStack for NetworkStackProxy<'a, S>
+where
+    S: embedded_nal::UdpClientStack,
+{
+    type UdpSocket = S::UdpSocket;
+    type Error = S::Error;
+
+    forward! {socket() -> Result<S::UdpSocket, S::Error>}
+    forward! {connect(socket: &mut S::UdpSocket, remote: embedded_nal::SocketAddr) -> Result<(), S::Error>}
+
+    forward! {send(socket: &mut S::UdpSocket, buffer: &[u8]) -> embedded_nal::nb::Result<(), S::Error>}
+    forward! {receive(socket: &mut S::UdpSocket, buffer: &mut [u8]) -> embedded_nal::nb::Result<(usize, embedded_nal::SocketAddr), S::Error>}
+    forward! {close(socket: S::UdpSocket) -> Result<(), S::Error>}
+}
+
 impl NetworkManager {
     /// Construct a new manager for a shared network stack
     ///