generic tcp proxy

pipe.go

package main

import (
	"flag"
	"io"
	"log"
	"net"
	"runtime"
)

var listenAddr = flag.String("listen", ":8025", "Address to listen on.")
var connectAddr = flag.String("connect", "mail:25", "Address to connect to.")

func myCopy(dst io.Writer, src io.Reader, ch chan bool) {
	defer func() { ch <- true }()
	io.Copy(dst, src)
}

func maybeFatal(msg string, e error) {
	if e != nil {
		log.Printf(msg, e)
		runtime.Goexit()
	}
}

func handleConn(c net.Conn, destAddr string) {
	defer c.Close()
	client, err := net.Dial("tcp", destAddr)
	maybeFatal("Error connecting to the other side:  %v", err)
	defer client.Close()

	log.Printf("Connected a new conn")

	ch := make(chan bool)

	go myCopy(c, client, ch)
	go myCopy(client, c, ch)

	<-ch
	c.Close()
	client.Close()
	<-ch

	log.Printf("Lost a connection")
}

func main() {
	flag.Parse()

	addr, err := net.ResolveTCPAddr("tcp", *listenAddr)
	maybeFatal("Error resolving listen address:  %v", err)
	l, err := net.ListenTCP("tcp", addr)
	maybeFatal("Error listening: %v", err)

	for {
		c, err := l.Accept()
		maybeFatal("Error connecting:  %v", err)
		go handleConn(c, *connectAddr)
	}
}