package com.atlassian.bitbucket.internal.ssh.server;

import org.apache.sshd.common.PropertyResolverUtils;
import org.apache.sshd.common.future.CloseFuture;
import org.apache.sshd.common.future.SshFutureListener;
import org.apache.sshd.server.channel.ChannelSession;

import java.io.IOException;
import java.util.TimerTask;
import java.util.concurrent.TimeUnit;

/**
 * A subclass of {@link ChannelSession} which ensures that the channels future is always resolved
 * <p>
 * Currently, sshd requires the client to acknowledge the channel close message before it closes the connection. This
 * can lead to a memory leak when using an unusually configured SSH client
 * (e.g. http://interrobeng.com/2013/08/25/speed-up-git-5x-to-50x/). This subclass automatically closes the channel
 * if no message is received after a configured period of time
 */
public class SelfClosingChannelSession extends ChannelSession {

    private static final long DEFAULT_CLOSE_TIMEOUT = 30000L;
    private static final String CLOSE_TIMEOUT = "channel-close-timeout";

    @Override
    public CloseFuture close(boolean immediately) {
        final CloseFuture closeFuture = super.close(immediately);

        if (!closeFuture.isDone()) {
            TimerTask task = new TimerTask() {
                @Override
                public void run() {
                    try {
                        // internally handleClose will check if the channel is currently closing
                        // and if it is, really close the channel. We are guaranteed to get this
                        // code path as we called super.close() above
                        handleClose();
                    } catch (IOException e) {
                        log.warn("Failed to force close channel {}. Marking it as closed anyway", getId(), e);
                        closeFuture.setClosed();
                    }
                }
            };
            getSession().getFactoryManager().getScheduledExecutorService()
                    .schedule(task, getCloseTimeoutInMs(), TimeUnit.MILLISECONDS);

            closeFuture.addListener(new CancelTimerTask(task));
        }

        return closeFuture;
    }

    private long getCloseTimeoutInMs() {
        long timeout = DEFAULT_CLOSE_TIMEOUT;

        try {
            Long configuredTimeout = PropertyResolverUtils.getLong(getSession(), CLOSE_TIMEOUT);
            if (configuredTimeout != null) {
                timeout = configuredTimeout;
            }
        } catch (NumberFormatException e) {
            // Ignore. Use the default
        }

        return timeout;
    }

    /**
     * Cancel the given timer task when the future completes
     */
    private static class CancelTimerTask implements SshFutureListener<CloseFuture> {

        private final TimerTask task;

        private CancelTimerTask(TimerTask task) {
            this.task = task;
        }

        @Override
        public void operationComplete(CloseFuture future) {
            task.cancel();
        }

    }

}
