Wormhole/internal/store/apikey.go

package store

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
)

func hashAPIKey(key string) string {
	sum := sha256.Sum256([]byte(key))
	return hex.EncodeToString(sum[:])
}

func generateAPIKeyPlain() (string, string, error) {
	b := make([]byte, 24)
	if _, err := rand.Read(b); err != nil {
		return "", "", err
	}
	plain := "wh_" + hex.EncodeToString(b)
	prefix := plain
	if len(prefix) > 12 {
		prefix = prefix[:12]
	}
	return plain, prefix, nil
}

func (s *Store) ListAPIKeys(userID uint, isAdmin bool) ([]APIKey, error) {
	var keys []APIKey
	q := s.db.Order("id asc")
	if !isAdmin {
		q = q.Where("user_id = ?", userID)
	}
	return keys, q.Find(&keys).Error
}

func (s *Store) CreateAPIKey(name string, userID uint) (*APIKey, string, error) {
	if name == "" {
		return nil, "", errors.New("name required")
	}
	plain, prefix, err := generateAPIKeyPlain()
	if err != nil {
		return nil, "", err
	}
	key := &APIKey{
		Name:      name,
		KeyPrefix: prefix,
		KeyHash:   hashAPIKey(plain),
		UserID:    userID,
		Status:    true,
	}
	if err := s.db.Create(key).Error; err != nil {
		return nil, "", err
	}
	return key, plain, nil
}

func (s *Store) DeleteAPIKey(id uint, userID uint, isAdmin bool) error {
	var key APIKey
	if err := s.db.First(&key, id).Error; err != nil {
		return err
	}
	if !isAdmin && key.UserID != userID {
		return fmt.Errorf("permission denied")
	}
	return s.db.Delete(&APIKey{}, id).Error
}

func (s *Store) ValidateAPIKey(plain string) (*User, error) {
	if plain == "" {
		return nil, errors.New("empty api key")
	}
	hash := hashAPIKey(plain)
	var key APIKey
	if err := s.db.Where("key_hash = ? AND status = ?", hash, true).First(&key).Error; err != nil {
		return nil, errors.New("invalid api key")
	}
	user, err := s.GetUserByID(key.UserID)
	if err != nil {
		return nil, err
	}
	if !user.Status {
		return nil, errors.New("user disabled")
	}
	return user, nil
}