path: root/src/org/gnunet/util/Connection.java

/*
 This file is part of GNUnet.
 (C) 2011, 2012 Christian Grothoff (and other contributing authors)

 GNUnet is free software; you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published
 by the Free Software Foundation; either version 3, or (at your
 option) any later version.

 GNUnet is distributed in the hope that it will be useful, but
 WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with GNUnet; see the file COPYING.  If not, write to the
 Free Software Foundation, Inc., 59 Temple Place - Suite 330,
 Boston, MA 02111-1307, USA.
 */

package org.gnunet.util;

import org.gnunet.construct.Construct;
import org.gnunet.construct.MessageLoader;
import org.gnunet.construct.ProtocolViolationException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOError;
import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SocketChannel;
import java.nio.channels.spi.SelectorProvider;
import java.util.LinkedList;
import java.util.List;

/**
 * Integrates sockets with the gnunet-java message loop / the scheduler.
 */
public class Connection {
    private static final Logger logger = LoggerFactory
            .getLogger(Connection.class);

    /**
     * The underlying socket the client is using to talk with the service.
     */
    private SocketChannel connectionChannel = null;

    private List<AddressProbe> addressProbes = null;

    /**
     * The task that is currently used by the resolve mechanism.
     */
    private Cancelable resolveHandle = null;

    /**
     * The task that is responsible for establishing the connection to the server.
     */
    private Cancelable connectHandle = null;

    /**
     * The ReceiveHelper responsible for receiving a whole message from the service
     * and calling the respective MessageReceiver.
     */
    private ReceiveHelper currentReceiveHelper = null;

    /**
     * The buffer with the (partial) message received from the service.
     * Initially, this buffer has the size of the smallest possible messages, but grows when
     * receiving larger messages.
     */
    private ByteBuffer recvBuffer = ByteBuffer.allocate(GnunetMessage.Header.SIZE);


    /**
     * The handle for the current transmission. Writes data to the socket.
     */
    private TransmitHelper currentTransmitHelper = null;
    /**
     * The handle for the next transmission. The next transmission will become the current
     * transmission once the current transmission has completed.
     * While nextTransmitHelper is not null, no new transmit requests may be scheduled.
     */
    private TransmitHelper nextTransmitHelper = null;

    /**
     * The transmitters passed to transmitReadyNotify(...) write to this buffer by calling
     * methods on the MessageSink passed to the Transmitter.transmit(MessageSink s) method.
     * Initially, this buffer has the size of the smallest possible messages, but grows when
     * transmitting larger messages.
     */
    private ByteBuffer transmitBuffer = ByteBuffer.allocate(GnunetMessage.Header.SIZE);
    private boolean disconnected = false;

    /**
     * Timeout task for the connect notify.
     */
    private Scheduler.TaskConfiguration notifyConnectedTimeout;

    /**
     * Continuation to call when connected
     */
    private Continuation notifyConnectedContinuation;


    /**
     * An address probe is a connection to a socket that may succeed or not.
     * The first address probe that succeeded is used for this connection.
     */
    private static class AddressProbe {
        Cancelable connectTask;
        SocketChannel channel;

        public void cancel() {
            if (connectTask != null) {
                connectTask.cancel();
            }
            if (channel != null) {
                try {
                    channel.close();
                } catch (IOException e) {
                    // nothing we can do here
                }
            }
        }
    }

    /**
     * Represents a request for transmission.
     */
    public interface TransmitHandle extends Cancelable {
        /**
         * Cancel a request for the transmit ready notification.
         * This does *not* cancel a transmission that already has been started.
         */
        public void cancel();
    }

    /**
     * An interface that allows the Transmitter.transmit method to deliver their messages
     * to the client, which sends them to the service.
     */
    public interface MessageSink {
        public void send(GnunetMessage.Body m);
    }

    /**
     * The ReceiveHelper is responsible for receiving a whole
     * GnunetMessage and call the respective MessageReceiver with the message on success,
     * and null on failure or timeout.
     */
    private class ReceiveHelper implements Scheduler.Task {
        private MessageReceiver receiver;
        private RelativeTime timeout;
        private GnunetMessage.Header msgh = null;
        private Scheduler.TaskConfiguration recvTask = null;
        private boolean finished = false;
        // is this receiver actively working? if not, the connection process has to kick off the receiver
        // (or select behaves badly)
        private boolean working = false;

        public ReceiveHelper(MessageReceiver receiver, RelativeTime timeout) {
            this.receiver = receiver;
            this.timeout = timeout;
        }

        public void dispatchMessage() {
            assert msgh != null;
            currentReceiveHelper = null;
            finished = true;
            recvBuffer.flip();

            boolean found = true;

            Class unionClass = null;

            try {
                unionClass = MessageLoader.getUnionClass(GnunetMessage.Body.class, msgh.messageType);
            } catch (ProtocolViolationException e) {
                found = false;
            }

            logger.debug("dispatching received message");
            if (found) {
                GnunetMessage msg;
                try {
                    msg = Construct.parseAs(recvBuffer, GnunetMessage.class);
                } catch (OutOfMemoryError e) {
                    throw new OutOfMemoryError("oom while parsing " + unionClass);
                }
                receiver.process(msg.body);
            } else {
                UnknownMessageBody b = new UnknownMessageBody();
                b.id = msgh.messageType;

                // may throw exception, doesn't matter as it's the last call
                receiver.process(b);
            }
        }

        @Override
        public void run(Scheduler.RunContext ctx) {
            recvTask = null;
            if (ctx.reasons.contains(Scheduler.Reason.TIMEOUT)) {
                currentReceiveHelper = null;
                receiver.handleError();
            } else if (ctx.reasons.contains(Scheduler.Reason.READ_READY)) {
                try {
                    int n = connectionChannel.read(recvBuffer);
                    if (n == -1) {
                        currentReceiveHelper = null;
                        logger.warn("lost connection to service");
                        connectionChannel.close();
                        connectionChannel = null;
                        if (Connection.this.currentTransmitHelper != null) {
                            Connection.this.currentTransmitHelper.cancel();
                            Connection.this.currentTransmitHelper = null;
                        }
                        try {
                            receiver.handleError();
                        } finally {
                            return;
                        }
                    }
                    logger.debug(String.format("read %s bytes from %s", n, connectionChannel.socket().toString()));
                } catch (IOException e) {
                    logger.error("read failed:", e);
                    try {
                        receiver.handleError();
                    } finally {
                        return;
                    }
                }
                if (recvBuffer.remaining() == 0) {
                    if (msgh != null) {
                        dispatchMessage();
                    } else {
                        recvBuffer.rewind();
                        msgh = Construct.parseAs(recvBuffer, GnunetMessage.Header.class);

                        logger.debug("expecting message of size {}, type {}", msgh.messageSize, msgh.messageType);
                        if (msgh.messageSize > GnunetMessage.Header.SIZE) {
                            if (recvBuffer.capacity() < msgh.messageSize) {
                                ByteBuffer buf = ByteBuffer.allocate(msgh.messageSize);
                                recvBuffer.flip();
                                buf.put(recvBuffer);
                                recvBuffer = buf;
                            }
                            recvBuffer.limit(msgh.messageSize);
                            schedule();
                        } else {
                            dispatchMessage();
                        }
                    }
                } else {
                    schedule();
                }
            } else if (ctx.reasons.contains(Scheduler.Reason.SHUTDOWN)) {
                // nothing to do!
            } else {
                // XXX: what to do here?
                throw new RuntimeException("receive failed");
            }
        }

        private void schedule() {
            working = true;
            recvTask = Scheduler.addRead(timeout, connectionChannel, this);
        }

        public void cancel() {
            if (finished) {
                throw new AssertionError("canceling finished receive");
            }
            if (recvTask != null) {
                recvTask.cancel();
                recvTask = null;
            }
        }
    }


    private class TransmitHelper implements Scheduler.Task, MessageSink {
        private final MessageTransmitter transmitter;

        private Cancelable notifyTimeoutTask;

        private Cancelable transmitTask = null;

        public TransmitHelper(final MessageTransmitter transmitter, RelativeTime notifyTimeout) {
            this.transmitter = transmitter;

            Scheduler.TaskConfiguration tc = new Scheduler.TaskConfiguration(notifyTimeout,
                    new Scheduler.Task() {
                        @Override
                        public void run(Scheduler.RunContext ctx) {
                            transmitter.handleError();
                        }
                    });

            notifyTimeoutTask = tc.schedule();
        }

        public void cancel() {
            if (transmitTask != null) {
                transmitTask.cancel();
                transmitTask = null;
            }
            if (notifyTimeoutTask != null) {
                notifyTimeoutTask.cancel();
                notifyTimeoutTask = null;
            }
        }

        @Override
        public void run(Scheduler.RunContext ctx) {
            this.transmitTask = null;
            if (connectionChannel == null) {
                logger.error("could not write to channel (null)");
                return;
            }
            try {
                int n = connectionChannel.write(transmitBuffer);
                // logger.debug("connectionChannel has written " + n + " bytes to " + connectionChannel.socket().toString());
            } catch (IOException e) {
                throw new IOError(e);
            }
            if (transmitBuffer.remaining() == 0) {
                //logger.debug("sent " + transmitBuffer.position() + "bytes complete message");
                if (nextTransmitHelper == null) {
                    currentTransmitHelper = null;
                } else {
                    currentTransmitHelper = nextTransmitHelper;
                    // we need to to this so the transmit callback can do notifyTransmitReady
                    TransmitHelper tmpTransmitHelper = nextTransmitHelper;
                    nextTransmitHelper = null;
                    tmpTransmitHelper.start();

                }
            } else {
                schedule();
            }
        }

        /**
         * called to notify when we are ready to put new messages in the transmit buffer
         */
        public void start() {
            notifyTimeoutTask.cancel();
            notifyTimeoutTask = null;
            transmitBuffer.clear();
            transmitter.transmit(TransmitHelper.this);
            transmitBuffer.flip();
            schedule();
        }

        private void schedule() {
            if (disconnected) {
                return;
            }
            // timeout is forever, because there is no way to directly limit the transmission time
            // of a message, only the max. wait time before transmission.
            // cancel must be called on the transmitTask if we disconnect
            Scheduler.TaskConfiguration tc = new Scheduler.TaskConfiguration(RelativeTime.FOREVER, this);
            tc.selectWrite(connectionChannel);
            this.transmitTask = tc.schedule();
        }

        @Override
        public void send(final GnunetMessage.Body m) {
            final GnunetMessage gm = new GnunetMessage();
            gm.header = new GnunetMessage.Header();
            gm.body = m;
            Construct.patch(gm);
            gm.header.messageSize = Construct.getSize(gm);
            byte[] b = Construct.toBinary(gm);
            if (b.length != gm.header.messageSize) {
                throw new AssertionError(
                        String.format("tried to send message with binary size %s but size in header %s",
                                b.length, gm.header.messageSize));
            }
            logger.debug("sending message (size={},type={})", b.length, gm.header.messageType);
            if (transmitBuffer.remaining() < b.length) {
                ByteBuffer buf = ByteBuffer.allocate(b.length + transmitBuffer.capacity());
                transmitBuffer.flip();
                buf.put(transmitBuffer);
                transmitBuffer = buf;
            }
            transmitBuffer.put(b);
        }
    }

    /**
     * Create a connection to the given hostname/port.
     *
     * @param hostname name of the host to connect to
     * @param port     port of the host to connect to
     */
    public Connection(String hostname, int port) {
        addressProbes = new LinkedList<AddressProbe>();
        ConnectionResolveHandler addressHandler = new ConnectionResolveHandler(port);
        resolveHandle = Resolver.getInstance().resolveHostname(hostname, RelativeTime.FOREVER, addressHandler);
    }

    public Connection(SocketChannel sock) {
        assert sock != null;
        this.connectionChannel = sock;
    }


    class ConnectionResolveHandler implements Resolver.AddressCallback {
        private final int port;

        public ConnectionResolveHandler(int port) {
            this.port = port;
        }

        @Override
        public void onAddress(InetAddress addr) {
            final SocketChannel channel = createChannel();
            try {
                channel.connect(new InetSocketAddress(addr, port));
            } catch (IOException e) {
                logger.error("could not connect to host");
                return;
            }

            final AddressProbe addressProbe = new AddressProbe();
            addressProbe.channel = channel;
            Scheduler.TaskConfiguration tc = new Scheduler.TaskConfiguration(RelativeTime.FOREVER,
                    new Scheduler.Task() {
                        @Override
                        public void run(Scheduler.RunContext ctx) {
                            addressProbe.connectTask = null;
                            if (ctx.reasons.contains(Scheduler.Reason.SHUTDOWN)) {
                                return;
                            }
                            Connection.this.finishConnect(addressProbe);
                        }
                    });

            // our channel has already disconnected
            if (!channel.isOpen()) {
                return;
            }

            tc.selectConnect(channel);

            addressProbe.connectTask = tc.schedule();
        }

        @Override
        public void onFinished() {
            resolveHandle = null;
        }

        @Override
        public void onTimeout() {
            // do nothing
            // todo: is this correct?
        }
    }


    private void finishConnect(AddressProbe probe) {
        // can happen if the addres probe task was already scheduled
        if (connectionChannel != null) {
            try {
                probe.channel.close();
            } catch (IOException e) {
                logger.error("could not close channel", e);
            }
            return;
        }

        SocketChannel channel = probe.channel;
        boolean connected;
        try {
            connected = channel.finishConnect();
        } catch (IOException e) {
            logger.debug("finishConnect() was not successful: {}", (Object) e);
            return;
        }

        if (!connected) {
            logger.error("socket reported OP_CONNECT but is not connected");
            return;
        }

        for (AddressProbe addressProbe : addressProbes) {
            if (addressProbe != probe && addressProbe.connectTask != null) {
                addressProbe.connectTask.cancel();
                try {
                    addressProbe.channel.close();
                } catch (IOException e) {
                    logger.error("could not close channel", e);
                }
            }
        }

        addressProbes.clear();

        connectionChannel = channel;

        if (currentTransmitHelper != null) {
            currentTransmitHelper.start();
        }
        if (currentReceiveHelper != null && !currentReceiveHelper.working) {
            currentReceiveHelper.schedule();
        }
        Continuation c = notifyConnectedContinuation;
        notifyConnectedContinuation = null;
        if (notifyConnectedTimeout != null) {
            notifyConnectedTimeout.cancel();
            notifyConnectedTimeout = null;
        }
        if (c != null) {
            c.cont(true);
        }
    }

    /**
     * Open a channel for this connection in non-blocking mode
     */
    private SocketChannel createChannel() {
        try {
            SocketChannel channel = SelectorProvider.provider().openSocketChannel();
            channel.configureBlocking(false);
            return channel;
        } catch (IOException e) {
            // this is fatal, no retry necessary
            throw new IOError(e);
        }
    }

    public boolean isConnected() {
        return connectionChannel != null && connectionChannel.isConnected();
    }


    public interface ReceiveHandle extends Cancelable {
        public void cancel();
    }

    /**
     * Receive one message from the network.
     *
     * @param timeout  deadline after which receiver.onError() will be called
     * @param receiver MessageReceiver that is responsible for the received message
     */
    public ReceiveHandle receive(RelativeTime timeout, final MessageReceiver receiver) {
        if (currentReceiveHelper != null) {
            throw new AssertionError("receive must not be called while receiving");
        }

        if (!isConnected()) {
            throw new AssertionError("cannot receive if not connected");
        }

        recvBuffer.clear();
        recvBuffer.limit(GnunetMessage.Header.SIZE);
        final ReceiveHelper rh = new ReceiveHelper(receiver, timeout);
        currentReceiveHelper = rh;

        // we can only schedule the receive helper if we are sure the connection is made, otherwise
        // select will misbehave!
        if (connectionChannel.isConnected()) {
            currentReceiveHelper.schedule();
        }

        return new ReceiveHandle() {
            @Override
            public void cancel() {
                rh.cancel();
            }
        };
    }

    /**
     * Call the transmitter once the we are ready to transmit data.
     *
     * @param size        number of bytes to send
     * @param timeout     after how long should we give up (and call transmitter.transmit(null))
     * @param transmitter the MessageTransmitter object to call once the client is ready to transmit or
     *                    when the timeout is over. Guaranteed to be called *after* notifyTransmitReady has returned.
     * @return a handle that can be used to cancel the transmit request, null if request could be satisfied immediately
     */
    public TransmitHandle notifyTransmitReady(int size, RelativeTime timeout, final MessageTransmitter transmitter) {
        if (disconnected) {
            throw new AssertionError("notifyTransmitReady called on a closed connection");
        }
        if (nextTransmitHelper != null) {
            throw new AssertionError(
                    "previous transmit request must have completed before calling notifyTransmitReady again");
        }

        if (timeout.getMicroseconds() <= 0) {
            throw new AssertionError("notifyTransmitReady timeout must be positive");
        }

        if (!isConnected()) {
            throw new AssertionError("notifyTransmitHandle can only be called once connected");
        }

        final TransmitHelper transmit = new TransmitHelper(transmitter, timeout);

        if (currentTransmitHelper == null) {
            currentTransmitHelper = transmit;
            currentTransmitHelper.start();
            return null;
        }

        nextTransmitHelper = transmit;

        return new TransmitHandle() {
            @Override
            public void cancel() {
                transmit.cancel();
            }
        };
    }


    /**
     * Call cont after establishing the connection or when the timeout has occured.
     *
     * @param timeout timeout
     * @param cont    continuation to call
     * @return
     */
    /* package-protected */ Cancelable notifyConnected(RelativeTime timeout, final Continuation cont) {
        if (notifyConnectedTimeout != null) {
            throw new AssertionError();
        }
        this.notifyConnectedContinuation = cont;
        this.notifyConnectedTimeout = Scheduler.addDelayed(timeout, new Scheduler.Task() {
            @Override
            public void run(Scheduler.RunContext ctx) {
                Continuation c = notifyConnectedContinuation;
                notifyConnectedContinuation = null;
                Connection.this.notifyConnectedTimeout = null;
                if (c != null) {
                    c.cont(false);
                }
            }
        });
        return this.notifyConnectedTimeout;
    }

    /**
     * Disconnect. There must not be any pending transmit/receive requests.
     * Any buffered data scheduled for writing is discarded.
     */
    public void disconnect() {
        if (disconnected) {
            logger.error("disconnect called twice");
        }
        disconnected = true;

        if (currentTransmitHelper != null) {
            currentTransmitHelper.cancel();
            currentTransmitHelper = null;
        }

        if (nextTransmitHelper != null) {
            nextTransmitHelper.cancel();
            nextTransmitHelper = null;
        }

        if (currentReceiveHelper != null) {
            currentReceiveHelper.cancel();
            currentReceiveHelper = null;
        }

        if (resolveHandle != null) {
            resolveHandle.cancel();
            resolveHandle = null;
        }
        if (connectHandle != null) {
            connectHandle.cancel();
            connectHandle = null;
        }
        if (connectionChannel != null) {
            try {
                connectionChannel.close();
            } catch (IOException e) {
                throw new IOError(e);
            }
            connectionChannel = null;
        }
    }

}