golang network proxy

go-proxy.go

// netcat go version
// Usage:
// 	go-nc hostname port
//
// - connect to hostname:port
// - read stdin and send to socket
// - read socket and write to stdout
package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"sync"
	"time"
)

var (
	timeout = flag.Duration("w", time.Minute,
		"If a connection and stdin are idle for more than timeout seconds, then the connection is silently closed.")
)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(flag.CommandLine.Output(), "Usage %s: [options] hostname port\n", os.Args[0])
		fmt.Fprintf(flag.CommandLine.Output(), "Options:\n")
		flag.PrintDefaults()
	}
}

type timeoutReader struct {
	conn    net.Conn
	timeout time.Duration
}

type timeoutWriter struct {
	conn    net.Conn
	timeout time.Duration
}

func (r *timeoutReader) Read(p []byte) (n int, err error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.timeout)); err != nil {
		panic(err)
	}

	return r.conn.Read(p)
}

func (r *timeoutWriter) Write(p []byte) (n int, err error) {
	if err := r.conn.SetWriteDeadline(time.Now().Add(r.timeout)); err != nil {
		panic(err)
	}

	return r.conn.Write(p)
}

func proxy(conn net.Conn) {
	var wg sync.WaitGroup
	defer wg.Wait()

	pipe := func(src io.Reader, dst io.Writer) {
		defer wg.Done()

		io.Copy(dst, src)
	}

	for _, rw := range []struct {
		reader io.Reader
		writer io.Writer
	}{
		{
			os.Stdin,
			&timeoutWriter{conn, *timeout},
		},
		{
			&timeoutReader{conn, *timeout},
			os.Stdout,
		},
	} {
		wg.Add(1)
		go pipe(rw.reader, rw.writer)
	}
}

func main() {
	flag.Parse()

	if flag.NArg() < 2 {
		flag.Usage()
		return
	}

	hostname := flag.Arg(0)
	port := flag.Arg(1)

	conn, err := net.Dial("tcp", fmt.Sprintf("%s:%s", hostname, port))
	if err != nil {
		log.Fatalf("Connection failed: %s:%s", hostname, port)
	}
	defer conn.Close()

	proxy(conn)
}