Add TCP Healthcheck Listener To STUN Service

Dockerfile

@@ -0,0 +1,17 @@
+FROM golang:1.17.7-buster
+ENV STUN_USER STUN
+ENV STUN_HOME /go/src/stun
+ARG GROUP_ID
+ARG USER_ID
+RUN groupadd --gid $GROUP_ID STUN && useradd -m -l --uid $USER_ID --gid $GROUP_ID $STUN_USER
+RUN mkdir -p $STUN_HOME && chown -R $STUN_USER:$STUN_USER $STUN_HOME
+USER $STUN_USER
+WORKDIR $STUN_HOME
+COPY main.go main.go
+EXPOSE 3478/udp
+EXPOSE 8888/tcp
+RUN GO111MODULE=off CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build .
+CMD ["./stun", "-hc"]
+
+# docker build --build-arg USER_ID=1234 --build-arg GROUP_ID=1234 -t stun-server .
+# docker run -it --rm -p 8888:8888/tcp -p 3478:3478/udp stun-server

main.go

@@ -1,12 +1,27 @@
 package main
 
 import (
+	"encoding/binary"
+	"flag"
 	"fmt"
 	"net"
-	"encoding/binary"
+	"os"
+	"time"
 )
 
-func isv4(ip []byte) (bool) {
+func tcpHealthCheckListener(l *net.TCPListener) {
+	for {
+		conn, err := l.AcceptTCP()
+		conn.SetKeepAlive(true)
+		conn.SetKeepAlivePeriod(1 * time.Minute)
+		if err != nil {
+			fmt.Println(err)
+		}
+		defer conn.Close()
+	}
+}
+
+func isv4(ip []byte) bool {
 	return binary.BigEndian.Uint32(ip[0:4]) == 0 && binary.BigEndian.Uint64(ip[4:12]) == 0xffff
 }
 
@@ -20,11 +35,11 @@ func xorv6(ip []byte, messageCookie []byte, transactionID []byte) {
 	xorv4(ip, messageCookie)
 
 	for x := 4; x < 16; x++ {
-		ip[x] ^= transactionID[x - 4]
+		ip[x] ^= transactionID[x-4]
 	}
 }
 
-func validateRequest(packet []byte, n int) (error) {
+func validateRequest(packet []byte, n int) error {
 	if n != 20 {
 		return fmt.Errorf("Request is %d bytes, should be 20", n)
 	}
@@ -38,7 +53,7 @@ func validateRequest(packet []byte, n int) (error) {
 	return nil
 }
 
-func makeResponse(packet []byte, addr *net.UDPAddr) ([]byte) {
+func makeResponse(packet []byte, addr *net.UDPAddr) []byte {
 	messageCookie := packet[4:8]
 	magicCookie := uint16(binary.BigEndian.Uint32(messageCookie) >> 16)
 	transactionID := packet[8:20]
@@ -60,25 +75,47 @@ func makeResponse(packet []byte, addr *net.UDPAddr) ([]byte) {
 
 	//header
 	binary.BigEndian.PutUint16(response[0:2], 0x0101)
-	binary.BigEndian.PutUint16(response[2:4], size - 20)
+	binary.BigEndian.PutUint16(response[2:4], size-20)
 	copy(response[4:8], messageCookie)
 	copy(response[8:20], transactionID)
 
 	//XOR-MAPPED-ADDRESS
 	binary.BigEndian.PutUint16(response[20:22], 0x0020)
-	binary.BigEndian.PutUint16(response[22:24], size - 20 - 4)
-	binary.BigEndian.PutUint16(response[26:28], uint16(addr.Port) ^ magicCookie)
+	binary.BigEndian.PutUint16(response[22:24], size-20-4)
+	binary.BigEndian.PutUint16(response[26:28], uint16(addr.Port)^magicCookie)
 
 	return response[0:size]
 }
 
 func main() {
-	addr := net.UDPAddr {
+
+	listenHealthcheck := flag.Bool("hc", false, "If this is specified, spawn a TCP health-check listener")
+	flag.Parse()
+
+	// Healthcheck
+	if *listenHealthcheck == true {
+		fmt.Println("Starting TCP healthcheck service.")
+		tcpAddr := net.TCPAddr{
+			Port: 8888,
+			IP:   net.ParseIP("::"),
+		}
+
+		hcServ, err := net.ListenTCP("tcp", &tcpAddr)
+		if err != nil {
+			fmt.Println(err)
+			os.Exit(-1)
+		}
+		go tcpHealthCheckListener(hcServ)
+	}
+
+	// STUN
+	fmt.Println("Starting UDP STUN service.")
+	udpAddr := net.UDPAddr{
 		Port: 3478,
-		IP: net.ParseIP("::"),
+		IP:   net.ParseIP("::"),
 	}
 
-	srv, err := net.ListenUDP("udp", &addr)
+	srv, err := net.ListenUDP("udp", &udpAddr)
 
 	if err == nil {
 		for {