Untitled

package it.danilotallarico.proxy;

import com.mongodb.client.*;
import org.bson.Document;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.net.*;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.*;

public class MultiPortReverseProxy {
    private static final Logger log = LoggerFactory.getLogger(MultiPortReverseProxy.class);

    // --- ENV / Config ---
    private final String mongoUri = env("MONGO_URI", "mongodb://127.0.0.1:27017");
    private final String dbName = env("DB_NAME", "reverse_proxy");
    private final String collName = env("COLLECTION", "routes");
    private final int refreshSeconds = Integer.parseInt(env("ROUTES_REFRESH_SECONDS", "10"));
    private final int headerLimitBytes = Integer.parseInt(env("HEADER_LIMIT_BYTES", "65536"));
    private final boolean forceConnectionClose = Boolean.parseBoolean(env("FORCE_CONNECTION_CLOSE", "true"));
    private final String defaultBackendHost = env("DEFAULT_BACKEND", ""); // optional fallback host
    private final boolean enableProxyProtocol = Boolean.parseBoolean(env("ENABLE_PROXY_PROTOCOL", "true"));
    private final int maxConns = Integer.parseInt(env("MAX_CONNS", "20000"));

    private final Set<Integer> httpPorts = parsePorts(env("HTTP_PORTS", "80,8080,8880,2052,2082,2086,2095"));
    private final Set<Integer> tlsPorts  = parsePorts(env("TLS_PORTS",  "443,2053,2083,2087,2096,8443"));

    private final ExecutorService vThreads = Executors.newVirtualThreadPerTaskExecutor();
    private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
    private final List<ServerSocket> listeners = new CopyOnWriteArrayList<>();
    private final Semaphore connLimiter = new Semaphore(maxConns);
    private volatile boolean running = true;

    private MongoClient mongo;
    private MongoCollection<Document> routesColl;
    private Map<String, String> routes = new ConcurrentHashMap<>(); // domain -> backendHost

    public static void main(String[] args) throws Exception {
        new MultiPortReverseProxy().start();
    }

    private void start() throws Exception {
        Runtime.getRuntime().addShutdownHook(new Thread(this::stop));

        mongo = MongoClients.create(mongoUri);
        routesColl = mongo.getDatabase(dbName).getCollection(collName);
        refreshRoutes();
        scheduler.scheduleAtFixedRate(this::safeRefreshRoutes, refreshSeconds, refreshSeconds, TimeUnit.SECONDS);

        // Start listeners
        for (int p : httpPorts) startListener(p, true);
        for (int p : tlsPorts)  startListener(p, false);

        log.info("Listeners avviati. HTTP: {} | TLS: {}", httpPorts, tlsPorts);
        // keep alive
        while (running) Thread.sleep(1000);
    }

    private void stop() {
        running = false;
        listeners.forEach(this::closeQuietly);
        closeQuietly(mongo);
        scheduler.shutdownNow();
        vThreads.shutdownNow();
        log.info("Proxy arrestato");
    }

    private void startListener(int port, boolean isHttp) {
        vThreads.submit(() -> {
            try (ServerSocket ss = new ServerSocket()) {
                ss.setReuseAddress(true);
                ss.bind(new InetSocketAddress("0.0.0.0", port), 65535);
                listeners.add(ss);
                log.info("In ascolto su :{} ({})", port, isHttp ? "HTTP" : "TLS");
                while (running) {
                    try {
                        Socket client = ss.accept();
                        if (!connLimiter.tryAcquire()) { quick503(client, isHttp); continue; }
                        vThreads.submit(() -> handle(client, isHttp));
                    } catch (IOException e) { if (running) log.warn("Accept {}: {}", port, e.toString()); }
                }
            } catch (IOException e) {
                log.error("Listener {} fallito: {}", port, e.toString());
            }
        });
    }

    private void handle(Socket client, boolean isHttp) {
        try (client) {
            final int localPort = ((InetSocketAddress) client.getLocalSocketAddress()).getPort();
            final String clientIp = ((InetAddress) ((InetSocketAddress) client.getRemoteSocketAddress()).getAddress()).getHostAddress();
            client.setTcpNoDelay(true);
            client.setSoTimeout(15000);

            InputStream cin = new BufferedInputStream(client.getInputStream());
            OutputStream cout = new BufferedOutputStream(client.getOutputStream());

            if (isHttp) {
                // Read headers
                byte[] hdr = readHttpHeaders(cin, headerLimitBytes);
                if (hdr == null) return;
                ParsedRequest req = parseRequest(hdr);
                if (req == null || req.host() == null) { badRequest(cout); return; }

                String domain = req.host().toLowerCase(Locale.ROOT);
                String backendHost = routes.get(domain);
                if (backendHost == null || backendHost.isBlank()) backendHost = defaultBackendHost;
                if (backendHost == null || backendHost.isBlank()) { notFound(cout, domain); return; }

                try (Socket be = new Socket()) {
                    be.setTcpNoDelay(true);
                    be.connect(new InetSocketAddress(backendHost, localPort), 3000);
                    InputStream bin = new BufferedInputStream(be.getInputStream());
                    OutputStream bout = new BufferedOutputStream(be.getOutputStream());

                    // rewrite + forward
                    byte[] rewritten = rewriteRequestHeaders(req, clientIp);
                    bout.write(rewritten); bout.flush();

                    Future<?> up = vThreads.submit(() -> pipe(cin, bout));
                    Future<?> dn = vThreads.submit(() -> pipe(bin, cout));
                    try { up.get(); } catch (Exception ignored) {}
                    try { dn.get(); } catch (Exception ignored) {}
                } catch (IOException ioe) { badGateway(cout); }

            } else {
                // TLS passthrough with SNI routing and optional PROXY v1
                // read initial bytes to parse ClientHello (up to 8KB)
                byte[] buf = readSome(cin, 8192, 1500);
                String sni = parseSNI(buf);
                String backendHost = (sni != null) ? routes.get(sni.toLowerCase(Locale.ROOT)) : null;
                if (backendHost == null || backendHost.isBlank()) backendHost = defaultBackendHost;
                if (backendHost == null || backendHost.isBlank()) { // no route → silent close
                    return;
                }

                try (Socket be = new Socket()) {
                    be.setTcpNoDelay(true);
                    be.connect(new InetSocketAddress(backendHost, localPort), 3000);
                    InputStream bin = new BufferedInputStream(be.getInputStream());
                    OutputStream bout = new BufferedOutputStream(be.getOutputStream());

                    // Optional PROXY v1 line BEFORE TLS bytes
                    if (enableProxyProtocol) {
                        String fam = (client.getInetAddress() instanceof Inet4Address) ? "TCP4" : "TCP6";
                        int cPort = ((InetSocketAddress) client.getRemoteSocketAddress()).getPort();
                        int dPort = ((InetSocketAddress) be.getLocalSocketAddress()).getPort();
                        String line = String.format("PROXY %s %s %s %d %d\r\n", fam, clientIp, dstIp, srcPort, dstPort);
                        bout.write(line.getBytes(StandardCharsets.US_ASCII));
                        bout.write(line.getBytes(StandardCharsets.US_ASCII));
                        bout.write(line.getBytes(StandardCharsets.US_ASCII));
                    }
                    // write the bytes already read, then pipe
                    if (buf != null && buf.length > 0) { bout.write(buf); }
                    bout.flush();

                    Future<?> up = vThreads.submit(() -> pipe(cin, bout));
                    Future<?> dn = vThreads.submit(() -> pipe(bin, cout));
                    try { up.get(); } catch (Exception ignored) {}
                    try { dn.get(); } catch (Exception ignored) {}
                } catch (IOException ioe) { /* drop */ }
            }
        } catch (IOException ignored) {
        } finally {
            connLimiter.release();
        }
    }

    // ---- Routes ----
    private void safeRefreshRoutes() { try { refreshRoutes(); } catch (Exception e) { log.warn("Refresh routes: {}", e.toString()); } }
    private void refreshRoutes() {
        Map<String,String> map = new ConcurrentHashMap<>();
        try (MongoCursor<Document> cur = routesColl.find().iterator()) {
            while (cur.hasNext()) {
                Document d = cur.next();
                String domain = opt(d, "domain");
                String host = opt(d, "backendHost");
                if (domain != null && !domain.isBlank() && host != null && !host.isBlank()) {
                    map.put(domain.toLowerCase(Locale.ROOT), host);
                }
            }
        }
        routes = map;
        log.info("Routes caricate: {}", routes.size());
    }

    // ---- HTTP helpers ----
    private static byte[] readHttpHeaders(InputStream in, int limit) throws IOException {
        ByteArrayOutputStream buf = new ByteArrayOutputStream();
        int b, count = 0, state = 0;
        while ((b = in.read()) != -1) {
            buf.write(b); count++;
            if (count > limit) return null;
            switch (state) {
                case 0 -> state = (b == '\r') ? 1 : 0;
                case 1 -> state = (b == '\n') ? 2 : 0;
                case 2 -> state = (b == '\r') ? 3 : 0;
                case 3 -> { if (b == '\n') return buf.toByteArray(); else state = 0; }
            }
        }
        return null;
    }

    record ParsedRequest(String method, String uri, String version, Map<String,String> headers, String host) {}

    private ParsedRequest parseRequest(byte[] headerBytes) {
        String head = new String(headerBytes, StandardCharsets.ISO_8859_1);
        String[] lines = head.split("\r\n");
        if (lines.length == 0) return null;
        String[] reqLine = lines[0].split(" ", 3);
        if (reqLine.length < 3) return null;
        String method = reqLine[0], uri = reqLine[1], version = reqLine[2];
        Map<String,String> headers = new LinkedHashMap<>();
        String host = null;
        for (int i=1; i<lines.length; i++) {
            String ln = lines[i];
            if (ln.isEmpty()) break;
            int idx = ln.indexOf(':'); if (idx <= 0) continue;
            String k = ln.substring(0, idx).trim();
            String v = ln.substring(idx+1).trim();
            headers.put(k, v);
            if (k.equalsIgnoreCase("Host")) host = v.split("\\s+")[0];
        }
        return new ParsedRequest(method, uri, version, headers, host);
    }

    private byte[] rewriteRequestHeaders(ParsedRequest req, String clientIp) {
        Map<String,String> h = new LinkedHashMap<>();
        req.headers().forEach((k,v) -> {
            if (k.equalsIgnoreCase("Connection") && FORCE_CONNECTION_CLOSE) return; // overwrite later
            h.put(k, v);
        });
        String xff = firstHeaderIgnoreCase(h, "X-Forwarded-For");
        h.put("X-Forwarded-For", (xff == null || xff.isBlank()) ? clientIp : xff+", "+clientIp);
        h.put("X-Real-IP", clientIp);
        h.put("CF-Connecting-IP", clientIp);
        h.put("X-Forwarded-Proto", "http");
        if (FORCE_CONNECTION_CLOSE) h.put("Connection", "close");

        StringBuilder sb = new StringBuilder();
        sb.append(req.method()).append(' ').append(req.uri()).append(' ').append(req.version()).append("\r\n");
        h.forEach((k,v)-> sb.append(k).append(": ").append(v).append("\r\n"));
        sb.append("\r\n");
        return sb.toString().getBytes(StandardCharsets.ISO_8859_1);
    }

    private byte[] rewriteRequestHeaders(ParsedRequest req, String clientIp) {
        Map<String,String> h = new LinkedHashMap<>();
        req.headers().forEach((k,v) -> {
            if (k.equalsIgnoreCase("Connection") && forceConnectionClose) return; // overwrite later
            h.put(k, v);
        });
        String xff = firstHeaderIgnoreCase(h, "X-Forwarded-For");
        h.put("X-Forwarded-For", (xff == null || xff.isBlank()) ? clientIp : xff+", "+clientIp);
        h.put("X-Real-IP", clientIp);
        h.put("CF-Connecting-IP", clientIp);
        h.put("X-Forwarded-Proto", "http");
        if (forceConnectionClose) h.put("Connection", "close");

        StringBuilder sb = new StringBuilder();
        sb.append(req.method()).append(' ').append(req.uri()).append(' ').append(req.version()).append("\r\n");
        h.forEach((k,v)-> sb.append(k).append(": ").append(v).append("\r\n"));
        sb.append("\r\n");
        return sb.toString().getBytes(StandardCharsets.ISO_8859_1);
    }

    private static String firstHeaderIgnoreCase(Map<String,String> headers, String name) {
        for (String k: headers.keySet()) if (k.equalsIgnoreCase(name)) return headers.get(k);
        return null;
    }

    // ---- TLS helpers ----
    private static byte[] readSome(InputStream in, int maxBytes, int timeoutMs) throws IOException {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        long end = System.currentTimeMillis() + timeoutMs;
        while (out.size() < maxBytes && System.currentTimeMillis() < end) {
            if (in.available() > 0) {
                byte[] buf = new byte[Math.min(in.available(), maxBytes - out.size())];
                int r = in.read(buf);
                if (r == -1) break; out.write(buf, 0, r);
                // stop early if we already have a full TLS record header + length
                if (out.size() >= 5) {
                    int len = ((out.toByteArray()[3]&0xFF)<<8) | (out.toByteArray()[4]&0xFF);
                    if (out.size() >= 5 + len) break;
                }
            } else {
                try { Thread.sleep(10); } catch (InterruptedException ignored) {}
            }
        }
        return out.toByteArray();
    }

    private static String parseSNI(byte[] data) {
        try {
            if (data == null || data.length < 5) return null;
            // TLS record header
            int pos = 0;
            if ((data[pos++] & 0xFF) != 0x16) return null; // Handshake
            pos += 2; // version
            int recLen = ((data[pos++] & 0xFF) << 8) | (data[pos++] & 0xFF);
            if (data.length < 5 + recLen) return null;
            // Handshake header
            if ((data[pos++] & 0xFF) != 0x01) return null; // ClientHello
            int hsLen = ((data[pos++] & 0xFF) << 16) | ((data[pos++] & 0xFF) << 8) | (data[pos++] & 0xFF);
            pos += 2; // client_version
            pos += 32; // random
            int sidLen = data[pos++] & 0xFF; pos += sidLen;
            int csLen = ((data[pos++] & 0xFF) << 8) | (data[pos++] & 0xFF); pos += csLen;
            int cmLen = data[pos++] & 0xFF; pos += cmLen;
            if (pos + 2 > data.length) return null;
            int extLen = ((data[pos++] & 0xFF) << 8) | (data[pos++] & 0xFF);
            int end = pos + extLen;
            while (pos + 4 <= end && end <= data.length) {
                int type = ((data[pos++] & 0xFF) << 8) | (data[pos++] & 0xFF);
                int len = ((data[pos++] & 0xFF) << 8) | (data[pos++] & 0xFF);
                if (type == 0x00) { // server_name
                    int listLen = ((data[pos++] & 0xFF) << 8) | (data[pos++] & 0xFF);
                    int listEnd = pos + listLen;
                    while (pos + 3 <= listEnd) {
                        int nameType = data[pos++] & 0xFF;
                        int nameLen = ((data[pos++] & 0xFF) << 8) | (data[pos++] & 0xFF);
                        if (nameType == 0 && pos + nameLen <= listEnd) {
                            String sni = new String(data, pos, nameLen, StandardCharsets.US_ASCII).toLowerCase(Locale.ROOT);
                            return sni;
                        } else { pos += nameLen; }
                    }
                    return null;
                } else { pos += len; }
            }
        } catch (Throwable ignored) {}
        return null;
    }

    // ---- Utilities ----
    private static void pipe(InputStream in, OutputStream out) {
        byte[] buf = new byte[32 * 1024];
        int r;
        try { while ((r = in.read(buf)) != -1) { out.write(buf, 0, r); out.flush(); } }
        catch (IOException ignored) {}
        finally { try { out.flush(); } catch (Exception ignored) {} }
    }

    private static void quick503(Socket c, boolean isHttp) {
        try (c) {
            if (isHttp) {
                OutputStream o = c.getOutputStream();
                o.write("HTTP/1.1 503 Service Unavailable\r\nConnection: close\r\nContent-Length: 0\r\n\r\n".getBytes(StandardCharsets.ISO_8859_1));
                o.flush();
            }
        } catch (IOException ignored) {}
    } catch (IOException ignored) {}
    }

    private static void badRequest(OutputStream out) throws IOException {
        out.write("HTTP/1.1 400 Bad Request\r\nConnection: close\r\nContent-Length: 0\r\n\r\n".getBytes(StandardCharsets.ISO_8859_1));
        out.flush();
    }

    private static void notFound(OutputStream out, String host) throws IOException {
        byte[] b = ("No route for host: " + host).getBytes(StandardCharsets.UTF_8);
        String h = "HTTP/1.1 404 Not Found\r\nConnection: close\r\nContent-Type: text/plain; charset=utf-8\r\nContent-Length: "+b.length+"\r\n\r\n";
        out.write(h.getBytes(StandardCharsets.ISO_8859_1)); out.write(b); out.flush();
    }

    private static void badGateway(OutputStream out) throws IOException {
        out.write("HTTP/1.1 502 Bad Gateway\r\nConnection: close\r\nContent-Length: 0\r\n\r\n".getBytes(StandardCharsets.ISO_8859_1));
        out.flush();
    }

    private static String env(String k, String def) {
        String v = System.getenv(k); return (v == null || v.isBlank()) ? def : v.trim();
    }

    private static Set<Integer> parsePorts(String s) {
        Set<Integer> out = new LinkedHashSet<>();
        for (String p : s.split(",")) {
            p = p.trim(); if (p.isEmpty()) continue; out.add(Integer.parseInt(p));
        }
        return out;
    }

    private static String opt(Document d, String k) { return d.getString(k); }

    private void closeQuietly(AutoCloseable c) { if (c!=null) try { c.close(); } catch (Exception ignored) {} }
}