protohackers/pkg/conn/conn.go

package conn

import (
	"bufio"
	"fmt"
	"net"
)

const (
	Port = ":3030"
)

func StartSimpleTCP(handler func(conn net.Conn)) error {
	l, err := net.Listen("tcp4", Port)
	if err != nil {
		return err
	}
	defer l.Close()

	fmt.Println(fmt.Sprintf("Listening TCP on port %s", Port))

	for {
		c, err := l.Accept()
		if err != nil {
			return err
		}
		go handler(c)
	}
}

func StartSimpleUDP(handler func(*net.UDPConn, *net.UDPAddr, []byte), buffSize int) error {
	srv, err := net.ResolveUDPAddr("udp4", Port)
	if err != nil {
		return err
	}

	c, err := net.ListenUDP("udp4", srv)
	if err != nil {
		return err
	}
	defer c.Close()

	fmt.Println(fmt.Sprintf("Listening UDP on port %s", Port))

	buf := make([]byte, buffSize)
	for {
		n, addr, err := c.ReadFromUDP(buf)
		if err != nil {
			return err
		}
		go handler(c, addr, buf[:n])
	}
}

func ReadAndRespond(c net.Conn, process func([]byte) []byte) error {
	defer c.Close()
	sc := bufio.NewScanner(c)
	for sc.Scan() {
		response := process(sc.Bytes())
		_, err := c.Write(response)
		if err != nil {
			return err
		}
	}
	if err := sc.Err(); err != nil {
		return err
	}
	return nil
}