How to test TCP/UDP connections in Go - Part 2

This was originally posted on my blog: devtheweb.io
Before continuing please ensure you've read the first installment of this series - How to test TCP/UDP connections in Go - Part 1

Introduction

In the last installment of this mini-series we took a look at testing TCP connections in Golang and how we can verify expected outputs.
In Part 2 we're going to look at taking the application one-step further by adding UDP capability. There are several things to be aware of when working with UDP or other stream-based protocols, all of which I will elaborate on below:

1. UDP connections use the net.PacketConn interface, instead of net.Listener
2. A UDP client connection in golang is just a net.UDPAddr instead of net.Conn, meaning we can't conn.Close()
3. Since net.PacketConn does not implement io.writer, we can't simply write to a connected client.

As someone with little network programming experience in any language, I came across quite a few “gotchas” trying to test the UPD connections. I started by creating an interface I assumed each protocol would satisfy - how wrong I was! Now that I’ve had more experience, utilised some fantastic documentation on the std lib, and learned from my mistakes, I will demonstrate how I solved the problem of meeting the same API and how to best communicate with a UDP service in hopes of saving you from having the same “gotchas” moments I had.

Update the tests

In true test-driven manner, there is one thing we need to do before we do anything else - update the tests from Part 1. Let’s update the tests to reflect the desired addition of the UDP.

var tcp, udp Server

func init() {
    // Start the new server
    tcp, err := NewServer("tcp", ":1123")
    if err != nil {
        log.Println("error starting TCP server")
        return
    }

    udp, err := NewServer("udp", ":6250")
    if err != nil {
        log.Println("error starting UDP server")
        return
    }

    // Run the servers in goroutines to stop blocking
    go func() {
        tcp.Run()
    }()
    go func() {
        udp.Run()
    }()
}

func TestNETServer_Running(t *testing.T) {
    // Simply check that the server is up and can
    // accept connections.
    servers := []struct {
        protocol string
        addr     string
    }{
        {"tcp", ":1123"},
        {"udp", ":6250"},
    }
    for _, serv := range servers {
        conn, err := net.Dial(serv.protocol, serv.addr)
        if err != nil {
            t.Error("could not connect to server: ", err)
        }
        defer conn.Close()
    }
}

func TestNETServer_Request(t *testing.T) {
    servers := []struct {
        protocol string
        addr     string
    }{
        {"tcp", ":1123"},
        {"udp", ":6250"},
    }

    tt := []struct {
        test    string
        payload []byte
        want    []byte
    }{
        {"Sending a simple request returns result", []byte("hello world\n"), []byte("Request received: hello world")},
        {"Sending another simple request works", []byte("goodbye world\n"), []byte("Request received: goodbye world")},
    }

    for _, serv := range servers {
        for _, tc := range tt {
            t.Run(tc.test, func(t *testing.T) {
                conn, err := net.Dial(serv.protocol, serv.addr)
                if err != nil {
                    t.Error("could not connect to server: ", err)
                }
                defer conn.Close()

                if _, err := conn.Write(tc.payload); err != nil {
                    t.Error("could not write payload to server:", err)
                }

                out := make([]byte, 1024)
                if _, err := conn.Read(out); err == nil {
                    if bytes.Compare(out, tc.want) == 0 {
                        t.Error("response did match expected output")
                    }
                } else {
                    t.Error("could not read from connection")
                }
            })
        }
    }
}

All we've done here is bootstrap the new UDPServer and then add a slice of servers to each test. This allows us to run the tests for each type of server connection with unique protocol and network address. If you run go test -v now, you should see the tests failing or erroring, but don't worry about that for now, we're going to fix it.

Next, at the bottom of net.go, we're going to write the minimum amount code we need to create our new type UDPServer and to implement the Server interface we defined in Part 1:

// UDPServer holds the necessary structure for our
// UDP server.
type UDPServer struct {
    addr   string
    server *net.UDPConn
}

// Run starts the UDP server.
func (u *UDPServer) Run() (err error) { return nil }

// Close ensures that the UDPServer is shut down gracefully.
func (u *UDPServer) Close() error { return nil }

The tests will still fail, however, we now have our foundations to build the UDPServer on. The next step is to implement the above two methods, like so:

// Run starts the UDP server.
func (u *UDPServer) Run() (err error) {
    laddr, err := net.ResolveUDPAddr("udp", u.addr)
    if err != nil {
        return errors.New("could not resolve UDP addr")
    }

    u.server, err = net.ListenUDP("udp", laddr)
    if err != nil {
        return errors.New("could not listen on UDP")
    }

    for {
        buf := make([]byte, 2048)
        n, conn, err := u.server.ReadFromUDP(buf)
        if err != nil {
            return errors.New("could not read from UDP")
        }
        if conn == nil {
            continue
        }

        u.server.WriteToUDP([]byte(fmt.Sprintf("Request recieved: %s", buf[:n])), conn)
    }
}

// Close ensures that the UDPServer is shut down gracefully.
func (u *UDPServer) Close() error {
    return u.server.Close()
}

That's it! Our tests should now all pass, and we have a working UDP connection Notice the differences here? First, we have to resolve the UDP address before we can start listening for connections and starting the server. Then we start accepting requests from the server. This is a little different with a UDP server. With a net.Listener, we can just .Accept() individual connections, however with UDP connections, we read from the server connection for requests and write each one to a buffer, we can then use the buffer to parse commands etc. ReadFromUDP returns three variables:

1. An integer representing the number of bytes written
2. UDP address of remote connection
3. Any errors encountered

We can use the first to parse only the number of written bytes, as shown in the example above with buf[:n]. Having the buffer sized to 2048 bytes allows us to listen for larger requests. It's also important to note that instead of writing to the connection, we have to use the server to write to the UDP addr of the connection. Since net.UDPConn doesn't implement io.Writer or io.Reader, we can't use the same approach we did with TCP by using the bufio package. Although I found this approach to be successful, I would love to hear your suggestions on how to overcome this problem, as there are bound to be cleaner ways of solving it.

Okay, so we've achieved the functionality we want, but can we abstract away some functionality? Yes, we can. Using the same approach as we did for TCP with a handleConnections method, we can reduce the responsibility of the Run method and ensure both servers have the same internal API. This benefits any other developers of the package as it provides a consistent way of working with different network protocols. Let's add the following:

func (u *UDPServer) Run() (err error) {
    laddr, err := net.ResolveUDPAddr("udp", u.addr)
    if err != nil {
        return errors.New("could not resolve UDP addr")
    }

    u.server, err = net.ListenUDP("udp", laddr)
    if err != nil {
        return errors.New("could not listen on UDP")
    }

    return u.handleConnections()
}

func (u *UDPServer) handleConnections() error {
    var err error
    for {
        buf := make([]byte, 2048)
        n, conn, err := u.server.ReadFromUDP(buf)
        if err != nil {
            log.Println(err)
            break
        }
        if conn == nil {
            continue
        }

        go u.handleConnection(conn, buf[:n])
    }
    return err
}

func (u *UDPServer) handleConnection(addr *net.UDPAddr, cmd []byte) {
    u.server.WriteToUDP([]byte(fmt.Sprintf("Request recieved: %s", cmd)), addr)
}

Voila! We've broken down the original Run method into two more functions that will achieve the same results and make the tests pass, but in a much cleaner way.

Conclusion

That's all for this series folks! I hope over the course of both posts I've clarified some of the differences and problems you may face when attempting to offer a tested, reliable, and simple service via TCP & UDP. Thanks again for reading, please consider sharing or reaching out to me on twitter @whg_codes.

The source code for this example can be found on my github here: github.com/williamhgough/devtheweb-source