proxy-pool/internal/adapters/redisactivity/extract.go

package redisactivity

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/url"
	"sort"
	"strconv"
	"time"

	extractionDomain "proxy-pool/internal/domain/extraction"
)

const defaultRedisIdempotencyTTL = 5 * time.Minute

type extractionDigestInput struct {
	Requested   int                          `json:"requested"`
	Fulfillment extractionDomain.Fulfillment `json:"fulfillment"`
	Protocols   []string                     `json:"protocols"`
	Regions     []string                     `json:"regions"`
	Carriers    []string                     `json:"carriers"`
	Upstreams   []string                     `json:"upstreams"`
}

type extractionFilterWire struct {
	Protocols []string `json:"protocols"`
	Regions   []string `json:"regions"`
	Carriers  []string `json:"carriers"`
	Upstreams []string `json:"upstreams"`
}

var _ extractionDomain.Store = (*Adapter)(nil)

func (a *Adapter) Extract(ctx context.Context, command extractionDomain.Command) (extractionDomain.Result, error) {
	result := extractionDomain.Result{Requested: command.Requested}
	if ctx == nil {
		return result, extractionDomain.ErrInvalidCommand
	}
	if err := ctx.Err(); err != nil {
		return result, err
	}
	if a == nil || command.Now.IsZero() || command.Requested < 0 || command.ReserveForGateway < 0 ||
		command.MinRemainingTTL < 0 || command.MaxHealthCheckAge < 0 || command.IdempotencyTTL < 0 ||
		(command.IdempotencyKey != "" && command.ClientID == "") ||
		(command.Fulfillment != extractionDomain.Partial && command.Fulfillment != extractionDomain.AllOrNothing) {
		return result, extractionDomain.ErrInvalidCommand
	}

	digestInput := extractionDigestInput{
		Requested: command.Requested, Fulfillment: command.Fulfillment,
		Protocols: canonicalFilter(command.Protocols), Regions: canonicalFilter(command.Regions),
		Carriers: canonicalFilter(command.Carriers), Upstreams: canonicalFilter(command.Upstreams),
	}
	requestDigest, err := extractionRequestDigest(digestInput)
	if err != nil {
		return result, err
	}
	if command.Requested == 0 && command.IdempotencyKey == "" {
		return result, nil
	}
	operationID := command.RequestID
	if operationID == "" {
		operationID, err = newOperationID()
		if err != nil {
			return result, err
		}
	}
	operationKey := a.keys.operation(digestParts(command.ClientID, operationID))
	idempotencyKey := operationKey
	hasIdempotency := 0
	if command.IdempotencyKey != "" {
		hasIdempotency = 1
		idempotencyKey = a.keys.idempotency(command.ClientID, command.IdempotencyKey)
	}
	filterPayload, err := json.Marshal(extractionFilterWire{
		Protocols: digestInput.Protocols, Regions: digestInput.Regions,
		Carriers: digestInput.Carriers, Upstreams: digestInput.Upstreams,
	})
	if err != nil {
		return result, fmt.Errorf("encode Redis extraction filters: %w", err)
	}
	keys := []string{
		a.keys.records, a.keys.unique, a.keys.idkeys, a.keys.expiry, a.keys.available,
		a.keys.inventory, a.keys.stateInventory, a.keys.owners, a.keys.ownerExpiry,
		operationKey, idempotencyKey,
	}
	keys = append(keys, a.extractionDriverKeys(digestInput)...)
	keys = append(keys, a.keys.healthDue, a.keys.healthQueued, a.keys.healthLeases,
		a.keys.healthTasks, a.keys.healthTaskExpiry, a.keys.healthRefTask, a.keys.healthUnhealthy)
	idempotencyTTL := command.IdempotencyTTL
	if idempotencyTTL == 0 {
		idempotencyTTL = defaultRedisIdempotencyTTL
	}
	scriptResult, err := runScript(ctx, a.client, extractScript, keys,
		command.Now.UnixMilli(), command.Requested, string(command.Fulfillment), command.ReserveForGateway,
		durationMillis(command.MinRemainingTTL), durationMillis(command.MaxHealthCheckAge),
		a.options.MaxCandidateScan, a.options.CleanupLimit, durationMillis(idempotencyTTL),
		operationTTLMillis(a.options.OperationTTL), requestDigest, hasIdempotency, string(filterPayload))
	if err != nil {
		return result, err
	}
	var reply extractScriptReply
	if err := decodeScriptResult(scriptResult, &reply); err != nil {
		return result, err
	}
	if reply.RequestDigest != requestDigest {
		return result, invalidScriptReply("extraction reply digest mismatch")
	}
	switch reply.Status {
	case scriptConflict:
		return result, extractionDomain.ErrIdempotencyConflict
	case scriptInsufficient:
		return result, extractionDomain.ErrInsufficientProxies
	case scriptUnavailable:
		return result, extractionDomain.ErrStoreUnavailable
	case scriptInvalid:
		return result, extractionDomain.ErrInvalidCommand
	case scriptOK:
		if reply.Record == "" {
			return result, invalidScriptReply("extraction reply omitted record")
		}
		committed, err := decodeIdempotencyRecord(reply.Record)
		if err != nil {
			return result, errors.Join(
				invalidScriptReply("extraction reply contained an invalid record"),
				fmt.Errorf("decode extraction record: %w", err),
			)
		}
		if committed.RequestDigest != requestDigest {
			return result, invalidScriptReply("extraction reply contained an invalid record")
		}
		return buildExtractionResult(committed.Result), nil
	default:
		return result, invalidScriptReply("unexpected extraction status")
	}
}

func (a *Adapter) extractionDriverKeys(input extractionDigestInput) []string {
	keys := make([]string, 0, 4)
	if len(input.Protocols) == 1 {
		keys = append(keys, a.keys.protocol(input.Protocols[0]))
	}
	if len(input.Regions) == 1 {
		keys = append(keys, a.keys.region(input.Regions[0]))
	}
	if len(input.Carriers) == 1 {
		keys = append(keys, a.keys.carrier(input.Carriers[0]))
	}
	if len(input.Upstreams) == 1 {
		keys = append(keys, a.keys.upstream(input.Upstreams[0]))
	}
	return keys
}

func extractionRequestDigest(input extractionDigestInput) (string, error) {
	payload, err := json.Marshal(input)
	if err != nil {
		return "", fmt.Errorf("encode extraction request digest: %w", err)
	}
	digest := sha256.Sum256(payload)
	return hex.EncodeToString(digest[:]), nil
}

func canonicalFilter(values []string) []string {
	if len(values) == 0 {
		return []string{}
	}
	unique := make(map[string]struct{}, len(values))
	for _, value := range values {
		unique[value] = struct{}{}
	}
	result := make([]string, 0, len(unique))
	for value := range unique {
		result = append(result, value)
	}
	sort.Strings(result)
	return result
}

func durationMillis(duration time.Duration) int64 {
	if duration <= 0 {
		return 0
	}
	return operationTTLMillis(duration)
}

func buildExtractionResult(result extractionDomain.Result) extractionDomain.Result {
	result.Items = append([]extractionDomain.Candidate(nil), result.Items...)
	for index := range result.Items {
		result.Items[index].URL = proxyURL(result.Items[index])
	}
	return result
}

func proxyURL(candidate extractionDomain.Candidate) string {
	parsed := url.URL{
		Scheme: candidate.Protocol,
		Host:   net.JoinHostPort(candidate.Host, strconv.FormatUint(uint64(candidate.Port), 10)),
	}
	if candidate.Password != "" {
		parsed.User = url.UserPassword(candidate.Username, candidate.Password)
	} else if candidate.Username != "" {
		parsed.User = url.User(candidate.Username)
	}
	return parsed.String()
}