Logging TCP proxy, with SSL/TLS support

logproxy.go

package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
)

var raddr string

var (
	useTLS   = flag.Bool("tls", false, "use TLS/SSL for both connections")
	certFile = flag.String("cert", "", "X.509 certificate file")
	keyFile  = flag.String("keyfile", "", "X.509 key file")
)

func usage() {
	fmt.Fprintf(os.Stderr, "usage: %s [options] <listen-addr> <remote-addr>\n", os.Args[0])
	flag.PrintDefaults()
	os.Exit(1)
}

func fatal(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format, args...)
	os.Exit(1)
}

func logproxy(client net.Conn) {
	var (
		remote net.Conn
		err    error
	)
	if *useTLS {
		remote, err = tls.Dial("tcp", raddr, nil)
	} else {
		remote, err = net.Dial("tcp", raddr)
	}
	if err != nil {
		log.Println("dial remote:", err)
		return
	}
	defer remote.Close()
	ctee := io.TeeReader(client, remote)
	go io.Copy(os.Stdout, ctee)
	rtee := io.TeeReader(remote, client)
	io.Copy(os.Stdout, rtee)
}

func cleanup() chan net.Conn {
	ch := make(chan net.Conn)
	go func() {
		for conn := range ch {
			log.Println("closing connection")
			conn.Close()
		}
	}()
	return ch
}

func handle(done chan<- net.Conn) chan net.Conn {
	ch := make(chan net.Conn)
	go func() {
		for c := range ch {
			log.Println("new connection")
			logproxy(c)
			done <- c
		}
	}()
	return ch
}

func main() {
	flag.Parse()
	if flag.NArg() != 2 {
		usage()
	}

	if *useTLS && !(exists(*certFile) && exists(*keyFile)) {
		fatal("cert and key file required with tls\n")
	}

	var err error

	laddr, err := net.ResolveTCPAddr("tcp", flag.Arg(0))
	if err != nil {
		fatal("listen addr:", err, "\n")
	}

	raddr = flag.Arg(1)

	var ln net.Listener

	ln, err = net.ListenTCP("tcp", laddr)
	if err != nil {
		fatal("listen:", err, "\n")
	}

	if *useTLS {
		config := tls.Config{}
		config.Certificates = make([]tls.Certificate, 1)
		config.Certificates[0], err = tls.LoadX509KeyPair(*certFile, *keyFile)
		if err != nil {
			fatal("loading cert:", err)
		}

		ln = tls.NewListener(ln, &config)
	}

	log.Println("listening on", laddr)

	done := cleanup()
	conns := handle(done)

	for {
		c, err := ln.Accept()
		if err != nil {
			log.Println("accept:", err)
			continue
		}
		conns <- c
	}
}

func exists(filename string) bool {
	_, err := os.Stat(filename)
	return !os.IsNotExist(err)
}