Implement context propagation through to SessionDetail

Author: imogenkraak

gateway/auth_manager.go

@@ -1,6 +1,7 @@
 package gateway
 
 import (
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"strings"
@@ -21,6 +22,7 @@ type SessionHandler interface {
 	UpdateSession(keyName string, session *user.SessionState, resetTTLTo int64, hashed bool) error
 	RemoveSession(orgID string, keyName string, hashed bool) bool
 	SessionDetail(orgID string, keyName string, hashed bool) (user.SessionState, bool)
+	SessionDetailContext(ctx context.Context, orgID string, keyName string, hashed bool) (user.SessionState, bool)
 	KeyExpired(newSession *user.SessionState) bool
 	Sessions(filter string) []string
 	ResetQuota(string, *user.SessionState, bool)
@@ -215,6 +217,84 @@ func (b *DefaultSessionManager) SessionDetail(orgID string, keyName string, hash
 	return session.Clone(), true
 }
 
+// SessionDetailContext returns the session detail using the storage engine with context support for cancellation
+func (b *DefaultSessionManager) SessionDetailContext(ctx context.Context, orgID string, keyName string, hashed bool) (user.SessionState, bool) {
+	select {
+	case <-ctx.Done():
+		log.WithFields(logrus.Fields{
+			"prefix":      "auth-mgr",
+			"inbound-key": b.Gw.obfuscateKey(keyName),
+		}).Debug("Context cancelled before session detail fetch")
+		return user.SessionState{}, false
+	default:
+	}
+
+	var jsonKeyVal string
+	var err error
+	keyId := keyName
+
+	// get session by key
+	if hashed {
+		jsonKeyVal, err = b.store.GetRawKey(b.store.GetKeyPrefix() + keyName)
+	} else {
+		if storage.TokenOrg(keyName) != orgID {
+			// try to get legacy and new format key at once
+			toSearchList := []string{}
+			if !b.Gw.GetConfig().DisableKeyActionsByUsername {
+				toSearchList = append(toSearchList, b.Gw.generateToken(orgID, keyName))
+			}
+
+			toSearchList = append(toSearchList, keyName)
+			for _, fallback := range b.Gw.GetConfig().HashKeyFunctionFallback {
+				if !b.Gw.GetConfig().DisableKeyActionsByUsername {
+					toSearchList = append(toSearchList, b.Gw.generateToken(orgID, keyName, fallback))
+				}
+			}
+
+			var jsonKeyValList []string
+
+			jsonKeyValList, err = b.store.GetMultiKey(toSearchList)
+			// pick the 1st non empty from the returned list
+			for idx, val := range jsonKeyValList {
+				if val != "" {
+					jsonKeyVal = val
+					keyId = toSearchList[idx]
+					break
+				}
+			}
+		} else {
+			// key is not an imported one
+			jsonKeyVal, err = b.store.GetKey(keyName)
+		}
+	}
+
+	select {
+	case <-ctx.Done():
+		log.WithFields(logrus.Fields{
+			"prefix":      "auth-mgr",
+			"inbound-key": b.Gw.obfuscateKey(keyName),
+		}).Debug("Context cancelled during session detail fetch")
+		return user.SessionState{}, false
+	default:
+	}
+
+	if err != nil {
+		log.WithFields(logrus.Fields{
+			"prefix":      "auth-mgr",
+			"inbound-key": b.Gw.obfuscateKey(keyName),
+			"err":         err,
+		}).Debug("Could not get session detail, key not found")
+		return user.SessionState{}, false
+	}
+	session := &user.SessionState{}
+	if err := json.Unmarshal([]byte(jsonKeyVal), &session); err != nil {
+		log.Error("Couldn't unmarshal session object (may be cache miss): ", err)
+		return user.SessionState{}, false
+	}
+	session.KeyID = keyId
+	return session.Clone(), true
+}
+
 func (b *DefaultSessionManager) Stop() {}
 
 // Sessions returns all sessions in the key store that match a filter key (a prefix)

gateway/middleware.go

@@ -368,7 +368,7 @@ func (t *BaseMiddleware) fetchOrgSessionWithTimeout(orgID string) (user.SessionS
 			}
 		}()
 
-		session, found := t.Spec.OrgSessionManager.SessionDetail(orgID, orgID, false)
+		session, found := t.Spec.OrgSessionManager.SessionDetailContext(timeoutCtx, orgID, orgID, false)
 		if found && t.Spec.GlobalConfig.EnforceOrgDataAge {
 			t.Logger().Debug("Setting data expiry: ", orgID)
 			t.Gw.ExpiryCache.Set(session.OrgID, session.DataExpires, cache.DefaultExpiration)

gateway/middleware_test.go

@@ -1,6 +1,7 @@
 package gateway
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -40,6 +41,17 @@ func (m mockStore) SessionDetail(orgID string, keyName string, hashed bool) (use
 	return sess.Clone(), !m.DetailNotFound
 }
 
+func (m mockStore) SessionDetailContext(ctx context.Context, orgID string, keyName string, hashed bool) (user.SessionState, bool) {
+	if m.Delay > 0 {
+		select {
+		case <-time.After(m.Delay):
+		case <-ctx.Done():
+			return user.SessionState{}, false
+		}
+	}
+	return sess.Clone(), !m.DetailNotFound
+}
+
 func TestBaseMiddleware_OrgSessionExpiry(t *testing.T) {
 	ts := StartTest(nil)
 	defer ts.Close()