diff --git a/net/lux_router.go b/net/lux_router.go index 3ecc262..adc7811 100644 --- a/net/lux_router.go +++ b/net/lux_router.go @@ -21,7 +21,7 @@ type LuxRouter struct { thisKey crypto.LuxKey keyStore crypto.LuxKeyStore - routes []LuxRoute + routes map[proto.LuxID]*LuxRoute channelLock sync.RWMutex outbound []LuxChannel @@ -34,7 +34,7 @@ func NewLuxRouter(key crypto.LuxKey, ks crypto.LuxKeyStore) LuxRouter { return LuxRouter{ thisKey: key, keyStore: ks, - routes: make([]LuxRoute, 0), + routes: make(map[proto.LuxID]*LuxRoute, 0), outbound: make([]LuxChannel, 0), inbound: make([]LuxChannel, 0), dgramChan: make(chan LuxDatagram), @@ -80,12 +80,12 @@ func (r *LuxRouter) CreateOutboundRoute(id proto.LuxID, chType LuxChannelType, u return err } - r.routes = append(r.routes, LuxRoute{ + r.routes[key.Id] = &LuxRoute{ Key: key, Destination: channel.Address, Associated: r.addOutboundChannel(channel), Nonces: NewLuxNonceList(), - }) + } return nil } @@ -95,12 +95,12 @@ func (r *LuxRouter) CreateInboundChannel(chType LuxChannelType, udpAddr string) return err } - r.routes = append(r.routes, LuxRoute{ + r.routes[r.thisKey.Id] = &LuxRoute{ Key: r.thisKey, Destination: channel.Address, Associated: r.addInboundChannel(channel), Nonces: NewLuxNonceList(), - }) + } return nil } @@ -164,7 +164,7 @@ func (r *LuxRouter) Stop() { r.outbound = r.outbound[0:] r.channelLock.Unlock() - r.routes = r.routes[0:] + r.routes = make(map[proto.LuxID]*LuxRoute) } func (r *LuxRouter) RecvDgram() LuxDatagram { @@ -175,42 +175,27 @@ func udpAddrEqual(addr *net.UDPAddr, other *net.UDPAddr) bool { return addr.IP.Equal(other.IP) && addr.Port == other.Port } -func (r *LuxRouter) getRouteIndex(udpAddr *net.UDPAddr) (LuxRoute, int) { - for idx, route := range r.routes { - if udpAddrEqual(route.Destination, udpAddr) { - return route, idx - } - } - return LuxRoute{}, -1 -} - -func (r *LuxRouter) getRouteIndexByID(id proto.LuxID) (LuxRoute, int) { - for idx, route := range r.routes { - if bytes.Equal(route.Key.Id.UUID[:], id.UUID[:]) { - return route, idx - } - } - - return LuxRoute{}, -1 -} - -func (r *LuxRouter) GetRoutes() []LuxRoute { +func (r *LuxRouter) GetRoutes() map[proto.LuxID]*LuxRoute { return r.routes } -func (r *LuxRouter) GetRoute(udpAddr *net.UDPAddr) (LuxRoute, bool) { - route, idx := r.getRouteIndex(udpAddr) - return route, idx != -1 +func (r *LuxRouter) GetRoute(udpAddr *net.UDPAddr) (*LuxRoute, bool) { + for _, route := range r.routes { + if udpAddrEqual(route.Destination, udpAddr) { + return route, true + } + } + + return nil, false } func (r *LuxRouter) DeleteRoute(route *LuxRoute) { - _, idx := r.getRouteIndex(route.Destination) - if idx == -1 { + if _, ok := r.routes[route.Key.Id]; !ok { return } r.CloseChannel(route.Associated, false) - r.routes = append(r.routes[idx:], r.routes[:idx+1]...) + delete(r.routes, route.Key.Id) } func (r *LuxRouter) Recv() (LuxPacket, error) { @@ -235,7 +220,7 @@ func (r *LuxRouter) Recv() (LuxPacket, error) { // check if LuxID matches if !bytes.Equal(packet.Target.UUID[:], route.Key.Id.UUID[:]) { // not matches.. we discard route and throw away packet - r.DeleteRoute(&route) + r.DeleteRoute(route) return packet, errors.New("bogus packet from established route") /* NOTE: * there may be rare situation where multiple clients behind NAT/CGNAT @@ -258,21 +243,21 @@ func (r *LuxRouter) Recv() (LuxPacket, error) { if bytes.Equal(packet.Target.UUID[:], key.Id.UUID[:]) { // key UUID and decrypted UUID matching - create OR update route and return packet + var ok bool var route *LuxRoute - if _, idx := r.getRouteIndexByID(packet.Target); idx != -1 { - route = &r.routes[idx] + if route, ok = r.routes[packet.Target]; ok { route.Destination = dgram.Target route.Associated = dgram.Channel // since packet arrived from different transport, we flush nonces route.Nonces.Flush() } else { - r.routes = append(r.routes, LuxRoute{ + r.routes[key.Id] = &LuxRoute{ Key: key, Destination: dgram.Target, Associated: dgram.Channel, Nonces: NewLuxNonceList(), - }) - route = &r.routes[len(r.routes)-1] + } + route = r.routes[key.Id] } // rotate nonce @@ -291,8 +276,8 @@ func (r *LuxRouter) Recv() (LuxPacket, error) { } func (r *LuxRouter) Send(packet LuxPacket) error { - route, idx := r.getRouteIndexByID(packet.Target) - if idx == -1 { + route, ok := r.routes[packet.Target] + if !ok { return errors.New("no route to peer") }