user-tools/sockio

   1 #!/usr/bin/env perl
   2
   3 # Takes a command and arguments as arguments.
   4 # Connects the command stdio to a tcp/udp/unix socket.
   5
   6 use Socket qw/AF_UNIX AF_INET SOCK_STREAM pack_sockaddr_in inet_aton sockaddr_un/;
   7 use POSIX ':sys_wait_h';
   8 use Data::Dumper;
   9 use Fcntl qw/F_GETFL F_SETFL O_NONBLOCK O_RDONLY O_WRONLY/;
  10 no if ($] >= 5.018), 'warnings' => 'experimental::smartmatch';
  11
  12
  13 sub blockio
  14 {
  15   flag($_[0], ~O_NONBLOCK);
  16 }
  17 sub nonblockio
  18 {
  19   flag($_[0], O_NONBLOCK);
  20 }
  21 sub flag
  22 {
  23   fcntl($_[0], F_SETFL, fcntl($_[0], F_GETFL, 0) | $_[1]);
  24 }
  25 sub sigchld
  26 {
  27   # This handler makes select(2) interrupted, and child process reaped,
  28   # and stdin excluded from watched fds.
  29   1;
  30 }
  31
  32
  33 if(defined $ENV{'SOCKIO_AF'})
  34 {
  35   $af = $ENV{'SOCKIO_AF'};
  36 }
  37 else
  38 {
  39   print STDERR "SOCKIO_AF not given, fallback to 'tcp'\n";
  40   $af = 'tcp';
  41 }
  42 unless($af ~~ [qw/file abstract tcp/])
  43 {
  44   die "invalid SOCKIO_AF given, choose one of file, abstract, tcp\n";
  45 }
  46
  47 if($af ~~ [qw/file abstract/])
  48 {
  49   socket $srv, AF_UNIX, SOCK_STREAM, 0 or die $!;
  50   setsockopt($srv, SOL_SOCKET, SO_PASSCRED, 1);
  51   if($af eq 'abstract')
  52   {
  53     $addr = sockaddr_un "\x00sockio-pid-$$";
  54   }
  55   else
  56   {
  57     $addr = sockaddr_un "/tmp/sockio-pid-$$";
  58   }
  59   bind $srv, $addr or die $!;
  60 }
  61 elsif($af eq 'tcp')
  62 {
  63   socket $srv, AF_INET, SOCK_STREAM, 0 or die $!;
  64   BIND_TCP:
  65   $port = int(rand() * (65_535 - 10_000)) + 10_000;
  66   $addr = pack_sockaddr_in($port, inet_aton('localhost'));
  67   print STDERR "bind to tcp/$port ... ";
  68   if(bind $srv, $addr)
  69   {
  70     print STDERR "ok\n";
  71   }
  72   else
  73   {
  74     print STDERR "$!\n";
  75     if($!{'EADDRINUSE'}) { goto BIND_TCP; }
  76     else { exit $!; }
  77   }
  78 }
  79
  80 listen $srv, 1 or die $!;
  81
  82
  83 $pid = fork;
  84 if($pid == 0)
  85 {
  86   if($af ~~ [qw/file abstract/])
  87   {
  88     socket $cln, AF_UNIX, SOCK_STREAM, 0 or die $!;
  89   }
  90   elsif($af == 'tcp')
  91   {
  92     socket $cln, AF_INET, SOCK_STREAM, 0 or die $!;
  93   }
  94   CONNECT:
  95   if(not connect $cln, $addr)
  96   {
  97     if($!{'ECONNREFUSED'}) { sleep 0.1; goto CONNECT; }
  98     else { die $!; }
  99   }
 100   open STDIN, '<&=', fileno $cln or die $!;
 101   open STDOUT, '>&=', fileno $cln or die $!;
 102   flag(STDIN, O_RDONLY);
 103   flag(STDOUT, O_WRONLY);
 104   exec {$ARGV[0]} @ARGV;
 105   my ($errno, $errstr) = (int $!, $!);
 106   warn "$0: ${ARGV[0]}: $errstr\n";
 107   exit 125+$errno;
 108 }
 109
 110
 111 accept $clnt, $srv;
 112 select $clnt;
 113 $|++;
 114 select STDOUT;
 115 $|++;
 116
 117 @watch_in = (\*STDIN, $clnt);
 118 undef $exitcode;
 119 nonblockio(STDIN);
 120 nonblockio($clnt);
 121 $SIG{'CHLD'} = \&sigchld;
 122
 123
 124
 125 while(1)
 126 {
 127   if(scalar @watch_in == 0)
 128   {
 129     last;
 130   }
 131
 132   $fds_in = '';
 133   for my $fd (@watch_in)
 134   {
 135     vec($fds_in, fileno($fd), 1) = 1;
 136   }
 137
 138   select($fds_in, undef, undef, undef);
 139
 140   if(vec($fds_in, fileno(STDIN), 1) == 1)
 141   {
 142     READ_STDIN:
 143     $stdin_bytes = sysread STDIN, $buf, 4096;
 144     if($stdin_bytes > 0)
 145     {
 146       print {$clnt} $buf;
 147       goto READ_STDIN;
 148     }
 149
 150     if($stdin_bytes eq 0)
 151     {
 152       shutdown $clnt, 1;
 153       close STDIN;
 154       @watch_in = grep {$_ ne \*STDIN} @watch_in;
 155     }
 156   }
 157   if(vec($fds_in, fileno($clnt), 1) == 1)
 158   {
 159     READ_CLNT:
 160     $clnt_bytes = sysread $clnt, $buf, 4096;
 161     if($clnt_bytes > 0)
 162     {
 163       print {STDOUT} $buf;
 164       goto READ_CLNT;
 165     }
 166
 167     if($clnt_bytes eq 0)
 168     {
 169       shutdown $clnt, 0;
 170       close STDOUT;
 171       @watch_in = grep {$_ ne $clnt} @watch_in;
 172     }
 173   }
 174
 175   if(not defined $exitcode and waitpid $pid, WNOHANG)
 176   {
 177     $exitcode = $? >> 8;
 178
 179     shutdown $clnt, 1;
 180     close STDIN;
 181     @watch_in = grep {$_ ne \*STDIN} @watch_in;
 182   }
 183 }
 184
 185
 186 if(not defined $exitcode)
 187 {
 188   waitpid $pid, 0;
 189   $exitcode = $? >> 8;
 190 }
 191 exit $exitcode;