Skip to content

Commit

Permalink
feat: specify the max TCP connections manually
Browse files Browse the repository at this point in the history 
support specifying the max TCP connections manually
  • Loading branch information
@X1r0z
X1r0z committed Feb 2, 2025
1 parent 352e442 commit 47b0011
Show file tree
Hide file tree
Showing 5 changed files with 59 additions and 28 deletions.
28 changes: 18 additions & 10 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,11 +64,12 @@ Port forwarding mode
Usage: pivot fwd [OPTIONS]

Options:
-l, --locals <LOCALS> Local listen IP address, format: [+][IP:]PORT
-r, --remotes <REMOTES> Remote connect IP address, format: [+]IP:PORT
-s, --socket <SOCKET> Unix domain socket path
-p, --protocol <PROTOCOL> Forward Protocol [default: tcp] [possible values: tcp, udp]
-h, --help Print help (see more with '--help')
-l, --locals <LOCALS> Local listen IP address, format: [+][IP:]PORT
-r, --remotes <REMOTES> Remote connect IP address, format: [+]IP:PORT
-s, --socket <SOCKET> Unix domain socket path
-p, --protocol <PROTOCOL> Forward Protocol [default: tcp] [possible values: tcp, udp]
-c, --connections <CONNECTIONS> Maximum connections [default: 32]
-h, --help Print help (see more with '--help')
```

Socks proxy mode
Expand All @@ -81,10 +82,11 @@ Socks proxy mode
Usage: pivot proxy [OPTIONS]

Options:
-l, --locals <LOCALS> Local listen IP address, format: [+][IP:]PORT
-r, --remote <REMOTE> Reverse server IP address, format: [+]IP:PORT
-a, --auth <AUTH> Authentication info, format: user:pass (other for random)
-h, --help Print help
-l, --locals <LOCALS> Local listen IP address, format: [+][IP:]PORT
-r, --remote <REMOTE> Reverse server IP address, format: [+]IP:PORT
-a, --auth <AUTH> Authentication info, format: user:pass (other for random)
-c, --connections <CONNECTIONS> Maximum connections [default: 32]
-h, --help Print help
```

Port reuse mode
Expand Down Expand Up @@ -127,6 +129,8 @@ Connect `10.0.0.1:8888` and `10.0.0.2:9999`, forward traffic between them.
./pivot fwd -r 10.0.0.1:8888 -r 10.0.0.1:9999
```

In this mode, specifying `-c` can set the maximum number of TCP connections (default is 32)

A basic example of accessing an intranet address through port forwarding.

```bash
Expand Down Expand Up @@ -230,7 +234,7 @@ curl http://vps:5555/version

### Socks Proxy

`pivot-rs` supports socks5 protocol (no/with authentication)
`pivot-rs` supports socks5 proxy (no/with authentication)

Forward socks proxy

Expand All @@ -256,6 +260,10 @@ The port 7777 in the above example is called the control port, which uses TCP mu

Therefore, the order of ports 7777 and 8888 **cannot be reversed**

In addition, in this scenario, the victim machine can specify the `-c` parameter to set the maximum number of connections (the default is 32)

*The maximum number of connections here refers to the maximum number of streams processed simultaneously in the TCP multiplexing scenario*

To enable authentication, simply add `user:pass` after the `-a` flag.

```bash
Expand Down
26 changes: 17 additions & 9 deletions README.zh.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,11 +64,12 @@ Port forwarding mode
Usage: pivot fwd [OPTIONS]

Options:
-l, --locals <LOCALS> Local listen IP address, format: [+][IP:]PORT
-r, --remotes <REMOTES> Remote connect IP address, format: [+]IP:PORT
-s, --socket <SOCKET> Unix domain socket path
-p, --protocol <PROTOCOL> Forward Protocol [default: tcp] [possible values: tcp, udp]
-h, --help Print help (see more with '--help')
-l, --locals <LOCALS> Local listen IP address, format: [+][IP:]PORT
-r, --remotes <REMOTES> Remote connect IP address, format: [+]IP:PORT
-s, --socket <SOCKET> Unix domain socket path
-p, --protocol <PROTOCOL> Forward Protocol [default: tcp] [possible values: tcp, udp]
-c, --connections <CONNECTIONS> Maximum connections [default: 32]
-h, --help Print help (see more with '--help')
```

Socks 代理模式
Expand All @@ -81,10 +82,11 @@ Socks proxy mode
Usage: pivot proxy [OPTIONS]

Options:
-l, --locals <LOCALS> Local listen IP address, format: [+][IP:]PORT
-r, --remote <REMOTE> Reverse server IP address, format: [+]IP:PORT
-a, --auth <AUTH> Authentication info, format: user:pass (other for random)
-h, --help Print help
-l, --locals <LOCALS> Local listen IP address, format: [+][IP:]PORT
-r, --remote <REMOTE> Reverse server IP address, format: [+]IP:PORT
-a, --auth <AUTH> Authentication info, format: user:pass (other for random)
-c, --connections <CONNECTIONS> Maximum connections [default: 32]
-h, --help Print help
```

端口复用模式
Expand Down Expand Up @@ -127,6 +129,8 @@ Options:
./pivot fwd -r 10.0.0.1:8888 -r 10.0.0.1:9999
```

在这种模式下, 指定 `-c` 可以设置最大 TCP 连接数 (默认为 32)

一个简单的内网端口转发的示例.

```bash
Expand Down Expand Up @@ -256,6 +260,10 @@ curl http://vps:5555/version

因此, 7777 和 8888 端口的顺序**不能颠倒**

另外, 在这种场景下, 受害者机器可以指定 `-c` 参数以设置最大连接数 (默认为 32)

*这里的最大连接数指的是在 TCP 多路复用场景下同时处理的最大 Stream 数*

如果需要开启身份验证, 只需要在 `-a` 参数后添加 `user:pass`

```bash
Expand Down
9 changes: 6 additions & 3 deletions src/forward.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,14 +11,15 @@ use tracing::{error, info};
#[cfg(target_family = "unix")]
use tokio::net::UnixStream;

use crate::{crypto, tcp, udp, Protocol, MAX_CONNECTIONS};
use crate::{crypto, tcp, udp, Protocol};

pub struct Forward {
locals: Vec<(String, bool)>,
remotes: Vec<(String, bool)>,
#[cfg(target_family = "unix")]
socket: Option<String>,
protocol: Protocol,
connections: usize,
}

impl Forward {
Expand All @@ -27,13 +28,15 @@ impl Forward {
remotes: Vec<(String, bool)>,
#[cfg(target_family = "unix")] socket: Option<String>,
protocol: Protocol,
connections: usize,
) -> Self {
Self {
locals,
remotes,
#[cfg(target_family = "unix")]
socket,
protocol,
connections,
}
}

Expand Down Expand Up @@ -160,7 +163,7 @@ impl Forward {
let connector2 = Arc::new(ssl2.then(|| crypto::get_tls_connector()));

// limit the number of concurrent connections
let semaphore = Arc::new(sync::Semaphore::new(MAX_CONNECTIONS));
let semaphore = Arc::new(sync::Semaphore::new(self.connections));

loop {
let permit = Arc::clone(&semaphore).acquire_owned().await.unwrap();
Expand Down Expand Up @@ -233,7 +236,7 @@ impl Forward {
let connector = Arc::new(ssl.then(|| crypto::get_tls_connector()));

// limit the number of concurrent connections
let semaphore = Arc::new(sync::Semaphore::new(MAX_CONNECTIONS));
let semaphore = Arc::new(sync::Semaphore::new(self.connections));

loop {
let permit = Arc::clone(&semaphore).acquire_owned().await.unwrap();
Expand Down
18 changes: 14 additions & 4 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,6 @@ mod tcp;
mod udp;
mod util;

pub const MAX_CONNECTIONS: usize = 32;
pub const MAX_MUX_CONNECTIONS: usize = 128;

#[derive(Parser)]
#[command(author, version, about = "Pivot: Port-Forwarding and Proxy Tool")]
pub struct Cli {
Expand Down Expand Up @@ -46,6 +43,11 @@ enum Mode {
#[arg(short, long)]
#[clap(value_enum, default_value = "tcp")]
protocol: Protocol,

/// Maximum connections
#[arg(short, long)]
#[clap(default_value = "32")]
connections: usize,
},

/// Socks proxy mode
Expand All @@ -61,6 +63,11 @@ enum Mode {
/// Authentication info, format: user:pass (other for random)
#[arg(short, long)]
auth: Option<String>,

/// Maximum connections
#[arg(short, long)]
#[clap(default_value = "32")]
connections: usize,
},

/// Port reuse mode
Expand Down Expand Up @@ -103,6 +110,7 @@ pub async fn run(cli: Cli) -> Result<()> {
#[cfg(target_family = "unix")]
socket,
protocol,
connections,
} => {
info!("Starting forward mode");

Expand All @@ -120,6 +128,7 @@ pub async fn run(cli: Cli) -> Result<()> {
#[cfg(target_family = "unix")]
socket,
protocol,
connections,
);

forward.start().await?;
Expand All @@ -128,14 +137,15 @@ pub async fn run(cli: Cli) -> Result<()> {
locals,
remote,
auth,
connections,
} => {
info!("Starting proxy mode");

let locals = util::parse_addrs(locals);
let remote = util::parse_addr(remote);
let auth = auth.map(|v| socks::UserPassAuth::new(v));

let proxy = Proxy::new(locals, remote, auth);
let proxy = Proxy::new(locals, remote, auth, connections);
proxy.start().await?;
}
Mode::Reuse {
Expand Down
6 changes: 4 additions & 2 deletions src/proxy.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,25 +9,27 @@ use crate::{
crypto,
socks::{self, UserPassAuth},
tcp::{self},
MAX_MUX_CONNECTIONS,
};

pub struct Proxy {
locals: Vec<(String, bool)>,
remote: Option<(String, bool)>,
auth: Option<UserPassAuth>,
connections: usize,
}

impl Proxy {
pub fn new(
locals: Vec<(String, bool)>,
remote: Option<(String, bool)>,
auth: Option<UserPassAuth>,
connections: usize,
) -> Self {
Self {
locals,
remote,
auth,
connections,
}
}

Expand Down Expand Up @@ -83,7 +85,7 @@ impl Proxy {
let auth = Arc::new(self.auth.clone());

// limit the number of concurrent connections
let semaphore = Arc::new(tokio::sync::Semaphore::new(MAX_MUX_CONNECTIONS));
let semaphore = Arc::new(tokio::sync::Semaphore::new(self.connections));

loop {
let permit = Arc::clone(&semaphore).acquire_owned().await;
Expand Down

0 comments on commit 47b0011

Please sign in to comment.