feature_torcontrol.py raw

   1  #!/usr/bin/env python3
   2  # Copyright (c) The Bitcoin Core developers
   3  # Distributed under the MIT software license, see the accompanying
   4  # file COPYING or http://www.opensource.org/licenses/mit-license.php.
   5  """Test torcontrol functionality with a mock Tor control server."""
   6  from contextlib import contextmanager
   7  import socket
   8  import threading
   9  from test_framework.test_framework import BitcoinTestFramework
  10  from test_framework.util import (
  11      assert_equal,
  12      ensure_for,
  13      p2p_port,
  14  )
  15  
  16  
  17  class MockTorControlServer:
  18      def __init__(self, port, manual_mode=False):
  19          self.port = port
  20          self.sock = None
  21          self.conn = None
  22          self.running = False
  23          self.thread = None
  24          self.received_commands = []
  25          self.manual_mode = manual_mode
  26          self.conn_ready = threading.Event()
  27  
  28      def start(self):
  29          self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  30          self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  31          self.sock.settimeout(1.0)
  32          self.sock.bind(('127.0.0.1', self.port))
  33          self.sock.listen(1)
  34          self.running = True
  35          self.thread = threading.Thread(target=self._serve)
  36          self.thread.daemon = True
  37          self.thread.start()
  38  
  39      def stop(self):
  40          self.running = False
  41          if self.conn:
  42              self.conn.close()
  43          if self.sock:
  44              self.sock.close()
  45          if self.thread:
  46              self.thread.join(timeout=5)
  47  
  48      def _serve(self):
  49          while self.running:
  50              try:
  51                  self.conn, _ = self.sock.accept()
  52                  self.conn.settimeout(1.0)
  53                  self.conn_ready.set()
  54                  self._handle_connection(self.conn)
  55              except socket.timeout:
  56                  continue
  57              except OSError:
  58                  break
  59  
  60      def _handle_connection(self, conn):
  61          try:
  62              buf = b""
  63              while self.running:
  64                  try:
  65                      data = conn.recv(1024)
  66                      if not data:
  67                          break
  68                      buf += data
  69                      while b"\r\n" in buf:
  70                          line, buf = buf.split(b"\r\n", 1)
  71                          command = line.decode('utf-8').strip()
  72                          if command:
  73                              self.received_commands.append(command)
  74                              if not self.manual_mode:
  75                                  response = self._get_response(command)
  76                                  conn.sendall(response.encode('utf-8'))
  77                  except socket.timeout:
  78                      continue
  79          finally:
  80              conn.close()
  81  
  82      def send_raw(self, data):
  83          if self.conn:
  84              self.conn.sendall(data.encode('utf-8'))
  85  
  86      def _get_response(self, command):
  87          if command == "PROTOCOLINFO 1":
  88              return (
  89                  "250-PROTOCOLINFO 1\r\n"
  90                  "250-AUTH METHODS=NULL\r\n"
  91                  "250-VERSION Tor=\"0.1.2.3\"\r\n"
  92                  "250 OK\r\n"
  93              )
  94          elif command == "AUTHENTICATE":
  95              return "250 OK\r\n"
  96          elif command.startswith("ADD_ONION"):
  97              return (
  98                  "250-ServiceID=testserviceid1234567890123456789012345678901234567890123456\r\n"
  99                  "250 OK\r\n"
 100              )
 101          elif command.startswith("GETINFO"):
 102              return "250-net/listeners/socks=\"127.0.0.1:9050\"\r\n250 OK\r\n"
 103          else:
 104              return "510 Unrecognized command\r\n"
 105  
 106  
 107  class TorControlTest(BitcoinTestFramework):
 108      def set_test_params(self):
 109          self.num_nodes = 1
 110  
 111      def next_port(self):
 112          self._port_counter = getattr(self, '_port_counter', 0) + 1
 113          return p2p_port(self.num_nodes + self._port_counter)
 114  
 115      def restart_with_mock(self, mock_tor):
 116          mock_tor.start()
 117          self.restart_node(0, extra_args=[
 118              f"-torcontrol=127.0.0.1:{mock_tor.port}",
 119              "-listenonion=1",
 120              "-debug=tor",
 121          ])
 122  
 123          # Wait for connection and PROTOCOLINFO command
 124          mock_tor.conn_ready.wait(timeout=10)
 125          self.wait_until(lambda: len(mock_tor.received_commands) >= 1, timeout=10)
 126          assert_equal(mock_tor.received_commands[0], "PROTOCOLINFO 1")
 127  
 128      @contextmanager
 129      def expect_disconnect(self, expect, mock_tor):
 130          initial_len = len(mock_tor.received_commands)
 131          yield
 132  
 133          if expect:
 134              # Expect to receive a PROTOCOLINFO 1 on reconnect, bumping the received
 135              # commands length.
 136              self.wait_until(lambda: len(mock_tor.received_commands) == initial_len + 1)
 137              assert_equal(mock_tor.received_commands[initial_len], "PROTOCOLINFO 1")
 138          else:
 139              # No disconnect, so no reconnect message
 140              ensure_for(duration=2, f=lambda: len(mock_tor.received_commands) == initial_len)
 141  
 142      def test_basic(self):
 143          self.log.info("Test Tor control basic functionality")
 144  
 145          mock_tor = MockTorControlServer(self.next_port())
 146          self.restart_with_mock(mock_tor)
 147  
 148          # Waiting for Tor control commands
 149          self.wait_until(lambda: len(mock_tor.received_commands) >= 4, timeout=10)
 150  
 151          # Verify expected protocol sequence
 152          assert_equal(mock_tor.received_commands[0], "PROTOCOLINFO 1")
 153          assert_equal(mock_tor.received_commands[1], "AUTHENTICATE")
 154          assert_equal(mock_tor.received_commands[2], "GETINFO net/listeners/socks")
 155          assert mock_tor.received_commands[3].startswith("ADD_ONION ")
 156          assert "PoWDefensesEnabled=1" in mock_tor.received_commands[3]
 157  
 158          # Clean up
 159          mock_tor.stop()
 160  
 161      def test_partial_data(self):
 162          self.log.info("Test that partial Tor control responses are buffered until complete")
 163  
 164          mock_tor = MockTorControlServer(self.next_port(), manual_mode=True)
 165          self.restart_with_mock(mock_tor)
 166  
 167          # Send partial response (no \r\n on last line)
 168          mock_tor.send_raw(
 169              "250-PROTOCOLINFO 1\r\n"
 170              "250-AUTH METHODS=NULL\r\n"
 171              "250 OK"
 172          )
 173  
 174          # Verify AUTHENTICATE is not sent
 175          ensure_for(duration=2, f=lambda: len(mock_tor.received_commands) == 1)
 176  
 177          # Complete the response
 178          mock_tor.send_raw("\r\n")
 179  
 180          # Should now process the complete response and send AUTHENTICATE
 181          self.wait_until(lambda: len(mock_tor.received_commands) >= 2, timeout=5)
 182          assert_equal(mock_tor.received_commands[1], "AUTHENTICATE")
 183  
 184          # Clean up
 185          mock_tor.stop()
 186  
 187      def test_pow_fallback(self):
 188          self.log.info("Test that ADD_ONION retries without PoW on 512 error")
 189  
 190          class NoPowServer(MockTorControlServer):
 191              def _get_response(self, command):
 192                  if command.startswith("ADD_ONION"):
 193                      if "PoWDefensesEnabled=1" in command:
 194                          return "512 Unrecognized option\r\n"
 195                      else:
 196                          return (
 197                              "250-ServiceID=testserviceid1234567890123456789012345678901234567890123456\r\n"
 198                              "250 OK\r\n"
 199                          )
 200                  return super()._get_response(command)
 201  
 202          mock_tor = NoPowServer(self.next_port())
 203          self.restart_with_mock(mock_tor)
 204  
 205          # Expect: PROTOCOLINFO, AUTHENTICATE, GETINFO, ADD_ONION (with PoW), ADD_ONION (without PoW)
 206          self.wait_until(lambda: len(mock_tor.received_commands) >= 5, timeout=10)
 207  
 208          # First ADD_ONION should have PoW enabled
 209          assert mock_tor.received_commands[3].startswith("ADD_ONION ")
 210          assert "PoWDefensesEnabled=1" in mock_tor.received_commands[3]
 211  
 212          # Retry should be ADD_ONION without PoW
 213          assert mock_tor.received_commands[4].startswith("ADD_ONION ")
 214          assert "PoWDefensesEnabled=1" not in mock_tor.received_commands[4]
 215  
 216          # Clean up
 217          mock_tor.stop()
 218  
 219      def test_oversized_line(self):
 220          mock_tor = MockTorControlServer(self.next_port(), manual_mode=True)
 221          self.restart_with_mock(mock_tor)
 222  
 223          MAX_LINE_LENGTH = 100000
 224  
 225          self.log.info("Test that Tor control does not disconnect with a MAX_LINE_LENGTH line.")
 226          with self.expect_disconnect(False, mock_tor):
 227              msg = "250-" + ("A" * (MAX_LINE_LENGTH - 5)) + "\r"
 228              assert_equal(len(msg), MAX_LINE_LENGTH)
 229              # The \n is not counted in line length.
 230              mock_tor.send_raw(msg + "\n")
 231  
 232          self.log.info("Test that Tor control disconnects with a MAX_LINE_LENGTH + 1 line")
 233          with self.expect_disconnect(True, mock_tor):
 234              msg = "250-" + ("A" * (MAX_LINE_LENGTH - 4)) + "\r"
 235              assert_equal(len(msg), MAX_LINE_LENGTH + 1)
 236              mock_tor.send_raw(msg + "\n")
 237  
 238          mock_tor.stop()
 239  
 240      def test_overmany_lines(self):
 241          mock_tor = MockTorControlServer(self.next_port(), manual_mode=True)
 242          self.restart_with_mock(mock_tor)
 243  
 244          MAX_LINE_COUNT = 1000
 245  
 246          self.log.info("Test that Tor control does not disconnect on receiving MAX_LINE_COUNT lines.")
 247          with self.expect_disconnect(False, mock_tor):
 248              for _ in range(MAX_LINE_COUNT - 1):
 249                  mock_tor.send_raw("250-Continuing\r\n")
 250              mock_tor.send_raw("250 OK\r\n")
 251  
 252          self.log.info("Test that Tor control disconnects on receiving MAX_LINE_COUNT + 1 lines.")
 253          with self.expect_disconnect(True, mock_tor):
 254              for _ in range(MAX_LINE_COUNT + 1):
 255                  mock_tor.send_raw("250-Continuing\r\n")
 256  
 257          mock_tor.stop()
 258  
 259      def run_test(self):
 260          self.test_basic()
 261          self.test_partial_data()
 262          self.test_pow_fallback()
 263          self.test_oversized_line()
 264          self.test_overmany_lines()
 265  
 266  
 267  if __name__ == '__main__':
 268      TorControlTest(__file__).main()
 269