From d7115fb8971ae1983a1114ccac3c843d1d8df29b Mon Sep 17 00:00:00 2001
From: dushu <shu.du@canonical.com>
Date: Thu, 11 Apr 2024 21:29:41 -0600
Subject: [PATCH] fix: update OpenID Connect session after user consent

---
 oauth2/handler.go                   | 10 ++++++++++
 persistence/sql/persister_oauth2.go | 24 ++++++++++++++++++++++++
 x/fosite_storer.go                  |  2 ++
 3 files changed, 36 insertions(+)

diff --git a/oauth2/handler.go b/oauth2/handler.go
index 01ec8d1fc6f..ba266804cfd 100644
--- a/oauth2/handler.go
+++ b/oauth2/handler.go
@@ -760,6 +760,16 @@ func (h *Handler) performOAuth2DeviceVerificationFlow(w http.ResponseWriter, r *
 		return
 	}
 
+	// Update the OpenID Connect session if "openid" scope is granted
+	if req.GetGrantedScopes().Has("openid") {
+		err = h.r.OAuth2Storage().UpdateOpenIDConnectSessionByRequestID(ctx, f.DeviceCodeRequestID.String(), req)
+		if err != nil {
+			x.LogError(r, err, h.r.Logger())
+			h.r.Writer().WriteError(w, r, err)
+			return
+		}
+	}
+
 	redirectURL := urlx.SetQuery(h.c.DeviceDoneURL(ctx), url.Values{"consent_verifier": {string(f.ConsentVerifier)}}).String()
 	http.Redirect(w, r, redirectURL, http.StatusFound)
 }
diff --git a/persistence/sql/persister_oauth2.go b/persistence/sql/persister_oauth2.go
index 340e30392df..808c48faeb2 100644
--- a/persistence/sql/persister_oauth2.go
+++ b/persistence/sql/persister_oauth2.go
@@ -446,6 +446,30 @@ func (p *Persister) CreateOpenIDConnectSession(ctx context.Context, signature st
 	return p.createSession(ctx, signature, requester, sqlTableOpenID)
 }
 
+// UpdateOpenIDConnectSessionByRequestID updates an OpenID session by requestID
+func (p *Persister) UpdateOpenIDConnectSessionByRequestID(ctx context.Context, requestID string, requester fosite.Requester) (err error) {
+	ctx, span := p.r.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.UpdateOpenIDConnectSessionByRequestID")
+	defer otelx.End(span, &err)
+
+	req, err := p.sqlSchemaFromRequest(ctx, requestID, requester, sqlTableOpenID)
+	if err != nil {
+		return err
+	}
+
+	stmt := fmt.Sprintf(
+		"UPDATE %s SET granted_scope=?, granted_audience=?, session_data=? WHERE request_id=? AND nid = ?",
+		OAuth2RequestSQL{Table: sqlTableOpenID}.TableName(),
+	)
+
+	/* #nosec G201 table is static */
+	err = p.Connection(ctx).RawQuery(stmt, req.GrantedScope, req.GrantedAudience, req.Session, requestID, p.NetworkID(ctx)).Exec()
+	if err != nil {
+		return sqlcon.HandleError(err)
+	}
+
+	return nil
+}
+
 func (p *Persister) GetOpenIDConnectSession(ctx context.Context, signature string, requester fosite.Requester) (_ fosite.Requester, err error) {
 	ctx, span := p.r.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.GetOpenIDConnectSession")
 	defer otelx.End(span, &err)
diff --git a/x/fosite_storer.go b/x/fosite_storer.go
index 2d4d33f7de4..088a222baca 100644
--- a/x/fosite_storer.go
+++ b/x/fosite_storer.go
@@ -43,6 +43,8 @@ type FositeStorer interface {
 
 	FlushInactiveRefreshTokens(ctx context.Context, notAfter time.Time, limit int, batchSize int) error
 
+	UpdateOpenIDConnectSessionByRequestID(ctx context.Context, requestID string, requester fosite.Requester) error
+
 	// DeleteOpenIDConnectSession deletes an OpenID Connect session.
 	// This is duplicated from Ory Fosite to help against deprecation linting errors.
 	DeleteOpenIDConnectSession(ctx context.Context, authorizeCode string) error