1 files changed, 142 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
+}