From 5cde417ad532a4e571796256f9b1753072fcc872 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E6=80=9D=E6=95=8F?= <huangsimin@yuandian.com>
Date: Fri, 26 Aug 2022 09:52:49 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B5=8B=E8=AF=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 config.go | 24 +++++++++++++++++++++++-
 tap.go    | 30 +++++++++++++++++++++---------
 2 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/config.go b/config.go
index 29f9c5d..8e9a7c8 100644
--- a/config.go
+++ b/config.go
@@ -1,8 +1,12 @@
 package main
 
 import (
+	"encoding/binary"
 	"log"
+	"net"
 	"os"
+	"strings"
+	"sync"
 
 	"gopkg.in/yaml.v3"
 )
@@ -17,12 +21,15 @@ type Config struct {
 		Self  Node   `yaml:"self"`
 		Nodes []Node `yaml:"nodes"`
 	} `yaml:"network"`
+
+	IPv4Nodes map[uint32]string
+	Lock      sync.Mutex
 }
 
 var config = NewConfig()
 
 func NewConfig() *Config {
-	cnf := &Config{}
+	cnf := &Config{IPv4Nodes: make(map[uint32]string)}
 	f, err := os.Open("config.yaml")
 	if err != nil {
 		log.Panic(err)
@@ -32,5 +39,20 @@ func NewConfig() *Config {
 	if err != nil {
 		log.Panic(err)
 	}
+
+	for _, node := range cnf.Network.Nodes {
+		values := strings.Split(node.Virt, "/")
+		if len(values) != 2 {
+			log.Panic("virt must have '/'. eg 1.1.1.1/24")
+		}
+		inet, net, err := net.ParseCIDR(node.Virt)
+		if err != nil {
+			log.Panic(err)
+		}
+
+		cnf.IPv4Nodes[binary.LittleEndian.Uint32([]byte(inet))] = node.Real
+		log.Printf("real %v virt %v %v", node.Real, inet, net)
+	}
+
 	return cnf
 }
diff --git a/tap.go b/tap.go
index 33cde84..dc5f124 100644
--- a/tap.go
+++ b/tap.go
@@ -3,8 +3,8 @@ package main
 import (
 	"bytes"
 	"context"
+	"encoding/binary"
 	"encoding/gob"
-	"flag"
 	"fmt"
 	"log"
 	"os/exec"
@@ -34,11 +34,11 @@ type RPCClient struct {
 	conn      *grpc.ClientConn
 }
 
-func (cli *RPCClient) connect(addr string) {
+func (cli *RPCClient) connect(realAddr string) {
 	log.Println("rpcclient start")
 	// defer log.Println("rpcclient exit")
 
-	conn, err := grpc.Dial(addr,
+	conn, err := grpc.Dial(realAddr,
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
 	)
 	if err != nil {
@@ -115,7 +115,6 @@ func (nc *NetCard) Run() {
 }
 
 func NewNetCard() *NetCard {
-	flag.Parse()
 
 	config := water.Config{
 		DeviceType: water.TUN,
@@ -147,11 +146,14 @@ func (nc *NetCard) runRead() {
 
 		var ok bool
 		var isLoop int32 = 1
+
+		// 20 微秒后停止收集数据
 		go func() {
 			var after = time.NewTimer(time.Millisecond * 20)
 			<-after.C
 			atomic.StoreInt32(&isLoop, 0)
 		}()
+
 		for atomic.LoadInt32(&isLoop) > 0 {
 			var rframe ethernet.Frame
 			rframe.Resize(1500)
@@ -161,19 +163,29 @@ func (nc *NetCard) runRead() {
 				log.Fatal(err)
 			}
 			rframe = rframe[:n]
+			if !waterutil.IsIPv4(rframe) {
+				continue
+			}
+
+			var realAddr string
+			config.Lock.Lock()
+			if _, ok = config.IPv4Nodes[binary.LittleEndian.Uint32([]byte(waterutil.IPv4Destination(rframe)))]; !ok {
+				continue
+			}
+			config.Lock.Unlock()
 
 			// log.Printf("Payload: % x\n", rframe.Payload())
 
-			log.Printf("Dst: %s %s\n", waterutil.IPv4Destination(rframe), rframe.Destination())
-			log.Printf("Src: %s %s\n", waterutil.IPv4Source(rframe), rframe.Source())
+			log.Printf("Dst: %s Src %s\n", waterutil.IPv4Destination(rframe), waterutil.IPv4Source(rframe))
+
 			log.Printf("Ethertype: % x %v\n", rframe.Ethertype(), waterutil.IsIPv4(rframe))
 
 			var buffer *[][]byte
-			var key = waterutil.IPv4Destination(rframe).String()
-			if buffer, ok = bytesMap[key]; !ok {
+
+			if buffer, ok = bytesMap[realAddr]; !ok {
 				mbuffer := make([][]byte, 0, 100)
 				buffer = &mbuffer
-				bytesMap[key] = buffer
+				bytesMap[realAddr] = buffer
 			}
 
 			*buffer = append(*buffer, []byte(rframe))