2 files changed, 98 insertions, 18 deletions

diff --git a/services/auth/source/oauth2/init.go b/services/auth/source/oauth2/init.go
index 343b24cf6f..edbbb8969e 100644
--- a/services/auth/source/oauth2/init.go
+++ b/services/auth/source/oauth2/init.go
@@ -5,23 +5,21 @@
 package oauth2
 
 import (
+	"encoding/gob"
 	"net/http"
 	"sync"
 
-	"code.gitea.io/gitea/models/db"
 	"code.gitea.io/gitea/models/login"
 	"code.gitea.io/gitea/modules/log"
 	"code.gitea.io/gitea/modules/setting"
 
 	"github.com/google/uuid"
+	"github.com/gorilla/sessions"
 	"github.com/markbates/goth/gothic"
 )
 
 var gothRWMutex = sync.RWMutex{}
 
-// SessionTableName is the table name that OAuth2 will use to store things
-const SessionTableName = "oauth2_session"
-
 // UsersStoreKey is the key for the store
 const UsersStoreKey = "gitea-oauth2-sessions"
 
@@ -34,23 +32,14 @@ func Init() error {
 		return err
 	}
 
-	store, err := db.CreateStore(SessionTableName, UsersStoreKey)
-	if err != nil {
-		return err
-	}
-
-	// according to the Goth lib:
-	// set the maxLength of the cookies stored on the disk to a larger number to prevent issues with:
-	// securecookie: the value is too long
-	// when using OpenID Connect , since this can contain a large amount of extra information in the id_token
-
-	// Note, when using the FilesystemStore only the session.ID is written to a browser cookie, so this is explicit for the storage on disk
-	store.MaxLength(setting.OAuth2.MaxTokenLength)
-
 	// Lock our mutex
 	gothRWMutex.Lock()
 
-	gothic.Store = store
+	gob.Register(&sessions.Session{})
+
+	gothic.Store = &SessionsStore{
+		maxLength: int64(setting.OAuth2.MaxTokenLength),
+	}
 
 	gothic.SetState = func(req *http.Request) string {
 		return uuid.New().String()
diff --git a/services/auth/source/oauth2/store.go b/services/auth/source/oauth2/store.go
new file mode 100644
index 0000000000..fc29056c24
--- /dev/null
+++ b/services/auth/source/oauth2/store.go
@@ -0,0 +1,91 @@
+// Copyright 2021 The Gitea Authors. All rights reserved.
+// Use of this source code is governed by a MIT-style
+// license that can be found in the LICENSE file.
+
+package oauth2
+
+import (
+	"encoding/gob"
+	"fmt"
+	"net/http"
+
+	"code.gitea.io/gitea/modules/log"
+	chiSession "gitea.com/go-chi/session"
+	"github.com/gorilla/sessions"
+)
+
+// SessionsStore creates a gothic store from our session
+type SessionsStore struct {
+	maxLength int64
+}
+
+// Get should return a cached session.
+func (st *SessionsStore) Get(r *http.Request, name string) (*sessions.Session, error) {
+	return st.getOrNew(r, name, false)
+}
+
+// New should create and return a new session.
+//
+// Note that New should never return a nil session, even in the case of
+// an error if using the Registry infrastructure to cache the session.
+func (st *SessionsStore) New(r *http.Request, name string) (*sessions.Session, error) {
+	return st.getOrNew(r, name, true)
+}
+
+// getOrNew gets the session from the chi-session if it exists. Override permits the overriding of an unexpected object.
+func (st *SessionsStore) getOrNew(r *http.Request, name string, override bool) (*sessions.Session, error) {
+	chiStore := chiSession.GetSession(r)
+
+	session := sessions.NewSession(st, name)
+
+	rawData := chiStore.Get(name)
+	if rawData != nil {
+		oldSession, ok := rawData.(*sessions.Session)
+		if ok {
+			session.ID = oldSession.ID
+			session.IsNew = oldSession.IsNew
+			session.Options = oldSession.Options
+			session.Values = oldSession.Values
+
+			return session, nil
+		} else if !override {
+			log.Error("Unexpected object in session at name: %s: %v", name, rawData)
+			return nil, fmt.Errorf("unexpected object in session at name: %s", name)
+		}
+	}
+
+	session.ID = chiStore.ID() // Simply copy the session id from the chi store
+
+	return session, chiStore.Set(name, session)
+}
+
+// Save should persist session to the underlying store implementation.
+func (st *SessionsStore) Save(r *http.Request, w http.ResponseWriter, session *sessions.Session) error {
+	chiStore := chiSession.GetSession(r)
+
+	if err := chiStore.Set(session.Name(), session); err != nil {
+		return err
+	}
+
+	if st.maxLength > 0 {
+		sizeWriter := &sizeWriter{}
+
+		_ = gob.NewEncoder(sizeWriter).Encode(session)
+		if sizeWriter.size > st.maxLength {
+			return fmt.Errorf("encode session: Data too long: %d > %d", sizeWriter.size, st.maxLength)
+		}
+	}
+
+	return chiStore.Release()
+}
+
+type sizeWriter struct {
+	size int64
+}
+
+func (s *sizeWriter) Write(data []byte) (int, error) {
+	s.size += int64(len(data))
+	return len(data), nil
+}
+
+var _ (sessions.Store) = &SessionsStore{}