Dashboard: Mix initials and custom gravatars in recent viewers list (#47212)

* Use Wiring to initialize Avatar Cache Server

Create AvatarCacheServer Provider function and pass it in as an
argument to HTTPServer. Also convert CacheServer to a singleton
so that we keep all cached Avatar info in one place for easier access

* Refactor avatar cache server and add 'isCustom' check

Avatar cache server needs to perform two similar fetches
back-to-back; break up functions to allow for easy reuse.
Then add handling to see if a user has a custom avatar.

* Add additional accessors so that /recents api can easily use the cache

* Minor mods to avatar server to facilitiate unit testing

* add unit tests for avatar fetching

* add error handling in case we somehow fetch gravatars while they are disabled

* linting: read error return value in unit test

* Use http package status codes

Co-authored-by: Ezequiel Victorero <evictorero@gmail.com>

* Use http package status codes

Co-authored-by: Ezequiel Victorero <evictorero@gmail.com>

* Use http package status codes

Co-authored-by: Ezequiel Victorero <evictorero@gmail.com>

* Incorporate suggestions from PR
-avoid mutating arguments
-change error handler function to private and make name more descriptive

Co-authored-by: Ezequiel Victorero <evictorero@gmail.com>

pkg/api/api.go

@@ -4,7 +4,6 @@ package api
 import (
 	"time"
 
-	"github.com/grafana/grafana/pkg/api/avatar"
 	"github.com/grafana/grafana/pkg/api/frontendlogging"
 	"github.com/grafana/grafana/pkg/api/routing"
 	"github.com/grafana/grafana/pkg/infra/log"
@@ -542,8 +541,7 @@ func (hs *HTTPServer) registerRoutes() {
 	r.Any("/api/gnet/*", reqSignedIn, hs.ProxyGnetRequest)
 
 	// Gravatar service.
-	avatarCacheServer := avatar.NewCacheServer(hs.Cfg)
-	r.Get("/avatar/:hash", avatarCacheServer.Handler)
+	r.Get("/avatar/:hash", hs.AvatarCacheServer.Handler)
 
 	// Snapshots
 	r.Post("/api/snapshots/", reqSnapshotPublicModeOrSignedIn, hs.CreateDashboardSnapshot)

pkg/api/avatar/avatar.go

@@ -34,22 +34,31 @@ const (
 // Avatar represents the avatar object.
 type Avatar struct {
 	hash      string
-	reqParams string
 	data      *bytes.Buffer
 	notFound  bool
+	isCustom  bool
 	timestamp time.Time
 }
 
-var alog = log.New("avatar")
+var (
+	alog = log.New("avatar")
+	// Represents a singleton AvatarCacheServer instance
+	csi *AvatarCacheServer
+	// Paremeters needed to fetch Gravatar with a retro fallback
+	gravatarReqParams = url.Values{
+		"d":    {"retro"},
+		"size": {"200"},
+		"r":    {"pg"},
+	}.Encode()
+	// Parameters needed to see if a Gravatar is custom
+	hasCustomReqParams = url.Values{
+		"d": {"404"},
+	}.Encode()
+	cacheInitOnce sync.Once
+)
 
 func New(hash string) *Avatar {
-	return &Avatar{
-		hash: hash,
-		reqParams: url.Values{
-			"d":    {"retro"},
-			"size": {"200"},
-			"r":    {"pg"}}.Encode(),
-	}
+	return &Avatar{hash: hash}
 }
 
 func (a *Avatar) Expired() bool {
@@ -61,16 +70,27 @@ func (a *Avatar) Encode(wr io.Writer) error {
 	return err
 }
 
-func (a *Avatar) Update() (err error) {
+func (a *Avatar) update(baseUrl string) (err error) {
+	customUrl := baseUrl + a.hash + "?"
 	select {
 	case <-time.After(time.Second * 3):
 		err = fmt.Errorf("get gravatar image %s timeout", a.hash)
-	case err = <-thunder.GoFetch(gravatarSource+a.hash+"?"+a.reqParams, a):
+	case err = <-thunder.GoFetch(customUrl, a):
 	}
 	return err
 }
 
-type CacheServer struct {
+func (a *Avatar) GetIsCustom() bool {
+	return a.isCustom
+}
+
+// Quick error handler to avoid multiple copy pastes
+func (a *Avatar) setAvatarNotFound() {
+	a.notFound = true
+	a.isCustom = false
+}
+
+type AvatarCacheServer struct {
 	cfg      *setting.Cfg
 	notFound *Avatar
 	cache    *gocache.Cache
@@ -78,7 +98,7 @@ type CacheServer struct {
 
 var validMD5 = regexp.MustCompile("^[a-fA-F0-9]{32}$")
 
-func (a *CacheServer) Handler(ctx *models.ReqContext) {
+func (a *AvatarCacheServer) Handler(ctx *models.ReqContext) {
 	hash := web.Params(ctx.Req)[":hash"]
 
 	if len(hash) != 32 || !validMD5.MatchString(hash) {
@@ -86,29 +106,7 @@ func (a *CacheServer) Handler(ctx *models.ReqContext) {
 		return
 	}
 
-	var avatar *Avatar
-	obj, exists := a.cache.Get(hash)
-	if exists {
-		avatar = obj.(*Avatar)
-	} else {
-		avatar = New(hash)
-	}
-
-	if avatar.Expired() {
-		// The cache item is either expired or newly created, update it from the server
-		if err := avatar.Update(); err != nil {
-			ctx.Logger.Debug("avatar update", "err", err)
-			avatar = a.notFound
-		}
-	}
-
-	if avatar.notFound {
-		avatar = a.notFound
-	} else if !exists {
-		if err := a.cache.Add(hash, avatar, gocache.DefaultExpiration); err != nil {
-			ctx.Logger.Debug("add avatar to cache", "err", err)
-		}
-	}
+	avatar := a.GetAvatarForHash(hash)
 
 	ctx.Resp.Header().Set("Content-Type", "image/jpeg")
 
@@ -120,12 +118,56 @@ func (a *CacheServer) Handler(ctx *models.ReqContext) {
 
 	if err := avatar.Encode(ctx.Resp); err != nil {
 		ctx.Logger.Warn("avatar encode error:", "err", err)
-		ctx.Resp.WriteHeader(500)
+		ctx.Resp.WriteHeader(http.StatusInternalServerError)
 	}
 }
 
-func NewCacheServer(cfg *setting.Cfg) *CacheServer {
-	return &CacheServer{
+func (a *AvatarCacheServer) GetAvatarForHash(hash string) *Avatar {
+	if setting.DisableGravatar {
+		alog.Warn("'GetGravatarForHash' called despite gravatars being disabled; returning default profile image")
+		return a.notFound
+	}
+	return a.getAvatarForHash(hash, gravatarSource)
+}
+
+func (a *AvatarCacheServer) getAvatarForHash(hash string, baseUrl string) *Avatar {
+	var avatar *Avatar
+	obj, exists := a.cache.Get(hash)
+	if exists {
+		avatar = obj.(*Avatar)
+	} else {
+		avatar = New(hash)
+	}
+
+	if avatar.Expired() {
+		// The cache item is either expired or newly created, update it from the server
+		if err := avatar.update(baseUrl); err != nil {
+			alog.Debug("avatar update", "err", err)
+			avatar = a.notFound
+		}
+	}
+
+	if avatar.notFound {
+		avatar = a.notFound
+	} else if !exists {
+		if err := a.cache.Add(hash, avatar, gocache.DefaultExpiration); err != nil {
+			alog.Debug("add avatar to cache", "err", err)
+		}
+	}
+	return avatar
+}
+
+// Access cache server singleton instance
+func ProvideAvatarCacheServer(cfg *setting.Cfg) *AvatarCacheServer {
+	cacheInitOnce.Do(func() {
+		csi = newCacheServer(cfg)
+	})
+
+	return csi
+}
+
+func newCacheServer(cfg *setting.Cfg) *AvatarCacheServer {
+	return &AvatarCacheServer{
 		cfg:      cfg,
 		notFound: newNotFound(cfg),
 		cache:    gocache.New(time.Hour, time.Hour*2),
@@ -133,7 +175,10 @@ func NewCacheServer(cfg *setting.Cfg) *CacheServer {
 }
 
 func newNotFound(cfg *setting.Cfg) *Avatar {
-	avatar := &Avatar{notFound: true}
+	avatar := &Avatar{
+		notFound: true,
+		isCustom: false,
+	}
 
 	// load user_profile png into buffer
 	// It's safe to ignore gosec warning G304 since the variable part of the file path comes from a configuration
@@ -176,11 +221,11 @@ func (t *Thunder) init() {
 	}
 }
 
-func (t *Thunder) Fetch(url string, avatar *Avatar) error {
+func (t *Thunder) Fetch(baseUrl string, avatar *Avatar) error {
 	t.once.Do(t.init)
 	task := &thunderTask{
-		Url:    url,
-		Avatar: avatar,
+		BaseUrl: baseUrl,
+		Avatar:  avatar,
 	}
 	task.Add(1)
 	t.q <- task
@@ -188,18 +233,18 @@ func (t *Thunder) Fetch(url string, avatar *Avatar) error {
 	return task.err
 }
 
-func (t *Thunder) GoFetch(url string, avatar *Avatar) chan error {
+func (t *Thunder) GoFetch(baseUrl string, avatar *Avatar) chan error {
 	c := make(chan error)
 	go func() {
-		c <- t.Fetch(url, avatar)
+		c <- t.Fetch(baseUrl, avatar)
 	}()
 	return c
 }
 
 // thunder download
 type thunderTask struct {
-	Url    string
-	Avatar *Avatar
+	BaseUrl string
+	Avatar  *Avatar
 	sync.WaitGroup
 	err error
 }
@@ -214,11 +259,48 @@ var client = &http.Client{
 	Transport: &http.Transport{Proxy: http.ProxyFromEnvironment},
 }
 
+// We fetch the same url with param tweaks twice in a row
+// Break out the fetch function in a way that makes each
+// Portion highly reusable
 func (a *thunderTask) fetch() error {
 	a.Avatar.timestamp = time.Now()
 
-	alog.Debug("avatar.fetch(fetch new avatar)", "url", a.Url)
-	req, err := http.NewRequest("GET", a.Url, nil)
+	alog.Debug("avatar.fetch(fetch new avatar)", "url", a.BaseUrl)
+	// First do the fetch to get the Gravatar with a retro icon fallback
+	err := performGet(a.BaseUrl+gravatarReqParams, a.Avatar, getGravatarHandler)
+
+	if err == nil {
+		// Next do a fetch with a 404 fallback to see if it's a custom gravatar
+		return performGet(a.BaseUrl+hasCustomReqParams, a.Avatar, checkIsCustomHandler)
+	}
+	return err
+}
+
+type ResponseHandler func(av *Avatar, resp *http.Response) error
+
+// Verifies the Gravatar response code was 200, then stores the image byte slice
+func getGravatarHandler(av *Avatar, resp *http.Response) error {
+	if resp.StatusCode != http.StatusOK {
+		av.setAvatarNotFound()
+		return fmt.Errorf("status code: %d", resp.StatusCode)
+	}
+
+	av.data = &bytes.Buffer{}
+	writer := bufio.NewWriter(av.data)
+
+	_, err := io.Copy(writer, resp.Body)
+	return err
+}
+
+// Uses the d=404 fallback to see if the gravatar we got back is custom
+func checkIsCustomHandler(av *Avatar, resp *http.Response) error {
+	av.isCustom = resp.StatusCode != http.StatusNotFound
+	return nil
+}
+
+// Reusable Get helper that allows us to pass in custom handling depending on the endpoint
+func performGet(url string, av *Avatar, handler ResponseHandler) error {
+	req, err := http.NewRequest("GET", url, nil)
 	if err != nil {
 		return err
 	}
@@ -227,9 +309,10 @@ func (a *thunderTask) fetch() error {
 	req.Header.Set("Accept-Language", "zh-CN,zh;q=0.8")
 	req.Header.Set("Cache-Control", "no-cache")
 	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/33.0.1750.154 Safari/537.36")
+	alog.Debug("Fetching avatar url with parameters", "url", url)
 	resp, err := client.Do(req)
 	if err != nil {
-		a.Avatar.notFound = true
+		av.setAvatarNotFound()
 		return fmt.Errorf("gravatar unreachable: %w", err)
 	}
 	defer func() {
@@ -238,14 +321,6 @@ func (a *thunderTask) fetch() error {
 		}
 	}()
 
-	if resp.StatusCode != 200 {
-		a.Avatar.notFound = true
-		return fmt.Errorf("status code: %d", resp.StatusCode)
-	}
-
-	a.Avatar.data = &bytes.Buffer{}
-	writer := bufio.NewWriter(a.Avatar.data)
-
-	_, err = io.Copy(writer, resp.Body)
+	err = handler(av, resp)
 	return err
 }

pkg/api/avatar/avatar_test.go

@@ -0,0 +1,120 @@
+package avatar
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/stretchr/testify/require"
+)
+
+const DEFAULT_NONSENSE_HASH string = "9e107d9d372bb6826bd81d3542a419d6"
+const CUSTOM_NONSENSE_HASH string = "d2a9116d4a63304733ca0f3471e57d16"
+
+var NONSENSE_BODY []byte = []byte("Bogus API response")
+
+func TestAvatar_AvatarRetrieval(t *testing.T) {
+	avc := ProvideAvatarCacheServer(setting.NewCfg())
+	callCounter := 0
+	mockServer := setupMockGravatarServer(&callCounter, false)
+
+	t.Cleanup(func() {
+		avc.cache.Flush()
+		mockServer.Close()
+	})
+
+	av := avc.getAvatarForHash(DEFAULT_NONSENSE_HASH, mockServer.URL+"/avatar/")
+	// verify there was a call to get the image and a call to the 404 fallback
+	require.Equal(t, callCounter, 2)
+	require.Equal(t, av.data.Bytes(), NONSENSE_BODY)
+
+	avc.getAvatarForHash(DEFAULT_NONSENSE_HASH, mockServer.URL+"/avatar/")
+	//since the avatar is cached, there should not have been anymore REST calls
+	require.Equal(t, callCounter, 2)
+}
+
+func TestAvatar_CheckCustom(t *testing.T) {
+	avc := ProvideAvatarCacheServer(setting.NewCfg())
+	callCounter := 0
+	mockServer := setupMockGravatarServer(&callCounter, false)
+
+	t.Cleanup(func() {
+		avc.cache.Flush()
+		mockServer.Close()
+	})
+
+	av := avc.getAvatarForHash(DEFAULT_NONSENSE_HASH, mockServer.URL+"/avatar/")
+	// verify this avatar is not marked custom
+	require.False(t, av.isCustom)
+
+	av2 := avc.getAvatarForHash(CUSTOM_NONSENSE_HASH, mockServer.URL+"/avatar/")
+	// verify this avatar is marked custom
+	require.True(t, av2.isCustom)
+}
+
+func TestAvatar_FallbackCase(t *testing.T) {
+	avc := ProvideAvatarCacheServer(setting.NewCfg())
+	callCounter := 0
+	mockServer := setupMockGravatarServer(&callCounter, true)
+
+	t.Cleanup(func() {
+		avc.cache.Flush()
+		mockServer.Close()
+	})
+
+	av := avc.getAvatarForHash(DEFAULT_NONSENSE_HASH, mockServer.URL+"/avatar/")
+	// the client should not have gotten a valid response back from the first call
+	// there should only be one REST call, and the avatar url should be the default
+	require.Equal(t, callCounter, 1)
+	require.False(t, av.isCustom)
+	require.True(t, av.notFound)
+	require.Equal(t, av, avc.notFound)
+}
+
+func TestAvatar_ExpirationHandler(t *testing.T) {
+	avc := ProvideAvatarCacheServer(setting.NewCfg())
+	callCounter := 0
+	mockServer := setupMockGravatarServer(&callCounter, false)
+
+	t.Cleanup(func() {
+		avc.cache.Flush()
+		mockServer.Close()
+	})
+
+	av := avc.getAvatarForHash(DEFAULT_NONSENSE_HASH, mockServer.URL+"/avatar/")
+	// verify there was a call to get the image and a call to the 404 fallback
+	require.Equal(t, callCounter, 2)
+	require.Equal(t, av.data.Bytes(), NONSENSE_BODY)
+
+	// manually expire the avatar in the cache
+	av.timestamp = av.timestamp.Add(-time.Minute * 15)
+	avc.getAvatarForHash(DEFAULT_NONSENSE_HASH, mockServer.URL+"/avatar/")
+	//since the avatar is expired, there should be two more REST calls
+	require.Equal(t, callCounter, 4)
+}
+
+func setupMockGravatarServer(counter *int, simulateError bool) *httptest.Server {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		(*counter)++
+		splitUri := strings.Split(r.RequestURI, "?")
+		urlHash := splitUri[0][len("/avatar/"):]
+		params := splitUri[1]
+		if params == "d=404" {
+			if urlHash == DEFAULT_NONSENSE_HASH {
+				w.WriteHeader(404)
+			} else {
+				_, _ = w.Write(NONSENSE_BODY)
+			}
+		} else {
+			if simulateError {
+				w.WriteHeader(500)
+			} else {
+				_, _ = w.Write(NONSENSE_BODY)
+			}
+		}
+	}))
+	return server
+}

pkg/api/dtos/models.go

@@ -79,11 +79,20 @@ func GetGravatarUrl(text string) string {
 		return ""
 	}
 
+	hash, _ := GetGravatarHash(text)
+	return fmt.Sprintf(setting.AppSubUrl+"/avatar/%x", hash)
+}
+
+func GetGravatarHash(text string) ([]byte, bool) {
+	if text == "" {
+		return make([]byte, 0), false
+	}
+
 	hasher := md5.New()
 	if _, err := hasher.Write([]byte(strings.ToLower(text))); err != nil {
 		mlog.Warn("Failed to hash text", "err", err)
 	}
-	return fmt.Sprintf(setting.AppSubUrl+"/avatar/%x", hasher.Sum(nil))
+	return hasher.Sum(nil), true
 }
 
 func GetGravatarUrlWithDefault(text string, defaultText string) string {

pkg/api/http_server.go

@@ -13,6 +13,7 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/grafana/grafana/pkg/api/avatar"
 	"github.com/grafana/grafana/pkg/api/routing"
 	httpstatic "github.com/grafana/grafana/pkg/api/static"
 	"github.com/grafana/grafana/pkg/bus"
@@ -147,6 +148,7 @@ type HTTPServer struct {
 	AlertNotificationService     *alerting.AlertNotificationService
 	DashboardsnapshotsService    *dashboardsnapshots.Service
 	PluginSettings               *pluginSettings.Service
+	AvatarCacheServer            *avatar.AvatarCacheServer
 }
 
 type ServerOptions struct {
@@ -178,6 +180,7 @@ func ProvideHTTPServer(opts ServerOptions, cfg *setting.Cfg, routeRegister routi
 	dashboardProvisioningService dashboards.DashboardProvisioningService, folderService dashboards.FolderService,
 	datasourcePermissionsService permissions.DatasourcePermissionsService, alertNotificationService *alerting.AlertNotificationService,
 	dashboardsnapshotsService *dashboardsnapshots.Service, commentsService *comments.Service, pluginSettings *pluginSettings.Service,
+	avatarCacheServer *avatar.AvatarCacheServer,
 ) (*HTTPServer, error) {
 	web.Env = cfg.Env
 	m := web.New()
@@ -250,6 +253,7 @@ func ProvideHTTPServer(opts ServerOptions, cfg *setting.Cfg, routeRegister routi
 		DashboardsnapshotsService:    dashboardsnapshotsService,
 		PluginSettings:               pluginSettings,
 		permissionServices:           permissionsServices,
+		AvatarCacheServer:            avatarCacheServer,
 	}
 	if hs.Listener != nil {
 		hs.log.Debug("Using provided listener")

pkg/server/wire.go

@@ -7,6 +7,7 @@ import (
 	"github.com/google/wire"
 	sdkhttpclient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
 	"github.com/grafana/grafana/pkg/api"
+	"github.com/grafana/grafana/pkg/api/avatar"
 	"github.com/grafana/grafana/pkg/api/routing"
 	"github.com/grafana/grafana/pkg/bus"
 	"github.com/grafana/grafana/pkg/expr"
@@ -232,6 +233,7 @@ var wireBasicSet = wire.NewSet(
 	wire.Bind(new(alerting.DashAlertExtractor), new(*alerting.DashAlertExtractorService)),
 	comments.ProvideService,
 	guardian.ProvideService,
+	avatar.ProvideAvatarCacheServer,
 	authproxy.ProvideAuthProxy,
 )