default_token_whitelist_checker.go

package auth

import (
	"context"
	"errors"
	"strconv"
	"time"
)

type DefaultTokenWhitelistChecker struct {
	Secret      string
	TokenIp     string
	TokenPrefix string
	VerifyToken func(tokenString string, secret string) (map[string]interface{}, int64, int64, error)
	CachePort   CachePort
	Level       int
}

func NewTokenWhitelistChecker(secret string, tokenIp string, keyPrefix string, verifyToken func(tokenString string, secret string) (map[string]interface{}, int64, int64, error), cacheService CachePort, level int) *DefaultTokenWhitelistChecker {
	return &DefaultTokenWhitelistChecker{secret, tokenIp, keyPrefix, verifyToken, cacheService, level}
}

func (b *DefaultTokenWhitelistChecker) generateKey(token string) string {
	return b.TokenPrefix + "::token::" + token
}

func (b *DefaultTokenWhitelistChecker) generateKeyForId(id string) string {
	return b.TokenPrefix + "::token::" + id
}

func (b *DefaultTokenWhitelistChecker) Add(ctx context.Context, id string, token string) error {
	_, _, eta, err := b.VerifyToken(token, b.Secret)
	if err != nil {
		return err
	}
	now := time.Now()

	if eta <= now.Unix() {
		return errors.New("token expired")
	}
	expire := time.Unix(eta, 0)
	dur := expire.Sub(now)

	key := b.generateKeyForId(id)
	return b.CachePort.Put(ctx, key, token, dur)
}

func (b *DefaultTokenWhitelistChecker) Check(ctx context.Context, id string, token string) bool {
	key := b.generateKeyForId(id)

	value, err := b.CachePort.Get(ctx, key)
	if err != nil {
		return false
	}
	if len(value) > 0 {
		tokenStore, _ := strconv.Unquote(value)
		if b.Level != 0 {
			if tokenStore != token {
				return false
			}
			return true
		}

		payloadStore, _, _, err1 := b.VerifyToken(tokenStore, b.Secret)
		payload, _, _, err2 := b.VerifyToken(token, b.Secret)
		if err1 != nil || err2 != nil {
			return false
		}
		ipStore, ok1 := payloadStore[b.TokenIp]
		ip, ok2 := payload[b.TokenIp]
		if ok1 && ok2 {
			if ip == ipStore {
				return true
			}
		}
	}
	return false
}