ping_service/manager/proxy.go

package main

import (
	"crypto/tls"
	"fmt"
	"io"
	"net/http"
	"sync/atomic"
	"time"
)

// Backend represents a backend service that can handle proxied requests
type Backend struct {
	WorkerID string
	URL      string
	Healthy  bool
}

// BackendPool manages a pool of backend services for load balancing
type BackendPool struct {
	workerType WorkerType
	store      *WorkerStore
	current    atomic.Uint64 // For round-robin
}

// NewBackendPool creates a new backend pool for a specific worker type
func NewBackendPool(workerType WorkerType, store *WorkerStore) *BackendPool {
	return &BackendPool{
		workerType: workerType,
		store:      store,
	}
}

// GetBackends returns all healthy backends of this pool's type
func (bp *BackendPool) GetBackends() []Backend {
	workers := bp.store.List()
	backends := make([]Backend, 0)

	for _, worker := range workers {
		if worker.Type == bp.workerType && worker.Healthy {
			backends = append(backends, Backend{
				WorkerID: worker.ID,
				URL:      worker.URL,
				Healthy:  worker.Healthy,
			})
		}
	}

	return backends
}

// NextBackend returns the next healthy backend using round-robin
func (bp *BackendPool) NextBackend() (*Backend, error) {
	backends := bp.GetBackends()

	if len(backends) == 0 {
		return nil, fmt.Errorf("no healthy %s backends available", bp.workerType)
	}

	// Round-robin selection
	idx := bp.current.Add(1) % uint64(len(backends))
	return &backends[idx], nil
}

// ProxyManager manages multiple backend pools
type ProxyManager struct {
	inputPool  *BackendPool
	outputPool *BackendPool
	client     *http.Client
}

// NewProxyManager creates a new proxy manager
func NewProxyManager(store *WorkerStore) *ProxyManager {
	// Create HTTP client that accepts self-signed certs (for internal services)
	transport := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		MaxIdleConns:    100,
		IdleConnTimeout: 90 * time.Second,
	}

	return &ProxyManager{
		inputPool:  NewBackendPool(WorkerTypeInput, store),
		outputPool: NewBackendPool(WorkerTypeOutput, store),
		client: &http.Client{
			Timeout:   30 * time.Second,
			Transport: transport,
		},
	}
}

// ProxyGetTarget forwards a GET request to an input service to get next target IP
func (pm *ProxyManager) ProxyGetTarget(w http.ResponseWriter, r *http.Request) error {
	backend, err := pm.inputPool.NextBackend()
	if err != nil {
		return err
	}

	// Forward GET /target request
	targetURL := fmt.Sprintf("%s/target", backend.URL)
	req, err := http.NewRequest("GET", targetURL, nil)
	if err != nil {
		return err
	}

	// Copy headers if needed
	req.Header.Set("User-Agent", "PingServiceManager-Gateway/1.0")

	resp, err := pm.client.Do(req)
	if err != nil {
		return fmt.Errorf("backend request failed: %v", err)
	}
	defer resp.Body.Close()

	// Copy response status and headers
	w.WriteHeader(resp.StatusCode)
	for key, values := range resp.Header {
		for _, value := range values {
			w.Header().Add(key, value)
		}
	}

	// Copy response body
	_, err = io.Copy(w, resp.Body)
	return err
}

// ProxyPostResult forwards a POST request to an output service to submit results
func (pm *ProxyManager) ProxyPostResult(w http.ResponseWriter, r *http.Request) error {
	backend, err := pm.outputPool.NextBackend()
	if err != nil {
		return err
	}

	// Forward POST /result request
	targetURL := fmt.Sprintf("%s/result", backend.URL)
	req, err := http.NewRequest("POST", targetURL, r.Body)
	if err != nil {
		return err
	}

	// Copy content type
	req.Header.Set("Content-Type", r.Header.Get("Content-Type"))
	req.Header.Set("User-Agent", "PingServiceManager-Gateway/1.0")

	resp, err := pm.client.Do(req)
	if err != nil {
		return fmt.Errorf("backend request failed: %v", err)
	}
	defer resp.Body.Close()

	// Copy response status and headers
	w.WriteHeader(resp.StatusCode)
	for key, values := range resp.Header {
		for _, value := range values {
			w.Header().Add(key, value)
		}
	}

	// Copy response body
	_, err = io.Copy(w, resp.Body)
	return err
}

// GetPoolStats returns statistics about backend pools
func (pm *ProxyManager) GetPoolStats() map[string]interface{} {
	inputBackends := pm.inputPool.GetBackends()
	outputBackends := pm.outputPool.GetBackends()

	return map[string]interface{}{
		"input_backends":  len(inputBackends),
		"output_backends": len(outputBackends),
		"total_backends":  len(inputBackends) + len(outputBackends),
	}
}