Fixed few memory leaks. Implement testing of the functionality.

input_service/http_input_service_test.go

@@ -0,0 +1,313 @@
+package main
+
+import (
+	"encoding/json"
+	"net/netip"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+)
+
+// TestIPParsingDoesNotPanic verifies that invalid IPs don't cause panics
+func TestIPParsingDoesNotPanic(t *testing.T) {
+	testCases := []string{
+		"not-an-ip",
+		"999.999.999.999",
+		"192.168.1",
+		"",
+		"192.168.1.1.1",
+		"hello world",
+		"2001:db8::1", // IPv6 (should be filtered)
+	}
+
+	// This test passes if it doesn't panic
+	for _, testIP := range testCases {
+		func() {
+			defer func() {
+				if r := recover(); r != nil {
+					t.Errorf("Parsing %q caused panic: %v", testIP, r)
+				}
+			}()
+
+			// Test the safe parsing logic
+			addr, err := netip.ParseAddr(testIP)
+			if err == nil && addr.IsValid() {
+				// Valid IP, this is fine
+			}
+		}()
+	}
+}
+
+// TestStateSerializationPreservesActiveGens verifies activeGens are saved/restored
+func TestStateSerializationPreservesActiveGens(t *testing.T) {
+	// Create a temporary server for testing
+	s := &Server{
+		allCIDRs:   []string{"192.0.2.0/24", "198.51.100.0/24", "203.0.113.0/24"},
+		globalSeen: make(map[string]bool),
+	}
+
+	// Create a generator with activeGens
+	gen, err := newIPGenerator(s, "test-consumer")
+	if err != nil {
+		t.Fatalf("Failed to create generator: %v", err)
+	}
+
+	// Generate some IPs to populate activeGens
+	for i := 0; i < 15; i++ {
+		_, err := gen.Next()
+		if err != nil {
+			break
+		}
+	}
+
+	// Verify we have activeGens
+	if len(gen.activeGens) == 0 {
+		t.Log("Warning: No activeGens created, test may not be comprehensive")
+	}
+
+	originalActiveGensCount := len(gen.activeGens)
+
+	// Build state
+	gen.mu.Lock()
+	state := gen.buildState()
+	gen.mu.Unlock()
+
+	// Verify activeGens were serialized
+	if len(state.ActiveGens) == 0 && originalActiveGensCount > 0 {
+		t.Errorf("ActiveGens not serialized: had %d active gens but state has 0", originalActiveGensCount)
+	}
+
+	// Create new generator and restore state
+	gen2, err := newIPGenerator(s, "test-consumer-2")
+	if err != nil {
+		t.Fatalf("Failed to create second generator: %v", err)
+	}
+
+	// Manually restore state (simulating loadState)
+	gen2.mu.Lock()
+	gen2.remainingCIDRs = state.RemainingCIDRs
+	gen2.totalCIDRsCount = state.TotalCIDRs
+
+	// Restore activeGens
+	if len(state.ActiveGens) > 0 {
+		gen2.activeGens = make([]*hostGenerator, 0, len(state.ActiveGens))
+		for _, genState := range state.ActiveGens {
+			hg, err := newHostGenerator(genState.CIDR)
+			if err != nil {
+				continue
+			}
+			hg.current, err = netip.ParseAddr(genState.Current)
+			if err != nil {
+				continue
+			}
+			hg.done = genState.Done
+			gen2.activeGens = append(gen2.activeGens, hg)
+		}
+	}
+	gen2.mu.Unlock()
+
+	// Verify activeGens were restored
+	if len(gen2.activeGens) != len(state.ActiveGens) {
+		t.Errorf("ActiveGens restoration failed: expected %d, got %d", len(state.ActiveGens), len(gen2.activeGens))
+	}
+}
+
+// TestGeneratorStateJSONSerialization verifies state can be marshaled/unmarshaled
+func TestGeneratorStateJSONSerialization(t *testing.T) {
+	state := GeneratorState{
+		RemainingCIDRs: []string{"192.0.2.0/24", "198.51.100.0/24"},
+		CurrentGen: &HostGenState{
+			CIDR:    "203.0.113.0/24",
+			Current: "203.0.113.10",
+			Done:    false,
+		},
+		ActiveGens: []HostGenState{
+			{CIDR: "192.0.2.0/24", Current: "192.0.2.5", Done: false},
+			{CIDR: "198.51.100.0/24", Current: "198.51.100.20", Done: false},
+		},
+		TotalCIDRs: 10,
+	}
+
+	// Marshal
+	data, err := json.Marshal(state)
+	if err != nil {
+		t.Fatalf("Failed to marshal state: %v", err)
+	}
+
+	// Unmarshal
+	var restored GeneratorState
+	if err := json.Unmarshal(data, &restored); err != nil {
+		t.Fatalf("Failed to unmarshal state: %v", err)
+	}
+
+	// Verify
+	if len(restored.RemainingCIDRs) != len(state.RemainingCIDRs) {
+		t.Error("RemainingCIDRs count mismatch")
+	}
+
+	if len(restored.ActiveGens) != len(state.ActiveGens) {
+		t.Errorf("ActiveGens count mismatch: expected %d, got %d", len(state.ActiveGens), len(restored.ActiveGens))
+	}
+
+	if restored.TotalCIDRs != state.TotalCIDRs {
+		t.Error("TotalCIDRs mismatch")
+	}
+
+	if restored.CurrentGen == nil {
+		t.Error("CurrentGen was not restored")
+	} else if restored.CurrentGen.CIDR != state.CurrentGen.CIDR {
+		t.Error("CurrentGen CIDR mismatch")
+	}
+}
+
+// TestHostGeneratorBasic verifies basic IP generation
+func TestHostGeneratorBasic(t *testing.T) {
+	gen, err := newHostGenerator("192.0.2.0/30")
+	if err != nil {
+		t.Fatalf("Failed to create host generator: %v", err)
+	}
+
+	// /30 network has 4 addresses: .0 (network), .1 and .2 (hosts), .3 (broadcast)
+	// We should get .1 and .2
+	ips := make([]string, 0)
+	for {
+		ip, ok := gen.next()
+		if !ok {
+			break
+		}
+		ips = append(ips, ip)
+	}
+
+	expectedCount := 2
+	if len(ips) != expectedCount {
+		t.Errorf("Expected %d IPs from /30 network, got %d: %v", expectedCount, len(ips), ips)
+	}
+
+	// Verify we got valid IPs
+	for _, ip := range ips {
+		addr, err := netip.ParseAddr(ip)
+		if err != nil || !addr.IsValid() {
+			t.Errorf("Generated invalid IP: %s", ip)
+		}
+	}
+}
+
+// TestGlobalDeduplication verifies that globalSeen prevents duplicates
+func TestGlobalDeduplication(t *testing.T) {
+	s := &Server{
+		allCIDRs:   []string{"192.0.2.0/29"},
+		globalSeen: make(map[string]bool),
+	}
+
+	// Mark some IPs as seen
+	s.globalSeen["192.0.2.1"] = true
+	s.globalSeen["192.0.2.2"] = true
+
+	if !s.globalSeen["192.0.2.1"] {
+		t.Error("IP should be marked as seen")
+	}
+
+	if s.globalSeen["192.0.2.100"] {
+		t.Error("Unseen IP should not be in globalSeen")
+	}
+}
+
+// TestIPGeneratorConcurrency verifies thread-safe generator access
+func TestIPGeneratorConcurrency(t *testing.T) {
+	s := &Server{
+		allCIDRs:   []string{"192.0.2.0/24", "198.51.100.0/24"},
+		globalSeen: make(map[string]bool),
+	}
+
+	gen, err := newIPGenerator(s, "test-consumer")
+	if err != nil {
+		t.Fatalf("Failed to create generator: %v", err)
+	}
+
+	done := make(chan bool)
+	errors := make(chan error, 10)
+
+	// Spawn multiple goroutines calling Next() concurrently
+	for i := 0; i < 10; i++ {
+		go func() {
+			for j := 0; j < 50; j++ {
+				_, err := gen.Next()
+				if err != nil {
+					errors <- err
+					break
+				}
+			}
+			done <- true
+		}()
+	}
+
+	// Wait for all goroutines
+	for i := 0; i < 10; i++ {
+		<-done
+	}
+
+	close(errors)
+	if len(errors) > 0 {
+		for err := range errors {
+			t.Errorf("Concurrent access error: %v", err)
+		}
+	}
+}
+
+// TestStatePersistence verifies state can be saved and loaded from disk
+func TestStatePersistence(t *testing.T) {
+	// Use default stateDir (progress_state) for this test
+	// Ensure it exists
+	if err := os.MkdirAll(stateDir, 0755); err != nil {
+		t.Fatalf("Failed to create state dir: %v", err)
+	}
+
+	s := &Server{
+		allCIDRs:   []string{"192.0.2.0/24"},
+		globalSeen: make(map[string]bool),
+	}
+
+	gen, err := newIPGenerator(s, "test-persistence-"+time.Now().Format("20060102150405"))
+	if err != nil {
+		t.Fatalf("Failed to create generator: %v", err)
+	}
+
+	// Generate some IPs
+	for i := 0; i < 10; i++ {
+		_, err := gen.Next()
+		if err != nil {
+			break
+		}
+	}
+
+	// Save state
+	if err := gen.saveState(); err != nil {
+		t.Fatalf("Failed to save state: %v", err)
+	}
+
+	// Verify state file was created
+	files, err := filepath.Glob(filepath.Join(stateDir, "*.json"))
+	if err != nil {
+		t.Fatalf("Failed to list state files: %v", err)
+	}
+
+	if len(files) == 0 {
+		t.Error("No state file was created")
+	}
+
+	// Create new generator and load state
+	gen2, err := newIPGenerator(s, gen.consumer)
+	if err != nil {
+		t.Fatalf("Failed to create second generator: %v", err)
+	}
+
+	if err := gen2.loadState(); err != nil {
+		t.Fatalf("Failed to load state: %v", err)
+	}
+
+	// Verify state was loaded (should have remaining CIDRs and progress)
+	if len(gen2.remainingCIDRs) == 0 && len(gen.remainingCIDRs) > 0 {
+		t.Error("State was not properly restored")
+	}
+}