a3s-gateway 0.2.1

A3S Gateway - AI-native API gateway with reverse proxy, routing, and agent orchestration
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

//! TCP proxy — bidirectional byte stream relay
//!
//! Handles raw TCP proxying by establishing a connection to the upstream
//! backend and relaying bytes in both directions.

use crate::error::{GatewayError, Result};
use tokio::io::AsyncWriteExt;
use tokio::net::TcpStream;

/// Relay bytes bidirectionally between two TCP streams
///
/// Copies data from client→upstream and upstream→client concurrently
/// until either side closes the connection.
pub async fn relay_tcp(mut client: TcpStream, mut upstream: TcpStream) -> Result<(u64, u64)> {
    let (mut client_read, mut client_write) = client.split();
    let (mut upstream_read, mut upstream_write) = upstream.split();

    let client_to_upstream = tokio::io::copy(&mut client_read, &mut upstream_write);
    let upstream_to_client = tokio::io::copy(&mut upstream_read, &mut client_write);

    let result = tokio::select! {
        result = client_to_upstream => {
            let bytes_sent = result.map_err(|e| {
                GatewayError::Other(format!("TCP relay client→upstream error: {}", e))
            })?;
            // Client closed, shutdown upstream write side
            let _ = upstream_write.shutdown().await;
            (bytes_sent, 0u64)
        }
        result = upstream_to_client => {
            let bytes_received = result.map_err(|e| {
                GatewayError::Other(format!("TCP relay upstream→client error: {}", e))
            })?;
            // Upstream closed, shutdown client write side
            let _ = client_write.shutdown().await;
            (0u64, bytes_received)
        }
    };

    Ok(result)
}

/// Connect to an upstream TCP server
pub async fn connect_upstream(address: &str) -> Result<TcpStream> {
    TcpStream::connect(address).await.map_err(|e| {
        GatewayError::ServiceUnavailable(format!(
            "TCP upstream connection to {} failed: {}",
            address, e
        ))
    })
}

/// Extract the host:port from a backend URL
///
/// Strips the protocol prefix if present:
/// - "http://127.0.0.1:8001" → "127.0.0.1:8001"
/// - "h2c://127.0.0.1:50051" → "127.0.0.1:50051"
/// - "127.0.0.1:9000" → "127.0.0.1:9000"
pub fn extract_address(url: &str) -> &str {
    if let Some(rest) = url.strip_prefix("http://") {
        rest.trim_end_matches('/')
    } else if let Some(rest) = url.strip_prefix("https://") {
        rest.trim_end_matches('/')
    } else if let Some(rest) = url.strip_prefix("h2c://") {
        rest.trim_end_matches('/')
    } else if let Some(rest) = url.strip_prefix("tcp://") {
        rest.trim_end_matches('/')
    } else {
        url.trim_end_matches('/')
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_extract_address_http() {
        assert_eq!(extract_address("http://127.0.0.1:8001"), "127.0.0.1:8001");
    }

    #[test]
    fn test_extract_address_https() {
        assert_eq!(
            extract_address("https://backend.example.com:443"),
            "backend.example.com:443"
        );
    }

    #[test]
    fn test_extract_address_h2c() {
        assert_eq!(extract_address("h2c://127.0.0.1:50051"), "127.0.0.1:50051");
    }

    #[test]
    fn test_extract_address_tcp() {
        assert_eq!(extract_address("tcp://127.0.0.1:9000"), "127.0.0.1:9000");
    }

    #[test]
    fn test_extract_address_bare() {
        assert_eq!(extract_address("127.0.0.1:9000"), "127.0.0.1:9000");
    }

    #[test]
    fn test_extract_address_trailing_slash() {
        assert_eq!(extract_address("http://127.0.0.1:8001/"), "127.0.0.1:8001");
    }

    #[tokio::test]
    async fn test_connect_upstream_invalid() {
        let result = connect_upstream("127.0.0.1:1").await;
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("TCP upstream connection"));
    }
}