astrogate/src/server/server.vala

using Invercargill;

namespace Astrogate {

    private static Configuration config;
    private static Tunnel? tunnel;
    private static Gee.LinkedList<ProxyConnection> closing_connections;

    public static int main(string[] args) {

        var config_path = args.length > 1 ? args[1] : "/etc/astrogate.conf";
        run.begin(config_path);

        new MainLoop().run();
        return 0;
    }

    private async void run(string config_path) {
        config = new Configuration();
        closing_connections = new Gee.LinkedList<ProxyConnection>();
        print("Reading config...\n");
        yield config.read_file(File.new_for_commandline_arg(config_path));

        var service = new SocketService();
        var service_ports = config.entries.group_by<uint16>(s => s.source_port, (a, b) => a == b);

        foreach (var address in config.bind_to) {
            foreach (var port in service_ports) {
                print(@"Binding to $(address):$(port.key) for hosts: $(port.items.to_string(i => i.hostname, ", "))\n");
                var socket_addr = new InetSocketAddress(address, port.key);
                service.add_address(socket_addr, SocketType.STREAM, SocketProtocol.TCP, null, null);
            }
        }

        if(config.tunnel_enabled) {
            tunnel = new Tunnel(config.tunnel_interface, 1500, config.tunnel_bind, config.tunnel_network, config.tunnel_netmask);
            tunnel.listen_threaded();
            var tunnel_dests = config.entries
                .select_many<ConfigurationDestination>(e => e.destinations)
                .where(d => d.mode == DestinationType.TUNNEL);

            foreach (var dest in tunnel_dests) {
                print(@"Whitelisting $(dest.socket_address.address) in the tunnel server\n");
                tunnel.whitelist_client(dest.socket_address.address);
            }
        }

        service.incoming.connect((c, o) => {
            handle_connection.begin(c);
            return false;
        });
        service.start ();
    }

    private async void handle_connection(SocketConnection connection) {
        try {
            var port = ((InetSocketAddress)connection.get_local_address()).port;
            var buffer = yield HeaderReader.read_sample(connection.input_stream);
            var type = HeaderReader.determine_type(buffer);
            if(type == ConnectionType.INVALID) {
                print(@"Could not determine type of incoming connection on port $port.\n");
                yield connection.close_async();
                return;
            }
            print(@"New $type connection on port $(port)\n");
            var name = HeaderReader.read_name(buffer, type);
            if(name != null) {
                var entry = get_entry(type, (uint16)port, name);

                if(entry == null) {
                    print(@"No matching configuration entry found matching $(type) $(name) $(port).\n");
                    yield connection.close_async();
                    return;
                }

                yield forward_connection(connection, entry, buffer);
            }
            else {
                print(@"Could not determine destination hostname\n");
                yield connection.close_async();
                return;
            }
        }
        catch(Error e) {
            warning(@"Error servicing connection: $(e.message)");
        }
    }

    private async void forward_connection(SocketConnection connection, ConfigurationItem entry, BinaryData initial_buffer) throws Error {
        ProxyConnection proxy;
        var remote_address = (InetSocketAddress)connection.get_remote_address();
        print(@"Proxying connection from $(remote_address) intended for $(entry.hostname):$(entry.source_port).\n");

        switch (entry.forward_mode) {
            case ForwardMode.ROUND_ROBIN:
                proxy = yield connect_round_robin(entry, remote_address);
                break;
            case ForwardMode.RANDOM:
                proxy = yield connect_random(entry, remote_address);
                break;
            case ForwardMode.FIRST_RESPONDER:
                proxy = yield connect_first_responder(entry, remote_address);
                break;
            default:
                assert_not_reached();
        }

        print(@"Connection from $(connection.get_remote_address().to_string()) forwarded to $(proxy.connection.get_remote_address()).\n");
        
        yield proxy.connection.output_stream.write_async(initial_buffer.to_array(), 0);
        yield proxy.connection.output_stream.flush_async(0);

        proxy.connection.output_stream.splice_async.begin(connection.input_stream, OutputStreamSpliceFlags.CLOSE_SOURCE | OutputStreamSpliceFlags.CLOSE_TARGET, 0, proxy.cancellation_token, () => proxy.cleanup());
        connection.output_stream.splice_async.begin(proxy.connection.input_stream, OutputStreamSpliceFlags.CLOSE_SOURCE | OutputStreamSpliceFlags.CLOSE_TARGET, 0, proxy.cancellation_token, () => proxy.cleanup());
    }

    private async ProxyConnection connect_round_robin(ConfigurationItem entry, InetSocketAddress source) throws Error {
        ProxyConnection? connection = null;
        foreach (var address in entry.destinations) {
            try{
                connection = yield connect_to_destination(source, address);
            }
            catch (Error e) {
                print(@"Failure during round-robin: $(e.message)\n");
            }
        }
        if(connection == null) {
            throw new IOError.HOST_UNREACHABLE("Could not reach any specified forwarding host");
        }
        return connection;
    }

    private async ProxyConnection connect_random(ConfigurationItem entry, InetSocketAddress source) throws Error {
        var items = entry.destinations.to_array();
        var i = Random.int_range(0, items.length);
        var address = items[i];

        var client = new SocketClient();
        return yield connect_to_destination(source, address);
    }

    private async ProxyConnection connect_first_responder(ConfigurationItem entry, InetSocketAddress source) throws Error {
        var cancel_token = new Cancellable();
        ProxyConnection? connection = null;
        var in_flight = 0;

        AsyncReadyCallback connection_cb = (obj, res) => {
            ProxyConnection conn;
            try {
                conn = connect_to_destination.end(res);
                if(connection == null){
                    connection = conn;
                }
            }
            catch(IOError.CANCELLED e) {
                print("Cancelled late connection\n");
                return;
            }
            catch(Error e) {
                print(@"Failure during connect-first (aka. fastest): $(e.message)\n");
            }
            in_flight--;
            connect_first_responder.callback();
        };

        foreach (var address in entry.destinations) {
            in_flight++;
            connect_to_destination.begin(source, address, cancel_token, connection_cb);
        }
        while(in_flight > 0 && connection == null) {
            print(@"In flight requests: $(in_flight)\n");
            yield;
        }
        if(connection == null) {
            throw new IOError.HOST_UNREACHABLE("Could not reach any specified forwarding host");
        }
        cancel_token.cancel();
        return connection;
    }

    private ConfigurationItem get_entry(ConnectionType type, uint16 port, string name) {
        var entry = config.entries
            .where(e => e.protocol == type)
            .where(e => e.source_port == port)
            .where(e => e.hostname == name)
            .first_or_default();

        // If there was no HTTP_WELL_KNOWN entry, try fallback to an HTTP entry
        if(entry == null && type == ConnectionType.HTTP_WELL_KNOWN) {
            return get_entry(ConnectionType.HTTP, port, name);
        }
        return entry;
    }

    private async ProxyConnection connect_to_destination(InetSocketAddress source, ConfigurationDestination dest, Cancellable? cancellation_token = null) throws Error {
        var client = new SocketClient();
        if(dest.mode == DestinationType.INTERNET) {
            var connection = yield client.connect_async(dest.socket_address, cancellation_token);
            return new ProxyConnection(connection);
        }
        if(tunnel == null) {
            throw new TunnelError.NOT_CONFIGURED(@"Cannot connect to $(dest.socket_address) via tunnel, tunnel is not configured");
        }
        
        var mapping = tunnel.map(source.address, dest.socket_address.address, (uint16)dest.socket_address.port);
        var local_addr = new InetSocketAddress(tunnel.proxy_internal_ip, mapping.source_port);
        client.set_local_address(local_addr);
        print(@"TUNNEL CONNECTION: $(local_addr) -> $(new InetSocketAddress(mapping.internal_address, mapping.target_port))");
        var connection = yield client.connect_async(new InetSocketAddress(mapping.internal_address, mapping.target_port), cancellation_token);
        connection.notify["closed"].connect(() => tunnel.unmap(mapping));
        return new ProxyConnection(connection, mapping);
    }

    private class ProxyConnection {
        public SocketConnection connection { get; private set; }
        public Cancellable cancellation_token { get; private set; }
        private TunnelAddressMapping? mapping;

        public ProxyConnection(SocketConnection conn, TunnelAddressMapping? map = null) {
            connection = conn;
            mapping = map;
            cancellation_token = new Cancellable();
        }

        public void cleanup() {
            closing_connections.add(this);
            print("Cleanup connection!\n");
            cancellation_token.cancel();
            connection = null;
            GLib.Timeout.add_seconds_once(10, () => {
                if(mapping != null) {
                    tunnel.unmap(mapping);
                }
                closing_connections.remove(this);
            });
        }

    }

}