golcg

- Unnamed repository; edit this file 'description' to name the repository.
git clone git://git.acid.vegas/-c.git
golcg.go (2954B)
      1 package golcg
      2 
      3 import (
      4 	"errors"
      5 	"fmt"
      6 	"math/rand"
      7 	"net"
      8 	"os"
      9 	"path/filepath"
     10 	"strings"
     11 	"time"
     12 )
     13 
     14 type LCG struct {
     15 	M       uint32
     16 	A       uint32
     17 	C       uint32
     18 	Current uint32
     19 }
     20 
     21 func NewLCG(seed int, m uint32) *LCG {
     22 	return &LCG{
     23 		M:       1<<32 - 1,
     24 		A:       1664525,
     25 		C:       1013904223,
     26 		Current: uint32(seed),
     27 	}
     28 }
     29 
     30 func (l *LCG) Next() uint32 {
     31 	l.Current = (l.A*l.Current + l.C) % l.M
     32 	return l.Current
     33 }
     34 
     35 type IPRange struct {
     36 	Start uint32
     37 	Total uint32
     38 }
     39 
     40 func NewIPRange(cidr string) (*IPRange, error) {
     41 	_, network, err := net.ParseCIDR(cidr)
     42 	if err != nil {
     43 		return nil, err
     44 	}
     45 
     46 	start := ipToUint32(network.IP)
     47 	ones, bits := network.Mask.Size()
     48 	hostBits := uint(bits - ones)
     49 
     50 	var total uint32
     51 	if hostBits == 32 {
     52 		total = 0
     53 	} else {
     54 		total = 1 << hostBits
     55 	}
     56 
     57 	return &IPRange{
     58 		Start: start,
     59 		Total: total,
     60 	}, nil
     61 }
     62 
     63 func (r *IPRange) GetIPAtIndex(index uint32) (string, error) {
     64 	if r.Total > 0 && index >= r.Total {
     65 		return "", errors.New("IP index out of range")
     66 	}
     67 
     68 	ip := uint32ToIP(r.Start + index)
     69 	return ip.String(), nil
     70 }
     71 
     72 func ipToUint32(ip net.IP) uint32 {
     73 	ip = ip.To4()
     74 	return uint32(ip[0])<<24 | uint32(ip[1])<<16 | uint32(ip[2])<<8 | uint32(ip[3])
     75 }
     76 
     77 func uint32ToIP(n uint32) net.IP {
     78 	ip := make(net.IP, 4)
     79 	ip[0] = byte(n >> 24)
     80 	ip[1] = byte(n >> 16)
     81 	ip[2] = byte(n >> 8)
     82 	ip[3] = byte(n)
     83 	return ip
     84 }
     85 
     86 func SaveState(seed int, cidr string, shard int, total int, lcgCurrent uint32) error {
     87 	fileName := fmt.Sprintf("golcg_%d_%s_%d_%d.state", seed, strings.Replace(cidr, "/", "_", -1), shard, total)
     88 	stateFile := filepath.Join(os.TempDir(), fileName)
     89 
     90 	return os.WriteFile(stateFile, []byte(fmt.Sprintf("%d", lcgCurrent)), 0644)
     91 }
     92 
     93 func IPStream(cidr string, shardNum, totalShards, seed int, state *uint32) (<-chan string, error) {
     94 	ipRange, err := NewIPRange(cidr)
     95 	if err != nil {
     96 		return nil, err
     97 	}
     98 
     99 	shardIndex := shardNum - 1
    100 
    101 	if seed == 0 {
    102 		rand.Seed(time.Now().UnixNano())
    103 		seed = rand.Intn(1<<32 - 1)
    104 	}
    105 
    106 	lcg := NewLCG(seed+shardIndex, 1<<32-1)
    107 	if state != nil {
    108 		lcg.Current = *state
    109 	}
    110 
    111 	var shardSize uint32
    112 	if ipRange.Total == 0 {
    113 		shardSize = (1<<32 - 1) / uint32(totalShards)
    114 		if uint32(shardIndex) < uint32(totalShards-1) {
    115 			shardSize++
    116 		}
    117 	} else {
    118 		shardSize = ipRange.Total / uint32(totalShards)
    119 		if uint32(shardIndex) < ipRange.Total%uint32(totalShards) {
    120 			shardSize++
    121 		}
    122 	}
    123 
    124 	out := make(chan string, 1000)
    125 	go func() {
    126 		defer close(out)
    127 		remaining := shardSize
    128 
    129 		for remaining > 0 {
    130 			next := lcg.Next()
    131 			var index uint32
    132 			if ipRange.Total > 0 {
    133 				index = next % ipRange.Total
    134 			} else {
    135 				index = next
    136 			}
    137 
    138 			if totalShards == 1 || index%uint32(totalShards) == uint32(shardIndex) {
    139 				ip, err := ipRange.GetIPAtIndex(index)
    140 				if err != nil {
    141 					continue
    142 				}
    143 				out <- ip
    144 				remaining--
    145 
    146 				if remaining%1000 == 0 {
    147 					SaveState(seed, cidr, shardNum, totalShards, lcg.Current)
    148 				}
    149 			}
    150 		}
    151 	}()
    152 
    153 	return out, nil
    154 }