package main

import (
	"fmt"
	"net"
)

func main() {
	tcpListener, err := net.Listen("tcp", "0.0.0.0:3333")
	checkError(err)
	defer tcpListener.Close()
	fmt.Println("TCP-Listening on 0.0.0.0:3333")

	udpListener, err := net.ListenPacket("udp", "0.0.0.0:3333")
	checkError(err)
	defer udpListener.Close()
	fmt.Println("UDP-Listening on 0.0.0.0:3333")

	go listenUDP(udpListener)

	for {
		tcpConn, err := tcpListener.Accept()
		checkError(err)
		defer tcpConn.Close()
		fmt.Println("TCP-Connected")

		go handleTCPConn(tcpConn)
	}
}

func handleTCPConn(conn net.Conn) {
	// Make a buffer to hold incoming data.
	buf := make([]byte, 1024)
	i := 0

	for {
		// Read the incoming connection into the buffer.
		_, err := conn.Read(buf)
		if err != nil {
			fmt.Println("TCP connection error ", err.Error())
			break
		}

		// Log what received
		fmt.Printf("%d - TCP:%s", i, buf)
		i++
	}
}

func listenUDP(udpListener net.PacketConn) {
	buffer := make([]byte, 1024)
	i := 0

	for {
		//simple read
		udpListener.ReadFrom(buffer)

		// Log what received
		fmt.Printf("%d - UDP:%s", i, buffer)
		i++
	}
}

func checkError(err error) {
	if err != nil {
		fmt.Println("Fatal error ", err.Error())
		panic(err)
	}
}

package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"syscall"
	"time"
)

func handleError(err error) {
	if err != nil {
		log.Fatalf("Error with the statsd connection %s", err.Error())
		panic(err)
	}
}

// From https://github.com/go-sql-driver/mysql/blob/master/conncheck.go
// Note that it works only on Unix, and it caused several allocations.
// I really want this is implemented in stdlib in efficient (and probably cross platform) way.
var errUnexpectedRead = errors.New("unexpected read from socket")

func connCheck(c net.Conn) error {
	var (
		n    int
		err  error
		buff [1]byte
	)

	sconn, ok := c.(syscall.Conn)
	if !ok {
		return nil
	}
	rc, err := sconn.SyscallConn()
	if err != nil {
		return err
	}
	rerr := rc.Read(func(fd uintptr) bool {
		n, err = syscall.Read(int(fd), buff[:])
		return true
	})
	switch {
	case rerr != nil:
		return rerr
	case n == 0 && err == nil:
		return io.EOF
	case n > 0:
		return errUnexpectedRead
	case err == syscall.EAGAIN || err == syscall.EWOULDBLOCK:
		return nil
	default:
		return err
	}
}

func runClient(server *string, port *string, protocol *string, sleep int) {
	address := fmt.Sprintf("%s:%s", *server, *port)
	prot := *protocol
	conn, err := net.Dial(prot, address)
	handleError(err)

	for {
		msg := fmt.Sprintf("%sTEST\n", prot)
		log.Printf("Sending message: %s", msg)
		_, err := conn.Write([]byte(msg))
		handleError(err)

		if prot == "tcp" {
			err := connCheck(conn)
			handleError(err)
		}

		time.Sleep(time.Duration(sleep) * time.Second)
	}
}

func main() {
	// Command flag: https://gobyexample.com/command-line-flags
	var srv, port, protocol string
	var sleep int
	flag.StringVar(&srv, "server", "127.0.0.1", "the target server")
	flag.StringVar(&port, "port", "3333", "the target port")
	flag.StringVar(&protocol, "protocol", "tcp", "protocol used [tcp|udp]")
	flag.IntVar(&sleep, "sleep", 1, "sleep seconds")

	flag.Parse()

	fmt.Printf("using server: %s, port: %s, protocol: %s, sleep seconds:%d\n", srv, port, protocol, sleep)
	runClient(&srv, &port, &protocol, sleep)
}