slimming/netcard.go

package main

import (
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	gen "slimming/proto/gen"
	"strings"
	"sync"
	"time"

	"github.com/474420502/check"
	"github.com/474420502/perfectshutdown"
	"github.com/songgao/packets/ethernet"
	"github.com/songgao/water"
	"google.golang.org/grpc"
)

type NetTunnel struct {
	ifce     *water.Interface
	shutdown *perfectshutdown.PerfectShutdown
	writer   chan []byte
	reader   chan []byte

	clients sync.Map
	ipv4key any

	gen.UnimplementedFrameServiceServer
}

func NewNetTunnel() *NetTunnel {
	var err error
	nt := &NetTunnel{
		writer:   make(chan []byte, 1000),
		reader:   make(chan []byte, 1000),
		shutdown: perfectshutdown.New(),
		clients:  sync.Map{},
		ipv4key:  config.CurrentNodeKey,
	}

	nt.shutdown.OnClose(func() {
		time.Sleep(time.Second * 1)
		log.Println("OnClose")
		os.Exit(0)
	})

	for ipv4, address := range config.IPv4Nodes {
		nt.clients.Store(ipv4, NewRPCClient(address))
	}

	// log.Println(config)
	nt.ifce, err = water.New(water.Config{
		DeviceType: water.TAP,
		PlatformSpecificParams: water.PlatformSpecificParams{
			Name: "stun",
		},
	})
	if err != nil {
		log.Panic(err)
	}

	cmdstr := fmt.Sprintf("ip addr add %s dev stun", config.Network.Self.Virt)
	log.Println(cmdstr)
	cmd := strings.Split(cmdstr, " ")
	err = exec.Command(cmd[0], cmd[1:]...).Run()
	if err != nil {
		log.Panic(err)
	}

	cmdstr = "ip link set dev stun up"
	log.Println(cmdstr)
	cmd = strings.Split(cmdstr, " ")
	err = exec.Command(cmd[0], cmd[1:]...).Run()
	if err != nil {
		log.Panic(err)
	}

	go func() {

		for rbuf := range nt.reader {

			var frame ethernet.Frame = rbuf

			log.Println(len(frame))
			switch frame.Ethertype() {
			case ethernet.ARP:
				log.Printf("bytes len: %d type(arp): % x", len(rbuf), frame.Ethertype()) // 数据长度
				nt.clients.Range(func(key, value any) bool {
					if key == nt.ipv4key {
						return true
					}

					client := value.(*RPCClient)
					client.CheckConnect()
					client.Frame <- frame
					log.Println(client.realAddr)
					return true
				})
			case ethernet.IPv4:
				log.Printf("bytes len: %d type(ipv4): % x", len(rbuf), frame.Ethertype())
				log.Println(len(frame))
			}

			// TODO: 判断地址 转入对应的客户端发送数据

			if nt.shutdown.IsClose() {
				break
			}
		}
	}()

	go func() {
		for wbuf := range nt.writer {
			var frame ethernet.Frame = wbuf
			log.Printf("ifce write %d", len(frame))
			_, err := nt.ifce.Write(frame)
			if err != nil {
				log.Println(err)
			}

			if nt.shutdown.IsClose() {
				break
			}
		}
	}()

	return nt
}

func (nt *NetTunnel) Run() {
	var (
		frame ethernet.Frame
		err   error
		n     int
	)

	lis, err := net.Listen("tcp", config.Network.Self.Real)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	server := grpc.NewServer()
	gen.RegisterFrameServiceServer(server, nt)
	log.Printf("server listening at %v", lis.Addr())
	defer server.Stop()

	go func() {
		if err := server.Serve(lis); err != nil {
			log.Fatalf("failed to serve: %v", err)
		}
	}()

	nt.shutdown.Loop(func(index int, ps *perfectshutdown.PerfectShutdown) {
		frame.Resize(1500)
		n, err = nt.ifce.Read(frame)
		check.CheckPanic(err)
		nt.reader <- frame[0:n] // 从网卡里读数据
	})
}

func (nt *NetTunnel) SendFrames(stream gen.FrameService_SendFramesServer) error {
	log.Printf("Start: %v", stream)

	nt.shutdown.Loop(func(index int, ps *perfectshutdown.PerfectShutdown) {
		request, err := stream.Recv()
		if err != nil {
			log.Panic(err)
		}

		frames := Decompress(request.GetFrames())
		if len(frames) > 0 {
			for _, frame := range frames {
				nt.writer <- frame
			}
		}
	})

	return nil
}