mirror of
https://github.com/netbirdio/netbird.git
synced 2024-12-14 10:50:45 +01:00
173 lines
4.0 KiB
Go
173 lines
4.0 KiB
Go
package routeselector
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"slices"
|
|
"sync"
|
|
|
|
"github.com/hashicorp/go-multierror"
|
|
"golang.org/x/exp/maps"
|
|
|
|
"github.com/netbirdio/netbird/client/errors"
|
|
route "github.com/netbirdio/netbird/route"
|
|
)
|
|
|
|
type RouteSelector struct {
|
|
mu sync.RWMutex
|
|
selectedRoutes map[route.NetID]struct{}
|
|
selectAll bool
|
|
}
|
|
|
|
func NewRouteSelector() *RouteSelector {
|
|
return &RouteSelector{
|
|
selectedRoutes: map[route.NetID]struct{}{},
|
|
// default selects all routes
|
|
selectAll: true,
|
|
}
|
|
}
|
|
|
|
// 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.selectedRoutes = map[route.NetID]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
|
|
}
|
|
|
|
rs.selectedRoutes[route] = struct{}{}
|
|
}
|
|
rs.selectAll = 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.selectAll = true
|
|
rs.selectedRoutes = map[route.NetID]struct{}{}
|
|
}
|
|
|
|
// DeselectRoutes removes specific routes from the selection.
|
|
// If the selector is in "select all" mode, it will transition to "select specific" mode.
|
|
func (rs *RouteSelector) DeselectRoutes(routes []route.NetID, allRoutes []route.NetID) error {
|
|
rs.mu.Lock()
|
|
defer rs.mu.Unlock()
|
|
|
|
if rs.selectAll {
|
|
rs.selectAll = false
|
|
rs.selectedRoutes = map[route.NetID]struct{}{}
|
|
for _, route := range allRoutes {
|
|
rs.selectedRoutes[route] = 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.selectedRoutes, route)
|
|
}
|
|
|
|
return errors.FormatErrorOrNil(err)
|
|
}
|
|
|
|
// DeselectAllRoutes deselects all routes, effectively disabling route selection.
|
|
func (rs *RouteSelector) DeselectAllRoutes() {
|
|
rs.mu.Lock()
|
|
defer rs.mu.Unlock()
|
|
|
|
rs.selectAll = false
|
|
rs.selectedRoutes = map[route.NetID]struct{}{}
|
|
}
|
|
|
|
// 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.selectAll {
|
|
return true
|
|
}
|
|
_, selected := rs.selectedRoutes[routeID]
|
|
return selected
|
|
}
|
|
|
|
// 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.selectAll {
|
|
return maps.Clone(routes)
|
|
}
|
|
|
|
filtered := route.HAMap{}
|
|
for id, rt := range routes {
|
|
if rs.IsSelected(id.NetID()) {
|
|
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 {
|
|
SelectedRoutes map[route.NetID]struct{} `json:"selected_routes"`
|
|
SelectAll bool `json:"select_all"`
|
|
}{
|
|
SelectAll: rs.selectAll,
|
|
SelectedRoutes: rs.selectedRoutes,
|
|
})
|
|
}
|
|
|
|
// 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.selectedRoutes = map[route.NetID]struct{}{}
|
|
rs.selectAll = true
|
|
return nil
|
|
}
|
|
|
|
var temp struct {
|
|
SelectedRoutes map[route.NetID]struct{} `json:"selected_routes"`
|
|
SelectAll bool `json:"select_all"`
|
|
}
|
|
|
|
if err := json.Unmarshal(data, &temp); err != nil {
|
|
return err
|
|
}
|
|
|
|
rs.selectedRoutes = temp.SelectedRoutes
|
|
rs.selectAll = temp.SelectAll
|
|
|
|
if rs.selectedRoutes == nil {
|
|
rs.selectedRoutes = map[route.NetID]struct{}{}
|
|
}
|
|
|
|
return nil
|
|
}
|