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   exit 127;
 106 }
 107
 108
 109 accept $clnt, $srv;
 110 select $clnt;
 111 $|++;
 112 select STDOUT;
 113 $|++;
 114
 115 @watch_in = (\*STDIN, $clnt);
 116 undef $exitcode;
 117 nonblockio(STDIN);
 118 nonblockio($clnt);
 119 $SIG{'CHLD'} = \&sigchld;
 120
 121
 122
 123 while(1)
 124 {
 125   if(scalar @watch_in == 0)
 126   {
 127     last;
 128   }
 129
 130   $fds_in = '';
 131   for my $fd (@watch_in)
 132   {
 133     vec($fds_in, fileno($fd), 1) = 1;
 134   }
 135
 136   select($fds_in, undef, undef, undef);
 137
 138   if(vec($fds_in, fileno(STDIN), 1) == 1)
 139   {
 140     READ_STDIN:
 141     $stdin_bytes = sysread STDIN, $buf, 4096;
 142     if($stdin_bytes > 0)
 143     {
 144       print {$clnt} $buf;
 145       goto READ_STDIN;
 146     }
 147
 148     if($stdin_bytes eq 0)
 149     {
 150       shutdown $clnt, 1;
 151       close STDIN;
 152       @watch_in = grep {$_ ne \*STDIN} @watch_in;
 153     }
 154   }
 155   if(vec($fds_in, fileno($clnt), 1) == 1)
 156   {
 157     READ_CLNT:
 158     $clnt_bytes = sysread $clnt, $buf, 4096;
 159     if($clnt_bytes > 0)
 160     {
 161       print {STDOUT} $buf;
 162       goto READ_CLNT;
 163     }
 164
 165     if($clnt_bytes eq 0)
 166     {
 167       shutdown $clnt, 0;
 168       close STDOUT;
 169       @watch_in = grep {$_ ne $clnt} @watch_in;
 170     }
 171   }
 172
 173   if(not defined $exitcode and waitpid $pid, WNOHANG)
 174   {
 175     $exitcode = $? >> 8;
 176
 177     shutdown $clnt, 1;
 178     close STDIN;
 179     @watch_in = grep {$_ ne \*STDIN} @watch_in;
 180   }
 181 }
 182
 183
 184 if(not defined $exitcode)
 185 {
 186   waitpid $pid, 0;
 187   $exitcode = $? >> 8;
 188 }
 189 exit $exitcode;