be.ems/lib/session/session.go

package session

import (
	"crypto/rand"
	"encoding/base64"
	"errors"
	"io"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"ems.agt/lib/log"
	"ems.agt/lib/oauth"
	"ems.agt/restagent/config"
)

// SessionMgr session manager
type SessManager struct {
	name     string
	expires  int64
	lock     sync.RWMutex
	sessions map[string]*Session
}

// Session
type Session struct {
	token      string
	time       time.Time
	permission []bool
	values     map[interface{}]interface{}
}

// NewSessionMgr create session manager
func NewSessManager(name string) *SessManager {
	smgr := &SessManager{name: name, expires: (int64)(config.GetExpiresFromConfig()), sessions: make(map[string]*Session)}
	go smgr.SessionGC()
	return smgr
}

// NewSession create session
func (smgr *SessManager) NewSession(w http.ResponseWriter, r *http.Request, plist []bool) string {
	smgr.lock.Lock()
	defer smgr.lock.Unlock()
	token := oauth.GenRandToken("omc") // Generate new token to session ID
	session := &Session{token: token, time: time.Now(), permission: plist, values: make(map[interface{}]interface{})}
	smgr.sessions[token] = session

	return token
}

// EndSession
func (smgr *SessManager) EndSession(w http.ResponseWriter, r *http.Request) {
	token := smgr.GetTokenFromHttpRequest(r)

	smgr.lock.Lock()
	defer smgr.lock.Unlock()
	delete(smgr.sessions, token)
}

// Handshake session, restart session
func (smgr *SessManager) ShakeSession(token string) bool {

	smgr.lock.Lock()
	defer smgr.lock.Unlock()
	for _, s := range smgr.sessions {
		if token == s.token {
			log.Debug("session time:", s.time)
			s.time = time.Now()
			return true
		}
	}
	return false
}

// EndSessionByID end the session by session ID
func (smgr *SessManager) DeleteSession(token string) {
	smgr.lock.Lock()
	defer smgr.lock.Unlock()
	delete(smgr.sessions, token)
}

// SetSessionValue set value fo session
func (smgr *SessManager) SetSessionValue(token string, key interface{}, value interface{}) error {
	smgr.lock.Lock()
	defer smgr.lock.Unlock()
	if session, ok := smgr.sessions[token]; ok {
		session.values[key] = value
		return nil
	}
	return errors.New("invalid session ID")
}

// GetSessionValue get value fo session
func (smgr *SessManager) GetSessionValue(token string, key interface{}) (interface{}, error) {
	smgr.lock.RLock()
	defer smgr.lock.RUnlock()
	if session, ok := smgr.sessions[token]; ok {
		if val, ok := session.values[key]; ok {
			return val, nil
		}
	}
	return nil, errors.New("invalid session ID")
}

func (smgr *SessManager) GetTokenFromHttpRequest(r *http.Request) string {
	for k, v := range r.Header {
		if strings.ToLower(k) == "accesstoken" && len(v) != 0 {
			log.Debug("AccessToken:", v[0])
			return v[0]
		}
	}

	return ""
}

// IsValidToken check token is valid or not
func (smgr *SessManager) IsValidToken(token string) bool {

	smgr.lock.Lock()
	defer smgr.lock.Unlock()
	if _, ok := smgr.sessions[token]; ok {
		return true
	}
	return false
}

// IsCarriedToken check token is carried
func (smgr *SessManager) IsCarriedToken(r *http.Request) (string, bool) {

	token := smgr.GetTokenFromHttpRequest(r)
	if token == "" {
		return "", false
	}
	return token, true
}

// GetPermissionFromSession get permission from session by token
func (smgr *SessManager) GetPermissionFromSession(token string) []bool {

	if s, ok := smgr.sessions[token]; ok {
		return s.permission
	}
	return nil
}

// SessionGC maintain session
func (smgr *SessManager) SessionGC() {
	smgr.lock.Lock()
	defer smgr.lock.Unlock()
	for token, session := range smgr.sessions {
		if session.time.Unix()+smgr.expires < time.Now().Unix() {
			delete(smgr.sessions, token)
		}
	}

	time.AfterFunc(time.Duration(smgr.expires)*time.Second, func() { smgr.SessionGC() })
}

// NewSessionID generate unique ID
func (smgr *SessManager) NewSessionID() string {
	b := make([]byte, 32)
	if _, err := io.ReadFull(rand.Reader, b); err != nil {
		nano := time.Now().UnixNano()
		return strconv.FormatInt(nano, 10)
	}
	return base64.URLEncoding.EncodeToString(b)
}