From 87f93c2872ececa0cb7f8c8175b51119d8d87da5 Mon Sep 17 00:00:00 2001
From: Yousong Zhou <zhouyousong@yunionyun.com>
Date: Thu, 18 Feb 2021 18:29:59 +0800
Subject: [PATCH] mcclient: auth: add SessionCache

---
 pkg/mcclient/auth/session_cache.go | 72 ++++++++++++++++++++++++++++++
 1 file changed, 72 insertions(+)
 create mode 100644 pkg/mcclient/auth/session_cache.go

diff --git a/pkg/mcclient/auth/session_cache.go b/pkg/mcclient/auth/session_cache.go
new file mode 100644
index 0000000000..e902dc5b61
--- /dev/null
+++ b/pkg/mcclient/auth/session_cache.go
@@ -0,0 +1,72 @@
+// Copyright 2019 Yunion
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package auth
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"yunion.io/x/onecloud/pkg/mcclient"
+)
+
+type SessionCache struct {
+	mu      sync.RWMutex
+	session *mcclient.ClientSession
+
+	Region     string
+	APIVersion string
+
+	// EarlyRefresh tells the cache how early to fetch a new session before
+	// actual expiration of the old
+	EarlyRefresh time.Duration
+
+	Token         mcclient.TokenCredential
+	UseAdminToken bool
+}
+
+func (sc *SessionCache) getToken() mcclient.TokenCredential {
+	if sc.Token != nil {
+		return sc.Token
+	}
+	if sc.UseAdminToken {
+		return manager.adminCredential
+	}
+	return nil
+}
+
+func (sc *SessionCache) Get(ctx context.Context) *mcclient.ClientSession {
+	mu := &sc.mu
+
+	{
+		mu.RLock()
+		s := sc.session
+		mu.RUnlock()
+
+		if s != nil {
+			token := s.GetToken()
+			expires := token.GetExpires()
+			if time.Now().Add(sc.EarlyRefresh).After(expires) {
+				return s
+			}
+		}
+	}
+
+	mu.Lock()
+	defer mu.Unlock()
+	token := sc.getToken()
+	sc.session = GetSession(ctx, token, sc.Region, sc.APIVersion)
+	return sc.session
+}