A simple STUN server in Go

stunned.go

package main

import (
	"flag"
	"fmt"
	"net"
	"os"
)

const udpConnType = "udp4"

// STUNMessageType defines possible values
// of STUN message type header field
type STUNMessageType uint16

const (
	// BindingRequest message type
	// means the message is a Binding Request
	BindingRequest STUNMessageType = 0x0001
	// BindingResponse message type
	// means the message is a Binding Response
	BindingResponse = 0x0101
	// BindingErrorResponse message type
	// means the message is a Binding Error Response
	BindingErrorResponse = 0x0111
	// SharedSecretRequest message type
	// means the message is a Shared Secret Request
	SharedSecretRequest = 0x0002
	// SharedSecretResponse message type
	// means the message is a Shared Secret Response
	SharedSecretResponse = 0x0102
	// SharedSecretErrorResponse message type
	// means the message is a Shared Secret Error Response
	SharedSecretErrorResponse = 0x0112
)

type incomingSTUNRequest struct {
	data          []byte
	remoteAddress *net.UDPAddr
}

func main() {
	listenPort := flag.Int(
		"port",
		3478,
		"The port for the STUN server to listen to",
	)
	workersCount := flag.Int(
		"workers",
		4,
		"The number of workers to spawn",
	)
	flag.Parse()

	listenAddressString := fmt.Sprintf("localhost:%d", listenPort)
	listenAddress, err := net.ResolveUDPAddr(udpConnType, listenAddressString)

	if err != nil {
		fmt.Println("Failed to resolve own address:", err.Error())
		os.Exit(1)
	}

	// for simplicity, will only work over IPv4
	udpConnection, err := net.ListenUDP(udpConnType, listenAddress)

	if err != nil {
		fmt.Println("Failed to setup UDP connection:", err.Error())
		os.Exit(1)
	}

	defer udpConnection.Close()

	fmt.Println("Listening on ", listenAddress)

	incomingPackets := make(chan incomingSTUNRequest)

	for workerIndex := 0; workerIndex < *workersCount; workerIndex++ {
		go stunnedWorker(incomingPackets, udpConnection)
	}

	for {
		// huge buffer to handle max possible UDP packet
		dataBuffer := make([]byte, 65507)
		packetSize, remoteAddress, err := udpConnection.ReadFromUDP(dataBuffer)
		if err != nil {
			fmt.Println("Failed to read UDP packet: ", err.Error())
			os.Exit(1)
		}
		incomingPackets <- incomingSTUNRequest{dataBuffer[:packetSize], remoteAddress}
	}
}

func stunnedWorker(packets <-chan incomingSTUNRequest, connection *net.UDPConn) {
	for packet := range packets {
		// 32 bytes is quite enough for header + IPv4 payload
		responseBuffer := make([]byte, 32)
		messageType := STUNMessageType(uint16(packet.data[0])<<8 + uint16(packet.data[1]))

		switch messageType {
		case BindingRequest:
			// don't need message length
			// messageLength := uint16(packet.data[2]) << 8 + uint16(packet.data[3])
			txID := packet.data[4:20]

			// set msg type to BindingResponse
			responseBuffer[0] = BindingResponse >> 8
			responseBuffer[1] = byte(BindingResponse % 256)

			// msg length is always 20 because it's an IPv4 (4 bytes) and a port (uint16)
			responseBuffer[2] = byte(0)
			responseBuffer[3] = byte(16)

			// set TX ID
			for byteNum := 4; byteNum < 20; byteNum++ {
				responseBuffer[byteNum] = txID[byteNum]
			}

			// write family to buffer (0x0001 for IPv4, 0x0002 for IPv6)
			responseBuffer[21] = 0x00
			responseBuffer[22] = 0x01
			// write port to buffer
			responseBuffer[23] = uint8(packet.remoteAddress.Port >> 8)
			responseBuffer[24] = uint8(packet.remoteAddress.Port & 0xff)
			// write IP to buffer
			copy(responseBuffer[25:28], packet.remoteAddress.IP.To4())

			// ideologically, this should go to the bottom of the function,
			// but since we are implementing this case only, that will
			// result in responses with undefined content
			connection.WriteToUDP(responseBuffer, packet.remoteAddress)

		case BindingResponse:
			fallthrough
		case BindingErrorResponse:
			fallthrough
		case SharedSecretRequest:
			fallthrough
		case SharedSecretResponse:
			fallthrough
		case SharedSecretErrorResponse:
			fmt.Println("Not implemented")
		default:
			fmt.Println("Unknown STUN message type")
		}
	}
}