pkg/a2s: Implement server probe

2 files changed, 240 insertions, 1 deletions

diff --git a/pkg/a2s/a2s.go b/pkg/a2s/a2s.go
index 2d68725..9d1d65d 100644
--- a/pkg/a2s/a2s.go
+++ b/pkg/a2s/a2s.go
@@ -1,3 +1,144 @@
 // Package a2s implements a small portion of the r2 netchannel used by Northstar
-// to test servers.
+// to probe servers.
 package a2s
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"encoding/binary"
+	"fmt"
+	"net"
+	"net/netip"
+	"time"
+)
+
+const ProbeUID uint64 = 1000000001337
+
+func Probe(addr netip.AddrPort, timeout time.Duration) error {
+	conn, err := net.DialUDP("udp", nil, net.UDPAddrFromAddrPort(addr))
+	if err != nil {
+		return fmt.Errorf("connect to server: %w", err)
+	}
+	defer conn.Close()
+
+	t := time.Now().Add(timeout)
+	conn.SetWriteDeadline(t)
+	conn.SetReadDeadline(t)
+
+	pkt, err := r2cryptoEncrypt(r2encodeGetChallenge(ProbeUID))
+	if err != nil {
+		return fmt.Errorf("encrypt connection packet: %w", err)
+	}
+	if _, err := conn.Write(pkt); err != nil {
+		return fmt.Errorf("send connection packet: %w", err)
+	}
+
+	resp := make([]byte, 1500)
+	if n, err := conn.Read(resp); err != nil {
+		return fmt.Errorf("receive packet: %w", err)
+	} else {
+		resp = resp[:n]
+	}
+
+	dec, err := r2cryptoDecrypt(resp)
+	if err != nil {
+		return fmt.Errorf("failed to decrypt received packet: %w", err)
+	}
+
+	uid, _, err := r2decodeChallenge(dec)
+	if err != nil {
+		return fmt.Errorf("invalid challenge: %w", err)
+	}
+	if uid != ProbeUID {
+		return fmt.Errorf("invalid challenge")
+	}
+	return nil
+}
+
+const (
+	r2cryptoNonceSize = 12
+	r2cryptoTagSize   = 16
+	r2cryptoKey       = "X3V.bXCfe3EhN'wb"
+	r2cryptoAAD       = "\x01\x02\x03\x04\x05\x06\x07\x08\t\n\x0b\x0c\r\x0e\x0f\x10"
+)
+
+func r2cryptoEncrypt(b []byte) ([]byte, error) {
+	pkt := make([]byte, r2cryptoNonceSize+r2cryptoTagSize+len(b))
+
+	nonce := pkt[:r2cryptoNonceSize]
+	if _, err := rand.Read(nonce); err != nil {
+		return nil, err
+	}
+
+	c, err := aes.NewCipher([]byte(r2cryptoKey))
+	if err != nil {
+		panic(err)
+	}
+
+	a, err := cipher.NewGCMWithTagSize(c, r2cryptoTagSize)
+	if err != nil {
+		panic(err)
+	}
+
+	// Go appends the ciphertext, then the tag to the dest (nonce)
+	tmp := a.Seal(nil, nonce, b, []byte(r2cryptoAAD))
+	copy(pkt[r2cryptoNonceSize:], tmp[len(b):])
+	copy(pkt[r2cryptoNonceSize+r2cryptoTagSize:], tmp)
+	return pkt, nil
+}
+
+func r2cryptoDecrypt(b []byte) ([]byte, error) {
+	if len(b) < r2cryptoNonceSize+r2cryptoTagSize+1 {
+		return nil, fmt.Errorf("packet too small")
+	}
+
+	c, err := aes.NewCipher([]byte(r2cryptoKey))
+	if err != nil {
+		panic(err)
+	}
+
+	a, err := cipher.NewGCMWithTagSize(c, r2cryptoTagSize)
+	if err != nil {
+		panic(err)
+	}
+
+	tmp := make([]byte, len(b)-r2cryptoNonceSize)
+	copy(tmp, b[r2cryptoNonceSize+r2cryptoTagSize:])
+	copy(tmp[len(tmp)-r2cryptoTagSize:], b[r2cryptoNonceSize:])
+
+	if tmp, err = a.Open(tmp[:0], b[:r2cryptoNonceSize], tmp, []byte(r2cryptoAAD)); err != nil {
+		return nil, err
+	}
+	return tmp, nil
+}
+
+func r2encodeGetChallenge(uid uint64) []byte {
+	var b bytes.Buffer
+	binary.Write(&b, binary.LittleEndian, int32(-1))
+	binary.Write(&b, binary.LittleEndian, uint8(72))
+	binary.Write(&b, binary.LittleEndian, []byte("connect\x00"))
+	binary.Write(&b, binary.LittleEndian, uint64(uid))
+	binary.Write(&b, binary.LittleEndian, uint8(2))
+	return b.Bytes()
+}
+
+func r2decodeChallenge(b []byte) (uint64, int32, error) {
+	var pkt struct {
+		Seq       int32
+		Type      uint8
+		Challenge int32
+		UID       uint64
+	}
+	if err := binary.Read(bytes.NewReader(b), binary.LittleEndian, &pkt); err != nil {
+		return 0, 0, err
+	}
+	if pkt.Seq != -1 {
+		return 0, 0, fmt.Errorf("not a connectionless packet")
+	}
+	if pkt.Type != 73 {
+		return 0, 0, fmt.Errorf("not a challenge response")
+	}
+	return pkt.UID, pkt.Challenge, nil
+}
diff --git a/pkg/a2s/a2s_test.go b/pkg/a2s/a2s_test.go
new file mode 100644
index 0000000..72eba3a
--- /dev/null
+++ b/pkg/a2s/a2s_test.go
@@ -0,0 +1,98 @@
+package a2s
+
+import (
+	"bytes"
+	"encoding/hex"
+	"fmt"
+	"testing"
+)
+
+func TestPacketRoundTrip(t *testing.T) {
+	b := r2encodeGetChallenge(1000000001337)
+	e := mustDecodeHex("ffffffff48636f6e6e656374003915a5d4e800000002")
+
+	if !bytes.Equal(b, e) {
+		t.Error("incorrect getchallenge encoding")
+	}
+
+	be, err := r2cryptoEncrypt(b)
+	if err != nil {
+		t.Errorf("failed to encrypt packet: %v", err)
+	}
+
+	bd, err := r2cryptoDecrypt(be)
+	if err != nil {
+		t.Errorf("failed to decrypt packet: %v", err)
+	}
+
+	if !bytes.Equal(bd, b) {
+		t.Error("incorrect decryption result")
+	}
+}
+
+func TestDecodeChallenge(t *testing.T) {
+	b := mustDecodeHex("f4ca55b7f53a2f9c19b563010d6964869648a23be1db9edce9f55ee3f9a02451be86ba56447740d1d893c34f3a854f6efbd47605ebf3211e05")
+
+	bd, err := r2cryptoDecrypt(b)
+	if err != nil {
+		t.Errorf("failed to decrypt packet: %v", err)
+	}
+
+	uid, challenge, err := r2decodeChallenge(bd)
+	if err != nil {
+		t.Errorf("failed to parse packet: %v", err)
+	}
+
+	if uid != 1000000001337 {
+		t.Errorf("incorrect uid")
+	}
+
+	if challenge != 81930672 {
+		t.Errorf("incorrect challenge")
+	}
+}
+func FuzzGetChallenge(f *testing.F) {
+	f.Add(uint64(0))
+	f.Add(uint64(1000000001337))
+
+	f.Fuzz(func(t *testing.T, uid uint64) {
+		b := r2encodeGetChallenge(uid)
+
+		be, err := r2cryptoEncrypt(b)
+		if err != nil {
+			t.Errorf("failed to encrypt packet: %v", err)
+		}
+
+		bd, err := r2cryptoDecrypt(be)
+		if err != nil {
+			t.Errorf("failed to decrypt packet: %v", err)
+		}
+
+		if !bytes.Equal(bd, b) {
+			t.Error("incorrect decryption result")
+		}
+	})
+}
+
+func FuzzChallenge(f *testing.F) {
+	f.Add(mustDecodeHex("aa"))
+	f.Add(mustDecodeHex("aaaaaaaaaaaaaa"))
+	f.Add(mustDecodeHex("00000000000000"))
+	f.Add(mustDecodeHex("09f7b6c1f41d91ecb41f370e9fd085610e5ee98827ba7aa9789557e18ddb2a28587f635a008aa71b9cb7b3f38b3ccd8d1ff0"))
+	f.Add(mustDecodeHex("edf3552e5d364fb3ab5505822c45c107208251b836022ad94698d920cfec348c469a861d14b5af2d8ca12702d09a7d91796e"))
+	f.Add(mustDecodeHex("f4ca55b7f53a2f9c19b563010d6964869648a23be1db9edce9f55ee3f9a02451be86ba56447740d1d893c34f3a854f6efbd47605ebf3211e05"))
+	f.Add(mustDecodeHex("bb8aaeed936b6dea21ba8bf4db5ca22a823a122307d5c6bc4124994581eb07b7996575acbbafe28ea4aee8bb58c681e33528470900007b012a"))
+
+	f.Fuzz(func(_ *testing.T, pkt []byte) {
+		// ensure this doesn't panic
+		r2cryptoDecrypt(pkt)
+	})
+}
+
+func mustDecodeHex(s string) []byte {
+	b, err := hex.DecodeString(s)
+	if err != nil {
+		panic(fmt.Errorf("decode %q: %w", s, err))
+	}
+	return b
+}