breakpilot-compliance/pca-platform/heuristic-service/internal/stepup/pow.go

package stepup

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/breakpilot/pca-platform/heuristic-service/internal/config"
)

// PoWService handles Proof-of-Work challenges
type PoWService struct {
	config     *config.PoWConfig
	challenges map[string]*PoWChallenge
	mu         sync.RWMutex
}

// PoWChallenge represents a Proof-of-Work challenge
type PoWChallenge struct {
	ID         string    `json:"id"`
	SessionID  string    `json:"session_id"`
	Challenge  string    `json:"challenge"`
	Difficulty int       `json:"difficulty"`
	CreatedAt  time.Time `json:"created_at"`
	ExpiresAt  time.Time `json:"expires_at"`
	Solved     bool      `json:"solved"`
}

// PoWChallengeResponse is sent to the client
type PoWChallengeResponse struct {
	ChallengeID string `json:"challenge_id"`
	Challenge   string `json:"challenge"`
	Difficulty  int    `json:"difficulty"`
	MaxTimeMs   int    `json:"max_time_ms"`
	Hint        string `json:"hint"`
}

// PoWVerifyRequest for verifying a solved challenge
type PoWVerifyRequest struct {
	SessionID   string `json:"session_id"`
	ChallengeID string `json:"challenge_id"`
	Challenge   string `json:"challenge"`
	Nonce       int64  `json:"nonce"`
}

// NewPoWService creates a new Proof-of-Work service
func NewPoWService(cfg *config.PoWConfig) *PoWService {
	return &PoWService{
		config:     cfg,
		challenges: make(map[string]*PoWChallenge),
	}
}

// CreateChallenge generates a new PoW challenge
func (s *PoWService) CreateChallenge(sessionID string) (*PoWChallengeResponse, error) {
	// Generate random challenge
	challengeBytes := make([]byte, 16)
	if _, err := rand.Read(challengeBytes); err != nil {
		return nil, err
	}
	challengeStr := hex.EncodeToString(challengeBytes)

	// Generate challenge ID
	idBytes := make([]byte, 8)
	rand.Read(idBytes)
	challengeID := hex.EncodeToString(idBytes)

	// Create challenge
	challenge := &PoWChallenge{
		ID:         challengeID,
		SessionID:  sessionID,
		Challenge:  challengeStr,
		Difficulty: s.config.Difficulty,
		CreatedAt:  time.Now(),
		ExpiresAt:  time.Now().Add(time.Duration(s.config.MaxDurationMs*2) * time.Millisecond),
		Solved:     false,
	}

	// Store challenge
	s.mu.Lock()
	s.challenges[challengeID] = challenge
	s.mu.Unlock()

	// Build response
	prefix := strings.Repeat("0", s.config.Difficulty)
	response := &PoWChallengeResponse{
		ChallengeID: challengeID,
		Challenge:   challengeStr,
		Difficulty:  s.config.Difficulty,
		MaxTimeMs:   s.config.MaxDurationMs,
		Hint:        fmt.Sprintf("Find nonce where SHA256(challenge + nonce) starts with '%s'", prefix),
	}

	return response, nil
}

// VerifyChallenge verifies a PoW solution
func (s *PoWService) VerifyChallenge(req *PoWVerifyRequest) (bool, error) {
	s.mu.RLock()
	challenge, exists := s.challenges[req.ChallengeID]
	s.mu.RUnlock()

	if !exists {
		return false, nil
	}

	// Check expiration
	if time.Now().After(challenge.ExpiresAt) {
		s.mu.Lock()
		delete(s.challenges, req.ChallengeID)
		s.mu.Unlock()
		return false, nil
	}

	// Check session match
	if challenge.SessionID != req.SessionID {
		return false, nil
	}

	// Check challenge string match
	if challenge.Challenge != req.Challenge {
		return false, nil
	}

	// Verify the proof of work
	input := fmt.Sprintf("%s%d", req.Challenge, req.Nonce)
	hash := sha256.Sum256([]byte(input))
	hashHex := hex.EncodeToString(hash[:])

	// Check if hash has required number of leading zeros
	prefix := strings.Repeat("0", challenge.Difficulty)
	if !strings.HasPrefix(hashHex, prefix) {
		return false, nil
	}

	// Mark as solved
	s.mu.Lock()
	challenge.Solved = true
	s.mu.Unlock()

	return true, nil
}

// VerifyProof is a standalone verification without stored challenge
// Useful for quick verification
func (s *PoWService) VerifyProof(challenge string, nonce int64, difficulty int) bool {
	input := fmt.Sprintf("%s%d", challenge, nonce)
	hash := sha256.Sum256([]byte(input))
	hashHex := hex.EncodeToString(hash[:])

	prefix := strings.Repeat("0", difficulty)
	return strings.HasPrefix(hashHex, prefix)
}

// CleanupExpiredChallenges removes expired challenges
func (s *PoWService) CleanupExpiredChallenges() {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := time.Now()
	for id, challenge := range s.challenges {
		if now.After(challenge.ExpiresAt) {
			delete(s.challenges, id)
		}
	}
}

// IsEnabled returns whether PoW is enabled
func (s *PoWService) IsEnabled() bool {
	return s.config.Enabled
}

// GetDifficulty returns configured difficulty
func (s *PoWService) GetDifficulty() int {
	return s.config.Difficulty
}