projectdiscovery-subfinder/pkg/subscraping/sources/github/tokenmanager.go

package github

import "time"

// Token struct
type Token struct {
	Hash         string
	RetryAfter   int64
	ExceededTime time.Time
}

// Tokens is the internal struct to manage the current token
// and the pool
type Tokens struct {
	current int
	pool    []Token
}

// NewTokenManager initialize the tokens pool
func NewTokenManager(keys []string) *Tokens {
	pool := []Token{}
	for _, key := range keys {
		t := Token{Hash: key, ExceededTime: time.Time{}, RetryAfter: 0}
		pool = append(pool, t)
	}

	return &Tokens{
		current: 0,
		pool:    pool,
	}
}

func (r *Tokens) setCurrentTokenExceeded(retryAfter int64) {
	if r.current >= len(r.pool) {
		r.current %= len(r.pool)
	}
	if r.pool[r.current].RetryAfter == 0 {
		r.pool[r.current].ExceededTime = time.Now()
		r.pool[r.current].RetryAfter = retryAfter
	}
}

// Get returns a new token from the token pool
func (r *Tokens) Get() *Token {
	resetExceededTokens(r)

	if r.current >= len(r.pool) {
		r.current %= len(r.pool)
	}

	result := &r.pool[r.current]
	r.current++

	return result
}

func resetExceededTokens(r *Tokens) {
	for i, token := range r.pool {
		if token.RetryAfter > 0 {
			if int64(time.Since(token.ExceededTime)/time.Second) > token.RetryAfter {
				r.pool[i].ExceededTime = time.Time{}
				r.pool[i].RetryAfter = 0
			}
		}
	}
}