Rendering: Store render key in remote cache (#22031)

By storing render key in remote cache it will enable
image renderer to use public facing url or load
balancer url to render images and thereby remove
the requirement of image renderer having to use the
url of the originating Grafana instance when running
HA setup (multiple Grafana instances).

Fixes #17704
Ref grafana/grafana-image-renderer#91

pkg/api/common_test.go

@@ -143,7 +143,7 @@ func setupScenarioContext(url string) *scenarioContext {
 		Delims:    macaron.Delims{Left: "[[", Right: "]]"},
 	}))
 
-	sc.m.Use(middleware.GetContextHandler(nil, nil))
+	sc.m.Use(middleware.GetContextHandler(nil, nil, nil))
 
 	return sc
 }

pkg/api/http_server.go

@@ -316,6 +316,7 @@ func (hs *HTTPServer) addMiddlewaresAndStaticRoutes() {
 	m.Use(middleware.GetContextHandler(
 		hs.AuthTokenService,
 		hs.RemoteCacheService,
+		hs.RenderService,
 	))
 	m.Use(middleware.OrgRedirect())
 

pkg/middleware/middleware.go

@@ -15,6 +15,7 @@ import (
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/remotecache"
 	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/services/rendering"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/util"
 )
@@ -39,6 +40,7 @@ var (
 func GetContextHandler(
 	ats models.UserTokenService,
 	remoteCache *remotecache.RemoteCache,
+	renderService rendering.Service,
 ) macaron.Handler {
 	return func(c *macaron.Context) {
 		ctx := &models.ReqContext{
@@ -62,7 +64,7 @@ func GetContextHandler(
 		// then look for api key in session (special case for render calls via api)
 		// then test if anonymous access is enabled
 		switch {
-		case initContextWithRenderAuth(ctx):
+		case initContextWithRenderAuth(ctx, renderService):
 		case initContextWithApiKey(ctx):
 		case initContextWithBasicAuth(ctx, orgId):
 		case initContextWithAuthProxy(remoteCache, ctx, orgId):

pkg/middleware/middleware_test.go

@@ -557,7 +557,7 @@ func middlewareScenario(t *testing.T, desc string, fn scenarioFunc) {
 		sc.userAuthTokenService = auth.NewFakeUserAuthTokenService()
 		sc.remoteCacheService = remotecache.NewFakeStore(t)
 
-		sc.m.Use(GetContextHandler(sc.userAuthTokenService, sc.remoteCacheService))
+		sc.m.Use(GetContextHandler(sc.userAuthTokenService, sc.remoteCacheService, nil))
 
 		sc.m.Use(OrgRedirect())
 

pkg/middleware/recovery_test.go

@@ -68,7 +68,7 @@ func recoveryScenario(t *testing.T, desc string, url string, fn scenarioFunc) {
 		sc.userAuthTokenService = auth.NewFakeUserAuthTokenService()
 		sc.remoteCacheService = remotecache.NewFakeStore(t)
 
-		sc.m.Use(GetContextHandler(sc.userAuthTokenService, sc.remoteCacheService))
+		sc.m.Use(GetContextHandler(sc.userAuthTokenService, sc.remoteCacheService, nil))
 		// mock out gc goroutine
 		sc.m.Use(OrgRedirect())
 

pkg/middleware/render_auth.go

 package middleware
 
 import (
-	"sync"
 	"time"
 
+	"github.com/grafana/grafana/pkg/services/rendering"
+
 	m "github.com/grafana/grafana/pkg/models"
-	"github.com/grafana/grafana/pkg/util"
 )
 
-var renderKeysLock sync.Mutex
-var renderKeys map[string]*m.SignedInUser = make(map[string]*m.SignedInUser)
-
-func initContextWithRenderAuth(ctx *m.ReqContext) bool {
+func initContextWithRenderAuth(ctx *m.ReqContext, renderService rendering.Service) bool {
 	key := ctx.GetCookie("renderKey")
 	if key == "" {
 		return false
 	}
 
-	renderKeysLock.Lock()
-	defer renderKeysLock.Unlock()
-
-	renderUser, exists := renderKeys[key]
+	renderUser, exists := renderService.GetRenderUser(key)
 	if !exists {
 		ctx.JsonApiErr(401, "Invalid Render Key", nil)
 		return true
 	}
 
 	ctx.IsSignedIn = true
-	ctx.SignedInUser = renderUser
+	ctx.SignedInUser = &m.SignedInUser{
+		OrgId:   renderUser.OrgID,
+		UserId:  renderUser.UserID,
+		OrgRole: m.RoleType(renderUser.OrgRole),
+	}
 	ctx.IsRenderCall = true
 	ctx.LastSeenAt = time.Now()
 	return true
 }
-
-func AddRenderAuthKey(orgId int64, userId int64, orgRole m.RoleType) (string, error) {
-	renderKeysLock.Lock()
-	defer renderKeysLock.Unlock()
-
-	key, err := util.GetRandomString(32)
-	if err != nil {
-		return "", err
-	}
-
-	renderKeys[key] = &m.SignedInUser{
-		OrgId:   orgId,
-		OrgRole: orgRole,
-		UserId:  userId,
-	}
-
-	return key, nil
-}
-
-func RemoveRenderAuthKey(key string) {
-	renderKeysLock.Lock()
-	defer renderKeysLock.Unlock()
-
-	delete(renderKeys, key)
-}

pkg/services/alerting/notifier_test.go

@@ -304,6 +304,10 @@ func (s *testRenderService) RenderErrorImage(err error) (*rendering.RenderResult
 	return &rendering.RenderResult{FilePath: "image.png"}, nil
 }
 
+func (s *testRenderService) GetRenderUser(key string) (*rendering.RenderUser, bool) {
+	return nil, false
+}
+
 var _ rendering.Service = &testRenderService{}
 
 type testImageUploader struct {

pkg/services/rendering/http_mode.go

@@ -26,7 +26,7 @@ var netClient = &http.Client{
 	Transport: netTransport,
 }
 
-func (rs *RenderingService) renderViaHttp(ctx context.Context, opts Opts) (*RenderResult, error) {
+func (rs *RenderingService) renderViaHttp(ctx context.Context, renderKey string, opts Opts) (*RenderResult, error) {
 	filePath, err := rs.getFilePathForNewImage()
 	if err != nil {
 		return nil, err
@@ -37,11 +37,6 @@ func (rs *RenderingService) renderViaHttp(ctx context.Context, opts Opts) (*Rend
 		return nil, err
 	}
 
-	renderKey, err := rs.getRenderKey(opts.OrgId, opts.UserId, opts.OrgRole)
-	if err != nil {
-		return nil, err
-	}
-
 	queryParams := rendererUrl.Query()
 	queryParams.Add("url", rs.getURL(opts.Path))
 	queryParams.Add("renderKey", renderKey)

pkg/services/rendering/interface.go

@@ -29,9 +29,10 @@ type RenderResult struct {
 	FilePath string
 }
 
-type renderFunc func(ctx context.Context, options Opts) (*RenderResult, error)
+type renderFunc func(ctx context.Context, renderKey string, options Opts) (*RenderResult, error)
 
 type Service interface {
 	Render(ctx context.Context, opts Opts) (*RenderResult, error)
 	RenderErrorImage(error error) (*RenderResult, error)
+	GetRenderUser(key string) (*RenderUser, bool)
 }

pkg/services/rendering/phantomjs.go

@@ -11,10 +11,9 @@ import (
 	"time"
 
 	"github.com/grafana/grafana/pkg/infra/log"
-	"github.com/grafana/grafana/pkg/middleware"
 )
 
-func (rs *RenderingService) renderViaPhantomJS(ctx context.Context, opts Opts) (*RenderResult, error) {
+func (rs *RenderingService) renderViaPhantomJS(ctx context.Context, renderKey string, opts Opts) (*RenderResult, error) {
 	var executable = "phantomjs"
 	if runtime.GOOS == "windows" {
 		executable = executable + ".exe"
@@ -33,12 +32,6 @@ func (rs *RenderingService) renderViaPhantomJS(ctx context.Context, opts Opts) (
 		return nil, err
 	}
 
-	renderKey, err := middleware.AddRenderAuthKey(opts.OrgId, opts.UserId, opts.OrgRole)
-	if err != nil {
-		return nil, err
-	}
-	defer middleware.RemoveRenderAuthKey(renderKey)
-
 	phantomDebugArg := "--debug=false"
 	if log.GetLogLevelFor("rendering") >= log.LvlDebug {
 		phantomDebugArg = "--debug=true"

pkg/services/rendering/plugin_mode.go

@@ -12,17 +12,12 @@ func (rs *RenderingService) startPlugin(ctx context.Context) error {
 	return rs.pluginInfo.Start(ctx)
 }
 
-func (rs *RenderingService) renderViaPlugin(ctx context.Context, opts Opts) (*RenderResult, error) {
+func (rs *RenderingService) renderViaPlugin(ctx context.Context, renderKey string, opts Opts) (*RenderResult, error) {
 	pngPath, err := rs.getFilePathForNewImage()
 	if err != nil {
 		return nil, err
 	}
 
-	renderKey, err := rs.getRenderKey(opts.OrgId, opts.UserId, opts.OrgRole)
-	if err != nil {
-		return nil, err
-	}
-
 	// gives plugin some additional time to timeout and return possible errors.
 	ctx, cancel := context.WithTimeout(ctx, opts.Timeout+time.Second*2)
 	defer cancel()

pkg/services/rendering/rendering.go

@@ -6,9 +6,11 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
+	"time"
+
+	"github.com/grafana/grafana/pkg/infra/remotecache"
 
 	"github.com/grafana/grafana/pkg/infra/log"
-	"github.com/grafana/grafana/pkg/middleware"
 	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/plugins"
 	"github.com/grafana/grafana/pkg/registry"
@@ -17,11 +19,20 @@ import (
 )
 
 func init() {
+	remotecache.Register(&RenderUser{})
 	registry.RegisterService(&RenderingService{})
 }
 
 var IsPhantomJSEnabled = false
 
+const renderKeyPrefix = "render-%s"
+
+type RenderUser struct {
+	OrgID   int64
+	UserID  int64
+	OrgRole string
+}
+
 type RenderingService struct {
 	log             log.Logger
 	pluginInfo      *plugins.RendererPlugin
@@ -30,6 +41,7 @@ type RenderingService struct {
 	inProgressCount int
 
 	Cfg                *setting.Cfg             `inject:""`
+	RemoteCacheService *remotecache.RemoteCache `inject:""`
 }
 
 func (rs *RenderingService) Init() error {
@@ -103,19 +115,40 @@ func (rs *RenderingService) Render(ctx context.Context, opts Opts) (*RenderResul
 		}, nil
 	}
 
+	if rs.renderAction != nil {
+		rs.log.Info("Rendering", "path", opts.Path)
+		renderKey, err := rs.generateAndStoreRenderKey(opts.OrgId, opts.UserId, opts.OrgRole)
+		if err != nil {
+			return nil, err
+		}
+
+		defer rs.deleteRenderKey(renderKey)
+
 		defer func() {
-		rs.inProgressCount -= 1
+			rs.inProgressCount--
 		}()
 
-	rs.inProgressCount += 1
-
-	if rs.renderAction != nil {
-		rs.log.Info("Rendering", "path", opts.Path)
-		return rs.renderAction(ctx, opts)
+		rs.inProgressCount++
+		return rs.renderAction(ctx, renderKey, opts)
 	}
 	return nil, fmt.Errorf("No renderer found")
 }
 
+func (rs *RenderingService) GetRenderUser(key string) (*RenderUser, bool) {
+	val, err := rs.RemoteCacheService.Get(fmt.Sprintf(renderKeyPrefix, key))
+	if err != nil {
+		rs.log.Error("Failed to get render key from cache", "error", err)
+	}
+
+	if val != nil {
+		if user, ok := val.(*RenderUser); ok {
+			return user, true
+		}
+	}
+
+	return nil, false
+}
+
 func (rs *RenderingService) getFilePathForNewImage() (string, error) {
 	rand, err := util.GetRandomString(20)
 	if err != nil {
@@ -152,6 +185,27 @@ func (rs *RenderingService) getURL(path string) string {
 	return fmt.Sprintf("%s://%s:%s/%s&render=1", protocol, rs.domain, setting.HttpPort, path)
 }
 
-func (rs *RenderingService) getRenderKey(orgId, userId int64, orgRole models.RoleType) (string, error) {
-	return middleware.AddRenderAuthKey(orgId, userId, orgRole)
+func (rs *RenderingService) generateAndStoreRenderKey(orgId, userId int64, orgRole models.RoleType) (string, error) {
+	key, err := util.GetRandomString(32)
+	if err != nil {
+		return "", err
+	}
+
+	err = rs.RemoteCacheService.Set(fmt.Sprintf(renderKeyPrefix, key), &RenderUser{
+		OrgID:   orgId,
+		UserID:  userId,
+		OrgRole: string(orgRole),
+	}, 5*time.Minute)
+	if err != nil {
+		return "", err
+	}
+
+	return key, nil
+}
+
+func (rs *RenderingService) deleteRenderKey(key string) {
+	err := rs.RemoteCacheService.Delete(fmt.Sprintf(renderKeyPrefix, key))
+	if err != nil {
+		rs.log.Error("Failed to delete render key", "error", err)
+	}
 }