Created
September 4, 2019 20:31
-
-
Save ksurent/2df280e410c6442f0d90a97ee952e43f to your computer and use it in GitHub Desktop.
A simple netcat-like utility written as a coding exercise
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment