From a138e2cc4251e7c99c55f4f1a5f97d1182e3aab4 Mon Sep 17 00:00:00 2001
From: Qiaosen Huang <joesonw@gmail.com>
Date: Sat, 23 Oct 2021 15:23:43 +0800
Subject: [PATCH] add Context version to Store methods

---
 data.go | 38 ++++++++++++++++++++++++++++++++++----
 1 file changed, 34 insertions(+), 4 deletions(-)

diff --git a/data.go b/data.go
index 1460e1b..e3061d6 100644
--- a/data.go
+++ b/data.go
@@ -59,7 +59,7 @@ func (s *SessionManager) Load(ctx context.Context, token string) (context.Contex
 		return s.addSessionDataToContext(ctx, newSessionData(s.Lifetime)), nil
 	}
 
-	b, found, err := s.Store.Find(token)
+	b, found, err := s.doStoreFind(ctx, token)
 	if err != nil {
 		return nil, err
 	} else if !found {
@@ -115,7 +115,7 @@ func (s *SessionManager) Commit(ctx context.Context) (string, time.Time, error)
 		}
 	}
 
-	if err := s.Store.Commit(sd.token, b, expiry); err != nil {
+	if err := s.doStoreCommit(ctx, sd.token, b, expiry); err != nil {
 		return "", time.Time{}, err
 	}
 
@@ -131,7 +131,7 @@ func (s *SessionManager) Destroy(ctx context.Context) error {
 	sd.mu.Lock()
 	defer sd.mu.Unlock()
 
-	err := s.Store.Delete(sd.token)
+	err := s.doStoreDelete(ctx, sd.token)
 	if err != nil {
 		return err
 	}
@@ -284,7 +284,7 @@ func (s *SessionManager) RenewToken(ctx context.Context) error {
 	sd.mu.Lock()
 	defer sd.mu.Unlock()
 
-	err := s.Store.Delete(sd.token)
+	err := s.doStoreDelete(ctx, sd.token)
 	if err != nil {
 		return err
 	}
@@ -578,3 +578,33 @@ func generateContextKey() contextKey {
 	atomic.AddUint64(&contextKeyID, 1)
 	return contextKey(fmt.Sprintf("session.%d", contextKeyID))
 }
+
+func (s *SessionManager) doStoreDelete(ctx context.Context, token string) (err error) {
+	c, ok := s.Store.(interface {
+		DeleteCtx(context.Context, string) error
+	})
+	if ok {
+		return c.DeleteCtx(ctx, token)
+	}
+	return s.Store.Delete(token)
+}
+
+func (s *SessionManager) doStoreFind(ctx context.Context, token string) (b []byte, found bool, err error) {
+	c, ok := s.Store.(interface {
+		FindCtx(context.Context, string) ([]byte, bool, error)
+	})
+	if ok {
+		return c.FindCtx(ctx, token)
+	}
+	return s.Store.Find(token)
+}
+
+func (s *SessionManager) doStoreCommit(ctx context.Context, token string, b []byte, expiry time.Time) (err error) {
+	c, ok := s.Store.(interface {
+		CommitCtx(context.Context, string, []byte, time.Time) error
+	})
+	if ok {
+		return c.CommitCtx(ctx, token, b, expiry)
+	}
+	return s.Store.Commit(token, b, expiry)
+}