#! /usr/bin/perl
#
# $Id$
# Copyright (C) 2004-2007 Jan Kratochvil <project-udpovertcp@jankratochvil.net>
# 
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; exactly version 2 of June 1991 is required
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA


use strict;
use warnings;
use Getopt::Long;
require IO::Socket::INET;
use Fcntl;
use Carp qw(cluck confess);
use Socket;
require IPC::Open2;


$SIG{__DIE__}=sub { confess @_; };
my $V=1;
$|=1;

my $D;
die "$0: <server-addr>:<port> [<local-port>:<server-addr>:<server-port>...]" if !GetOptions(
	"d|debug+",\$D,
);

my $forward_arg=shift||"";
my($forward_hostuser,$forward_portuser)=($forward_arg=~/^([\S^:]+):([\S^:]+)$/);
my($forward_in,$forward_out);
my $listen_tcp;
if ($forward_arg=~/^\w+$/) {
	$listen_tcp=IO::Socket::INET->new(
		LocalPort=>$forward_arg,
		Proto=>"tcp",
		Listen=>1,
		ReuseAddr=>1,
	) or die "socket(listen-tcp-<$forward_arg>): $!";
	die "Excessive arguments: @ARGV" if @ARGV;
} elsif ($forward_hostuser) {
	$forward_in=$forward_out=IO::Socket::INET->new(
		Proto=>"tcp",
		PeerHost=>$forward_hostuser,
		PeerPort=>$forward_portuser,
	) or die "socket(connect-tcp-<$forward_hostuser>:<$forward_portuser>): $!";
} elsif ($forward_arg) {
	IPC::Open2::open2($forward_in,$forward_out,$forward_arg) or die "$forward_arg: $!";
} else {
	$forward_in=\*STDIN;
	$forward_out=\*STDOUT;
	die "Excessive arguments: @ARGV" if @ARGV;
}

my %listenport;
while (@ARGV) {
	my $triple=shift;
	my($localport,$serverhostuser,$serverportuser)=($triple=~/^([\S^:]+):([\S^:]+):([\S^:]+)$/);
	my $serverhost=gethostbyname $serverhostuser or die "resolving host $serverhostuser: $!";
	my $serverport=getservbyname($serverportuser,"udp")||$serverportuser;
	my $serveraddr=sockaddr_in($serverport,$serverhost) or die "assembling ".inet_ntoa($serverhost).":$serverport: $!";
	die if 16!=length $serveraddr;
	my $localsock=IO::Socket::INET->new(
		LocalPort=>$localport,
		Proto=>"udp",
		ReuseAddr=>1,
	) or die "socket(listen-udp-<$forward_arg>): $!";
	die $localport if $listenport{$localsock->sockname()};
	$listenport{$serveraddr}={
		"localsock"=>$localsock,
		"serverhostuser"=>$serverhostuser,
		"serverportuser"=>$serverportuser,
		"serveraddr"=>$serveraddr,
	};
}

my $TYPE_SEND=0;	# SEND,sockaddr_in-where,sockaddr_in-from,length,data

sub sockaddr_in_name($)
{
	my($sockaddr_in)=@_;
	my($port,$iaddr)=sockaddr_in($sockaddr_in);
	my $peer_addr=inet_ntoa($iaddr);
	return "$peer_addr:$port";
}

sub send_udp($$$)
{
	my($to,$from,$data)=@_;

	my $listenport=$listenport{$from}||={
		"localsock"=>do {
			my($from_port,$from_iaddr)=sockaddr_in $from;
			warn "UDP-new-local: port=$from_port\n" if $D;
			IO::Socket::INET->new(
				LocalPort=>$from_port,
				ReuseAddr=>1,
				Proto=>"udp",
			) or die "New local UDP socket";
		},
		"serveraddr"=>$from,
	};
	warn "UDP-send: to=<".sockaddr_in_name($to).">,from-stored-serveraddr=<".sockaddr_in_name($from).">\n" if $D;
	my $got=send $listenport->{"localsock"},$data,0,$to;
	cluck "sendmsg(): $!" if !$got || $got!=length $data;
}

sub got_udp($$$)
{
	my($listenport,$peer_addr,$data)=@_;
	die if 16!=length $peer_addr;
	warn "TCP-write: to=<".sockaddr_in_name($listenport->{"serveraddr"}).">,from=<".sockaddr_in_name($peer_addr).">\n" if $D;
	my $packet=pack "Ca16a16Sa*",$TYPE_SEND,$listenport->{"serveraddr"},$peer_addr,length($data),$data;
	my $got=syswrite $forward_out,$packet;
	cluck "syswrite(): $!" if !$got || $got!=length $packet;
}

warn "START\n" if $D;
my $forward_buf="";
for (;;) {
	my $rfds="";
	vec($rfds,fileno($forward_in),1)=1 if $forward_in;
	vec($rfds,fileno($listen_tcp),1)=1 if $listen_tcp;
	for my $listenport (values(%listenport)) {
		vec($rfds,fileno($listenport->{"localsock"}),1)=1;
	}
	my $got=select $rfds,undef(),undef(),undef();
	die "Invalid select(2): ".Dumper($got) if !defined $got || $got<0;
	if ($listen_tcp && vec($rfds,fileno($listen_tcp),1)) {
		if ($forward_in) {
			$forward_in eq $forward_out or close $forward_out or confess "Error closing old forward_out socket: $!";
			close $forward_in or confess "Error closing old forward_in socket: $!";
		}
		accept $forward_in,$listen_tcp or confess "Error accepting new TCP socket: $!";
		$forward_out=$forward_in;
	}
	if ($forward_in && vec($rfds,fileno($forward_in),1)) {
		my $buf;
		fcntl($forward_in,F_SETFL,O_NONBLOCK) or die "fnctl(forward_in,F_SETFL,O_NONBLOCK)";
		my $got=sysread $forward_in,$buf,0x100000;
		fcntl($forward_in,F_SETFL,0)          or die "fnctl(forward_in,F_SETFL,0)";
		if (!$got) {
			die "Got TCP EOF/error on forward_in";
		} elsif ($got==length $buf) {
			$forward_buf.=$buf;
		} else {
			confess "Invalid socket read return value: $got";
		}
		while (length($forward_buf)>=1+16+16+2) {
			my($type,$to,$from,$length,$data)=unpack "Ca16a16Sa*",$forward_buf;
			last if length($data)<$length;
			($data,$forward_buf)=unpack "a${length}a*",$data;
			die if $length!=length $data;
			die if $type!=$TYPE_SEND;
			send_udp $to,$from,$data;
		}
	}
	for my $listenport (values(%listenport)) {
		my $localsock=$listenport->{"localsock"};
		next if !vec($rfds,fileno($localsock),1);
		my $buf;
		my $peer_addr=recv $localsock,$buf,0x100000,0;
		if (!$peer_addr || !defined $buf) {
			die "Got TCP EOF/error on one of the listenports";
		} else {
			got_udp $listenport,$peer_addr,$buf;
		}
	}
}