From 85485bb1473337c63e15353e3ddd6a0f3effc9d0 Mon Sep 17 00:00:00 2001
From: Bohdan Horbeshko <bodqhrohro@gmail.com>
Date: Sun, 12 Jan 2025 22:05:24 -0500
Subject: [PATCH] Retrieve XMPP client features

---
 telegram/client.go | 32 ++++++++++++++-----------
 xmpp/handlers.go   | 59 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 75 insertions(+), 16 deletions(-)

diff --git a/telegram/client.go b/telegram/client.go
index 79f27d5..846e5a3 100644
--- a/telegram/client.go
+++ b/telegram/client.go
@@ -48,6 +48,9 @@ type Client struct {
 	lastMsgIds    map[int64]string
 	msgHashSeed   maphash.Seed
 
+	XmppClientFeatures     map[string]*[]string
+	XmppClientFeaturesLock sync.Mutex
+
 	locks           clientLocks
 	SendMessageLock sync.Mutex
 }
@@ -109,20 +112,21 @@ func NewClient(conf config.TelegramConfig, jid string, component *xmpp.Component
 	}
 
 	return &Client{
-		parameters:      &parameters,
-		xmpp:            component,
-		jid:             jid,
-		Session:         session,
-		resources:       make(map[string]bool),
-		content:         &conf.Content,
-		cache:           cache.NewCache(),
-		outbox:          make(map[string]string),
-		editOutbox:      make(map[string]string),
-		options:         options,
-		DelayedStatuses: make(map[int64]*DelayedStatus),
-		lastMsgHashes:   make(map[int64]uint64),
-		lastMsgIds:      make(map[int64]string),
-		msgHashSeed:     maphash.MakeSeed(),
+		parameters:         &parameters,
+		xmpp:               component,
+		jid:                jid,
+		Session:            session,
+		resources:          make(map[string]bool),
+		content:            &conf.Content,
+		cache:              cache.NewCache(),
+		outbox:             make(map[string]string),
+		editOutbox:         make(map[string]string),
+		options:            options,
+		DelayedStatuses:    make(map[int64]*DelayedStatus),
+		lastMsgHashes:      make(map[int64]uint64),
+		lastMsgIds:         make(map[int64]string),
+		msgHashSeed:        maphash.MakeSeed(),
+		XmppClientFeatures: make(map[string]*[]string),
 		locks: clientLocks{
 			chatMessageLocks: make(map[int64]*sync.Mutex),
 		},
diff --git a/xmpp/handlers.go b/xmpp/handlers.go
index 811cef6..d5ce6f0 100644
--- a/xmpp/handlers.go
+++ b/xmpp/handlers.go
@@ -15,6 +15,7 @@ import (
 	"dev.narayana.im/narayana/telegabber/xmpp/extensions"
 	"dev.narayana.im/narayana/telegabber/xmpp/gateway"
 
+	"github.com/google/uuid"
 	log "github.com/sirupsen/logrus"
 	"github.com/soheilhy/args"
 	"gosrc.io/xmpp"
@@ -40,7 +41,7 @@ func HandleIq(s xmpp.Sender, p stanza.Packet) {
 	}
 
 	log.Debugf("%#v", iq)
-	if iq.Type == "get" {
+	if iq.Type == stanza.IQTypeGet {
 		_, ok := iq.Payload.(*extensions.IqVcardTemp)
 		if ok {
 			go handleGetVcardIq(s, iq, TypeVCardTemp)
@@ -68,12 +69,18 @@ func HandleIq(s xmpp.Sender, p stanza.Packet) {
 			go handleGetQueryRegister(s, iq)
 			return
 		}
-	} else if iq.Type == "set" {
+	} else if iq.Type == stanza.IQTypeSet {
 		query, ok := iq.Payload.(*extensions.QueryRegister)
 		if ok {
 			go handleSetQueryRegister(s, iq, query)
 			return
 		}
+	} else if iq.Type == stanza.IQTypeResult {
+		discoInfo, ok := iq.Payload.(*stanza.DiscoInfo)
+		if ok {
+			go handleClientFeatures(iq, discoInfo)
+			return
+		}
 	}
 }
 
@@ -412,6 +419,7 @@ func handlePresence(s xmpp.Sender, p stanza.Presence) {
 						newArgs...,
 					)
 				}
+				probeClientFeatures(p.From, component)
 				session.UpdateChatNicknames()
 			}
 		}()
@@ -687,6 +695,53 @@ func iqAnswerSetError(answer *stanza.IQ, payload *extensions.QueryRegister, code
 	}
 }
 
+func probeClientFeatures(jid string, component *xmpp.Component) {
+	id, err := uuid.NewRandom()
+	if err != nil {
+		log.Error("Could not generate ID for a client features probe")
+		return
+	}
+
+	probe := stanza.IQ{
+		Attrs: stanza.Attrs{
+			From: gateway.Jid.Bare(),
+			To:   jid,
+			Id:   id.String(),
+			Type: stanza.IQTypeGet,
+		},
+		Payload: &stanza.DiscoInfo{},
+	}
+	log.Debugf("%#v", probe)
+
+	gateway.ResumableSend(component, &probe)
+}
+
+func handleClientFeatures(iq *stanza.IQ, discoInfo *stanza.DiscoInfo) {
+	fromJid, err := stanza.NewJid(iq.From)
+	if err != nil {
+		log.Error("Invalid from JID!")
+		return
+	}
+	bareFrom := fromJid.Bare()
+
+	session, ok := sessions[bareFrom]
+	if !ok {
+		log.Errorf("Got client features for unknown JID %v", bareFrom)
+		return
+	}
+
+	var features []string
+	for _, feature := range discoInfo.Features {
+		features = append(features, feature.Var)
+	}
+
+	session.XmppClientFeaturesLock.Lock()
+	session.XmppClientFeatures[fromJid.Resource] = &features
+	session.XmppClientFeaturesLock.Unlock()
+
+	log.Debugf("Features for %v: %#v", iq.From, features)
+}
+
 func toToID(to string) (int64, bool) {
 	toParts := strings.Split(to, "@")
 	if len(toParts) < 2 {