From 51b235b61ab50f91aa951b30c0bdc75fe72a94ba Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Thu, 4 Jan 2024 20:26:42 +0700
Subject: [PATCH] internal/clientinfo: implement ndp listen

So when new clients join the network, ctrld can really the event and
update client information to NDP table quickly.
---
 go.mod                             |  1 +
 go.sum                             |  2 +
 internal/clientinfo/client_info.go |  6 ++
 internal/clientinfo/ndp.go         | 93 ++++++++++++++++++++++++++++++
 4 files changed, 102 insertions(+)

diff --git a/go.mod b/go.mod
index fec32ef..4d38445 100644
--- a/go.mod
+++ b/go.mod
@@ -15,6 +15,7 @@ require (
 	github.com/jaytaylor/go-hostsfile v0.0.0-20220426042432-61485ac1fa6c
 	github.com/josharian/native v1.1.1-0.20230202152459-5c7d0dd6ab86
 	github.com/kardianos/service v1.2.1
+	github.com/mdlayher/ndp v1.0.1
 	github.com/miekg/dns v1.1.55
 	github.com/olekukonko/tablewriter v0.0.5
 	github.com/pelletier/go-toml/v2 v2.0.8
diff --git a/go.sum b/go.sum
index c792103..a64357a 100644
--- a/go.sum
+++ b/go.sum
@@ -201,6 +201,8 @@ github.com/mattn/go-runewidth v0.0.14 h1:+xnbZSEeDbOIg5/mE6JF0w6n9duR1l3/WmbinWV
 github.com/mattn/go-runewidth v0.0.14/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
 github.com/mdlayher/ethernet v0.0.0-20190606142754-0394541c37b7 h1:lez6TS6aAau+8wXUP3G9I3TGlmPFEq2CTxBaRqY6AGE=
 github.com/mdlayher/ethernet v0.0.0-20190606142754-0394541c37b7/go.mod h1:U6ZQobyTjI/tJyq2HG+i/dfSoFUt8/aZCM+GKtmFk/Y=
+github.com/mdlayher/ndp v1.0.1 h1:+yAD79/BWyFlvAoeG5ncPS0ItlHP/eVbH7bQ6/+LVA4=
+github.com/mdlayher/ndp v1.0.1/go.mod h1:rf3wKaWhAYJEXFKpgF8kQ2AxypxVbfNcZbqoAo6fVzk=
 github.com/mdlayher/netlink v1.7.2 h1:/UtM3ofJap7Vl4QWCPDGXY8d3GIY2UGSDbK+QWmY8/g=
 github.com/mdlayher/netlink v1.7.2/go.mod h1:xraEF7uJbxLhc5fpHL4cPe221LI2bdttWlU+ZGLfQSw=
 github.com/mdlayher/raw v0.0.0-20190606142536-fef19f00fc18/go.mod h1:7EpbotpCmVZcu+KCX4g9WaRNuu11uyhiW7+Le1dKawg=
diff --git a/internal/clientinfo/client_info.go b/internal/clientinfo/client_info.go
index 72ef971..1fe1083 100644
--- a/internal/clientinfo/client_info.go
+++ b/internal/clientinfo/client_info.go
@@ -209,6 +209,12 @@ func (t *Table) init() {
 				t.refreshers = append(t.refreshers, discover)
 			}
 		}
+		ctx, cancel := context.WithCancel(context.Background())
+		go func() {
+			<-t.quitCh
+			cancel()
+		}()
+		go t.ndp.listen(ctx)
 	}
 	// PTR lookup.
 	if t.discoverPTR() {
diff --git a/internal/clientinfo/ndp.go b/internal/clientinfo/ndp.go
index 337c414..600b54c 100644
--- a/internal/clientinfo/ndp.go
+++ b/internal/clientinfo/ndp.go
@@ -2,10 +2,19 @@ package clientinfo
 
 import (
 	"bufio"
+	"context"
+	"errors"
+	"fmt"
 	"io"
 	"net"
+	"net/netip"
 	"strings"
 	"sync"
+	"time"
+
+	"github.com/mdlayher/ndp"
+
+	"github.com/Control-D-Inc/ctrld"
 )
 
 // ndpDiscover provides client discovery functionality using NDP protocol.
@@ -60,6 +69,55 @@ func (nd *ndpDiscover) List() []string {
 	return ips
 }
 
+// listen listens on ipv6 link local for Neighbor Solicitation message
+// to update new neighbors information to ndp table.
+func (nd *ndpDiscover) listen(ctx context.Context) {
+	ifi, err := firstInterfaceWithV6LinkLocal()
+	if err != nil {
+		ctrld.ProxyLogger.Load().Debug().Err(err).Msg("failed to find valid ipv6")
+		return
+	}
+	c, ip, err := ndp.Listen(ifi, ndp.LinkLocal)
+	if err != nil {
+		ctrld.ProxyLogger.Load().Debug().Err(err).Msg("ndp listen failed")
+		return
+	}
+	defer c.Close()
+	ctrld.ProxyLogger.Load().Debug().Msgf("listening ndp on: %s", ip.String())
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
+		_ = c.SetReadDeadline(time.Now().Add(30 * time.Second))
+		msg, _, from, readErr := c.ReadFrom()
+		if readErr != nil {
+			var opErr *net.OpError
+			if errors.As(readErr, &opErr) && (opErr.Timeout() || opErr.Temporary()) {
+				continue
+			}
+			ctrld.ProxyLogger.Load().Debug().Err(readErr).Msg("ndp read loop error")
+			return
+		}
+
+		// Only looks for neighbor solicitation message, since new clients
+		// which join network will broadcast this message to us.
+		am, ok := msg.(*ndp.NeighborSolicitation)
+		if !ok {
+			continue
+		}
+		fromIP := from.String()
+		for _, opt := range am.Options {
+			if lla, ok := opt.(*ndp.LinkLayerAddress); ok {
+				mac := lla.Addr.String()
+				nd.mac.Store(fromIP, mac)
+				nd.ip.Store(mac, fromIP)
+			}
+		}
+	}
+}
+
 // scanWindows populates NDP table using information from "netsh" command.
 func (nd *ndpDiscover) scanWindows(r io.Reader) {
 	scanner := bufio.NewScanner(r)
@@ -124,3 +182,38 @@ func parseMAC(mac string) string {
 	hw, _ := net.ParseMAC(normalizeMac(mac))
 	return hw.String()
 }
+
+// firstInterfaceWithV6LinkLocal returns the first interface which is capable of using NDP.
+func firstInterfaceWithV6LinkLocal() (*net.Interface, error) {
+	ifis, err := net.Interfaces()
+	if err != nil {
+		return nil, err
+	}
+
+	for _, ifi := range ifis {
+		// Skip if iface is down/loopback/non-multicast.
+		if ifi.Flags&net.FlagUp == 0 || ifi.Flags&net.FlagLoopback != 0 || ifi.Flags&net.FlagMulticast == 0 {
+			continue
+		}
+
+		addrs, err := ifi.Addrs()
+		if err != nil {
+			return nil, err
+		}
+
+		for _, addr := range addrs {
+			ipNet, ok := addr.(*net.IPNet)
+			if !ok {
+				continue
+			}
+			ip, ok := netip.AddrFromSlice(ipNet.IP)
+			if !ok {
+				return nil, fmt.Errorf("invalid ip address: %s", ipNet.String())
+			}
+			if ip.Is6() && !ip.Is4In6() {
+				return &ifi, nil
+			}
+		}
+	}
+	return nil, errors.New("no interface can be used")
+}