From 575c792156e25e0f71cc3d4781489706e76a21eb Mon Sep 17 00:00:00 2001
From: Joel Speed <Joel.speed@hotmail.co.uk>
Date: Sun, 28 Jan 2018 22:37:07 +0000
Subject: [PATCH] Store most recent refresh token in offline sessions

---
 server/handlers.go | 33 +++++++++++++++++++++++++++++++++
 storage/storage.go |  3 +++
 2 files changed, 36 insertions(+)

diff --git a/server/handlers.go b/server/handlers.go
index b528918f..3be6f380 100644
--- a/server/handlers.go
+++ b/server/handlers.go
@@ -505,6 +505,39 @@ func (s *Server) finalizeLogin(identity connector.Identity, authReq storage.Auth
 	s.logger.Infof("login successful: connector %q, username=%q, preferred_username=%q, email=%q, groups=%q",
 		authReq.ConnectorID, claims.Username, claims.PreferredUsername, email, claims.Groups)
 
+	if _, ok := conn.(connector.RefreshConnector); ok {
+		// Try to retrieve an existing OfflineSession object for the corresponding user.
+		if session, err := s.storage.GetOfflineSessions(identity.UserID, authReq.ConnectorID); err != nil {
+			if err != storage.ErrNotFound {
+				s.logger.Errorf("failed to get offline session: %v", err)
+				return "", err
+			}
+			offlineSessions := storage.OfflineSessions{
+				UserID:        identity.UserID,
+				ConnID:        authReq.ConnectorID,
+				Refresh:       make(map[string]*storage.RefreshTokenRef),
+				ConnectorData: identity.ConnectorData,
+			}
+
+			// Create a new OfflineSession object for the user and add a reference object for
+			// the newly received refreshtoken.
+			if err := s.storage.CreateOfflineSessions(offlineSessions); err != nil {
+				s.logger.Errorf("failed to create offline session: %v", err)
+				return "", err
+			}
+		} else {
+			// Update existing OfflineSession obj with new RefreshTokenRef.
+			if err := s.storage.UpdateOfflineSessions(session.UserID, session.ConnID, func(old storage.OfflineSessions) (storage.OfflineSessions, error) {
+				old.ConnectorData = identity.ConnectorData
+				return old, nil
+			}); err != nil {
+				s.logger.Errorf("failed to update offline session: %v", err)
+				return "", err
+			}
+
+		}
+	}
+
 	return path.Join(s.issuerURL.Path, "/approval") + "?req=" + authReq.ID, nil
 }
 
diff --git a/storage/storage.go b/storage/storage.go
index 235f74e0..cb2a7e0c 100644
--- a/storage/storage.go
+++ b/storage/storage.go
@@ -273,6 +273,9 @@ type OfflineSessions struct {
 	// Refresh is a hash table of refresh token reference objects
 	// indexed by the ClientID of the refresh token.
 	Refresh map[string]*RefreshTokenRef
+
+	// Authentication data provided by an upstream source.
+	ConnectorData []byte
 }
 
 // Password is an email to password mapping managed by the storage.