<commit_before>/*
 * Copyright 2015-2025 the original author or authors.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software distributed under the License
 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 * or implied. See the License for the specific language governing permissions and limitations under
 * the License.
 */

package sockslib.server;

import java.util.concurrent.CountDownLatch;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import sockslib.client.SocksProxy;
import sockslib.client.SocksSocket;
import sockslib.common.ProtocolErrorException;
import sockslib.common.SocksException;
import sockslib.common.methods.SocksMethod;
import sockslib.server.io.Pipe;
import sockslib.server.io.PipeListener;
import sockslib.server.io.SocketPipe;
import sockslib.server.msg.CommandMessage;
import sockslib.server.msg.CommandResponseMessage;
import sockslib.server.msg.MethodSelectionMessage;
import sockslib.server.msg.MethodSelectionResponseMessage;
import sockslib.server.msg.ServerReply;

import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;

/**
 * The class <code>Socks5Handler</code> represents a handler that can handle SOCKS5 protocol.
 *
 * @author <NAME>
 * @version 1.0
 * @date Apr 16, 2015 11:03:49 AM
 */
public class Socks5Handler implements SocksHandler {

  /**
   * Logger that subclasses also can use.
   */
  protected static final Logger logger = LoggerFactory.getLogger(Socks5Handler.class);

  /**
   * Protocol version.
   */
  private static final int VERSION = 0x5;

  /**
   * Session.
   */
  private Session session;

  /**
   * Method selector.
   */
  private MethodSelector methodSelector;

  private int bufferSize;

  private SocksProxy proxy;

  private SocksProxyServer socksProxyServer;

  private SessionManager sessionManager;

  @Override
  public void handle(Session session) throws Exception {
    sessionManager = getSocksProxyServer().getSessionManager();
    sessionManager.sessionOnCreate(session);

    MethodSelectionMessage msg = new MethodSelectionMessage();
    session.read(msg);

    if (msg.getVersion() != VERSION) {
      throw new ProtocolErrorException();
    }
    SocksMethod selectedMethod = methodSelector.select(msg);

    logger.debug("SESSION[{}] Response client:{}", session.getId(), selectedMethod.getMethodName());
    // send select method.
    session.write(new MethodSelectionResponseMessage(VERSION, selectedMethod));

    // do method.
    selectedMethod.doMethod(session);

    CommandMessage commandMessage = new CommandMessage();
    session.read(commandMessage); // Read command request.

    //    logger.info("SESSION[{}] request:{}  {}:{}", session.getId(), commandMessage.getCommand(),
    //        commandMessage.getAddressType() != AddressType.DOMAIN_NAME ?
    //            commandMessage.getInetAddress() :
    //            commandMessage.getHost(), commandMessage.getPort());

    // If there is a SOCKS exception in command message, It will send a right response to client.
    if (commandMessage.hasSocksException()) {
      ServerReply serverReply = commandMessage.getSocksException().getServerReply();
      session.write(new CommandResponseMessage(serverReply));
      logger.info("SESSION[{}] will close, because {}", session.getId(), serverReply);
      return;
    }

    /**************************** DO COMMAND ******************************************/
    sessionManager.sessionOnCommand(session, commandMessage);
    switch (commandMessage.getCommand()) {
      case BIND:
        doBind(session, commandMessage);
        break;
      case CONNECT:
        doConnect(session, commandMessage);
        break;
      case UDP_ASSOCIATE:
        doUDPAssociate(session, commandMessage);
        break;
    }
  }

  @Override
  public void doConnect(Session session, CommandMessage commandMessage) throws SocksException,
      IOException {

    ServerReply reply = null;
    Socket socket = null;
    InetAddress bindAddress = null;
    int bindPort = 0;
    InetAddress remoteServerAddress = commandMessage.getInetAddress();
    int remoteServerPort = commandMessage.getPort();

    // set default bind address.
    byte[] defaultAddress = {0, 0, 0, 0};
    bindAddress = InetAddress.getByAddress(defaultAddress);
    // DO connect
    try {
      // Connect directly.
      if (proxy == null) {
        socket = new Socket(remoteServerAddress, remoteServerPort);
      } else {
        socket = new SocksSocket(proxy, remoteServerAddress, remoteServerPort);
      }
      bindAddress = socket.getLocalAddress();
      bindPort = socket.getLocalPort();
      reply = ServerReply.SUCCEEDED;

    } catch (IOException e) {
      if (e.getMessage().equals("Connection refused")) {
        reply = ServerReply.CONNECTION_REFUSED;
      } else if (e.getMessage().equals("Operation timed out")) {
        reply = ServerReply.TTL_EXPIRED;
      } else if (e.getMessage().equals("Network is unreachable")) {
        reply = ServerReply.NETWORK_UNREACHABLE;
      } else if (e.getMessage().equals("Connection timed out")) {
        reply = ServerReply.TTL_EXPIRED;
      } else {
        reply = ServerReply.GENERAL_SOCKS_SERVER_FAILURE;
      }
      logger.info("SESSION[{}] connect {} [{}] exception:{}", session.getId(), new
          InetSocketAddress(remoteServerAddress, remoteServerPort), reply, e.getMessage());
    }

    CommandResponseMessage responseMessage =
        new CommandResponseMessage(VERSION, reply, bindAddress, bindPort);
    session.write(responseMessage);
    if (reply != ServerReply.SUCCEEDED) { // 如果返回失败信息，则退出该方法。
      session.close();
      return;
    }

    Pipe pipe = new SocketPipe(session.getSocket(), socket);
    pipe.setName("SESSION[" + session.getId() + "]");
    pipe.setBufferSize(bufferSize);
    if(getSocksProxyServer().getPipeInitializer() != null){
      pipe = getSocksProxyServer().getPipeInitializer().initialize(pipe);
    }

    waitForPipe(pipe, session);
  }

  /**
   * This function starts the pipe and blocks until the pipe stops.
   * @param pipe A pipe that has not yet been started.
   * @throws InterruptedException if this thread is interrupted.
   */
  private static void runPipe(final Pipe pipe) throws InterruptedException {
    CountDownLatch latch = new CountDownLatch(1);
    pipe.addPipeListener(new PipeListener() {
      @Override
      public void onStart(Pipe pipe) {
      }

      @Override
      public void onStop(Pipe pipe) {
        latch.countDown();
      }

      @Override
      public void onTransfer(Pipe pipe, byte[] buffer, int bufferLength) {
      }

      @Override
      public void onError(Pipe pipe, Exception exception) {
      }
    });

    pipe.start(); // This method will build two threads to run two internal pipes.
    while (pipe.isRunning()) {
      latch.await();
    }
  }

  /**<commit_msg>Make proxy more robust against infinite loops<commit_after>/*
 * Copyright 2015-2025 the original author or authors.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software distributed under the License
 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 * or implied. See the License for the specific language governing permissions and limitations under
 * the License.
 */

package sockslib.server;

import java.util.concurrent.CountDownLatch;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import sockslib.client.SocksProxy;
import sockslib.client.SocksSocket;
import sockslib.common.ProtocolErrorException;
import sockslib.common.SocksException;
import sockslib.common.methods.SocksMethod;
import sockslib.server.io.Pipe;
import sockslib.server.io.PipeListener;
import sockslib.server.io.SocketPipe;
import sockslib.server.msg.CommandMessage;
import sockslib.server.msg.CommandResponseMessage;
import sockslib.server.msg.MethodSelectionMessage;
import sockslib.server.msg.MethodSelectionResponseMessage;
import sockslib.server.msg.ServerReply;

import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;

/**
 * The class <code>Socks5Handler</code> represents a handler that can handle SOCKS5 protocol.
 *
 * @author <NAME>
 * @version 1.0
 * @date Apr 16, 2015 11:03:49 AM
 */
public class Socks5Handler implements SocksHandler {

  /**
   * Logger that subclasses also can use.
   */
  protected static final Logger logger = LoggerFactory.getLogger(Socks5Handler.class);

  /**
   * Protocol version.
   */
  private static final int VERSION = 0x5;

  /**
   * Session.
   */
  private Session session;

  /**
   * Method selector.
   */
  private MethodSelector methodSelector;

  private int bufferSize;

  private SocksProxy proxy;

  private SocksProxyServer socksProxyServer;

  private SessionManager sessionManager;

  @Override
  public void handle(Session session) throws Exception {
    sessionManager = getSocksProxyServer().getSessionManager();
    sessionManager.sessionOnCreate(session);

    MethodSelectionMessage msg = new MethodSelectionMessage();
    session.read(msg);

    if (msg.getVersion() != VERSION) {
      throw new ProtocolErrorException();
    }
    SocksMethod selectedMethod = methodSelector.select(msg);

    logger.debug("SESSION[{}] Response client:{}", session.getId(), selectedMethod.getMethodName());
    // send select method.
    session.write(new MethodSelectionResponseMessage(VERSION, selectedMethod));

    // do method.
    selectedMethod.doMethod(session);

    CommandMessage commandMessage = new CommandMessage();
    session.read(commandMessage); // Read command request.

    //    logger.info("SESSION[{}] request:{}  {}:{}", session.getId(), commandMessage.getCommand(),
    //        commandMessage.getAddressType() != AddressType.DOMAIN_NAME ?
    //            commandMessage.getInetAddress() :
    //            commandMessage.getHost(), commandMessage.getPort());

    // If there is a SOCKS exception in command message, It will send a right response to client.
    if (commandMessage.hasSocksException()) {
      ServerReply serverReply = commandMessage.getSocksException().getServerReply();
      session.write(new CommandResponseMessage(serverReply));
      logger.info("SESSION[{}] will close, because {}", session.getId(), serverReply);
      return;
    }

    /**************************** DO COMMAND ******************************************/
    sessionManager.sessionOnCommand(session, commandMessage);
    switch (commandMessage.getCommand()) {
      case BIND:
        doBind(session, commandMessage);
        break;
      case CONNECT:
        doConnect(session, commandMessage);
        break;
      case UDP_ASSOCIATE:
        doUDPAssociate(session, commandMessage);
        break;
    }
  }

  @Override
  public void doConnect(Session session, CommandMessage commandMessage) throws SocksException,
      IOException {

    ServerReply reply = null;
    Socket socket = null;
    InetAddress bindAddress = null;
    int bindPort = 0;
    InetAddress remoteServerAddress = commandMessage.getInetAddress();
    int remoteServerPort = commandMessage.getPort();

    // set default bind address.
    byte[] defaultAddress = {0, 0, 0, 0};
    bindAddress = InetAddress.getByAddress(defaultAddress);
    // DO connect
    try {
      // Connect directly.
      if (proxy == null) {
        socket = new Socket(remoteServerAddress, remoteServerPort);
      } else {
        socket = new SocksSocket(proxy, remoteServerAddress, remoteServerPort);
      }
      bindAddress = socket.getLocalAddress();
      bindPort = socket.getLocalPort();
      reply = ServerReply.SUCCEEDED;

    } catch (IOException e) {
      if (e.getMessage().equals("Connection refused")) {
        reply = ServerReply.CONNECTION_REFUSED;
      } else if (e.getMessage().equals("Operation timed out")) {
        reply = ServerReply.TTL_EXPIRED;
      } else if (e.getMessage().equals("Network is unreachable")) {
        reply = ServerReply.NETWORK_UNREACHABLE;
      } else if (e.getMessage().equals("Connection timed out")) {
        reply = ServerReply.TTL_EXPIRED;
      } else {
        reply = ServerReply.GENERAL_SOCKS_SERVER_FAILURE;
      }
      logger.info("SESSION[{}] connect {} [{}] exception:{}", session.getId(), new
          InetSocketAddress(remoteServerAddress, remoteServerPort), reply, e.getMessage());
    }

    CommandResponseMessage responseMessage =
        new CommandResponseMessage(VERSION, reply, bindAddress, bindPort);
    session.write(responseMessage);
    if (reply != ServerReply.SUCCEEDED) { // 如果返回失败信息，则退出该方法。
      session.close();
      return;
    }

    Pipe pipe = new SocketPipe(session.getSocket(), socket);
    pipe.setName("SESSION[" + session.getId() + "]");
    pipe.setBufferSize(bufferSize);
    if(getSocksProxyServer().getPipeInitializer() != null){
      pipe = getSocksProxyServer().getPipeInitializer().initialize(pipe);
    }

    waitForPipe(pipe, session);
  }

  /**
   * This PipeListener makes it possible to block until onStop is called.
   */
  private static class StopListener implements PipeListener {
    private final CountDownLatch latch = new CountDownLatch(1);
    private boolean stopped = false;

    /**
     * Blocks until onStop is called.
     * @throws InterruptedException if this thread is interrupted.
     */
    void await() throws InterruptedException {
      latch.await();
    }

    boolean wasStopped() {
      return stopped;
    }

    @Override
    public void onStart(Pipe pipe) {
    }

    @Override
    public void onStop(Pipe pipe) {
      stopped = true;
      latch.countDown();
    }

    @Override
    public void onTransfer(Pipe pipe, byte[] buffer, int bufferLength) {
    }

    @Override
    public void onError(Pipe pipe, Exception exception) {
    }
  }

  /**
   * This function starts the pipe and blocks until the pipe stops.
   * @param pipe A pipe that has not yet been started.
   * @throws InterruptedException if this thread is interrupted.
   */
  private static void runPipe(final Pipe pipe) throws InterruptedException {
    StopListener listener = new StopListener();
    pipe.addPipeListener(listener);

    pipe.start(); // This method will build two threads to run two internal pipes.

    // In normal operation, pipe.isRunning() and !listener.wasStopped() should be
    // equal.  However, they can be different if start() failed, or if there is a
    // bug in SocketPipe or StreamPipe.  Checking both ensures that there is no
    // possibility of a hang or busy-loop here.
    while (pipe.isRunning() && !listener.wasStopped()) {
      listener.await();
    }
  }

  /**
{'fraction_non_alphanumeric': 0.11135156403698603, 'fraction_numerical': 0.005246245655452816, 'mean_word_length': 3.472140762463343, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1893381', 'n_tokens_mistral': 4294, 'n_tokens_neox': 4154, 'n_words': 1139}