azukaar/Cosmos-Server
1
0
Fork 0
mirror of https://github.com/azukaar/Cosmos-Server.git synced 2026-10-12 02:11:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

[skip ci] lb

Browse Source
This commit is contained in:
Yann Stepienik
2026-03-03 18:02:38 +00:00
parent 3f2e12112e
commit 409fe81d4c
7 changed files with 331 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/constellation/DNS.go
+3 -3
View File
@@ -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)
src/constellation/tunnels.go
+61 -10
View File
@@ -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)
}
}
src/proxy/SocketProxy.go
+109 -74
View File
@@ -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()
src/proxy/buildFromConfig.go
+7 -4
View File
@@ -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))
}
}
src/proxy/routeTo.go
+77 -1
View File
@@ -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)
src/proxy/tunnel_lb.go
+65
View File
@@ -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{}
src/utils/types.go
+9 -2
View File
@@ -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 {