axum-forwarded-tcpstream.rs

use axum::{body::Body, http::StatusCode, response::IntoResponse, routing::get, Router};
use futures::Stream;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, BufReader, ReadBuf};
use tokio::net::{TcpListener, TcpStream};
use tower_http::cors::CorsLayer;

// This struct represent a stream and an amount of bytes to read from it
struct LimitedStream {
    reader: BufReader<TcpStream>,
    remaining: usize,
}

impl LimitedStream {
    const fn new(reader: BufReader<TcpStream>, remaining: usize) -> Self {
        Self { reader, remaining }
    }
}

impl Stream for LimitedStream {
    type Item = Result<Vec<u8>, std::io::Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();

        if this.remaining == 0 {
            return Poll::Ready(None);
        }

        // Create a buffer that is big enough to hold the necessary data but not bigger than 1024 bytes
        let buf_size = std::cmp::min(1024, this.remaining);
        let mut buf = vec![0; buf_size];
        let mut read_buf = ReadBuf::new(&mut buf);

        // Start reading
        match Pin::new(&mut this.reader).poll_read(cx, &mut read_buf) {
            Poll::Ready(Ok(())) => {
                // Get how many bytes were read
                let n = read_buf.filled().len();
                if n == 0 {
                    Poll::Ready(None)
                } else {
                    // Decrement the amount of bytes left by what was read
                    this.remaining -= n;
                    buf.truncate(n);
                    Poll::Ready(Some(Ok(buf)))
                }
            }
            Poll::Ready(Err(e)) => Poll::Ready(Some(Err(e))),
            Poll::Pending => Poll::Pending,
        }
    }
}

async fn handler() -> impl IntoResponse {
    // Open a TCP connection with the console
    let stream = TcpStream::connect("192.168.1.18:730").await.unwrap();

    // Send the command to get the contents of a file
    stream.writable().await.unwrap();
    stream
        .try_write(b"getfile name=\"HDD:\\xbdm.ini\"\r\nbye\r\n")
        .unwrap();

    // Read the "201- connected\r\n" message into a dummy buffer
    stream.readable().await.unwrap();
    let mut reader = BufReader::new(stream);
    let mut connected_message = String::new();
    reader.read_line(&mut connected_message).await.unwrap();
    if connected_message != "201- connected\r\n" {
        return (
            StatusCode::INTERNAL_SERVER_ERROR,
            Body::from("Not connected"),
        );
    }

    // Read the "203- binary response follows\r\n" message
    let mut binary_response_message = String::new();
    reader
        .read_line(&mut binary_response_message)
        .await
        .unwrap();
    if binary_response_message != "203- binary response follows\r\n" {
        return (StatusCode::NOT_FOUND, Body::from("Not found"));
    }

    // Read the content length from the response (first 4 bytes)
    let mut length_buffer = [0_u8; 4];
    reader.read_exact(&mut length_buffer).await.unwrap();
    let content_length = u32::from_le_bytes(length_buffer);

    // Create a stream that will read the content_length bytes
    let limited_stream = LimitedStream::new(reader, content_length as usize);
    let body = Body::from_stream(limited_stream);

    (StatusCode::OK, body)
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/", get(handler))
        .layer(CorsLayer::permissive());

    let listener = TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}