A simple netcat-like utility written as a coding exercise

nc.go

package main

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

var (
	timeout = flag.Duration("timeout", 0, "I/O timeout (TCP only).")
	listen  = flag.Bool("listen", false, "Server mode.")
	verbose = flag.Bool("verbose", false, "Enable debug output.")
	udp     = flag.Bool("udp", false, "Use UDP.")
	inet4   = flag.Bool("4", false, "Force IPv4 only.")
	inet6   = flag.Bool("6", false, "Force IPv6 only.")
	stay    = flag.Bool("stay", false, "Remain listening after a connection is closed (TCP server only).")
)

func main() {
	log.SetOutput(os.Stderr)

	flag.Parse()

	addr := flag.Arg(0)
	if addr == "" {
		fmt.Fprintln(os.Stderr, "Usage: netcat [-flags] host:port")
		os.Exit(1)
	}

	var err error
	if *listen {
		err = listener(addr, *udp, *inet4, *inet6, *stay, *timeout)
	} else {
		err = sender(addr, *udp, *inet4, *inet6, *timeout)
	}
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func listener(addr string, udp, inet4, inet6, stay bool, timeout time.Duration) error {
	var network string

	if udp {
		network = "udp"
		if inet4 {
			network = "udp4"
		} else if inet6 {
			network = "udp6"
		}

		return udpListener(network, addr)
	}

	network = "tcp"
	if inet4 {
		network = "tcp4"
	} else if inet6 {
		network = "tcp6"
	}

	return tcpListener(network, addr, stay, timeout)
}

func sender(addr string, udp, inet4, inet6 bool, timeout time.Duration) error {
	var network string

	if udp {
		network = "udp"
		if inet4 {
			network = "udp4"
		} else if inet6 {
			network = "udp6"
		}

		return udpSender(network, addr)
	}

	network = "tcp"
	if inet4 {
		network = "tcp4"
	} else if inet6 {
		network = "tcp6"
	}

	return tcpSender(network, addr, timeout)
}

func tcpListener(network, addr string, stay bool, timeout time.Duration) error {
	l, err := net.Listen(network, addr)
	if err != nil {
		return err
	}
	defer func() {
		l.Close()
		logf("stopped listening on %s", l.Addr())
	}()
	logf("listening on %s", l.Addr())

	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}
		logf("accepted connection from %s", conn.RemoteAddr())

		err = handleConn(conn, timeout, true, true)
		conn.Close()
		logf("connection from %s closed", conn.RemoteAddr())
		if err != nil {
			return err
		}

		if !stay {
			break
		}
	}

	return nil
}

func tcpSender(network, addr string, timeout time.Duration) error {
	conn, err := net.DialTimeout(network, addr, timeout)
	if err != nil {
		return err
	}
	defer func() {
		conn.Close()
		logf("disconnected from %s", conn.RemoteAddr())
	}()
	logf("connected to %s", conn.RemoteAddr())

	return handleConn(conn, timeout, true, true)
}

func handleConn(conn net.Conn, timeout time.Duration, r, w bool) error {
	if !r && !w {
		panic("connection must have a reader or a writer")
	}

	setDeadline := func() error {
		if timeout > 0 {
			return conn.SetDeadline(time.Now().Add(timeout))
		}

		return nil
	}

	errCh := make(chan error)
	if r {
		go ioLoop(os.Stdout, conn, setDeadline, errCh)
	}
	if w {
		go ioLoop(conn, os.Stdin, setDeadline, errCh)
	}

	return <-errCh
}

func ioLoop(dst io.Writer, src io.Reader, deadline func() error, errCh chan<- error) {
	for {
		if err := deadline(); err != nil {
			errCh <- err
			return
		}

		if n, err := io.Copy(dst, src); err != nil || n == 0 {
			errCh <- err
			return
		}
	}
}

func logf(format string, args ...interface{}) {
	if *verbose {
		log.Println(fmt.Sprintf(format, args...))
	}
}

func udpListener(network, addr string) error {
	uaddr, err := net.ResolveUDPAddr(network, addr)
	if err != nil {
		return err
	}

	conn, err := net.ListenUDP(network, uaddr)
	if err != nil {
		return err
	}
	defer func() {
		conn.Close()
		logf("stopped listening for datagrams on %s", conn.LocalAddr())
	}()
	logf("listening for datagrams on %s", conn.LocalAddr())

	return handleConn(conn, 0, true, false)
}

func udpSender(network, addr string) error {
	conn, err := net.Dial(network, addr)
	if err != nil {
		return err
	}
	defer func() {
		conn.Close()
		logf("no more datagrams to %s", conn.RemoteAddr())
	}()
	logf("ready to send datagrams to %s", conn.RemoteAddr())

	return handleConn(conn, 0, false, true)
}