diff --git a/clients/clients.go b/clients/clients.go index 3e3ec04..1e47606 100644 --- a/clients/clients.go +++ b/clients/clients.go @@ -29,9 +29,9 @@ import ( ) type WgConfig struct { - IpAddress string `json:"ipAddress"` - Peers []Peer `json:"peers"` - // Targets []Target `json:"targets"` + IpAddress string `json:"ipAddress"` + Peers []Peer `json:"peers"` + Targets []Target `json:"targets"` } type Target struct { @@ -181,6 +181,7 @@ func NewWireGuardService(interfaceName string, mtu int, generateAndSaveKeyTo str wsClient.RegisterHandler("newt/wg/peer/update", service.handleUpdatePeer) wsClient.RegisterHandler("newt/wg/target/add", service.handleAddTarget) wsClient.RegisterHandler("newt/wg/target/remove", service.handleRemoveTarget) + wsClient.RegisterHandler("newt/wg/target/update", service.handleUpdateTarget) return service, nil } @@ -866,6 +867,60 @@ func (s *WireGuardService) handleAddTarget(msg websocket.WSMessage) { logger.Info("Added target subnet %s with port ranges: %v", target.CIDR, target.PortRange) } +func (s *WireGuardService) handleUpdateTarget(msg websocket.WSMessage) { + logger.Debug("Received message: %v", msg.Data) + + // you are going to get a oldTarget and a newTarget in the message + type UpdateTargetRequest struct { + OldTarget Target `json:"oldTarget"` + NewTarget Target `json:"newTarget"` + } + + jsonData, err := json.Marshal(msg.Data) + if err != nil { + logger.Info("Error marshaling data: %v", err) + return + } + + var request UpdateTargetRequest + if err := json.Unmarshal(jsonData, &request); err != nil { + logger.Info("Error unmarshaling data: %v", err) + return + } + + if s.tnet == nil { + logger.Info("Netstack not initialized") + return + } + + prefix, err := netip.ParsePrefix(request.OldTarget.CIDR) + if err != nil { + logger.Info("Invalid CIDR %s: %v", request.OldTarget.CIDR, err) + return + } + + s.tnet.RemoveProxySubnetRule(prefix) + + // Now add the new target + newPrefix, err := netip.ParsePrefix(request.NewTarget.CIDR) + if err != nil { + logger.Info("Invalid CIDR %s: %v", request.NewTarget.CIDR, err) + return + } + + var portRanges []netstack2.PortRange + for _, pr := range request.NewTarget.PortRange { + portRanges = append(portRanges, netstack2.PortRange{ + Min: pr.Min, + Max: pr.Max, + }) + } + + s.tnet.AddProxySubnetRule(newPrefix, portRanges) + + logger.Info("Updated target subnet from %s to %s", request.OldTarget.CIDR, request.NewTarget.CIDR) +} + func (s *WireGuardService) handleRemoveTarget(msg websocket.WSMessage) { logger.Debug("Received message: %v", msg.Data) diff --git a/docker/client.go b/docker/docker.go similarity index 100% rename from docker/client.go rename to docker/docker.go diff --git a/logger/logger.go b/logger/logger.go index 50911ac..c647443 100644 --- a/logger/logger.go +++ b/logger/logger.go @@ -35,8 +35,12 @@ func NewLoggerWithWriter(writer LogWriter) *Logger { } // Init initializes the default logger -func Init() *Logger { +func Init(logger *Logger) *Logger { once.Do(func() { + if logger != nil { + defaultLogger = logger + return + } defaultLogger = NewLogger() }) return defaultLogger @@ -45,7 +49,7 @@ func Init() *Logger { // GetLogger returns the default logger instance func GetLogger() *Logger { if defaultLogger == nil { - Init() + Init(nil) } return defaultLogger }