968 lines
24 KiB
Go
968 lines
24 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"math/rand"
|
|
"net"
|
|
"net/http"
|
|
"net/netip"
|
|
"os"
|
|
"os/signal"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"syscall"
|
|
"time"
|
|
)
|
|
|
|
const (
|
|
repoDir = "cloud-provider-ip-addresses"
|
|
port = 8080
|
|
stateDir = "progress_state"
|
|
saveInterval = 30 * time.Second
|
|
cleanupInterval = 5 * time.Minute
|
|
generatorTTL = 24 * time.Hour
|
|
maxImportSize = 10 * 1024 * 1024 // 10MB
|
|
interleavedGens = 10 // Number of concurrent CIDR generators to interleave
|
|
)
|
|
|
|
// GeneratorState represents the serializable state of a generator
|
|
type GeneratorState struct {
|
|
RemainingCIDRs []string `json:"remaining_cidrs"`
|
|
CurrentGen *HostGenState `json:"current_gen,omitempty"`
|
|
ActiveGens []HostGenState `json:"active_gens,omitempty"`
|
|
TotalCIDRs int `json:"total_cidrs"`
|
|
}
|
|
|
|
type HostGenState struct {
|
|
CIDR string `json:"cidr"`
|
|
Current string `json:"current"`
|
|
Done bool `json:"done"`
|
|
}
|
|
|
|
// IPGenerator generates IPs from CIDR ranges lazily
|
|
type IPGenerator struct {
|
|
mu sync.Mutex
|
|
rng *rand.Rand
|
|
totalCIDRsCount int
|
|
remainingCIDRs []string
|
|
currentGen *hostGenerator
|
|
activeGens []*hostGenerator // Multiple active generators for interleaving
|
|
genRotationIdx int // Current rotation index
|
|
consumer string
|
|
dirty atomic.Bool
|
|
seenIPs map[string]bool // Deduplication map
|
|
}
|
|
|
|
type hostGenerator struct {
|
|
prefix netip.Prefix
|
|
current netip.Addr
|
|
last netip.Addr
|
|
done bool
|
|
}
|
|
|
|
func addrToUint32(a netip.Addr) uint32 {
|
|
b := a.As4()
|
|
return uint32(b[0])<<24 | uint32(b[1])<<16 | uint32(b[2])<<8 | uint32(b[3])
|
|
}
|
|
|
|
func uint32ToAddr(u uint32) netip.Addr {
|
|
return netip.AddrFrom4([4]byte{byte(u >> 24), byte(u >> 16), byte(u >> 8), byte(u)})
|
|
}
|
|
|
|
func newHostGenerator(cidr string) (*hostGenerator, error) {
|
|
prefix, err := netip.ParsePrefix(cidr)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
prefix = prefix.Masked()
|
|
if !prefix.IsValid() || !prefix.Addr().Is4() {
|
|
return nil, fmt.Errorf("invalid IPv4 prefix")
|
|
}
|
|
if prefix.Addr().IsMulticast() {
|
|
return nil, fmt.Errorf("multicast network")
|
|
}
|
|
|
|
ip := prefix.Addr()
|
|
maskLen := prefix.Bits()
|
|
|
|
var first, last netip.Addr
|
|
lastUint := addrToUint32(ip) | ((1 << (32 - uint(maskLen))) - 1)
|
|
last = uint32ToAddr(lastUint)
|
|
|
|
if maskLen == 32 {
|
|
first = ip
|
|
last = ip
|
|
} else if maskLen == 31 {
|
|
first = ip
|
|
// last already ip + 1
|
|
} else {
|
|
first = ip.Next()
|
|
last = last.Prev()
|
|
}
|
|
|
|
if !prefix.Contains(first) || !prefix.Contains(last) {
|
|
return nil, fmt.Errorf("invalid range")
|
|
}
|
|
|
|
return &hostGenerator{
|
|
prefix: prefix,
|
|
current: first,
|
|
last: last,
|
|
done: false,
|
|
}, nil
|
|
}
|
|
|
|
func (hg *hostGenerator) next() (string, bool) {
|
|
if hg.done {
|
|
return "", false
|
|
}
|
|
|
|
if !hg.prefix.Contains(hg.current) || addrToUint32(hg.current) > addrToUint32(hg.last) {
|
|
hg.done = true
|
|
return "", false
|
|
}
|
|
|
|
if hg.current.IsMulticast() {
|
|
hg.current = hg.current.Next()
|
|
return hg.next()
|
|
}
|
|
|
|
ip := hg.current.String()
|
|
hg.current = hg.current.Next()
|
|
|
|
return ip, true
|
|
}
|
|
|
|
func (hg *hostGenerator) getState() HostGenState {
|
|
return HostGenState{
|
|
CIDR: hg.prefix.String(),
|
|
Current: hg.current.String(),
|
|
Done: hg.done,
|
|
}
|
|
}
|
|
|
|
func newIPGenerator(s *Server, consumer string) (*IPGenerator, error) {
|
|
gen := &IPGenerator{
|
|
rng: rand.New(rand.NewSource(time.Now().UnixNano())),
|
|
consumer: consumer,
|
|
seenIPs: make(map[string]bool),
|
|
activeGens: make([]*hostGenerator, 0, interleavedGens),
|
|
}
|
|
|
|
// Try to load existing state
|
|
if err := gen.loadState(); err == nil {
|
|
log.Printf("📂 Loaded saved state for consumer: %s", consumer)
|
|
return gen, nil
|
|
}
|
|
|
|
// No saved state, initialize fresh
|
|
gen.remainingCIDRs = append([]string{}, s.allCIDRs...)
|
|
gen.rng.Shuffle(len(gen.remainingCIDRs), func(i, j int) {
|
|
gen.remainingCIDRs[i], gen.remainingCIDRs[j] = gen.remainingCIDRs[j], gen.remainingCIDRs[i]
|
|
})
|
|
gen.totalCIDRsCount = len(gen.remainingCIDRs)
|
|
gen.dirty.Store(true)
|
|
|
|
log.Printf("🆕 New generator for %s: %d total CIDRs", consumer, gen.totalCIDRsCount)
|
|
|
|
return gen, nil
|
|
}
|
|
|
|
func (g *IPGenerator) Next() (string, error) {
|
|
g.mu.Lock()
|
|
defer g.mu.Unlock()
|
|
|
|
// Ensure we have enough active generators for interleaving
|
|
for len(g.activeGens) < interleavedGens && len(g.remainingCIDRs) > 0 {
|
|
cidr := g.remainingCIDRs[0]
|
|
g.remainingCIDRs = g.remainingCIDRs[1:]
|
|
|
|
if !strings.Contains(cidr, "/") {
|
|
cidr += "/32"
|
|
}
|
|
|
|
newGen, err := newHostGenerator(cidr)
|
|
if err != nil {
|
|
g.dirty.Store(true)
|
|
continue
|
|
}
|
|
|
|
g.activeGens = append(g.activeGens, newGen)
|
|
g.dirty.Store(true)
|
|
}
|
|
|
|
// Try to get IP from rotating generators
|
|
maxAttempts := len(g.activeGens) * 100 // Avoid infinite loop
|
|
for attempt := 0; attempt < maxAttempts || len(g.activeGens) > 0; attempt++ {
|
|
if len(g.activeGens) == 0 {
|
|
if len(g.remainingCIDRs) == 0 {
|
|
return "", fmt.Errorf("no more IPs available")
|
|
}
|
|
// Refill active generators
|
|
for len(g.activeGens) < interleavedGens && len(g.remainingCIDRs) > 0 {
|
|
cidr := g.remainingCIDRs[0]
|
|
g.remainingCIDRs = g.remainingCIDRs[1:]
|
|
|
|
if !strings.Contains(cidr, "/") {
|
|
cidr += "/32"
|
|
}
|
|
|
|
newGen, err := newHostGenerator(cidr)
|
|
if err != nil {
|
|
g.dirty.Store(true)
|
|
continue
|
|
}
|
|
|
|
g.activeGens = append(g.activeGens, newGen)
|
|
g.dirty.Store(true)
|
|
}
|
|
if len(g.activeGens) == 0 {
|
|
return "", fmt.Errorf("no more IPs available")
|
|
}
|
|
}
|
|
|
|
// Round-robin through active generators
|
|
g.genRotationIdx = g.genRotationIdx % len(g.activeGens)
|
|
gen := g.activeGens[g.genRotationIdx]
|
|
|
|
ip, ok := gen.next()
|
|
if !ok {
|
|
// Remove exhausted generator
|
|
g.activeGens = append(g.activeGens[:g.genRotationIdx], g.activeGens[g.genRotationIdx+1:]...)
|
|
g.dirty.Store(true)
|
|
if g.genRotationIdx >= len(g.activeGens) && len(g.activeGens) > 0 {
|
|
g.genRotationIdx = 0
|
|
}
|
|
continue
|
|
}
|
|
|
|
// Check deduplication
|
|
if g.seenIPs[ip] {
|
|
g.genRotationIdx = (g.genRotationIdx + 1) % max(len(g.activeGens), 1)
|
|
continue
|
|
}
|
|
|
|
g.seenIPs[ip] = true
|
|
g.genRotationIdx = (g.genRotationIdx + 1) % max(len(g.activeGens), 1)
|
|
g.dirty.Store(true)
|
|
return ip, nil
|
|
}
|
|
|
|
return "", fmt.Errorf("no more unique IPs available")
|
|
}
|
|
|
|
func max(a, b int) int {
|
|
if a > b {
|
|
return a
|
|
}
|
|
return b
|
|
}
|
|
|
|
func (g *IPGenerator) buildState() GeneratorState {
|
|
// Assumes mu is held
|
|
state := GeneratorState{
|
|
RemainingCIDRs: append([]string{}, g.remainingCIDRs...),
|
|
TotalCIDRs: g.totalCIDRsCount,
|
|
}
|
|
if g.currentGen != nil && !g.currentGen.done {
|
|
state.CurrentGen = &HostGenState{
|
|
CIDR: g.currentGen.prefix.String(),
|
|
Current: g.currentGen.current.String(),
|
|
Done: false,
|
|
}
|
|
}
|
|
// Save activeGens to preserve interleaving state
|
|
if len(g.activeGens) > 0 {
|
|
state.ActiveGens = make([]HostGenState, 0, len(g.activeGens))
|
|
for _, gen := range g.activeGens {
|
|
if gen != nil && !gen.done {
|
|
state.ActiveGens = append(state.ActiveGens, HostGenState{
|
|
CIDR: gen.prefix.String(),
|
|
Current: gen.current.String(),
|
|
Done: false,
|
|
})
|
|
}
|
|
}
|
|
}
|
|
return state
|
|
}
|
|
|
|
func (g *IPGenerator) getState() GeneratorState {
|
|
g.mu.Lock()
|
|
defer g.mu.Unlock()
|
|
return g.buildState()
|
|
}
|
|
|
|
func (g *IPGenerator) saveState() error {
|
|
g.mu.Lock()
|
|
if !g.dirty.Load() {
|
|
g.mu.Unlock()
|
|
return nil
|
|
}
|
|
state := g.buildState()
|
|
g.dirty.Store(false)
|
|
g.mu.Unlock()
|
|
|
|
// Ensure state directory exists
|
|
if err := os.MkdirAll(stateDir, 0755); err != nil {
|
|
return fmt.Errorf("failed to create state directory: %w", err)
|
|
}
|
|
|
|
// Use hash of consumer as filename
|
|
hash := sha256.Sum256([]byte(g.consumer))
|
|
filename := hex.EncodeToString(hash[:])
|
|
filePath := filepath.Join(stateDir, filename+".json")
|
|
|
|
// Write to temp file first, then rename
|
|
tempPath := filePath + ".tmp"
|
|
file, err := os.Create(tempPath)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to create temp state file: %w", err)
|
|
}
|
|
defer file.Close()
|
|
|
|
encoder := json.NewEncoder(file)
|
|
encoder.SetIndent("", " ")
|
|
if err := encoder.Encode(state); err != nil {
|
|
os.Remove(tempPath)
|
|
return fmt.Errorf("failed to encode state: %w", err)
|
|
}
|
|
|
|
if err := os.Rename(tempPath, filePath); err != nil {
|
|
os.Remove(tempPath)
|
|
return fmt.Errorf("failed to rename state file: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (g *IPGenerator) loadState() error {
|
|
// Use hash of consumer as filename
|
|
hash := sha256.Sum256([]byte(g.consumer))
|
|
filename := hex.EncodeToString(hash[:])
|
|
filePath := filepath.Join(stateDir, filename+".json")
|
|
|
|
file, err := os.Open(filePath)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer file.Close()
|
|
|
|
var state GeneratorState
|
|
if err := json.NewDecoder(file).Decode(&state); err != nil {
|
|
return fmt.Errorf("failed to decode state: %w", err)
|
|
}
|
|
|
|
// Restore state
|
|
g.remainingCIDRs = state.RemainingCIDRs
|
|
g.totalCIDRsCount = state.TotalCIDRs
|
|
|
|
if state.CurrentGen != nil {
|
|
gen, err := newHostGenerator(state.CurrentGen.CIDR)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
gen.current, err = netip.ParseAddr(state.CurrentGen.Current)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
gen.done = state.CurrentGen.Done
|
|
g.currentGen = gen
|
|
}
|
|
|
|
// Restore activeGens to preserve interleaving state
|
|
if len(state.ActiveGens) > 0 {
|
|
g.activeGens = make([]*hostGenerator, 0, len(state.ActiveGens))
|
|
for _, genState := range state.ActiveGens {
|
|
gen, err := newHostGenerator(genState.CIDR)
|
|
if err != nil {
|
|
log.Printf("⚠️ Failed to restore activeGen %s: %v", genState.CIDR, err)
|
|
continue
|
|
}
|
|
gen.current, err = netip.ParseAddr(genState.Current)
|
|
if err != nil {
|
|
log.Printf("⚠️ Failed to parse current IP for activeGen %s: %v", genState.CIDR, err)
|
|
continue
|
|
}
|
|
gen.done = genState.Done
|
|
g.activeGens = append(g.activeGens, gen)
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// Server holds per-consumer generators
|
|
type Server struct {
|
|
generators map[string]*IPGenerator
|
|
lastAccess map[string]time.Time
|
|
allCIDRs []string
|
|
// MULTI-INSTANCE LIMITATION: globalSeen is instance-local, not shared across
|
|
// multiple input_service instances. In multi-instance deployments, either:
|
|
// 1. Use session affinity for ping workers (same worker always talks to same instance)
|
|
// 2. POST discovered hops to ALL input_service instances, or
|
|
// 3. Implement shared deduplication backend (Redis, database, etc.)
|
|
// Without this, different instances may serve duplicate hops.
|
|
globalSeen map[string]bool // Global deduplication across all sources (instance-local)
|
|
mu sync.RWMutex
|
|
stopSaver chan struct{}
|
|
stopCleanup chan struct{}
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
func newServer() *Server {
|
|
s := &Server{
|
|
generators: make(map[string]*IPGenerator),
|
|
lastAccess: make(map[string]time.Time),
|
|
globalSeen: make(map[string]bool),
|
|
stopSaver: make(chan struct{}),
|
|
stopCleanup: make(chan struct{}),
|
|
}
|
|
|
|
if err := s.loadAllCIDRs(); err != nil {
|
|
log.Fatalf("❌ Failed to load CIDRs: %v", err)
|
|
}
|
|
|
|
s.startPeriodicSaver()
|
|
s.startCleanup()
|
|
|
|
return s
|
|
}
|
|
|
|
func (s *Server) loadAllCIDRs() error {
|
|
// Find all IP files
|
|
var fileList []string
|
|
err := filepath.Walk(repoDir, func(path string, info os.FileInfo, err error) error {
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !info.IsDir() && strings.HasSuffix(path, ".txt") && strings.Contains(strings.ToLower(path), "ips") {
|
|
fileList = append(fileList, path)
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return fmt.Errorf("failed to scan repo directory: %w", err)
|
|
}
|
|
|
|
if len(fileList) == 0 {
|
|
return fmt.Errorf("no IP files found in %s", repoDir)
|
|
}
|
|
|
|
// Load all CIDRs
|
|
for _, path := range fileList {
|
|
file, err := os.Open(path)
|
|
if err != nil {
|
|
log.Printf("⚠️ Failed to open %s: %v", path, err)
|
|
continue
|
|
}
|
|
scanner := bufio.NewScanner(file)
|
|
for scanner.Scan() {
|
|
line := strings.TrimSpace(scanner.Text())
|
|
if line == "" || strings.HasPrefix(line, "#") {
|
|
continue
|
|
}
|
|
fields := strings.Fields(line)
|
|
for _, field := range fields {
|
|
if field != "" {
|
|
// Accept CIDRs (contains /) or valid IP addresses
|
|
isCIDR := strings.Contains(field, "/")
|
|
isValidIP := false
|
|
if !isCIDR {
|
|
if addr, err := netip.ParseAddr(field); err == nil && addr.IsValid() {
|
|
isValidIP = true
|
|
}
|
|
}
|
|
if isCIDR || isValidIP {
|
|
s.allCIDRs = append(s.allCIDRs, field)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
file.Close()
|
|
if err := scanner.Err(); err != nil {
|
|
log.Printf("⚠️ Error reading %s: %v", path, err)
|
|
}
|
|
}
|
|
|
|
log.Printf("📁 Found %d IP files", len(fileList))
|
|
log.Printf("📊 Total CIDRs discovered: %d", len(s.allCIDRs))
|
|
|
|
return nil
|
|
}
|
|
|
|
func (s *Server) startPeriodicSaver() {
|
|
s.wg.Add(1)
|
|
go func() {
|
|
defer s.wg.Done()
|
|
ticker := time.NewTicker(saveInterval)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
s.saveAllStates()
|
|
case <-s.stopSaver:
|
|
s.saveAllStates()
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
func (s *Server) startCleanup() {
|
|
s.wg.Add(1)
|
|
go func() {
|
|
defer s.wg.Done()
|
|
ticker := time.NewTicker(cleanupInterval)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
s.cleanupOldGenerators()
|
|
case <-s.stopCleanup:
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
func (s *Server) cleanupOldGenerators() {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
for consumer, t := range s.lastAccess {
|
|
if now.Sub(t) > generatorTTL {
|
|
delete(s.generators, consumer)
|
|
delete(s.lastAccess, consumer)
|
|
|
|
// Remove state file
|
|
hash := sha256.Sum256([]byte(consumer))
|
|
fn := hex.EncodeToString(hash[:]) + ".json"
|
|
p := filepath.Join(stateDir, fn)
|
|
if err := os.Remove(p); err != nil && !os.IsNotExist(err) {
|
|
log.Printf("⚠️ Failed to remove state file %s: %v", p, err)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *Server) saveAllStates() {
|
|
s.mu.RLock()
|
|
gens := make([]*IPGenerator, 0, len(s.generators))
|
|
for _, gen := range s.generators {
|
|
gens = append(gens, gen)
|
|
}
|
|
s.mu.RUnlock()
|
|
|
|
for _, gen := range gens {
|
|
if err := gen.saveState(); err != nil {
|
|
log.Printf("⚠️ Failed to save state for %s: %v", gen.consumer, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *Server) shutdown() {
|
|
close(s.stopSaver)
|
|
close(s.stopCleanup)
|
|
s.wg.Wait()
|
|
log.Println("💾 All states saved")
|
|
}
|
|
|
|
func (s *Server) getGenerator(consumer string) (*IPGenerator, error) {
|
|
s.mu.RLock()
|
|
gen, exists := s.generators[consumer]
|
|
s.mu.RUnlock()
|
|
|
|
if exists {
|
|
s.mu.Lock()
|
|
s.lastAccess[consumer] = time.Now()
|
|
s.mu.Unlock()
|
|
return gen, nil
|
|
}
|
|
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
// Double-check
|
|
if gen, exists := s.generators[consumer]; exists {
|
|
s.lastAccess[consumer] = time.Now()
|
|
return gen, nil
|
|
}
|
|
|
|
newGen, err := newIPGenerator(s, consumer)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
s.generators[consumer] = newGen
|
|
s.lastAccess[consumer] = time.Now()
|
|
log.Printf("🆕 New consumer: %s", consumer)
|
|
|
|
return newGen, nil
|
|
}
|
|
|
|
func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodGet {
|
|
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
|
|
addrPort, err := netip.ParseAddrPort(r.RemoteAddr)
|
|
consumerStr := r.RemoteAddr
|
|
if err == nil {
|
|
consumerStr = addrPort.Addr().String()
|
|
} else {
|
|
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
|
if err == nil {
|
|
consumerStr = host
|
|
}
|
|
}
|
|
|
|
gen, err := s.getGenerator(consumerStr)
|
|
if err != nil {
|
|
log.Printf("❌ Failed to get generator for %s: %v", consumerStr, err)
|
|
http.Error(w, "Internal server error", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
|
|
ip, err := gen.Next()
|
|
if err != nil {
|
|
log.Printf("❌ Failed to get IP for %s: %v", consumerStr, err)
|
|
http.Error(w, "No more IPs available", http.StatusServiceUnavailable)
|
|
return
|
|
}
|
|
|
|
log.Printf("📤 Serving IP to %s: %s", consumerStr, ip)
|
|
|
|
w.Header().Set("Content-Type", "text/plain")
|
|
fmt.Fprintf(w, "%s\n", ip)
|
|
}
|
|
|
|
type ConsumerStatus struct {
|
|
Consumer string `json:"consumer"`
|
|
RemainingCIDRs int `json:"remaining_cidrs"`
|
|
HasActiveGen bool `json:"has_active_gen"`
|
|
TotalCIDRs int `json:"total_cidrs"`
|
|
}
|
|
|
|
type StatusResponse struct {
|
|
TotalConsumers int `json:"total_consumers"`
|
|
Consumers []ConsumerStatus `json:"consumers"`
|
|
StateDirectory string `json:"state_directory"`
|
|
SaveInterval string `json:"save_interval"`
|
|
}
|
|
|
|
func (s *Server) handleStatus(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodGet {
|
|
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
|
|
response := StatusResponse{
|
|
TotalConsumers: len(s.generators),
|
|
Consumers: make([]ConsumerStatus, 0, len(s.generators)),
|
|
StateDirectory: stateDir,
|
|
SaveInterval: saveInterval.String(),
|
|
}
|
|
|
|
for consumer, gen := range s.generators {
|
|
gen.mu.Lock()
|
|
status := ConsumerStatus{
|
|
Consumer: consumer,
|
|
RemainingCIDRs: len(gen.remainingCIDRs),
|
|
HasActiveGen: gen.currentGen != nil,
|
|
TotalCIDRs: gen.totalCIDRsCount,
|
|
}
|
|
gen.mu.Unlock()
|
|
|
|
response.Consumers = append(response.Consumers, status)
|
|
}
|
|
|
|
w.Header().Set("Content-Type", "application/json")
|
|
encoder := json.NewEncoder(w)
|
|
encoder.SetIndent("", " ")
|
|
if err := encoder.Encode(response); err != nil {
|
|
log.Printf("❌ Failed to encode status response: %v", err)
|
|
}
|
|
}
|
|
|
|
type ExportResponse struct {
|
|
ExportedAt time.Time `json:"exported_at"`
|
|
States map[string]GeneratorState `json:"states"`
|
|
}
|
|
|
|
func (s *Server) handleExport(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodGet {
|
|
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
|
|
// Force save all current states first
|
|
s.saveAllStates()
|
|
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
|
|
response := ExportResponse{
|
|
ExportedAt: time.Now(),
|
|
States: make(map[string]GeneratorState, len(s.generators)),
|
|
}
|
|
|
|
for consumer, gen := range s.generators {
|
|
response.States[consumer] = gen.getState()
|
|
}
|
|
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=state-export-%s.json",
|
|
time.Now().Format("2006-01-02-150405")))
|
|
|
|
encoder := json.NewEncoder(w)
|
|
encoder.SetIndent("", " ")
|
|
if err := encoder.Encode(response); err != nil {
|
|
log.Printf("❌ Failed to encode export response: %v", err)
|
|
}
|
|
}
|
|
|
|
func (s *Server) handleImport(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodPost {
|
|
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
|
|
reader := http.MaxBytesReader(w, r.Body, maxImportSize)
|
|
defer r.Body.Close()
|
|
|
|
var exportData ExportResponse
|
|
if err := json.NewDecoder(reader).Decode(&exportData); err != nil {
|
|
if err == io.EOF || strings.Contains(err.Error(), "EOF") {
|
|
http.Error(w, "Invalid or empty request body", http.StatusBadRequest)
|
|
} else {
|
|
http.Error(w, fmt.Sprintf("Failed to decode import data: %v", err), http.StatusBadRequest)
|
|
}
|
|
return
|
|
}
|
|
|
|
// Ensure state directory exists
|
|
if err := os.MkdirAll(stateDir, 0755); err != nil {
|
|
http.Error(w, fmt.Sprintf("Failed to create state directory: %v", err), http.StatusInternalServerError)
|
|
return
|
|
}
|
|
|
|
imported := 0
|
|
failed := 0
|
|
|
|
for consumer, state := range exportData.States {
|
|
// Use hash for filename
|
|
hash := sha256.Sum256([]byte(consumer))
|
|
filename := hex.EncodeToString(hash[:])
|
|
filePath := filepath.Join(stateDir, filename+".json")
|
|
|
|
file, err := os.Create(filePath)
|
|
if err != nil {
|
|
log.Printf("⚠️ Failed to create state file for %s: %v", consumer, err)
|
|
failed++
|
|
continue
|
|
}
|
|
|
|
encoder := json.NewEncoder(file)
|
|
encoder.SetIndent("", " ")
|
|
if err := encoder.Encode(state); err != nil {
|
|
file.Close()
|
|
log.Printf("⚠️ Failed to encode state for %s: %v", consumer, err)
|
|
failed++
|
|
os.Remove(filePath)
|
|
continue
|
|
}
|
|
|
|
file.Close()
|
|
imported++
|
|
}
|
|
|
|
response := map[string]interface{}{
|
|
"imported": imported,
|
|
"failed": failed,
|
|
"total": len(exportData.States),
|
|
"message": fmt.Sprintf("Successfully imported %d consumer states", imported),
|
|
}
|
|
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(response)
|
|
|
|
log.Printf("📥 Imported %d consumer states (%d failed)", imported, failed)
|
|
}
|
|
|
|
// HopsRequest is the payload from output_service
|
|
type HopsRequest struct {
|
|
Hops []string `json:"hops"`
|
|
}
|
|
|
|
func (s *Server) handleHops(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodPost {
|
|
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
|
|
defer r.Body.Close()
|
|
|
|
var req HopsRequest
|
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
http.Error(w, "Invalid JSON", http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
added := 0
|
|
duplicates := 0
|
|
|
|
for _, hop := range req.Hops {
|
|
// Validate IP
|
|
addr, err := netip.ParseAddr(hop)
|
|
if err != nil {
|
|
log.Printf("⚠️ Invalid hop IP: %s", hop)
|
|
continue
|
|
}
|
|
|
|
// Skip if not IPv4
|
|
if !addr.Is4() {
|
|
continue
|
|
}
|
|
|
|
// Skip multicast, private, loopback
|
|
if addr.IsMulticast() || addr.IsLoopback() || addr.IsPrivate() {
|
|
continue
|
|
}
|
|
|
|
// Check global deduplication
|
|
if s.globalSeen[hop] {
|
|
duplicates++
|
|
continue
|
|
}
|
|
|
|
// Add to global pool
|
|
s.globalSeen[hop] = true
|
|
s.allCIDRs = append(s.allCIDRs, hop)
|
|
added++
|
|
}
|
|
|
|
log.Printf("🔍 Received %d hops: %d new, %d duplicates", len(req.Hops), added, duplicates)
|
|
|
|
response := map[string]interface{}{
|
|
"status": "ok",
|
|
"received": len(req.Hops),
|
|
"added": added,
|
|
"duplicates": duplicates,
|
|
}
|
|
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(response)
|
|
}
|
|
|
|
// ServiceInfo represents service metadata for discovery
|
|
type ServiceInfo struct {
|
|
ServiceType string `json:"service_type"`
|
|
Version string `json:"version"`
|
|
Name string `json:"name"`
|
|
InstanceID string `json:"instance_id"`
|
|
Capabilities []string `json:"capabilities"`
|
|
}
|
|
|
|
func (s *Server) handleServiceInfo(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodGet {
|
|
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
|
|
hostname, _ := os.Hostname()
|
|
if hostname == "" {
|
|
hostname = "unknown"
|
|
}
|
|
|
|
info := ServiceInfo{
|
|
ServiceType: "input",
|
|
Version: "1.0.0",
|
|
Name: "http_input_service",
|
|
InstanceID: hostname,
|
|
Capabilities: []string{"target_generation", "cidr_import", "hop_discovery"},
|
|
}
|
|
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(info)
|
|
}
|
|
|
|
func main() {
|
|
// Check if repo directory exists
|
|
if _, err := os.Stat(repoDir); os.IsNotExist(err) {
|
|
log.Fatalf("❌ Error: Directory '%s' not found", repoDir)
|
|
}
|
|
|
|
server := newServer()
|
|
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("/", server.handleRequest)
|
|
mux.HandleFunc("/status", server.handleStatus)
|
|
mux.HandleFunc("/service-info", server.handleServiceInfo)
|
|
mux.HandleFunc("/export", server.handleExport)
|
|
mux.HandleFunc("/import", server.handleImport)
|
|
mux.HandleFunc("/hops", server.handleHops)
|
|
|
|
httpServer := &http.Server{
|
|
Addr: fmt.Sprintf(":%d", port),
|
|
Handler: mux,
|
|
ReadTimeout: 10 * time.Second,
|
|
WriteTimeout: 10 * time.Second,
|
|
IdleTimeout: 60 * time.Second,
|
|
}
|
|
|
|
// Graceful shutdown handling
|
|
go func() {
|
|
sigChan := make(chan os.Signal, 1)
|
|
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
|
|
<-sigChan
|
|
|
|
log.Println("\n🛑 Shutting down gracefully...")
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
// Stop savers and save final state
|
|
server.shutdown()
|
|
|
|
if err := httpServer.Shutdown(ctx); err != nil {
|
|
log.Printf("❌ Error during shutdown: %v", err)
|
|
}
|
|
}()
|
|
|
|
log.Printf("🌐 HTTP Input Server running on http://localhost:%d", port)
|
|
log.Printf(" Serving individual IPv4 host addresses lazily")
|
|
log.Printf(" In highly mixed random order per consumer")
|
|
log.Printf(" 🔄 Interleaving %d CIDRs to avoid same-subnet consecutive IPs", interleavedGens)
|
|
log.Printf(" 💾 Progress saved every %v to '%s' directory", saveInterval, stateDir)
|
|
log.Printf(" 📊 Status endpoint: http://localhost:%d/status", port)
|
|
log.Printf(" 📤 Export endpoint: http://localhost:%d/export", port)
|
|
log.Printf(" 📥 Import endpoint: http://localhost:%d/import (POST)", port)
|
|
log.Printf(" 🔍 Hops endpoint: http://localhost:%d/hops (POST)", port)
|
|
log.Printf(" Press Ctrl+C to stop")
|
|
|
|
if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
|
log.Fatalf("❌ Server error: %v", err)
|
|
}
|
|
|
|
log.Println("✅ Server stopped cleanly")
|
|
} |