From 409fe81d4ceba4ce0496c3261b2a4afa584728eb Mon Sep 17 00:00:00 2001 From: Yann Stepienik Date: Tue, 3 Mar 2026 18:01:05 +0000 Subject: [PATCH] [skip ci] lb --- src/constellation/DNS.go | 6 +- src/constellation/tunnels.go | 71 ++++++++++++-- src/proxy/SocketProxy.go | 183 +++++++++++++++++++++-------------- src/proxy/buildFromConfig.go | 11 ++- src/proxy/routeTo.go | 78 ++++++++++++++- src/proxy/tunnel_lb.go | 65 +++++++++++++ src/utils/types.go | 11 ++- 7 files changed, 331 insertions(+), 94 deletions(-) create mode 100644 src/proxy/tunnel_lb.go diff --git a/src/constellation/DNS.go b/src/constellation/DNS.go index 6be0f804..7fd1152e 100644 --- a/src/constellation/DNS.go +++ b/src/constellation/DNS.go @@ -88,11 +88,11 @@ func handleDNSRequest(w dns.ResponseWriter, r *dns.Msg) { } else { for _, q := range r.Question { for _, tunnel := range remoteTunnels { - for _, from := range tunnel.From { - if from == currentName { + for _, target := range tunnel.Targets { + if target.DeviceName == currentName { continue } - destination := CachedDeviceNames[from] + destination := CachedDeviceNames[target.DeviceName] if destination != "" { if strings.HasSuffix(q.Name, tunnel.Route.Host + ".") && q.Qtype == dns.TypeA { utils.Debug("DNS Overwrite " + tunnel.Route.Host + " with " + destination) diff --git a/src/constellation/tunnels.go b/src/constellation/tunnels.go index 1687e114..11bb095f 100644 --- a/src/constellation/tunnels.go +++ b/src/constellation/tunnels.go @@ -267,18 +267,18 @@ func UpdateLocalTunnelCache() { } for _, tunnelRoute := range heartbeat.Tunnels { - // Skip tunnels from ourselves - if heartbeat.DeviceName == currentDeviceName { - continue - } if tunnelRoute.Tunnel == "_ANY_" || tunnelRoute.Tunnel == currentDeviceName { + target := utils.TunnelTarget{ + DeviceName: heartbeat.DeviceName, + TargetURL: tunnelRoute.Target, + } if existing, ok := byName[tunnelRoute.Name]; ok { - existing.From = append(existing.From, heartbeat.DeviceName) + existing.Targets = append(existing.Targets, target) } else { tunnelRoute.Const_IsTunneled = true byName[tunnelRoute.Name] = &utils.ConstellationTunnel{ - Route: tunnelRoute, - From: []string{heartbeat.DeviceName}, + Route: tunnelRoute, + Targets: []utils.TunnelTarget{target}, } } } @@ -296,9 +296,11 @@ func UpdateLocalTunnelCache() { copied := make([]utils.ConstellationTunnel, len(t)) for i, tunnel := range t { copied[i] = tunnel - copied[i].From = make([]string, len(tunnel.From)) - copy(copied[i].From, tunnel.From) - sort.Strings(copied[i].From) + copied[i].Targets = make([]utils.TunnelTarget, len(tunnel.Targets)) + copy(copied[i].Targets, tunnel.Targets) + sort.Slice(copied[i].Targets, func(a, b int) bool { + return copied[i].Targets[a].DeviceName < copied[i].Targets[b].DeviceName + }) } sort.Slice(copied, func(i, j int) bool { return copied[i].Route.Name < copied[j].Route.Name @@ -348,4 +350,53 @@ func GetLocalTunnelCache() []utils.ConstellationTunnel { func IsTunneled(route utils.ProxyRouteConfig) bool { return route.Const_IsTunneled +} + +func ensureStickyBucket() (nats.KeyValue, error) { + if js == nil { + return nil, nats.ErrConnectionClosed + } + kv, err := js.KeyValue("tunnel-sticky") + if err == nil { + return kv, nil + } + kv, err = js.CreateKeyValue(&nats.KeyValueConfig{ + Bucket: "tunnel-sticky", + TTL: 120 * time.Second, + Storage: nats.MemoryStorage, + Replicas: 1, + }) + return kv, err +} + +func GetStickyTarget(clientKey string) (string, bool) { + clientConfigLock.RLock() + defer clientConfigLock.RUnlock() + + kv, err := ensureStickyBucket() + if err != nil { + return "", false + } + + entry, err := kv.Get(clientKey) + if err != nil { + return "", false + } + return string(entry.Value()), true +} + +func SetStickyTarget(clientKey string, deviceName string) { + clientConfigLock.RLock() + defer clientConfigLock.RUnlock() + + kv, err := ensureStickyBucket() + if err != nil { + utils.Error("[NATS] Error accessing sticky KV bucket", err) + return + } + + _, err = kv.Put(clientKey, []byte(deviceName)) + if err != nil { + utils.Error("[NATS] Error setting sticky target", err) + } } \ No newline at end of file diff --git a/src/proxy/SocketProxy.go b/src/proxy/SocketProxy.go index 398473e2..9c79945f 100644 --- a/src/proxy/SocketProxy.go +++ b/src/proxy/SocketProxy.go @@ -20,6 +20,14 @@ var ( proxiesLock sync.Mutex ) +func stripTargetProtocol(target string) string { + parts := strings.Split(target, "://") + if len(parts) > 1 { + return parts[1] + } + return parts[0] +} + type ProxyInfo struct { stop chan bool stopped chan bool @@ -88,7 +96,7 @@ var mapToTCPUDP = map[string]string{ "quic": "udp", // QUIC is a transport protocol that runs over UDP } -func startProxy(listenAddr string, target string, proxyInfo *ProxyInfo, isHTTPProxy bool, route utils.ProxyRouteConfig) { +func startProxy(listenAddr string, target string, targets []utils.TunnelTarget, proxyInfo *ProxyInfo, isHTTPProxy bool, route utils.ProxyRouteConfig) { // remove any protocol destinationArr := strings.Split(target, "://") listenProtocol := "tcp" @@ -137,13 +145,13 @@ func startProxy(listenAddr string, target string, proxyInfo *ProxyInfo, isHTTPPr if listenProtocol == "tcp" { utils.Debug("[SocketProxy] Starting TCP proxy on " + listenAddr + " -> " + target) - handleTCPProxy(listener, target, proxyInfo, isHTTPProxy, route, listenAddr) + handleTCPProxy(listener, target, targets, proxyInfo, isHTTPProxy, route, listenAddr) } else { - handleUDPProxy(packetConn, target, proxyInfo, route, listenAddr) + handleUDPProxy(packetConn, target, targets, proxyInfo, route, listenAddr) } } -func handleTCPProxy(listener net.Listener, target string, proxyInfo *ProxyInfo, isHTTPProxy bool, route utils.ProxyRouteConfig, listenAddr string) { +func handleTCPProxy(listener net.Listener, target string, targets []utils.TunnelTarget, proxyInfo *ProxyInfo, isHTTPProxy bool, route utils.ProxyRouteConfig, listenAddr string) { for { acceptChan := make(chan net.Conn) acceptErrChan := make(chan error) @@ -175,7 +183,18 @@ func handleTCPProxy(listener net.Listener, target string, proxyInfo *ProxyInfo, utils.Debug("[SocketProxy] New TCP connection accepted on " + listenAddr) - server, err := net.Dial("tcp", target) + stickyKey := "" + if route.LBStickyMode { + ip, _ := utils.SplitIP(client.RemoteAddr().String()) + stickyKey = ip + } + selected := DefaultTunnelLB.SelectTarget(targets, route.Name, route.LBMode, route.LBStickyMode, stickyKey) + dialTarget := target + if selected != nil { + dialTarget = stripTargetProtocol(selected.TargetURL) + } + + server, err := net.Dial("tcp", dialTarget) if err != nil { utils.Error("[SocketProxy] Failed to connect to TCP server", err) shieldedClient.Close() @@ -186,7 +205,7 @@ func handleTCPProxy(listener net.Listener, target string, proxyInfo *ProxyInfo, } } -func handleUDPProxy(packetConn net.PacketConn, target string, proxyInfo *ProxyInfo, route utils.ProxyRouteConfig, listenAddr string) { +func handleUDPProxy(packetConn net.PacketConn, target string, targets []utils.TunnelTarget, proxyInfo *ProxyInfo, route utils.ProxyRouteConfig, listenAddr string) { targetAddr, err := net.ResolveUDPAddr("udp", target) if err != nil { utils.Error("[SocketProxy] Failed to resolve UDP target address", err) @@ -211,22 +230,34 @@ func handleUDPProxy(packetConn net.PacketConn, target string, proxyInfo *ProxyIn continue } - // utils.Debug("[SocketProxy] New UDP packet received on " + listenAddr) - // Apply UDP shield middleware if needed shieldedBuffer := UDPSmartShieldMiddleware("proxy-"+listenAddr, route)(buffer[:n], remoteAddr) if shieldedBuffer == nil { continue } - _, err = packetConn.WriteTo(buffer[:n], targetAddr) + stickyKey := "" + if route.LBStickyMode { + ip, _ := utils.SplitIP(remoteAddr.String()) + stickyKey = ip + } + selected := DefaultTunnelLB.SelectTarget(targets, route.Name, route.LBMode, route.LBStickyMode, stickyKey) + dialAddr := targetAddr + if selected != nil { + resolved, err := net.ResolveUDPAddr("udp", stripTargetProtocol(selected.TargetURL)) + if err == nil { + dialAddr = resolved + } + } + + _, err = packetConn.WriteTo(buffer[:n], dialAddr) if err != nil { utils.Error("[SocketProxy] Failed to forward UDP packet to server", err) continue } // Handle response from target - go handleUDPResponse(packetConn, remoteAddr, targetAddr, proxyInfo) + go handleUDPResponse(packetConn, remoteAddr, dialAddr, proxyInfo) } } } @@ -255,6 +286,7 @@ func handleUDPResponse(packetConn net.PacketConn, clientAddr net.Addr, targetAdd type PortsPair struct { From string To string + Targets []utils.TunnelTarget isHTTPProxy bool route utils.ProxyRouteConfig } @@ -271,7 +303,7 @@ func initInternalPortProxy(ports []PortsPair) { } activeProxies[port.From] = proxyInfo - go startProxy(port.From, port.To, proxyInfo, port.isHTTPProxy, port.route) + go startProxy(port.From, port.To, port.Targets, proxyInfo, port.isHTTPProxy, port.route) } } @@ -324,15 +356,8 @@ func InitInternalSocketProxy() { routes := config.HTTPConfig.ProxyConfig.Routes - tunnels := []utils.ProxyRouteConfig{} remoteTunnels := constellation.GetLocalTunnelCache() - for _, tunnel := range remoteTunnels { - routeConfig := tunnel.Route - if !routeConfig.Disabled { - tunnels = append(tunnels, routeConfig) - } - } - + remoteListRoutes := []utils.ProxyRouteConfig{} for _, shares := range config.RemoteStorage.Shares { route := shares.Route @@ -347,69 +372,79 @@ func InitInternalSocketProxy() { targetPort = HTTPSPort } - routesList := []utils.ProxyRouteConfig{} - routesList = append(routesList, tunnels...) - routesList = append(routesList, remoteListRoutes...) - routesList = append(routesList, routes...) + addPortPair := func(route utils.ProxyRouteConfig, targets []utils.TunnelTarget) { + if !route.UseHost || !strings.Contains(route.Host, ":") || route.Disabled { + return + } + hostname := route.Host + port := strings.Split(hostname, ":")[1] + if _, err := strconv.Atoi(port); err != nil { + return + } + destination := "" + isHTTPProxy := strings.HasPrefix(route.Target, "http://") || strings.HasPrefix(route.Target, "https://") + sourceHostname := hostname + var pairTargets []utils.TunnelTarget + if isHTTPProxy { + destination = "localhost:" + targetPort + sourceHostname = ":" + port + } else { + destination = route.Target + pairTargets = targets - for _, route := range routesList { - if route.UseHost && strings.Contains(route.Host, ":") && !route.Disabled { - hostname := route.Host - port := strings.Split(hostname, ":")[1] - // if port is a number - if _, err := strconv.Atoi(port); err == nil { - destination := "" - isHTTPProxy := strings.HasPrefix(route.Target, "http://") || strings.HasPrefix(route.Target, "https://") - sourceHostname := hostname - if isHTTPProxy { - destination = "localhost:" + targetPort - sourceHostname = ":" + port + if route.Mode == "SERVAPP" && (!utils.IsInsideContainer || utils.IsHostNetwork) { + url, err := url.Parse(destination) + if err != nil { + utils.Error("[SocketProxy] Create socket Route", err) } else { - destination = route.Target - - if route.Mode == "SERVAPP" && (!utils.IsInsideContainer || utils.IsHostNetwork) { - url, err := url.Parse(destination) - if err != nil { - utils.Error("[SocketProxy] Create socket Route", err) - } else { - targetHost := url.Hostname() - - targetIP, err := docker.GetContainerIPByName(targetHost) - if err != nil { - utils.Error("[SocketProxy] Can't find container", err) - } else { - utils.Debug("[SocketProxy] Dockerless socket Target IP: " + targetIP) - destination = targetIP + ":" + url.Port() - if url.Scheme != "" { - destination = url.Scheme + "://" + destination - } - } + targetHost := url.Hostname() + + targetIP, err := docker.GetContainerIPByName(targetHost) + if err != nil { + utils.Error("[SocketProxy] Can't find container", err) + } else { + utils.Debug("[SocketProxy] Dockerless socket Target IP: " + targetIP) + destination = targetIP + ":" + url.Port() + if url.Scheme != "" { + destination = url.Scheme + "://" + destination } } } - - portpair := PortsPair{sourceHostname, destination, isHTTPProxy, route} - - if isHTTPProxy && allowHTTPLocal && utils.IsLocalIP(route.Host) { - portpair.To = "localhost:" + HTTPPort - } - - // Check if this port is already in the list (avoid duplicates) - alreadyExists := false - for _, existing := range expectedPorts { - if existing.From == portpair.From { - alreadyExists = true - break - } - } - if alreadyExists { - utils.MajorError("[SocketProxy] Duplicate port detected: "+portpair.From+". Multiple routes are trying to use the same port.", nil) - } else { - expectedPorts = append(expectedPorts, portpair) - } } + } + + portpair := PortsPair{sourceHostname, destination, pairTargets, isHTTPProxy, route} + + if isHTTPProxy && allowHTTPLocal && utils.IsLocalIP(route.Host) { + portpair.To = "localhost:" + HTTPPort + } + + // Check if this port is already in the list (avoid duplicates) + alreadyExists := false + for _, existing := range expectedPorts { + if existing.From == portpair.From { + alreadyExists = true + break + } + } + if alreadyExists { + utils.MajorError("[SocketProxy] Duplicate port detected: "+portpair.From+". Multiple routes are trying to use the same port.", nil) + } else { + expectedPorts = append(expectedPorts, portpair) + } + } + + for _, tunnel := range remoteTunnels { + if !tunnel.Route.Disabled { + addPortPair(tunnel.Route, tunnel.Targets) } } + for _, route := range remoteListRoutes { + addPortPair(route, nil) + } + for _, route := range routes { + addPortPair(route, nil) + } StopAllProxies() diff --git a/src/proxy/buildFromConfig.go b/src/proxy/buildFromConfig.go index 0707c037..fc3ef547 100644 --- a/src/proxy/buildFromConfig.go +++ b/src/proxy/buildFromConfig.go @@ -3,6 +3,7 @@ package proxy import ( "github.com/gorilla/mux" "net/http" + "strings" "github.com/azukaar/cosmos-server/src/utils" "github.com/azukaar/cosmos-server/src/constellation" @@ -17,9 +18,8 @@ func BuildFromConfig(router *mux.Router, config utils.ProxyConfig) *mux.Router { remoteTunnels := constellation.GetLocalTunnelCache() for _, tunnel := range remoteTunnels { - routeConfig := tunnel.Route - if !routeConfig.Disabled { - RouterGen(routeConfig, router, RouteTo(routeConfig)) + if !tunnel.Route.Disabled && (strings.HasPrefix(tunnel.Route.Target, "http://") || strings.HasPrefix(tunnel.Route.Target, "https://")) { + RouterGen(tunnel.Route, router, TunnelRouteTo(tunnel, DefaultTunnelLB)) } } @@ -29,12 +29,15 @@ func BuildFromConfig(router *mux.Router, config utils.ProxyConfig) *mux.Router { if route.Disabled { continue } + if !strings.HasPrefix(route.Target, "http://") && !strings.HasPrefix(route.Target, "https://") { + continue + } RouterGen(route, router, RouteTo(route)) } for i := len(config.Routes)-1; i >= 0; i-- { routeConfig := config.Routes[i] - if !routeConfig.Disabled { + if !routeConfig.Disabled && (strings.HasPrefix(routeConfig.Target, "http://") || strings.HasPrefix(routeConfig.Target, "https://")) { RouterGen(routeConfig, router, RouteTo(routeConfig)) } } diff --git a/src/proxy/routeTo.go b/src/proxy/routeTo.go index e805bcd3..c0b6f097 100644 --- a/src/proxy/routeTo.go +++ b/src/proxy/routeTo.go @@ -235,6 +235,82 @@ func NewProxy(targetHost string, AcceptInsecureHTTPSTarget bool, DisableHeaderHa } +func TunnelRouteTo(tunnel utils.ConstellationTunnel, lb *TunnelLoadBalancer) http.Handler { + type targetProxy struct { + target utils.TunnelTarget + handler http.Handler + } + + proxies := make([]targetProxy, 0, len(tunnel.Targets)) + for _, t := range tunnel.Targets { + route := tunnel.Route + route.Target = t.TargetURL + proxy, err := NewProxy(t.TargetURL, route.AcceptInsecureHTTPSTarget, route.DisableHeaderHardening, route) + if err != nil { + utils.Error("Create Tunnel Route for "+t.DeviceName, err) + continue + } + proxies = append(proxies, targetProxy{t, proxy}) + } + + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if len(proxies) == 0 { + http.Error(w, "No tunnel targets available", http.StatusBadGateway) + return + } + + targets := make([]utils.TunnelTarget, len(proxies)) + for i, p := range proxies { + targets[i] = p.target + } + + sticky := tunnel.Route.LBStickyMode + stickyKey := "" + if sticky { + stickyKey = GetClientID(r, tunnel.Route) + } + + // For HTTP, check cookie first (survives IP changes on mobile) + var selected *utils.TunnelTarget + if sticky { + if cookie, err := r.Cookie("_cosmos_tunnel_lb"); err == nil && cookie.Value != "" { + for i, p := range proxies { + if p.target.DeviceName == cookie.Value { + selected = &proxies[i].target + break + } + } + // Cookie device gone — fall through to SelectTarget + } + } + + if selected == nil { + selected = lb.SelectTarget(targets, tunnel.Route.Name, tunnel.Route.LBMode, sticky, stickyKey) + } + if selected == nil { + http.Error(w, "No tunnel targets available", http.StatusBadGateway) + return + } + + if sticky { + http.SetCookie(w, &http.Cookie{ + Name: "_cosmos_tunnel_lb", + Value: selected.DeviceName, + HttpOnly: true, + SameSite: http.SameSiteLaxMode, + MaxAge: 3600, + }) + } + + for _, p := range proxies { + if p.target.DeviceName == selected.DeviceName { + p.handler.ServeHTTP(w, r) + return + } + } + }) +} + func RouteTo(route utils.ProxyRouteConfig) http.Handler { // initialize a reverse proxy and pass the actual backend server url here @@ -247,7 +323,7 @@ func RouteTo(route utils.ProxyRouteConfig) http.Handler { } } - if(routeType == "SERVAPP" || routeType == "PROXY") { + if(routeType == "SERVAPP" || routeType == "PROXY") { proxy, err := NewProxy(destination, route.AcceptInsecureHTTPSTarget, route.DisableHeaderHardening, route) if err != nil { utils.Error("Create Route", err) diff --git a/src/proxy/tunnel_lb.go b/src/proxy/tunnel_lb.go new file mode 100644 index 00000000..e597dceb --- /dev/null +++ b/src/proxy/tunnel_lb.go @@ -0,0 +1,65 @@ +package proxy + +import ( + "sync" + "sync/atomic" + + "github.com/azukaar/cosmos-server/src/constellation" + "github.com/azukaar/cosmos-server/src/utils" +) + +type routeCounter struct { + val atomic.Uint64 +} + +type TunnelLoadBalancer struct { + counters sync.Map // routeName -> *routeCounter +} + +func (lb *TunnelLoadBalancer) getCounter(routeName string) *routeCounter { + if v, ok := lb.counters.Load(routeName); ok { + return v.(*routeCounter) + } + c := &routeCounter{} + actual, _ := lb.counters.LoadOrStore(routeName, c) + return actual.(*routeCounter) +} + +func (lb *TunnelLoadBalancer) SelectTarget(targets []utils.TunnelTarget, routeName string, mode string, sticky bool, stickyKey string) *utils.TunnelTarget { + if len(targets) == 0 { + return nil + } + + // Check sticky assignment + if sticky && stickyKey != "" { + deviceName, ok := constellation.GetStickyTarget(stickyKey) + if ok { + for i := range targets { + if targets[i].DeviceName == deviceName { + return &targets[i] + } + } + // Sticky device gone from targets, fall through to re-assign + } + } + + // Mode selection + var selected *utils.TunnelTarget + switch mode { + case "round_robin": + c := lb.getCounter(routeName) + idx := c.val.Add(1) - 1 + selected = &targets[idx%uint64(len(targets))] + default: // "first" or "" + selected = &targets[0] + } + + // Store sticky assignment + if sticky && stickyKey != "" && selected != nil { + constellation.SetStickyTarget(stickyKey, selected.DeviceName) + } + + return selected +} + +var DefaultTunnelLB = &TunnelLoadBalancer{} diff --git a/src/utils/types.go b/src/utils/types.go index dfe18435..32a88bac 100644 --- a/src/utils/types.go +++ b/src/utils/types.go @@ -262,6 +262,8 @@ type ProxyRouteConfig struct { ExtraHeaders map[string]string `yaml:"extra_headers,omitempty"` SkipURLClean bool `yaml:"skip_url_clean"` UseH2C bool `yaml:"use_h2c"` + LBMode string `yaml:"lb_mode" json:"lb_mode,omitempty"` + LBStickyMode bool `yaml:"lb_sticky_mode" json:"lb_sticky_mode,omitempty"` Const_IsTunneled bool `yaml:"-", json:"-"` } @@ -309,9 +311,14 @@ type ConstellationConfig struct { IPRange string } +type TunnelTarget struct { + DeviceName string `json:"deviceName"` + TargetURL string `json:"targetURL"` +} + type ConstellationTunnel struct { - Route ProxyRouteConfig - From []string + Route ProxyRouteConfig + Targets []TunnelTarget } type ConstellationDNSEntry struct {