go.fuchsia.dev/infra@v0.0.0-20240507153436-9b593402251b/cmd/buildproxywrap/socket_sink.py (about)

     1  #!/usr/bin/env python3
     2  # Copyright 2023 The Fuchsia Authors. All rights reserved.
     3  # Use of this source code is governed by a BSD-style license that can
     4  # found in the LICENSE file.
     5  """Test server that just listens on a socket.
     6  
     7  This does the minimum amount of work to
     8  1) open a unix domain socket file, bind it
     9  2) start listening and consume any data
    10  
    11  The important thing is that it creates a socket
    12  after some delay.
    13  
    14  This is only used by relay_test.go.
    15  """
    16  
    17  import os
    18  import socket
    19  import sys
    20  import time
    21  from typing import Sequence
    22  
    23  
    24  def extract_socket_path(arg: str) -> str:
    25      colon = arg.find(":")
    26      if colon != -1:
    27          arg = arg[colon + 1 :]
    28      comma = arg.find(",")
    29      if comma != -1:
    30          arg = arg[:comma]
    31      return arg
    32  
    33  
    34  def sink(socket_arg: str) -> None:
    35      with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
    36          s.bind(socket_arg)
    37          s.listen()
    38          conn, addr = s.accept()
    39          with conn:
    40              print(f"Connected by {addr}")
    41              while True:
    42                  data = conn.recv(1024)
    43                  if not data:
    44                      break
    45                  # just sink the data, don't send it anywhere
    46  
    47  
    48  def main(argv: Sequence[str]) -> None:
    49      socket_arg = extract_socket_path(argv[0])
    50      print("Snoozing...")
    51      time.sleep(1)  # intentional delay
    52      print(f"Sinking data on socket {socket_arg}")
    53      if os.path.exists(socket_arg):
    54          os.remove(socket_arg)
    55  
    56      # Run until interrupted
    57      try:
    58          sink(socket_arg)
    59      finally:
    60          print("Shutting down and cleaning up")
    61          os.remove(socket_arg)
    62  
    63  
    64  if __name__ == "__main__":
    65      sys.exit(main(sys.argv[1:]))