package routeselector import ( "encoding/json" "fmt" "slices" "sync" "github.com/hashicorp/go-multierror" "golang.org/x/exp/maps" "github.com/netbirdio/netbird/client/errors" "github.com/netbirdio/netbird/route" ) type RouteSelector struct { mu sync.RWMutex deselectedRoutes map[route.NetID]struct{} deselectAll bool } func NewRouteSelector() *RouteSelector { return &RouteSelector{ deselectedRoutes: map[route.NetID]struct{}{}, deselectAll: false, } } // SelectRoutes updates the selected routes based on the provided route IDs. func (rs *RouteSelector) SelectRoutes(routes []route.NetID, appendRoute bool, allRoutes []route.NetID) error { rs.mu.Lock() defer rs.mu.Unlock() if !appendRoute || rs.deselectAll { maps.Clear(rs.deselectedRoutes) for _, r := range allRoutes { rs.deselectedRoutes[r] = struct{}{} } } var err *multierror.Error for _, route := range routes { if !slices.Contains(allRoutes, route) { err = multierror.Append(err, fmt.Errorf("route '%s' is not available", route)) continue } delete(rs.deselectedRoutes, route) } rs.deselectAll = false return errors.FormatErrorOrNil(err) } // SelectAllRoutes sets the selector to select all routes. func (rs *RouteSelector) SelectAllRoutes() { rs.mu.Lock() defer rs.mu.Unlock() rs.deselectAll = false maps.Clear(rs.deselectedRoutes) } // DeselectRoutes removes specific routes from the selection. func (rs *RouteSelector) DeselectRoutes(routes []route.NetID, allRoutes []route.NetID) error { rs.mu.Lock() defer rs.mu.Unlock() if rs.deselectAll { return nil } var err *multierror.Error for _, route := range routes { if !slices.Contains(allRoutes, route) { err = multierror.Append(err, fmt.Errorf("route '%s' is not available", route)) continue } rs.deselectedRoutes[route] = struct{}{} } return errors.FormatErrorOrNil(err) } // DeselectAllRoutes deselects all routes, effectively disabling route selection. func (rs *RouteSelector) DeselectAllRoutes() { rs.mu.Lock() defer rs.mu.Unlock() rs.deselectAll = true maps.Clear(rs.deselectedRoutes) } // IsSelected checks if a specific route is selected. func (rs *RouteSelector) IsSelected(routeID route.NetID) bool { rs.mu.RLock() defer rs.mu.RUnlock() if rs.deselectAll { return false } _, deselected := rs.deselectedRoutes[routeID] return !deselected } // FilterSelected removes unselected routes from the provided map. func (rs *RouteSelector) FilterSelected(routes route.HAMap) route.HAMap { rs.mu.RLock() defer rs.mu.RUnlock() if rs.deselectAll { return route.HAMap{} } filtered := route.HAMap{} for id, rt := range routes { netID := id.NetID() _, deselected := rs.deselectedRoutes[netID] if !deselected { filtered[id] = rt } } return filtered } // MarshalJSON implements the json.Marshaler interface func (rs *RouteSelector) MarshalJSON() ([]byte, error) { rs.mu.RLock() defer rs.mu.RUnlock() return json.Marshal(struct { DeselectedRoutes map[route.NetID]struct{} `json:"deselected_routes"` DeselectAll bool `json:"deselect_all"` }{ DeselectedRoutes: rs.deselectedRoutes, DeselectAll: rs.deselectAll, }) } // UnmarshalJSON implements the json.Unmarshaler interface // If the JSON is empty or null, it will initialize like a NewRouteSelector. func (rs *RouteSelector) UnmarshalJSON(data []byte) error { rs.mu.Lock() defer rs.mu.Unlock() // Check for null or empty JSON if len(data) == 0 || string(data) == "null" { rs.deselectedRoutes = map[route.NetID]struct{}{} rs.deselectAll = false return nil } var temp struct { DeselectedRoutes map[route.NetID]struct{} `json:"deselected_routes"` DeselectAll bool `json:"deselect_all"` } if err := json.Unmarshal(data, &temp); err != nil { return err } rs.deselectedRoutes = temp.DeselectedRoutes rs.deselectAll = temp.DeselectAll if rs.deselectedRoutes == nil { rs.deselectedRoutes = map[route.NetID]struct{}{} } return nil }