golang tcp proxy

proxy.go

package main

import (
	"fmt"
	"io"
	"net"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
)

func warn(err error) {
	fmt.Fprintln(os.Stderr, err)
}
func warnm(msg string, err error) {
	fmt.Fprintf(os.Stderr, "%s: %v\n", msg, err)
}

func simplexProxy(from, to *net.TCPConn, wg *sync.WaitGroup) {
	defer wg.Done()
	_, err := io.Copy(to, from)
	if err != nil && !strings.HasSuffix(err.Error(), "use of closed network connection") {
		warn(err)
	}
	to.Close()
	from.Close()
}

func duplexProxy(local *net.TCPConn, target string) {
	defer local.Close()
	remote, err := net.Dial("tcp", target)
	if err != nil {
		warn(err)
		return
	}
	defer remote.Close()

	wg := &sync.WaitGroup{}
	wg.Add(2)
	go simplexProxy(local, remote.(*net.TCPConn), wg)
	go simplexProxy(remote.(*net.TCPConn), local, wg)
	wg.Wait()
}

func main() {
	if len(os.Args) != 3 {
		fmt.Fprintf(os.Stderr, "Usage: %s listen-ip:listen-port target-ip:target-port\n", os.Args[0])
		os.Exit(1)
	}
	bindTo, sendTo := os.Args[1], os.Args[2]

	c := make(chan os.Signal, 2)
	signal.Notify(c, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		sig := <-c
		if sig != syscall.SIGTERM {
			os.Exit(1)
		}
		os.Exit(0)
	}()

	listener, err := net.Listen("tcp", bindTo)
	if err != nil {
		warnm(fmt.Sprintf("can't bind to %s", bindTo), err)
		return
	}
	defer listener.Close()

	for {
		connection, err := listener.Accept()
		if err == nil {
			go duplexProxy(connection.(*net.TCPConn), sendTo)
		} else {
			warnm("error accepting connection", err)
		}
	}
}