star/net/nats.go

package net

import (
	"b612.me/starlog"
	"context"
	"fmt"
	"net"
	"strings"
	"sync/atomic"
)

type NatTesterServer struct {
	MainPort string
	AltPort  string
	MainIP   string
	AltIP    string
	LogPath  string
	stopCtx  context.Context
	stopFn   context.CancelFunc
	maina    *net.UDPConn
	mainb    *net.UDPConn
	alt      *net.UDPConn
	running  int32
}

func (n *NatTesterServer) Run() error {
	if atomic.LoadInt32(&n.running) > 0 {
		starlog.Errorln("already running")
		return fmt.Errorf("already running")
	}
	atomic.StoreInt32(&n.running, 1)
	defer atomic.StoreInt32(&n.running, 0)
	if n.LogPath != "" {
		starlog.SetLogFile(n.LogPath, starlog.Std, true)
		starlog.Infof("Log file set to %s\n", n.LogPath)
	}
	starlog.Infof("MainPort: %s\n", n.MainPort)
	starlog.Infof("AltPort: %s\n", n.AltPort)
	tmp, err := net.Dial("udp", "8.8.8.8:53")
	if err != nil {
		return err
	}
	starlog.Infof("Current Output IP: %s\n", tmp.LocalAddr().(*net.UDPAddr).IP.String())
	tmp.Close()
	n.stopCtx, n.stopFn = context.WithCancel(context.Background())
	mainaaddr, err := net.ResolveUDPAddr("udp", n.MainIP+":"+n.MainPort)
	if err != nil {
		return err
	}
	mainbaddr, err := net.ResolveUDPAddr("udp", n.MainIP+":"+n.AltPort)
	if err != nil {
		return err
	}
	n.maina, err = net.ListenUDP("udp", mainaaddr)
	if err != nil {
		return err
	}
	starlog.Infof("UDP MainIP:MainPort Listening on %s\n", n.maina.LocalAddr().String())
	n.mainb, err = net.ListenUDP("udp", mainbaddr)
	if err != nil {
		return err
	}
	starlog.Infof("UDP MainIP:AltPort Listening on %s\n", n.mainb.LocalAddr().String())
	altaddr, err := net.ResolveUDPAddr("udp", n.AltIP+":"+n.AltPort)
	if err != nil {
		return err
	}
	n.alt, err = net.ListenUDP("udp", altaddr)
	if err != nil {
		return err
	}
	starlog.Infof("UDP AltIP:AltPort Listening on %s\n", n.alt.LocalAddr().String())
	go func() {
		for {
			select {
			case <-n.stopCtx.Done():
				starlog.Infoln("Stopping,Reason: Context Done")
				return
			default:
			}
			buf := make([]byte, 1024)
			num, r, e := n.alt.ReadFromUDP(buf)
			if e != nil {
				continue
			}
			go n.Analyse(n.alt, r, strings.Split(string(buf[:num]), "::"))
		}
	}()
	for {
		select {
		case <-n.stopCtx.Done():
			starlog.Infoln("Stopping,Reason: Context Done")
			n.maina.Close()
			n.mainb.Close()
			n.alt.Close()
			return nil
		default:
		}
		buf := make([]byte, 1024)
		num, r, e := n.maina.ReadFromUDP(buf)
		if e != nil {
			continue
		}
		go n.Analyse(n.maina, r, strings.Split(string(buf[:num]), "::"))
	}
}

func (n *NatTesterServer) Analyse(c *net.UDPConn, r *net.UDPAddr, cmds []string) error {
	switch cmds[0] {
	case "ip":
		c.WriteToUDP([]byte("ip::"+r.String()), r)
		starlog.Infof("Recv IP Request from %s,Local: %s\n", r.String(), c.LocalAddr().String())
	case "startnat1":
		n.alt.WriteToUDP([]byte("stage1"), r)
		starlog.Infof("Start NAT1 Test from %s,Recv Local:%s Send Local:%s\n", r.String(), c.LocalAddr().String(), n.alt.LocalAddr().String())
	case "stage1recv":
		n.maina.WriteToUDP([]byte("result::nat1"), r)
		starlog.Infof("Recv NAT1 Test from %s,Recv Local:%s Send Local:%s\n", r.String(), c.LocalAddr().String(), n.maina.LocalAddr().String())
	case "startnat2":
		n.mainb.WriteToUDP([]byte("stage2"), r)
		starlog.Infof("Start NAT2 Test from %s,Recv Local:%s Send Local:%s\n", r.String(), c.LocalAddr().String(), n.mainb.LocalAddr().String())
	case "stage2recv":
		n.maina.WriteToUDP([]byte("result::nat2"), r)
		starlog.Infof("Recv NAT2 Test from %s,Recv Local:%s Send Local:%s\n", r.String(), c.LocalAddr().String(), n.maina.LocalAddr().String())
	case "startnat3":
		n.maina.WriteToUDP([]byte("stage3"), r)
		starlog.Infof("Start NAT3 Test from %s,Recv Local:%s Send Local:%s\n", r.String(), c.LocalAddr().String(), n.maina.LocalAddr().String())
	case "stage3recv":
		n.maina.WriteToUDP([]byte("result::nat3"), r)
		starlog.Infof("Recv NAT3 Test from %s,Recv Local:%s Send Local:%s\n", r.String(), c.LocalAddr().String(), n.maina.LocalAddr().String())
	}
	return nil
}

func main() {
	server := NatTesterServer{
		MainIP:   "10.0.0.5",
		AltIP:    "10.0.0.2",
		MainPort: "41127",
		AltPort:  "46610",
	}
	fmt.Println(server.Run())
}