From 295e9c9a1020eafd6cad1320e0219bf0124762ce Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Mon, 5 Nov 2018 17:31:10 +0000
Subject: [PATCH] Cache crypto-key routes (until routing table changes)

---
 src/yggdrasil/ckr.go | 51 ++++++++++++++++++++++++++++++++------------
 1 file changed, 37 insertions(+), 14 deletions(-)

diff --git a/src/yggdrasil/ckr.go b/src/yggdrasil/ckr.go
index c4f6d697..f24cd9d5 100644
--- a/src/yggdrasil/ckr.go
+++ b/src/yggdrasil/ckr.go
@@ -16,6 +16,8 @@ type cryptokey struct {
 	enabled    bool
 	ipv4routes []cryptokey_route
 	ipv6routes []cryptokey_route
+	ipv4cache  map[address]cryptokey_route
+	ipv6cache  map[address]cryptokey_route
 }
 
 type cryptokey_route struct {
@@ -27,6 +29,8 @@ func (c *cryptokey) init(core *Core) {
 	c.core = core
 	c.ipv4routes = make([]cryptokey_route, 0)
 	c.ipv6routes = make([]cryptokey_route, 0)
+	c.ipv4cache = make(map[address]cryptokey_route, 0)
+	c.ipv6cache = make(map[address]cryptokey_route, 0)
 }
 
 func (c *cryptokey) isEnabled() bool {
@@ -68,12 +72,20 @@ func (c *cryptokey) addRoute(cidr string, dest string) error {
 				subnet:      *ipnet,
 				destination: boxPubKey,
 			})
+
 			// Sort so most specific routes are first
 			sort.Slice(c.ipv6routes, func(i, j int) bool {
 				im, _ := c.ipv6routes[i].subnet.Mask.Size()
 				jm, _ := c.ipv6routes[j].subnet.Mask.Size()
 				return im > jm
 			})
+
+			// Clear the cache as this route might change future routing
+			// Setting an empty slice keeps the memory whereas nil invokes GC
+			for k := range c.ipv6cache {
+				delete(c.ipv6cache, k)
+			}
+
 			return nil
 		}
 	} else if prefixsize == net.IPv4len*8 {
@@ -83,28 +95,39 @@ func (c *cryptokey) addRoute(cidr string, dest string) error {
 	return errors.New("Unspecified error")
 }
 
-func (c *cryptokey) getPublicKeyForAddress(addr string) (boxPubKey, error) {
-	ipaddr := net.ParseIP(addr)
+func (c *cryptokey) getPublicKeyForAddress(addr address) (boxPubKey, error) {
+	// Check if there's a cache entry for this addr
+	if route, ok := c.ipv6cache[addr]; ok {
+		var box boxPubKey
+		copy(box[:boxPubKeyLen], route.destination)
+		return box, nil
+	}
 
-	if ipaddr.To4() == nil {
-		// IPv6
+	// No cache was found - start by converting the address into a net.IP
+	ip := make(net.IP, 16)
+	copy(ip[:16], addr[:])
+
+	// Check whether it's an IPv4 or an IPv6 address
+	if ip.To4() == nil {
+		// Check if we have a route. At this point c.ipv6routes should be
+		// pre-sorted so that the most specific routes are first
 		for _, route := range c.ipv6routes {
-			if route.subnet.Contains(ipaddr) {
+			// Does this subnet match the given IP?
+			if route.subnet.Contains(ip) {
+				// Cache the entry for future packets to get a faster lookup
+				c.ipv6cache[addr] = route
+
+				// Return the boxPubKey
 				var box boxPubKey
 				copy(box[:boxPubKeyLen], route.destination)
 				return box, nil
 			}
 		}
 	} else {
-		// IPv4
+		// IPv4 isn't supported yet
 		return boxPubKey{}, errors.New("IPv4 not supported at this time")
-		/*
-			    for _, route := range c.ipv4routes {
-						if route.subnet.Contains(ipaddr) {
-							return route.destination, nil
-						}
-					}
-		*/
 	}
-	return boxPubKey{}, errors.New(fmt.Sprintf("No route to %s", addr))
+
+	// No route was found if we got to this point
+	return boxPubKey{}, errors.New(fmt.Sprintf("No route to %s", ip.String()))
 }