From 418281c1fffe83dbb2dfc73865444485863c66e6 Mon Sep 17 00:00:00 2001
From: Alexander PapaTutuWawa <papatutuwawa@polynom.me>
Date: Fri, 17 Apr 2020 22:46:05 +0200
Subject: [PATCH] Implement disco request for avatars

---
 bridge/xmpp/helpers.go | 21 +++++++++++++++++++++
 bridge/xmpp/xmpp.go    | 38 ++++++++++++++++++++++++++++++--------
 2 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/bridge/xmpp/helpers.go b/bridge/xmpp/helpers.go
index eb6a5366..28f0e65a 100644
--- a/bridge/xmpp/helpers.go
+++ b/bridge/xmpp/helpers.go
@@ -4,6 +4,7 @@ import (
 	"regexp"
 
 	"github.com/42wim/matterbridge/bridge/config"
+	"github.com/matterbridge/go-xmpp"
 )
 
 var pathRegex = regexp.MustCompile("[^a-zA-Z0-9]+")
@@ -28,3 +29,23 @@ func (b *Bxmpp) cacheAvatar(msg *config.Message) string {
 	}
 	return ""
 }
+
+func discoSupportsAvatar(items []*xmpp.DiscoItems) bool {
+	for _, item := range items {
+		if item.Node == xmpp.XMPPNS_AVATAR_PEP_DATA {
+			return true
+		}
+	}
+
+	return false
+}
+
+func (b *Bxmpp) handleDisco(items xmpp.DiscoItems) {
+	if discoSupportsAvatar(items) {
+		b.Log.Debugf("%s supports avatars", items.Jid)
+		b.avatarAvailability[items.Jid] = avatarAvailable
+	} else {
+		b.Log.Debugf("%s does not support avatars", items.Jid)
+		b.avatarAvailability[items.Jid] = avatarUnavailable
+	}
+}
diff --git a/bridge/xmpp/xmpp.go b/bridge/xmpp/xmpp.go
index 7122534f..6f9fea34 100644
--- a/bridge/xmpp/xmpp.go
+++ b/bridge/xmpp/xmpp.go
@@ -15,6 +15,14 @@ import (
 	"github.com/rs/xid"
 )
 
+type avatarAvailability int
+
+const (
+	avatarDiscoSent avatarAvailability = iota
+	avatarAvailable
+	avatarUnavailable
+)
+
 type Bxmpp struct {
 	*bridge.Config
 
@@ -24,14 +32,16 @@ type Bxmpp struct {
 	connected bool
 	sync.RWMutex
 
-	avatarMap map[string]string
+	avatarMap          map[string]string
+	avatarAvailability map[string]avatarAvailability
 }
 
 func New(cfg *bridge.Config) bridge.Bridger {
 	return &Bxmpp{
-		Config:    cfg,
-		xmppMap:   make(map[string]string),
-		avatarMap: make(map[string]string),
+		Config:             cfg,
+		xmppMap:            make(map[string]string),
+		avatarMap:          make(map[string]string),
+		avatarAvailability: make(map[string]avatarAvailability),
 	}
 }
 
@@ -237,10 +247,20 @@ func (b *Bxmpp) handleXMPP() error {
 					event = config.EventTopicChange
 				}
 
-				avatar := getAvatar(b.avatarMap, v.Remote, b.General)
-				if avatar == "" {
-					b.Log.Debugf("Requesting avatar data")
-					b.xc.AvatarRequestData(v.Remote)
+				// First check whether the user even supports
+				// avatars
+				avatar := ""
+				state, sok := b.avatarAvailability[v.Remote]
+				if !sok {
+					b.Log.Debugf("Sending disco to %s", v.Remote)
+					b.xc.DiscoverEntityItems(v.Remote)
+					b.avatarAvailability[v.Remote] = avatarDiscoSent
+				} else if state == avatarAvailable {
+					avatar = getAvatar(b.avatarMap, v.Remote, b.General)
+					if avatar == "" {
+						b.Log.Debugf("Requesting avatar data")
+						b.xc.AvatarRequestData(v.Remote)
+					}
 				}
 
 				msgID := v.ID
@@ -271,6 +291,8 @@ func (b *Bxmpp) handleXMPP() error {
 			}
 		case xmpp.AvatarData:
 			b.handleDownloadAvatar(v)
+		case xmpp.DiscoItems:
+			b.handleDisco(v)
 		case xmpp.Presence:
 			// Do nothing.
 		}