slimming/tap.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os/exec"
	gen "slimming/proto/gen"
	"strings"
	"sync/atomic"
	"time"

	"github.com/songgao/packets/ethernet"
	"github.com/songgao/water"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
)

type NetCard struct {
	FrameChan chan [][]byte
	ifce      *water.Interface
	cli       *RPCClient
	server    *RPCServer
}

type RPCClient struct {
	FrameChan chan [][]byte
	conn      *grpc.ClientConn
}

func (cli *RPCClient) connect() {
	log.Println("rpcclient start")
	// defer log.Println("rpcclient exit")

	conn, err := grpc.Dial(*othersAddr,
		grpc.WithTransportCredentials(insecure.NewCredentials()),
	)
	if err != nil {
		log.Printf("did not connect: %v", err)
		return
	}

	cli.conn = conn
	go cli.run()
	time.Sleep(time.Millisecond)
}

func (cli *RPCClient) run() {
	defer log.Println("rpcclient exit")
	defer func() {
		if err := recover(); err != nil {
			log.Println("recover")
			cli.conn = nil
		}
	}()

	c := gen.NewFrameServiceClient(cli.conn)

	stream, err := c.SendFrames(context.Background())
	if err != nil {
		log.Panic(err)
	}

	for {
		// Contact the server and print out its response.

		// 发到对面的网卡
		err = stream.Send(&gen.Request{
			Frames: <-cli.FrameChan,
		})

		if err != nil {
			log.Panic(err)
		}
	}
}

func (nc *NetCard) Run() {
	go nc.runRead()
	go nc.runWrite()
	// go nc.cli.run()
	time.Sleep(time.Second)
	log.Printf("ip addr add %s/24 dev stap", *cardIP)
	cmd := strings.Split(fmt.Sprintf("ip addr add %s/24 dev stap", *cardIP), " ")
	err := exec.Command(cmd[0], cmd[1:]...).Run()
	if err != nil {
		log.Panic(err)
	}
	cmd = strings.Split("ip link set dev stap up", " ")
	err = exec.Command(cmd[0], cmd[1:]...).Run()
	if err != nil {
		log.Panic(err)
	}
	nc.server.run()
}

func NewNetCard() *NetCard {
	flag.Parse()

	config := water.Config{
		DeviceType: water.TAP,
	}
	config.Name = "stap"

	ifce, err := water.New(config)
	if err != nil {
		log.Panic(err)
	}

	nc := &NetCard{
		FrameChan: make(chan [][]byte, 2000),
		ifce:      ifce,
		cli:       &RPCClient{FrameChan: make(chan [][]byte, 2000)},
	}
	nc.server = newRPCServer(nc)
	return nc
}

func (nc *NetCard) runRead() {
	log.Println("start netcard read")

	var ifce *water.Interface = nc.ifce

	for {

		var framesBytes [][]byte
		var isLoop int32 = 1
		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)
			n, err := ifce.Read([]byte(rframe))
			if err != nil {
				log.Fatal(err)
			}
			rframe = rframe[:n]
			framesBytes = append(framesBytes, []byte(rframe))
		}

		log.Println(len(framesBytes))
		if len(framesBytes) > 0 {
			if nc.cli.conn == nil {
				nc.cli.connect()
			}
			if nc.cli.conn != nil {
				nc.cli.FrameChan <- framesBytes // 网卡数据 发到对方
			}
		}

		// 写到grpc服务

		// log.Printf("Dst: %s\n", rframe.Destination()[0:4])
		// log.Printf("Src: %s\n", rframe.Source()[0:4])
		// log.Printf("Ethertype: % x\n", rframe.Ethertype())
		// log.Printf("Payload: % x\n", rframe.Payload())

	}
}

func (nc *NetCard) runWrite() {
	var ifce *water.Interface = nc.ifce

	for wframes := range nc.FrameChan {
		log.Printf("get wframes %d", len(wframes))
		for _, wframe := range wframes {
			_, err := ifce.Write(wframe)
			if err != nil {
				panic(err)
			}
		}
	}

}