pkg/a2s: Implement Atlas sigreq1 connectionless packet

2 files changed, 56 insertions, 0 deletions

diff --git a/pkg/a2s/a2s.go b/pkg/a2s/a2s.go
index 04648c8..11cfc7c 100644
--- a/pkg/a2s/a2s.go
+++ b/pkg/a2s/a2s.go
@@ -6,8 +6,11 @@ import (
 	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
+	"crypto/hmac"
 	"crypto/rand"
+	"crypto/sha256"
 	"encoding/binary"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"net"
@@ -67,6 +70,37 @@ func Probe(addr netip.AddrPort, timeout time.Duration) error {
 	return nil
 }
 
+func AtlasSigreq1(addr netip.AddrPort, timeout time.Duration, key string, obj any) error {
+	// TODO: unconnected udp for performance
+
+	a := net.UDPAddrFromAddrPort(addr)
+	conn, err := net.DialUDP("udp", nil, a)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+
+	t := time.Now().Add(timeout)
+	conn.SetWriteDeadline(t)
+
+	j, err := json.Marshal(obj)
+	if err != nil {
+		return err
+	}
+
+	pkt, err := r2cryptoEncrypt(r2encodeAtlasSigreq1([]byte(key), j))
+	if err != nil {
+		return fmt.Errorf("encrypt packet: %w", err)
+	}
+	if _, err := conn.Write(pkt); err != nil {
+		if errors.Is(err, os.ErrDeadlineExceeded) {
+			err = fmt.Errorf("%w: %v", ErrTimeout, err)
+		}
+		return fmt.Errorf("send packet: %w", err)
+	}
+	return nil
+}
+
 const (
 	r2cryptoNonceSize = 12
 	r2cryptoTagSize   = 16
@@ -134,6 +168,19 @@ func r2encodeGetChallenge(uid uint64) []byte {
 	return b.Bytes()
 }
 
+func r2encodeAtlasSigreq1(key, data []byte) []byte {
+	sig := hmac.New(sha256.New, key)
+	sig.Write(data)
+
+	var b []byte
+	b = append(b, '\xFF', '\xFF', '\xFF', '\xFF') // connectionless: int32(-1)
+	b = append(b, 'T')                            // packet kind
+	b = append(b, "sigreq1\x00"...)               // packet type
+	b = sig.Sum(b)                                // sigreq1 - signature
+	b = append(b, data...)                        // sigreq1 - data
+	return b
+}
+
 func r2decodeChallenge(b []byte) (uint64, int32, error) {
 	var pkt struct {
 		Seq       int32
diff --git a/pkg/a2s/a2s_test.go b/pkg/a2s/a2s_test.go
index 72eba3a..b97b5c4 100644
--- a/pkg/a2s/a2s_test.go
+++ b/pkg/a2s/a2s_test.go
@@ -51,6 +51,15 @@ func TestDecodeChallenge(t *testing.T) {
 		t.Errorf("incorrect challenge")
 	}
 }
+
+func TestAtlasSigreq1(t *testing.T) {
+	b := mustDecodeHex("ffffffff547369677265713100803dab964c7c71851c05de40f5bf4cf72743951c96f2f0b81139ca780203260674657374")
+	a := r2encodeAtlasSigreq1([]byte("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"), []byte("test"))
+	if !bytes.Equal(a, b) {
+		t.Errorf("incorrect encoding: expected %x, got %x", b, a)
+	}
+}
+
 func FuzzGetChallenge(f *testing.F) {
 	f.Add(uint64(0))
 	f.Add(uint64(1000000001337))