org.smallmind.web.reverse.http1_1.ReverseProxyService Maven / Gradle / Ivy

Go to download
/*
 * Copyright (c) 2007 through 2024 David Berkman
 *
 * This file is part of the SmallMind Code Project.
 *
 * The SmallMind Code Project is free software, you can redistribute
 * it and/or modify it under either, at your discretion...
 *
 * 1) The terms of GNU Affero General Public License as published by the
 * Free Software Foundation, either version 3 of the License, or (at
 * your option) any later version.
 *
 * ...or...
 *
 * 2) The terms of the Apache License, Version 2.0.
 *
 * The SmallMind Code Project 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 or Apache License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * and the Apache License along with the SmallMind Code Project. If not, see
 *  or .
 *
 * Additional permission under the GNU Affero GPL version 3 section 7
 * ------------------------------------------------------------------
 * If you modify this Program, or any covered work, by linking or
 * combining it with other code, such other code is not for that reason
 * alone subject to any of the requirements of the GNU Affero GPL
 * version 3.
 */
package org.smallmind.web.reverse.http1_1;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.StandardSocketOptions;
import java.nio.ByteBuffer;
import java.nio.channels.SelectableChannel;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.Iterator;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import org.smallmind.scribe.pen.LoggerManager;

public class ReverseProxyService {

  private final CountDownLatch terminationLatch = new CountDownLatch(1);
  private final AtomicBoolean closed = new AtomicBoolean(false);
  private final ServerSocketChannel serverSocketChannel;
  private final Selector selector;
  private final ProxyDictionary dictionary;
  private final ProxyExecutor proxyExecutor;
  private final Lock selectLock = new ReentrantLock(true);
  private final Lock loopLock = new ReentrantLock(true);
  private final ByteBuffer byteBuffer = ByteBuffer.allocate(8192);
  private final int connectTimeoutMillis;
  private EventLoop eventLoop;

  public ReverseProxyService (String host, int port, ProxyDictionary dictionary, int connectTimeoutMillis, int concurrencyLimit)
    throws IOException {

    this.dictionary = dictionary;
    this.connectTimeoutMillis = connectTimeoutMillis;

    proxyExecutor = new ProxyExecutor(concurrencyLimit);

    serverSocketChannel = ServerSocketChannel.open();
    serverSocketChannel.configureBlocking(false);
    serverSocketChannel.socket().bind(new InetSocketAddress(host, port));

    serverSocketChannel.register(selector = Selector.open(), SelectionKey.OP_ACCEPT);

    startEventLoop();
  }

  private void startEventLoop () {

    Thread eventThread;

    eventThread = new Thread(eventLoop = new EventLoop());
    eventThread.setDaemon(true);
    eventThread.start();
  }

  public void destroy ()
    throws IOException, InterruptedException {

    eventLoop.stop();

    if (closed.compareAndSet(false, true)) {
      selector.close();
      serverSocketChannel.close();
      proxyExecutor.shutdown();
    }
  }

  public SelectionKey keyFor (SelectableChannel selectableChannel) {

    loopLock.lock();
    try {
      selectLock.lock();
      try {

        return selectableChannel.keyFor(selector);
      } finally {
        selectLock.unlock();
      }
    } finally {
      loopLock.unlock();
    }
  }

  public void execute (SocketChannel sourceChannel, Runnable runnable) {

    proxyExecutor.execute(sourceChannel, runnable);
  }

  public ProxyTarget lookup (HttpRequestFrame httpRequestFrame)
    throws ProtocolException {

    ProxyTarget target;

    if ((target = dictionary.lookup(httpRequestFrame)) == null) {
      throw new ProtocolException(CannedResponse.NOT_FOUND);
    }

    return target;
  }

  public void connectDestination (final SocketChannel sourceSocketChannel, final HttpRequestFrameReader httpRequestFrameReader, final ProxyTarget target) {

    execute(sourceSocketChannel, () -> {

      SocketChannel destinationChannel = null;

      try {
        destinationChannel = SocketChannel.open();
        destinationChannel.socket().connect(new InetSocketAddress(target.getHost(), target.getPort()), connectTimeoutMillis);
        destinationChannel.setOption(StandardSocketOptions.SO_KEEPALIVE, true).setOption(StandardSocketOptions.TCP_NODELAY, true).configureBlocking(false);

        loopLock.lock();
        try {
          selectLock.lock();
          try {
            destinationChannel.register(selector, SelectionKey.OP_READ, new HttpResponseFrameReader(ReverseProxyService.this, sourceSocketChannel, destinationChannel));
          } finally {
            selectLock.unlock();
          }
        } finally {
          loopLock.unlock();
        }

        httpRequestFrameReader.registerDestination(target, destinationChannel);
      } catch (IOException ioException) {
        try {
          if (destinationChannel != null) {
            destinationChannel.close();
          }
        } catch (IOException closeException) {
          LoggerManager.getLogger(ReverseProxyService.class).error(closeException);
        }

        httpRequestFrameReader.fail(CannedResponse.NOT_FOUND, null);
      }
    });
  }

  private class EventLoop implements Runnable {

    private Thread runnableThread;
    private boolean stopped = false;

    private void stop ()
      throws InterruptedException {

      stopped = true;

      if (runnableThread != null) {
        runnableThread.interrupt();
      }

      terminationLatch.await();
    }

    @Override
    public void run () {

      try {
        runnableThread = Thread.currentThread();

        while (!stopped) {
          try {

            int selectedKeyCount = 0;

            selectLock.lock();
            try {
              selectedKeyCount = selector.select(100);
            } catch (IOException ioException) {
              LoggerManager.getLogger(ReverseProxyService.class).error(ioException);
            } finally {
              selectLock.unlock();
            }

            if ((!stopped) && (selectedKeyCount > 0)) {

              Iterator selectionKeyIter = selector.selectedKeys().iterator();

              loopLock.lock();
              try {
                while (selectionKeyIter.hasNext()) {

                  final SelectionKey selectionKey = selectionKeyIter.next();

                  try {
                    if (selectionKey.isValid()) {
                      if (selectionKey.isAcceptable()) {

                        SocketChannel sourceChannel = null;

                        try {
                          sourceChannel = (SocketChannel)((ServerSocketChannel)selectionKey.channel()).accept().setOption(StandardSocketOptions.SO_KEEPALIVE, true).setOption(StandardSocketOptions.TCP_NODELAY, true).configureBlocking(false);
                          sourceChannel.register(selector, SelectionKey.OP_READ, new HttpRequestFrameReader(ReverseProxyService.this, sourceChannel, connectTimeoutMillis));
                        } catch (IOException ioException) {
                          try {
                            if (sourceChannel != null) {
                              sourceChannel.close();
                            }
                          } catch (IOException closeException) {
                            LoggerManager.getLogger(ReverseProxyService.class).error(closeException);
                          }
                        }
                      } else if (selectionKey.isReadable() && selectionKey.channel().isOpen()) {

                        int bytesRead = 0;

                        try {
                          byteBuffer.clear();
                          bytesRead = ((SocketChannel)selectionKey.channel()).read(byteBuffer);
                        } catch (IOException ioException) {
                          ((FrameReader)selectionKey.attachment()).fail(CannedResponse.BAD_REQUEST, (SocketChannel)selectionKey.channel());
                        }

                        if (bytesRead > 0) {
                          byteBuffer.flip();
                          ((FrameReader)selectionKey.attachment()).processInput(selectionKey, byteBuffer);
                        } else if (bytesRead < 0) {
                          ((FrameReader)selectionKey.attachment()).closeChannels();
                        }
                      }
                    } else {
                      selectionKey.cancel();
                    }
                  } finally {
                    selectionKeyIter.remove();
                  }
                }
              } finally {
                loopLock.unlock();
              }
            }
          } catch (Exception exception) {
            LoggerManager.getLogger(ReverseProxyService.class).error(exception);
          }
        }
      } finally {
        terminationLatch.countDown();
      }
    }
  }
}