go tcp server/client

tcpclient.go

package main 

import (
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"net"
	"io"
	"time"
)

const (
	TIMEOUT = 10
)

func main() {
	address, timeout, ok := getConnectParam(os.Args)
	if !ok {
		return
	}

	conn, err := net.DialTimeout("tcp", address, timeout)
	if err != nil {
		fmt.Println("Error connecting: ", err.Error())
		return
	}

	// signal handle channel
	c := make(chan os.Signal, 1)
	// remote reading quit channel
	qr := make(chan bool, 1)
	// main thread quit channel
	qm := make(chan bool, 1)

	// handle ctrl-c
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
	go handleSignals(c, qr, func() {
		conn.Close()
	})
	// print out all received
	go outRemote(os.Stdout, conn, qr, qm)
	
	defer conn.Close()
	
	// send all stdin out
	if _, err := io.Copy(conn, os.Stdin); err != nil {
		fmt.Fprintln(os.Stderr, "Error reading: ", err)
	}
	
	// blocking until all has been done.
	select {
	case <- qm:
	}
}

func outRemote(out io.Writer, conn net.Conn, qr chan bool, qm chan bool) {
	if _, err := io.Copy(out, conn); err != nil {
		select {
		case <- qr:
			// error because the connection was closed in main thread.
			close(qm)
		default:
			fmt.Fprintln(os.Stderr, "Error remote reading: ", err)
			
		}
	}
	conn.Close()
	os.Exit(0)
}

func handleSignals(c chan os.Signal, q chan bool, exec func()) {
	for _ = range c {
		exec()
		// notify signal handled.
		close(q)
	}
}

func getConnectParam(args []string) (string, time.Duration, bool) {
	if len(args) < 2 {
		fmt.Println(usage(args[0]))
		return "", 0, false
	}
	
	var timeout time.Duration
	if len(args) > 2 {
		timeout, _ = time.ParseDuration(args[2])
	}
	if timeout == 0 {
		timeout = time.Duration(TIMEOUT * time.Second)
	}
	
	return args[1], timeout, true
}

func usage(filename string) string {
	return fmt.Sprintf("Usage: %s <address (ex. localhost:8016, google.com:http, [2001:db8::1]:http, etc.)> [timeout (ex. 10s, 300ms, etc.)]", filename)
}

tcpsvr.go

package main 

import (
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"net"
	"io"
)

func main() {
	if len(os.Args) < 2 {
		fmt.Println(usage(os.Args[0]))
		return
	}
	
	ln, err := net.Listen("tcp", fmt.Sprintf(":%s", os.Args[1]))
	
	if err != nil {
		fmt.Println("Error listening: ", err.Error())
		return
	}
	
	conn, err := ln.Accept()
	
	if err != nil {
		fmt.Println("Error accepting: ", err.Error())
		return
	}

	// signal handle channel
	c := make(chan os.Signal, 1)
	// remote reading quit channel
	qr := make(chan bool, 1)
	// main thread quit channel
	qm := make(chan bool, 1)

	// handle ctrl-c
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
	go handleSignals(c, qr, func() {
		conn.Close()
	})
	// print out all received
	go outRemote(os.Stdout, conn, qr, qm)
	
	defer conn.Close()
	
	// send all stdin out
	if _, err := io.Copy(conn, os.Stdin); err != nil {
		fmt.Fprintln(os.Stderr, "Error reading: ", err)
	}
	
	// blocking until all has been done.
	select {
	case <- qm:
	}
}

func outRemote(out io.Writer, conn net.Conn, qr chan bool, qm chan bool) {
	if _, err := io.Copy(out, conn); err != nil {
		select {
		case <- qr:
			// error because the connection was closed in main thread.
			close(qm)
		default:
			fmt.Fprintln(os.Stderr, "Error remote reading: ", err)
			
		}
	}
	conn.Close()
	os.Exit(0)
}

func handleSignals(c chan os.Signal, q chan bool, exec func()) {
	for _ = range c {
		exec()
		// notify signal handled.
		close(q)
	}
}

func usage(filename string) string {
	return fmt.Sprintf("%s <port number to listen>", filename)
}