Add TCP Healthcheck Listener To STUN Service

Dockerfile

@@ -0,0 +1,21 @@
+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
+RUN mkdir -p stun-server health-server
+COPY ./stun-server/main.go stun-server/main.go
+COPY ./health-server/main.go health-server/main.go
+COPY run_services.sh run_services.sh
+EXPOSE 3478/udp
+EXPOSE 8888/tcp
+RUN GO111MODULE=off CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -o stun-server/service stun-server/main.go 
+RUN GO111MODULE=off CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -o health-server/service health-server/main.go 
+CMD ["./run_services.sh"]
+
+# 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

health-server/main.go

@@ -0,0 +1,125 @@
+package main
+
+import (
+	"bufio"
+	"encoding/binary"
+	"fmt"
+	"log"
+	"net"
+	"net/http"
+	"sync"
+	"time"
+)
+
+const MagicCookie uint32 = 0x2112A442
+const MethodBinding uint16 = 0x001
+
+type StunHeader struct { // 20 bytes
+	Type          uint16   // 2 bytes
+	Length        uint16   // 2 bytes
+	MessageCookie uint32   // 4 bytes
+	TransactionID [12]byte // 12 bytes
+}
+
+var mu sync.RWMutex
+var stunServerIsHealthy = false
+
+type HttpHandler struct{}
+
+func (h HttpHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	if isHealthy() {
+		w.WriteHeader(http.StatusOK)
+	} else {
+		w.WriteHeader(http.StatusServiceUnavailable)
+	}
+}
+
+func isHealthy() bool {
+	mu.RLock()
+	defer mu.RUnlock()
+	return stunServerIsHealthy
+}
+
+func setHealthStatus(s bool) {
+	mu.Lock()
+	stunServerIsHealthy = s
+	mu.Unlock()
+}
+func testLocalStunService() {
+	log.Println("Running STUN test")
+	err := makeLocalStunBindingRequest()
+	if err != nil {
+		log.Printf("An error was observed while testing the STUN service %v", err)
+		setHealthStatus(false)
+		return
+	}
+	setHealthStatus(true)
+}
+
+func encodeToBytes(sh StunHeader) []byte {
+	buf := make([]byte, 20)
+	binary.BigEndian.PutUint16(buf[0:], sh.Type)
+	binary.BigEndian.PutUint16(buf[2:], sh.Length)
+	binary.BigEndian.PutUint32(buf[4:], sh.MessageCookie)
+	copy(buf[8:], sh.TransactionID[:])
+	return buf
+}
+
+func makeLocalStunBindingRequest() error {
+
+	// create a valid test binding request for the STUN server
+	var b = encodeToBytes(StunHeader{Type: MethodBinding,
+		Length:        0,
+		MessageCookie: MagicCookie,
+		TransactionID: [12]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}})
+
+	// send the STUN message binding datagram
+	conn, err := net.Dial("udp", "127.0.0.1:3478")
+	defer conn.Close()
+	if err != nil {
+		log.Printf("An error was triggered while sending STUN request %v", err)
+		return err
+	}
+	conn.Write(b)
+
+	// receive the resonse from STUN server, or timeout and fail if it is slow
+	rBuf := make([]byte, 128)
+	conn.SetReadDeadline(time.Now().Add(1 * time.Second))
+	_, err = bufio.NewReader(conn).Read(rBuf)
+	if err != nil {
+		log.Printf("An error was triggered while reading STUN response %v", err)
+		return err
+	}
+	log.Println("Success")
+	return nil
+}
+
+//
+
+func main() {
+
+	fmt.Println("Starting TCP healthcheck service.")
+	// create a new http request handler
+	handler := HttpHandler{}
+
+	// initialize the healthy state
+	testLocalStunService()
+
+	// run the check on a timer from now on
+	ticker := time.NewTicker(10 * time.Second)
+	quit := make(chan struct{})
+	go func() {
+		for {
+			select {
+			case <-ticker.C:
+				go testLocalStunService()
+			case <-quit:
+				ticker.Stop()
+				return
+			}
+		}
+	}()
+
+	// listen and serve the health service
+	http.ListenAndServe(":8888", handler)
+}

run_services.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+
+_term() {
+  echo "Caught SIGTERM signal!"
+  kill -TERM "$stun_service_pid" 2>/dev/null
+  kill -TERM "$health_service_pid" 2>/dev/null
+}
+
+trap _term SIGTERM
+
+# Start the stun server process
+./stun-server/service &
+stun_service_pid=$!
+
+# Start the health check side-car service process
+./health-server/service &
+health_service_pid=$!
+
+# Wait for any of the process to exit
+wait -n
+
+# Exit with status of process that exited first
+exit $?

main.go → stun-server/main.go

@@ -1,12 +1,12 @@
 package main
 
 import (
+	"encoding/binary"
 	"fmt"
 	"net"
-	"encoding/binary"
 )
 
-func isv4(ip []byte) (bool) {
+func isv4(ip []byte) bool {
 	return binary.BigEndian.Uint32(ip[0:4]) == 0 && binary.BigEndian.Uint64(ip[4:12]) == 0xffff
 }
 
@@ -20,11 +20,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 +38,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,22 +60,22 @@ 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 {
+	addr := net.UDPAddr{
 		Port: 3478,
-		IP: net.ParseIP("::"),
+		IP:   net.ParseIP("::"),
 	}
 
 	srv, err := net.ListenUDP("udp", &addr)

stun-server/main_test.go

@@ -0,0 +1,128 @@
+package main
+
+import (
+	"encoding/binary"
+	"fmt"
+	"net"
+	"reflect"
+	"testing"
+)
+
+// GO111MODULE=off CGO_ENABLED=0 go test
+
+const MagicCookie uint32 = 0x2112A442
+const MethodBinding uint16 = 0x001
+
+type StunHeader struct { // 20 bytes
+	Type          uint16   // 2 bytes
+	Length        uint16   // 2 bytes
+	MessageCookie uint32   // 4 bytes
+	TransactionID [12]byte // 12 bytes
+}
+
+func encodeStunHeaderToBytes(sh StunHeader) []byte {
+	buf := make([]byte, 20)
+	binary.BigEndian.PutUint16(buf[0:], sh.Type)
+	binary.BigEndian.PutUint16(buf[2:], sh.Length)
+	binary.BigEndian.PutUint32(buf[4:], sh.MessageCookie)
+	copy(buf[8:], sh.TransactionID[:])
+	return buf
+}
+
+func TestValidateRequest(t *testing.T) {
+	type test struct {
+		input []byte
+		want  error
+	}
+
+	tests := []test{
+		{input: make([]byte, 20), want: fmt.Errorf("Request is type 0, should be Binding Request (0x0001)")},
+		{input: make([]byte, 2), want: fmt.Errorf("Request is 2 bytes, should be 20")},
+		{input: encodeStunHeaderToBytes(StunHeader{Type: 0x1234,
+			Length:        0,
+			MessageCookie: MagicCookie,
+			TransactionID: [12]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}}),
+			want: fmt.Errorf("Request is type 1234, should be Binding Request (0x0001)")},
+		{input: encodeStunHeaderToBytes(StunHeader{Type: MethodBinding,
+			Length:        0,
+			MessageCookie: MagicCookie,
+			TransactionID: [12]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}}), want: nil},
+	}
+
+	for _, tc := range tests {
+		got := validateRequest(tc.input, len(tc.input))
+		if !reflect.DeepEqual(tc.want, got) {
+			t.Fatalf("expected: %v, got: %v", tc.want, got)
+		}
+	}
+}
+
+func printGolangBuffer(bt []byte) {
+	s := ""
+	for _, b := range bt {
+		s += fmt.Sprintf("%#.2x", b)
+		s += ","
+	}
+	fmt.Printf("[]byte{%s}\n", s)
+}
+
+func TestMakeResponse(t *testing.T) {
+	type test struct {
+		packet []byte
+		addr   *net.UDPAddr
+		want   []byte
+	}
+
+	// TODO: IPv6
+	addr1, _ := net.ResolveUDPAddr("udp", ":0")
+	addr2, _ := net.ResolveUDPAddr("udp", "127.0.0.1:22")
+	addr3, _ := net.ResolveUDPAddr("udp", "255.255.255.255:65535")
+
+	tests := []test{
+		{
+			packet: encodeStunHeaderToBytes(StunHeader{
+				Type:          0x1234,
+				Length:        0,
+				MessageCookie: MagicCookie,
+				TransactionID: [12]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12},
+			}),
+			addr: addr1,
+			want: []byte{0x01, 0x01, 0x00, 0x18, 0x21, 0x12, 0xa4, 0x42, 0x01, 0x02, 0x03,
+				0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x00, 0x20,
+				0x00, 0x14, 0x00, 0x02, 0x21, 0x12, 0x21, 0x12, 0xa4, 0x42, 0x01,
+				0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c},
+		},
+		{
+			packet: encodeStunHeaderToBytes(StunHeader{
+				Type:          0x1001,
+				Length:        0,
+				MessageCookie: MagicCookie,
+				TransactionID: [12]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12},
+			}),
+			addr: addr2,
+			want: []byte{0x01, 0x01, 0x00, 0x0c, 0x21, 0x12, 0xa4, 0x42, 0x01, 0x02, 0x03,
+				0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x00, 0x20,
+				0x00, 0x08, 0x00, 0x01, 0x21, 0x04, 0x5e, 0x12, 0xa4, 0x43},
+		},
+		{
+			packet: encodeStunHeaderToBytes(StunHeader{
+				Type:          0x1001,
+				Length:        65535,
+				MessageCookie: MagicCookie,
+				TransactionID: [12]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12},
+			}),
+			addr: addr3,
+			want: []byte{0x01, 0x01, 0x00, 0x0c, 0x21, 0x12, 0xa4, 0x42, 0x01, 0x02, 0x03, 0x04, 0x05,
+				0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x00, 0x20, 0x00, 0x08, 0x00, 0x01,
+				0xde, 0xed, 0xde, 0xed, 0x5b, 0xbd},
+		},
+	}
+
+	for _, tc := range tests {
+		got := makeResponse(tc.packet, tc.addr)
+		//printGolangBuffer(got)
+		if !reflect.DeepEqual(tc.want, got) {
+			t.Fatalf("expected: %v, got: %v", tc.want, got)
+		}
+	}
+}