sten.wtf / STEN Vault source

src/main/java/wtf/sten/vault/Proxy.java

package wtf.sten.vault;

import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
import io.netty.channel.ChannelDuplexHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelPromise;

import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.charset.StandardCharsets;

public final class Proxy extends ChannelDuplexHandler {
	private final InetSocketAddress proxy;
	private final String user;
	private final String pass;
	private InetSocketAddress dest;
	private ChannelPromise connectPromise;
	private int step;
	private ByteBuf pending;

	private Proxy(InetSocketAddress proxy, String user, String pass) {
		this.proxy = proxy;
		this.user = user == null ? "" : user;
		this.pass = pass == null ? "" : pass;
	}

	public static void install(Channel channel) {
		Parsed parsed = parse(VaultConfig.get().proxy);
		if (parsed == null) return;
		channel.pipeline().addFirst("sten-socks", new Proxy(parsed.address, parsed.user, parsed.pass));
	}

	public static Parsed parse(String raw) {
		if (raw == null) return null;
		String value = raw.trim();
		if (value.isEmpty()) return null;
		if (value.startsWith("socks5://") || value.startsWith("socks://")) {
			value = value.substring(value.indexOf("://") + 3);
		}
		String user = "";
		String pass = "";
		int at = value.lastIndexOf('@');
		if (at > 0) {
			String auth = value.substring(0, at);
			value = value.substring(at + 1);
			int colon = auth.indexOf(':');
			if (colon >= 0) {
				user = auth.substring(0, colon);
				pass = auth.substring(colon + 1);
			} else {
				user = auth;
			}
		}
		int colon = value.lastIndexOf(':');
		if (colon <= 0 || colon == value.length() - 1) return null;
		String host = value.substring(0, colon).replace("[", "").replace("]", "");
		int port;
		try {
			port = Integer.parseInt(value.substring(colon + 1));
		} catch (Exception ignored) {
			return null;
		}
		if (host.isBlank() || port <= 0 || port > 65535) return null;
		return new Parsed(new InetSocketAddress(host, port), user, pass);
	}

	@Override
	public void connect(ChannelHandlerContext ctx, SocketAddress remote, SocketAddress local, ChannelPromise promise) {
		if (!(remote instanceof InetSocketAddress address)) {
			ctx.connect(remote, local, promise);
			return;
		}
		dest = address;
		connectPromise = promise;
		ctx.connect(proxy, local, ctx.newPromise().addListener(future -> {
			if (!future.isSuccess() && connectPromise != null) {
				connectPromise.setFailure(future.cause());
			}
		}));
	}

	@Override
	public void channelActive(ChannelHandlerContext ctx) throws Exception {
		super.channelActive(ctx);
		ByteBuf hello = ctx.alloc().buffer(4);
		hello.writeByte(5);
		if (user.isEmpty()) {
			hello.writeByte(1);
			hello.writeByte(0);
		} else {
			hello.writeByte(2);
			hello.writeByte(0);
			hello.writeByte(2);
		}
		ctx.writeAndFlush(hello);
		step = 1;
	}

	@Override
	public void channelRead(ChannelHandlerContext ctx, Object msg) {
		if (step >= 4 || !(msg instanceof ByteBuf buf)) {
			ctx.fireChannelRead(msg);
			return;
		}
		if (pending == null) pending = ctx.alloc().buffer();
		pending.writeBytes(buf);
		buf.release();
		while (consume(ctx)) {
		}
	}

	private boolean consume(ChannelHandlerContext ctx) {
		if (step == 1) {
			if (pending.readableBytes() < 2) return false;
			pending.readByte();
			int method = pending.readByte() & 0xff;
			if (method == 2 && !user.isEmpty()) {
				byte[] u = user.getBytes(StandardCharsets.UTF_8);
				byte[] p = pass.getBytes(StandardCharsets.UTF_8);
				ByteBuf auth = ctx.alloc().buffer(3 + u.length + p.length);
				auth.writeByte(1);
				auth.writeByte(u.length);
				auth.writeBytes(u);
				auth.writeByte(p.length);
				auth.writeBytes(p);
				ctx.writeAndFlush(auth);
				step = 2;
			} else if (method == 0) {
				sendConnect(ctx);
				step = 3;
			} else {
				fail(ctx, "SOCKS5 method rejected");
			}
			return true;
		}
		if (step == 2) {
			if (pending.readableBytes() < 2) return false;
			pending.readByte();
			if ((pending.readByte() & 0xff) != 0) {
				fail(ctx, "SOCKS5 login rejected");
				return false;
			}
			sendConnect(ctx);
			step = 3;
			return true;
		}
		if (step == 3) {
			if (pending.readableBytes() < 4) return false;
			pending.markReaderIndex();
			pending.readByte();
			int status = pending.readByte() & 0xff;
			pending.readByte();
			int atyp = pending.readByte() & 0xff;
			int extra = atyp == 1 ? 6 : atyp == 4 ? 18 : atyp == 3 ? 1 : -1;
			if (extra < 0) {
				fail(ctx, "SOCKS5 bind failed");
				return false;
			}
			if (atyp == 3) {
				if (pending.readableBytes() < 1) {
					pending.resetReaderIndex();
					return false;
				}
				extra = (pending.readByte() & 0xff) + 2;
			}
			if (pending.readableBytes() < extra) {
				pending.resetReaderIndex();
				return false;
			}
			pending.skipBytes(extra);
			if (status != 0) {
				fail(ctx, "SOCKS5 connect failed (" + status + ")");
				return false;
			}
			step = 4;
			if (connectPromise != null) connectPromise.trySuccess();
			if (pending.readableBytes() > 0) ctx.fireChannelRead(pending.retain());
			pending.release();
			pending = null;
			return false;
		}
		return false;
	}

	private void sendConnect(ChannelHandlerContext ctx) {
		String host = dest.getHostString();
		byte[] name = host.getBytes(StandardCharsets.US_ASCII);
		ByteBuf req = ctx.alloc().buffer(7 + name.length);
		req.writeByte(5);
		req.writeByte(1);
		req.writeByte(0);
		req.writeByte(3);
		req.writeByte(name.length);
		req.writeBytes(name);
		req.writeShort(dest.getPort());
		ctx.writeAndFlush(req);
	}

	private void fail(ChannelHandlerContext ctx, String reason) {
		if (connectPromise != null) connectPromise.tryFailure(new IllegalStateException(reason));
		ctx.close();
	}

	public record Parsed(InetSocketAddress address, String user, String pass) {
	}
}