Reef/wg/wg_test.go

package wg

import (
	"context"
	"io"
	"net"
	"net/netip"
	"strconv"
	"testing"
	"time"
)

func freeUDP(t *testing.T) int {
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer pc.Close()
	return pc.LocalAddr().(*net.UDPAddr).Port
}

func TestTunnelCarriesTCPAndUDP(t *testing.T) {
	ka, _ := GenerateKey()
	kb, _ := GenerateKey()
	pa, _ := PublicKey(ka)
	pb, _ := PublicKey(kb)
	portA := freeUDP(t)

	a, err := New(ka, netip.MustParseAddr("100.89.0.1"), portA) // the "edge": has a fixed port
	if err != nil {
		t.Fatal(err)
	}
	defer a.Close()
	b, err := New(kb, netip.MustParseAddr("100.89.0.2"), 0) // the "site": behind NAT, dials out
	if err != nil {
		t.Fatal(err)
	}
	defer b.Close()
	if err := a.SetPeers([]Peer{{PublicKey: pb, AllowedIP: "100.89.0.2"}}); err != nil {
		t.Fatal(err)
	}
	if err := b.SetPeers([]Peer{{PublicKey: pa, AllowedIP: "100.89.0.1", Endpoint: net.JoinHostPort("127.0.0.1", strconv.Itoa(portA))}}); err != nil {
		t.Fatal(err)
	}

	// TCP: site listens inside the tunnel, edge dials it
	ln, err := b.ListenTCP(9111)
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			go func() { defer c.Close(); io.Copy(c, c) }()
		}
	}()
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	c, err := a.DialContext(ctx, "tcp", "100.89.0.2:9111")
	if err != nil {
		t.Fatalf("dial through tunnel: %v", err)
	}
	defer c.Close()
	c.Write([]byte("hello"))
	buf := make([]byte, 5)
	c.SetReadDeadline(time.Now().Add(5 * time.Second))
	if _, err := io.ReadFull(c, buf); err != nil || string(buf) != "hello" {
		t.Fatalf("echo: %q %v", buf, err)
	}

	// UDP both ways
	pc, err := b.ListenUDP(9112)
	if err != nil {
		t.Fatal(err)
	}
	defer pc.Close()
	go func() {
		p := make([]byte, 2048)
		for {
			n, from, err := pc.ReadFrom(p)
			if err != nil {
				return
			}
			pc.WriteTo(append([]byte("re:"), p[:n]...), from)
		}
	}()
	u, err := a.DialContext(ctx, "udp", "100.89.0.2:9112")
	if err != nil {
		t.Fatal(err)
	}
	defer u.Close()
	u.Write([]byte("ping"))
	u.SetReadDeadline(time.Now().Add(5 * time.Second))
	rb := make([]byte, 64)
	n, err := u.Read(rb)
	if err != nil || string(rb[:n]) != "re:ping" {
		t.Fatalf("udp: %q %v", rb[:n], err)
	}

	// the site shows a recent handshake with the edge
	if hs := b.Handshakes(); time.Since(hs[pa]) > time.Minute {
		t.Errorf("no handshake recorded: %v", hs)
	}

	// a second peer can be added and removed without disturbing the first
	kc, _ := GenerateKey()
	pcKey, _ := PublicKey(kc)
	if err := a.SetPeers([]Peer{{PublicKey: pb, AllowedIP: "100.89.0.2"}, {PublicKey: pcKey, AllowedIP: "100.89.0.3"}}); err != nil {
		t.Fatal(err)
	}
	c.Write([]byte("again"))
	buf2 := make([]byte, 5)
	c.SetReadDeadline(time.Now().Add(5 * time.Second))
	if _, err := io.ReadFull(c, buf2); err != nil || string(buf2) != "again" {
		t.Fatalf("existing connection broke when a peer was added: %q %v", buf2, err)
	}
	if err := a.SetPeers([]Peer{{PublicKey: pb, AllowedIP: "100.89.0.2"}}); err != nil {
		t.Fatal(err)
	}
}