Add support for linking by entering code on phone (tulir#429)

Co-authored-by: burstfreeze <[email protected]>

client.go

@@ -130,6 +130,8 @@ type Client struct {
 	// Should SubscribePresence return an error if no privacy token is stored for the user?
 	ErrorOnSubscribePresenceWithoutToken bool
 
+	phoneLinkingCache *phoneLinkingCache
+
 	uniqueID  string
 	idCounter uint32
 

mdtest/main.go

@@ -178,6 +178,16 @@ func parseJID(arg string) (types.JID, bool) {
 
 func handleCmd(cmd string, args []string) {
 	switch cmd {
+	case "pair-phone":
+		if len(args) < 1 {
+			log.Errorf("Usage: pair-phone <number>")
+			return
+		}
+		linkingCode, err := cli.PairPhone(args[0], true, whatsmeow.PairClientUnknown, "whatsmeow")
+		if err != nil {
+			panic(err)
+		}
+		fmt.Println("Linking code:", linkingCode)
 	case "reconnect":
 		cli.Disconnect()
 		err := cli.Connect()

notification.go

@@ -259,6 +259,8 @@ func (cli *Client) handleNotification(node *waBinary.Node) {
 		go cli.handleMediaRetryNotification(node)
 	case "privacy_token":
 		go cli.handlePrivacyTokenNotification(node)
+	case "link_code_companion_reg":
+		go cli.tryHandleCodePairNotification(node)
 	// Other types: business, disappearing_mode, server, status, pay, psa
 	default:
 		cli.Log.Debugf("Unhandled notification with type %s", notifType)

pair-code.go

@@ -0,0 +1,224 @@
+// Copyright (c) 2023 Tulir Asokan
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package whatsmeow
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"crypto/sha256"
+	"encoding/base32"
+	"fmt"
+	"regexp"
+	"strconv"
+
+	"golang.org/x/crypto/curve25519"
+	"golang.org/x/crypto/pbkdf2"
+
+	waBinary "go.mau.fi/whatsmeow/binary"
+	"go.mau.fi/whatsmeow/types"
+	"go.mau.fi/whatsmeow/util/hkdfutil"
+	"go.mau.fi/whatsmeow/util/keys"
+)
+
+type PairClientType int
+
+const (
+	PairClientUnknown PairClientType = iota
+	PairClientChrome
+	PairClientEdge
+	PairClientFirefox
+	PairClientIE
+	PairClientOpera
+	PairClientSafari
+	PairClientElectron
+	PairClientUWP
+	PairClientOtherWebClient
+)
+
+var notNumbers = regexp.MustCompile("[^0-9]")
+var linkingBase32 = base32.NewEncoding("123456789ABCDEFGHJKLMNPQRSTVWXYZ")
+
+type phoneLinkingCache struct {
+	jid         types.JID
+	keyPair     *keys.KeyPair
+	linkingCode string
+	pairingRef  string
+}
+
+func randomBytes(length int) []byte {
+	random := make([]byte, length)
+	_, err := rand.Read(random)
+	if err != nil {
+		panic(fmt.Errorf("failed to get random bytes: %w", err))
+	}
+	return random
+}
+
+func generateCompanionEphemeralKey() (ephemeralKeyPair *keys.KeyPair, ephemeralKey []byte, encodedLinkingCode string) {
+	ephemeralKeyPair = keys.NewKeyPair()
+	salt := randomBytes(32)
+	iv := randomBytes(16)
+	linkingCode := randomBytes(5)
+	encodedLinkingCode = linkingBase32.EncodeToString(linkingCode)
+	linkCodeKey := pbkdf2.Key([]byte(encodedLinkingCode), salt, 2<<16, 32, sha256.New)
+	linkCipherBlock, _ := aes.NewCipher(linkCodeKey)
+	encryptedPubkey := ephemeralKeyPair.Pub[:]
+	cipher.NewCTR(linkCipherBlock, iv).XORKeyStream(encryptedPubkey, encryptedPubkey)
+	ephemeralKey = make([]byte, 80)
+	copy(ephemeralKey[0:32], salt)
+	copy(ephemeralKey[32:48], iv)
+	copy(ephemeralKey[48:80], encryptedPubkey)
+	return
+}
+
+func (cli *Client) PairPhone(phone string, showPushNotification bool, clientType PairClientType, clientDisplayName string) (string, error) {
+	ephemeralKeyPair, ephemeralKey, encodedLinkingCode := generateCompanionEphemeralKey()
+	phone = notNumbers.ReplaceAllString(phone, "")
+	jid := types.NewJID(phone, types.DefaultUserServer)
+	resp, err := cli.sendIQ(infoQuery{
+		Namespace: "md",
+		Type:      iqSet,
+		To:        types.ServerJID,
+		Content: []waBinary.Node{{
+			Tag: "link_code_companion_reg",
+			Attrs: waBinary.Attrs{
+				"jid":   jid,
+				"stage": "companion_hello",
+
+				"should_show_push_notification": strconv.FormatBool(showPushNotification),
+			},
+			Content: []waBinary.Node{
+				{Tag: "link_code_pairing_wrapped_companion_ephemeral_pub", Content: ephemeralKey},
+				{Tag: "companion_server_auth_key_pub", Content: cli.Store.NoiseKey.Pub[:]},
+				{Tag: "companion_platform_id", Content: strconv.Itoa(int(clientType))},
+				{Tag: "companion_platform_display", Content: clientDisplayName},
+				{Tag: "link_code_pairing_nonce", Content: []byte{0}},
+			},
+		}},
+	})
+	if err != nil {
+		return "", err
+	}
+	pairingRefNode, ok := resp.GetOptionalChildByTag("link_code_companion_reg", "link_code_pairing_ref")
+	if !ok {
+		return "", &ElementMissingError{Tag: "link_code_pairing_ref", In: "code link registration response"}
+	}
+	pairingRef, ok := pairingRefNode.Content.([]byte)
+	if !ok {
+		return "", fmt.Errorf("unexpected type %T in content of link_code_pairing_ref tag", pairingRefNode.Content)
+	}
+	cli.phoneLinkingCache = &phoneLinkingCache{
+		jid:         jid,
+		keyPair:     ephemeralKeyPair,
+		linkingCode: encodedLinkingCode,
+		pairingRef:  string(pairingRef),
+	}
+	return encodedLinkingCode[0:4] + "-" + encodedLinkingCode[4:], nil
+}
+
+func (cli *Client) tryHandleCodePairNotification(parentNode *waBinary.Node) {
+	err := cli.handleCodePairNotification(parentNode)
+	if err != nil {
+		cli.Log.Errorf("Failed to handle code pair notification: %s", err)
+	}
+}
+
+func (cli *Client) handleCodePairNotification(parentNode *waBinary.Node) error {
+	node, ok := parentNode.GetOptionalChildByTag("link_code_companion_reg")
+	if !ok {
+		return &ElementMissingError{
+			Tag: "link_code_companion_reg",
+			In:  "notification",
+		}
+	}
+	linkCache := cli.phoneLinkingCache
+	if linkCache == nil {
+		return fmt.Errorf("received code pair notification without a pending pairing")
+	}
+	linkCodePairingRef, _ := node.GetChildByTag("link_code_pairing_ref").Content.([]byte)
+	if string(linkCodePairingRef) != linkCache.pairingRef {
+		return fmt.Errorf("pairing ref mismatch in code pair notification")
+	}
+	wrappedPrimaryEphemeralPub, ok := node.GetChildByTag("link_code_pairing_wrapped_primary_ephemeral_pub").Content.([]byte)
+	if !ok {
+		return &ElementMissingError{
+			Tag: "link_code_pairing_wrapped_primary_ephemeral_pub",
+			In:  "notification",
+		}
+	}
+	primaryIdentityPub, ok := node.GetChildByTag("primary_identity_pub").Content.([]byte)
+	if !ok {
+		return &ElementMissingError{
+			Tag: "primary_identity_pub",
+			In:  "notification",
+		}
+	}
+
+	advSecretRandom := randomBytes(32)
+	keyBundleSalt := randomBytes(32)
+	keyBundleNonce := randomBytes(12)
+
+	// Decrypt the primary device's ephemeral public key, which was encrypted with the 8-character pairing code,
+	// then compute the DH shared secret using our ephemeral private key we generated earlier.
+	primarySalt := wrappedPrimaryEphemeralPub[0:32]
+	primaryIV := wrappedPrimaryEphemeralPub[32:48]
+	primaryEncryptedPubkey := wrappedPrimaryEphemeralPub[48:80]
+	linkCodeKey := pbkdf2.Key([]byte(linkCache.linkingCode), primarySalt, 2<<16, 32, sha256.New)
+	linkCipherBlock, err := aes.NewCipher(linkCodeKey)
+	if err != nil {
+		return fmt.Errorf("failed to create link cipher: %w", err)
+	}
+	primaryDecryptedPubkey := make([]byte, 32)
+	cipher.NewCTR(linkCipherBlock, primaryIV).XORKeyStream(primaryDecryptedPubkey, primaryEncryptedPubkey)
+	ephemeralSharedSecret, err := curve25519.X25519(linkCache.keyPair.Priv[:], primaryDecryptedPubkey)
+	if err != nil {
+		return fmt.Errorf("failed to compute ephemeral shared secret: %w", err)
+	}
+
+	// Encrypt and wrap key bundle containing our identity key, the primary device's identity key and the randomness used for the adv key.
+	keyBundleEncryptionKey := hkdfutil.SHA256(ephemeralSharedSecret, keyBundleSalt, []byte("link_code_pairing_key_bundle_encryption_key"), 32)
+	keyBundleCipherBlock, err := aes.NewCipher(keyBundleEncryptionKey)
+	if err != nil {
+		return fmt.Errorf("failed to create key bundle cipher: %w", err)
+	}
+	keyBundleGCM, err := cipher.NewGCM(keyBundleCipherBlock)
+	if err != nil {
+		return fmt.Errorf("failed to create key bundle GCM: %w", err)
+	}
+	plaintextKeyBundle := concatBytes(cli.Store.IdentityKey.Pub[:], primaryIdentityPub, advSecretRandom)
+	encryptedKeyBundle := keyBundleGCM.Seal(nil, keyBundleNonce, plaintextKeyBundle, nil)
+	wrappedKeyBundle := concatBytes(keyBundleSalt, keyBundleNonce, encryptedKeyBundle)
+
+	// Compute the adv secret key (which is used to authenticate the pair-success event later)
+	identitySharedKey, err := curve25519.X25519(cli.Store.IdentityKey.Priv[:], primaryIdentityPub)
+	if err != nil {
+		return fmt.Errorf("failed to compute identity shared key: %w", err)
+	}
+	advSecretInput := append(append(ephemeralSharedSecret, identitySharedKey...), advSecretRandom...)
+	advSecret := hkdfutil.SHA256(advSecretInput, nil, []byte("adv_secret"), 32)
+	cli.Store.AdvSecretKey = advSecret
+
+	_, err = cli.sendIQ(infoQuery{
+		Namespace: "md",
+		Type:      iqSet,
+		To:        types.ServerJID,
+		Content: []waBinary.Node{{
+			Tag: "link_code_companion_reg",
+			Attrs: waBinary.Attrs{
+				"jid":   linkCache.jid,
+				"stage": "companion_finish",
+			},
+			Content: []waBinary.Node{
+				{Tag: "link_code_pairing_wrapped_key_bundle", Content: wrappedKeyBundle},
+				{Tag: "companion_identity_public", Content: cli.Store.IdentityKey.Pub[:]},
+				{Tag: "link_code_pairing_ref", Content: linkCodePairingRef},
+			},
+		}},
+	})
+	return err
+}