LFG

internal/proxy/loadbalancer.go

@@ -0,0 +1,217 @@
+package proxy
+
+import (
+	"math/rand"
+	"sync"
+	"time"
+
+	"github.com/iwasforcedtobehere/goRZ/internal/config"
+)
+
+// RoundRobinLoadBalancer implements round-robin load balancing
+type RoundRobinLoadBalancer struct {
+	targets []*config.TargetConfig
+	current int
+	mu      sync.Mutex
+}
+
+// NewRoundRobinLoadBalancer creates a new round-robin load balancer
+func NewRoundRobinLoadBalancer(targets []config.TargetConfig) *RoundRobinLoadBalancer {
+	t := make([]*config.TargetConfig, len(targets))
+	for i := range targets {
+		t[i] = &targets[i]
+	}
+	return &RoundRobinLoadBalancer{
+		targets: t,
+		current: 0,
+	}
+}
+
+// NextTarget returns the next target using round-robin algorithm
+func (lb *RoundRobinLoadBalancer) NextTarget() (*config.TargetConfig, error) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	// Filter healthy targets
+	healthyTargets := make([]*config.TargetConfig, 0)
+	for _, target := range lb.targets {
+		if target.Healthy {
+			healthyTargets = append(healthyTargets, target)
+		}
+	}
+
+	if len(healthyTargets) == 0 {
+		return nil, ErrNoHealthyTargets
+	}
+
+	// Get next target
+	target := healthyTargets[lb.current%len(healthyTargets)]
+	lb.current = (lb.current + 1) % len(healthyTargets)
+
+	return target, nil
+}
+
+// UpdateTargets updates the targets list
+func (lb *RoundRobinLoadBalancer) UpdateTargets(targets []config.TargetConfig) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	t := make([]*config.TargetConfig, len(targets))
+	for i := range targets {
+		t[i] = &targets[i]
+	}
+	lb.targets = t
+	lb.current = 0
+}
+
+// RandomLoadBalancer implements random load balancing
+type RandomLoadBalancer struct {
+	targets []*config.TargetConfig
+	rand    *rand.Rand
+	mu      sync.Mutex
+}
+
+// NewRandomLoadBalancer creates a new random load balancer
+func NewRandomLoadBalancer(targets []config.TargetConfig) *RandomLoadBalancer {
+	t := make([]*config.TargetConfig, len(targets))
+	for i := range targets {
+		t[i] = &targets[i]
+	}
+	return &RandomLoadBalancer{
+		targets: t,
+		rand:    rand.New(rand.NewSource(time.Now().UnixNano())),
+	}
+}
+
+// NextTarget returns a random target
+func (lb *RandomLoadBalancer) NextTarget() (*config.TargetConfig, error) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	// Filter healthy targets
+	healthyTargets := make([]*config.TargetConfig, 0)
+	for _, target := range lb.targets {
+		if target.Healthy {
+			healthyTargets = append(healthyTargets, target)
+		}
+	}
+
+	if len(healthyTargets) == 0 {
+		return nil, ErrNoHealthyTargets
+	}
+
+	// Get random target
+	index := lb.rand.Intn(len(healthyTargets))
+	return healthyTargets[index], nil
+}
+
+// UpdateTargets updates the targets list
+func (lb *RandomLoadBalancer) UpdateTargets(targets []config.TargetConfig) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	t := make([]*config.TargetConfig, len(targets))
+	for i := range targets {
+		t[i] = &targets[i]
+	}
+	lb.targets = t
+}
+
+// LeastConnectionsLoadBalancer implements least connections load balancing
+type LeastConnectionsLoadBalancer struct {
+	targets      []*config.TargetConfig
+	connections map[string]int
+	mu          sync.Mutex
+}
+
+// NewLeastConnectionsLoadBalancer creates a new least connections load balancer
+func NewLeastConnectionsLoadBalancer(targets []config.TargetConfig) *LeastConnectionsLoadBalancer {
+	t := make([]*config.TargetConfig, len(targets))
+	for i := range targets {
+		t[i] = &targets[i]
+	}
+
+	connections := make(map[string]int)
+	for _, target := range t {
+		connections[target.Name] = 0
+	}
+
+	return &LeastConnectionsLoadBalancer{
+		targets:      t,
+		connections: connections,
+	}
+}
+
+// NextTarget returns the target with the least connections
+func (lb *LeastConnectionsLoadBalancer) NextTarget() (*config.TargetConfig, error) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	// Filter healthy targets and find the one with least connections
+	var selectedTarget *config.TargetConfig
+	minConnections := -1
+
+	for _, target := range lb.targets {
+		if target.Healthy {
+			connections := lb.connections[target.Name]
+			if minConnections == -1 || connections < minConnections {
+				minConnections = connections
+				selectedTarget = target
+			}
+		}
+	}
+
+	if selectedTarget == nil {
+		return nil, ErrNoHealthyTargets
+	}
+
+	// Increment connection count
+	lb.connections[selectedTarget.Name]++
+
+	return selectedTarget, nil
+}
+
+// ReleaseConnection decrements the connection count for a target
+func (lb *LeastConnectionsLoadBalancer) ReleaseConnection(targetName string) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	if count, exists := lb.connections[targetName]; exists && count > 0 {
+		lb.connections[targetName] = count - 1
+	}
+}
+
+// UpdateTargets updates the targets list
+func (lb *LeastConnectionsLoadBalancer) UpdateTargets(targets []config.TargetConfig) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	t := make([]*config.TargetConfig, len(targets))
+	for i := range targets {
+		t[i] = &targets[i]
+	}
+
+	// Update connections map
+	connections := make(map[string]int)
+	for _, target := range t {
+		// Preserve existing connection count if target exists
+		if count, exists := lb.connections[target.Name]; exists {
+			connections[target.Name] = count
+		} else {
+			connections[target.Name] = 0
+		}
+	}
+
+	lb.targets = t
+	lb.connections = connections
+}
+
+// ErrNoHealthyTargets is returned when no healthy targets are available
+var ErrNoHealthyTargets = errorString("no healthy targets available")
+
+// errorString is a simple string-based error type
+type errorString string
+
+func (e errorString) Error() string {
+	return string(e)
+}

internal/proxy/loadbalancer_test.go

@@ -0,0 +1,270 @@
+package proxy
+
+import (
+	"testing"
+
+	"github.com/iwasforcedtobehere/goRZ/internal/config"
+)
+
+func TestRoundRobinLoadBalancer(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target3", Address: "http://localhost:8083", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewRoundRobinLoadBalancer(targets)
+
+	// Test that targets are selected in round-robin order
+	for i := 0; i < 6; i++ {
+		target, err := lb.NextTarget()
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		expectedTarget := targets[i%3]
+		if target.Name != expectedTarget.Name {
+			t.Errorf("Expected target %s, got %s", expectedTarget.Name, target.Name)
+		}
+	}
+}
+
+func TestRoundRobinLoadBalancerWithUnhealthyTargets(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: false},
+		{Name: "target3", Address: "http://localhost:8083", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewRoundRobinLoadBalancer(targets)
+
+	// Test that only healthy targets are selected
+	for i := 0; i < 4; i++ {
+		target, err := lb.NextTarget()
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		if target.Name == "target2" {
+			t.Errorf("Selected unhealthy target: %s", target.Name)
+		}
+
+		// Should alternate between target1 and target3
+		if i%2 == 0 && target.Name != "target1" {
+			t.Errorf("Expected target1, got %s", target.Name)
+		}
+		if i%2 == 1 && target.Name != "target3" {
+			t.Errorf("Expected target3, got %s", target.Name)
+		}
+	}
+}
+
+func TestRoundRobinLoadBalancerNoHealthyTargets(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: false},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: false},
+	}
+
+	lb := NewRoundRobinLoadBalancer(targets)
+
+	_, err := lb.NextTarget()
+	if err != ErrNoHealthyTargets {
+		t.Errorf("Expected ErrNoHealthyTargets, got %v", err)
+	}
+}
+
+func TestRoundRobinLoadBalancerUpdateTargets(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewRoundRobinLoadBalancer(targets)
+
+	// Get first target
+	target, err := lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target1" {
+		t.Errorf("Expected target1, got %s", target.Name)
+	}
+
+	// Update targets
+	newTargets := []config.TargetConfig{
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target3", Address: "http://localhost:8083", Protocol: "http", Weight: 1, Healthy: true},
+	}
+	lb.UpdateTargets(newTargets)
+
+	// Test that new targets are selected
+	for i := 0; i < 4; i++ {
+		target, err := lb.NextTarget()
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		expectedTarget := newTargets[i%2]
+		if target.Name != expectedTarget.Name {
+			t.Errorf("Expected target %s, got %s", expectedTarget.Name, target.Name)
+		}
+	}
+}
+
+func TestRandomLoadBalancer(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target3", Address: "http://localhost:8083", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewRandomLoadBalancer(targets)
+
+	// Test that targets are selected randomly
+	// We'll just check that we don't get errors and that all targets are eventually selected
+	selected := make(map[string]bool)
+	for i := 0; i < 100; i++ {
+		target, err := lb.NextTarget()
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		selected[target.Name] = true
+	}
+
+	// Check that all targets were selected at least once
+	for _, target := range targets {
+		if !selected[target.Name] {
+			t.Errorf("Target %s was never selected", target.Name)
+		}
+	}
+}
+
+func TestRandomLoadBalancerWithUnhealthyTargets(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: false},
+		{Name: "target3", Address: "http://localhost:8083", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewRandomLoadBalancer(targets)
+
+	// Test that only healthy targets are selected
+	for i := 0; i < 100; i++ {
+		target, err := lb.NextTarget()
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		if target.Name == "target2" {
+			t.Errorf("Selected unhealthy target: %s", target.Name)
+		}
+	}
+}
+
+func TestLeastConnectionsLoadBalancer(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewLeastConnectionsLoadBalancer(targets)
+
+	// Initially, both targets should have 0 connections
+	// The first target should be selected
+	target, err := lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target1" {
+		t.Errorf("Expected target1, got %s", target.Name)
+	}
+
+	// Now target1 should have 1 connection, target2 should have 0
+	// So target2 should be selected
+	target, err = lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target2" {
+		t.Errorf("Expected target2, got %s", target.Name)
+	}
+
+	// Now both targets should have 1 connection
+	// target1 should be selected again
+	target, err = lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target1" {
+		t.Errorf("Expected target1, got %s", target.Name)
+	}
+
+	// Release a connection from target1
+	lb.(*LeastConnectionsLoadBalancer).ReleaseConnection("target1")
+
+	// Now target1 should have 1 connection, target2 should have 1 connection
+	// But since we released from target1, it should be selected
+	target, err = lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target1" {
+		t.Errorf("Expected target1, got %s", target.Name)
+	}
+}
+
+func TestLeastConnectionsLoadBalancerUpdateTargets(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewLeastConnectionsLoadBalancer(targets)
+
+	// Get first target
+	target, err := lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target1" {
+		t.Errorf("Expected target1, got %s", target.Name)
+	}
+
+	// Update targets
+	newTargets := []config.TargetConfig{
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target3", Address: "http://localhost:8083", Protocol: "http", Weight: 1, Healthy: true},
+	}
+	lb.UpdateTargets(newTargets)
+
+	// Test that new targets are selected
+	// The first new target should be selected
+	target, err = lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target2" {
+		t.Errorf("Expected target2, got %s", target.Name)
+	}
+
+	// The second new target should be selected
+	target, err = lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target3" {
+		t.Errorf("Expected target3, got %s", target.Name)
+	}
+}
+
+func TestLeastConnectionsLoadBalancerNoHealthyTargets(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: false},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: false},
+	}
+
+	lb := NewLeastConnectionsLoadBalancer(targets)
+
+	_, err := lb.NextTarget()
+	if err != ErrNoHealthyTargets {
+		t.Errorf("Expected ErrNoHealthyTargets, got %v", err)
+	}
+}

internal/proxy/proxy.go

@@ -0,0 +1,238 @@
+package proxy
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"sync"
+	"time"
+
+	"github.com/iwasforcedtobehere/goRZ/internal/config"
+	"github.com/iwasforcedtobehere/goRZ/internal/logger"
+)
+
+// Server represents the reverse proxy server
+type Server struct {
+	config      *config.Config
+	logger      *logger.Logger
+	loadBalancer LoadBalancer
+	healthChecker *HealthChecker
+}
+
+// LoadBalancer defines the interface for load balancing strategies
+type LoadBalancer interface {
+	NextTarget() (*config.TargetConfig, error)
+	UpdateTargets(targets []config.TargetConfig)
+}
+
+// HealthChecker handles health checking for proxy targets
+type HealthChecker struct {
+	config     *config.Config
+	logger     *logger.Logger
+	targets    map[string]*config.TargetConfig
+	httpClient *http.Client
+	stopChan   chan struct{}
+	mu         sync.RWMutex
+}
+
+// NewServer creates a new reverse proxy server
+func NewServer(cfg *config.Config, logger *logger.Logger) (*Server, error) {
+	// Create load balancer based on configuration
+	var lb LoadBalancer
+	switch cfg.Proxy.LoadBalancer {
+	case "roundrobin":
+		lb = NewRoundRobinLoadBalancer(cfg.Proxy.Targets)
+	case "leastconn":
+		lb = NewLeastConnectionsLoadBalancer(cfg.Proxy.Targets)
+	case "random":
+		lb = NewRandomLoadBalancer(cfg.Proxy.Targets)
+	default:
+		return nil, fmt.Errorf("unsupported load balancer: %s", cfg.Proxy.LoadBalancer)
+	}
+
+	// Create health checker
+	healthChecker := NewHealthChecker(cfg, logger)
+
+	return &Server{
+		config:       cfg,
+		logger:       logger,
+		loadBalancer: lb,
+		healthChecker: healthChecker,
+	}, nil
+}
+
+// ServeHTTP handles HTTP requests
+func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	// Get next target based on load balancing strategy
+	target, err := s.loadBalancer.NextTarget()
+	if err != nil {
+		s.logger.Error("Failed to get target", logger.Error(err))
+		http.Error(w, "Service unavailable", http.StatusServiceUnavailable)
+		return
+	}
+
+	// Check if target is healthy
+	if !target.Healthy {
+		s.logger.Warn("Target is unhealthy", logger.String("target", target.Name))
+		http.Error(w, "Service unavailable", http.StatusServiceUnavailable)
+		return
+	}
+
+	// Create reverse proxy
+	targetURL, err := url.Parse(target.Address)
+	if err != nil {
+		s.logger.Error("Failed to parse target URL", 
+			logger.String("target", target.Name), 
+			logger.String("address", target.Address),
+			logger.Error(err))
+		http.Error(w, "Internal server error", http.StatusInternalServerError)
+		return
+	}
+
+	proxy := httputil.NewSingleHostReverseProxy(targetURL)
+
+	// Set up custom director to modify request
+	originalDirector := proxy.Director
+	proxy.Director = func(req *http.Request) {
+		originalDirector(req)
+		// Add custom headers
+		req.Header.Set("X-Forwarded-Host", req.Host)
+		req.Header.Set("X-Forwarded-Proto", "https")
+		if req.TLS != nil {
+			req.Header.Set("X-Forwarded-Ssl", "on")
+		}
+		// Add proxy identification
+		req.Header.Set("X-Proxy-Server", "goRZ")
+	}
+
+	// Set up error handler
+	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
+		s.logger.Error("Proxy error", logger.Error(err))
+		http.Error(w, "Proxy error", http.StatusBadGateway)
+	}
+
+	// Set up modify response handler
+	proxy.ModifyResponse = func(resp *http.Response) error {
+		// Add custom headers to response
+		resp.Header.Set("X-Served-By", "goRZ")
+		resp.Header.Set("X-Target", target.Name)
+		return nil
+	}
+
+	// Serve the request
+	s.logger.Debug("Proxying request", 
+		logger.String("method", r.Method),
+		logger.String("path", r.URL.Path),
+		logger.String("target", target.Name))
+
+	proxy.ServeHTTP(w, r)
+}
+
+// Start starts the proxy server
+func (s *Server) Start() error {
+	// Start health checker
+	if err := s.healthChecker.Start(); err != nil {
+		return fmt.Errorf("failed to start health checker: %w", err)
+	}
+	return nil
+}
+
+// Stop stops the proxy server
+func (s *Server) Stop() {
+	s.healthChecker.Stop()
+}
+
+// NewHealthChecker creates a new health checker
+func NewHealthChecker(cfg *config.Config, logger *logger.Logger) *HealthChecker {
+	targets := make(map[string]*config.TargetConfig)
+	for i := range cfg.Proxy.Targets {
+		targets[cfg.Proxy.Targets[i].Name] = &cfg.Proxy.Targets[i]
+	}
+
+	return &HealthChecker{
+		config:     cfg,
+		logger:     logger,
+		targets:    targets,
+		httpClient: &http.Client{Timeout: 5 * time.Second},
+		stopChan:   make(chan struct{}),
+	}
+}
+
+// Start starts the health checker
+func (h *HealthChecker) Start() error {
+	// Initial health check
+	h.checkAllTargets()
+
+	// Start periodic health checks
+	go func() {
+		ticker := time.NewTicker(time.Duration(h.config.Proxy.HealthCheckInterval) * time.Second)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ticker.C:
+				h.checkAllTargets()
+			case <-h.stopChan:
+				return
+			}
+		}
+	}()
+
+	return nil
+}
+
+// Stop stops the health checker
+func (h *HealthChecker) Stop() {
+	close(h.stopChan)
+}
+
+// checkAllTargets checks the health of all targets
+func (h *HealthChecker) checkAllTargets() {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	for name, target := range h.targets {
+		healthy := h.checkTargetHealth(target)
+		if target.Healthy != healthy {
+			h.logger.Info("Target health status changed", 
+				logger.String("target", name), 
+				logger.Bool("healthy", healthy))
+			target.Healthy = healthy
+		}
+	}
+}
+
+// checkTargetHealth checks the health of a single target
+func (h *HealthChecker) checkTargetHealth(target *config.TargetConfig) bool {
+	targetURL, err := url.Parse(target.Address)
+	if err != nil {
+		h.logger.Error("Failed to parse target URL for health check", 
+			logger.String("target", target.Name),
+			logger.Error(err))
+		return false
+	}
+
+	healthURL := *targetURL
+	healthURL.Path = h.config.Proxy.HealthCheckPath
+
+	req, err := http.NewRequest("GET", healthURL.String(), nil)
+	if err != nil {
+		h.logger.Error("Failed to create health check request", 
+			logger.String("target", target.Name),
+			logger.Error(err))
+		return false
+	}
+
+	resp, err := h.httpClient.Do(req)
+	if err != nil {
+		h.logger.Error("Health check request failed", 
+			logger.String("target", target.Name),
+			logger.Error(err))
+		return false
+	}
+	defer resp.Body.Close()
+
+	return resp.StatusCode == http.StatusOK
+}