star/net/natt.go

package net

import (
	"b612.me/apps/b612/netforward"
	"b612.me/starlog"
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"github.com/huin/goupnp/dcps/internetgateway2"
	"math/rand"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

type NatThroughs struct {
	Lists               []*NatThrough
	WebPort             int
	AutoUPnP            bool
	KeepAlivePeriod     int
	KeepAliveIdel       int
	KeepAliveCount      int
	STUN                string
	Remote              string
	HealthCheckInterval int
}

func (n *NatThroughs) Close() {
	for _, v := range n.Lists {
		v.Close()
	}
}

func (n *NatThroughs) Parse(reqs []string) error {
	if n.KeepAlivePeriod == 0 {
		n.KeepAlivePeriod = 10
	}
	if n.KeepAliveIdel == 0 {
		n.KeepAliveIdel = 30
	}
	if n.KeepAliveCount == 0 {
		n.KeepAliveCount = 5
	}
	if n.STUN == "" {
		n.STUN = "turn.b612.me:3478"
	}
	for _, v := range reqs {
		var req = NatThrough{
			Forward: netforward.NetForward{
				LocalAddr:       "0.0.0.0",
				DialTimeout:     time.Second * 5,
				UDPTimeout:      time.Second * 20,
				KeepAlivePeriod: n.KeepAlivePeriod,
				KeepAliveIdel:   n.KeepAliveIdel,
				KeepAliveCount:  n.KeepAliveCount,
				UsingKeepAlive:  true,
				EnableTCP:       true,
			},
			Type:                "tcp",
			STUN:                n.STUN,
			Remote:              n.Remote,
			KeepAlivePeriod:     n.KeepAlivePeriod,
			KeepAliveIdel:       n.KeepAliveIdel,
			KeepAliveCount:      n.KeepAliveCount,
			AutoUPnP:            n.AutoUPnP,
			HealthCheckInterval: n.HealthCheckInterval,
		}
		strs := strings.Split(v, ",")
		switch len(strs) {
		case 1:
			req.Type = "tcp"
			req.Forward.RemoteURI = strs[0]
		case 2:
			ipport := strings.Split(strs[0], ":")
			if len(ipport) == 1 {
				port, err := strconv.Atoi(ipport[0])
				if err != nil {
					return err
				}
				req.Forward.LocalPort = port
			} else {
				req.Forward.LocalAddr = ipport[0]
				port, err := strconv.Atoi(ipport[1])
				if err != nil {
					return err
				}
				req.Forward.LocalPort = port
			}
			req.Type = "tcp"
			req.Forward.RemoteURI = strs[1]
		case 3:
			ipport := strings.Split(strs[1], ":")
			if len(ipport) == 1 {
				port, err := strconv.Atoi(ipport[0])
				if err != nil {
					return err
				}
				req.Forward.LocalPort = port
			} else {
				req.Forward.LocalAddr = ipport[0]
				port, err := strconv.Atoi(ipport[1])
				if err != nil {
					return err
				}
				req.Forward.LocalPort = port
			}
			req.Type = "tcp"
			req.Forward.RemoteURI = strs[2]
			req.Name = strs[0]
		case 4:
			ipport := strings.Split(strs[2], ":")
			if len(ipport) == 1 {
				port, err := strconv.Atoi(ipport[0])
				if err != nil {
					return err
				}
				req.Forward.LocalPort = port
			} else {
				req.Forward.LocalAddr = ipport[0]
				port, err := strconv.Atoi(ipport[1])
				if err != nil {
					return err
				}
				req.Forward.LocalPort = port
			}
			req.Type = strings.ToLower(strs[0])
			req.Forward.RemoteURI = strs[3]
			req.Name = strs[1]
		}
		n.Lists = append(n.Lists, &req)
	}
	return nil
}

func (n *NatThroughs) Run() error {
	go n.WebService()
	wg := sync.WaitGroup{}
	for _, v := range n.Lists {
		wg.Add(1)
		go func(v *NatThrough) {
			defer wg.Done()
			if err := v.Run(); err != nil {
				starlog.Errorf("Failed to run forward: %v\n", err)
			}
			v.HealthCheck()
		}(v)
	}
	wg.Wait()
	return nil
}

type nattinfo struct {
	Id      int    `json:"id"`
	Name    string `json:"name"`
	Ext     string `json:"ext"`
	Local   string `json:"local"`
	Forward string `json:"forward"`
}

func (n *NatThroughs) WebService() error {
	if n.WebPort == 0 {
		return nil
	}
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", n.WebPort))
	if err != nil {
		return err
	}
	starlog.Infof("Web service listen on %d\n", n.WebPort)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		var str string
		for k, v := range n.Lists {
			str += fmt.Sprintf("id:%d name:%s : %s <----> %s <-----> %s\n", k, v.Name, v.ExtUrl, v.localipport, v.Forward.RemoteURI)
		}
		w.Write([]byte(str))
	})
	http.HandleFunc("/json", func(w http.ResponseWriter, r *http.Request) {
		var res []nattinfo
		for k, v := range n.Lists {
			res = append(res, nattinfo{
				Id:      k,
				Name:    v.Name,
				Ext:     v.ExtUrl,
				Local:   v.localipport,
				Forward: v.Forward.RemoteURI,
			})
		}
		w.Header().Set("Content-Type", "application/json")
		data, _ := json.Marshal(res)
		w.Write(data)
		return
	})
	http.HandleFunc("/jump", func(w http.ResponseWriter, r *http.Request) {
		types := "https://"
		name := r.URL.Query().Get("name")
		if name == "" {
			w.Write([]byte("id is empty"))
			return
		}
		if r.URL.Query().Get("type") == "http" {
			types = "http://"
		}
		for _, v := range n.Lists {
			if v.Name == name {
				http.Redirect(w, r, types+v.ExtUrl, http.StatusFound)
				return
			}
		}
	})
	return http.Serve(listener, nil)
}

// NatThrough 类似于natter.py 是一个用于Full Cone NAT直接穿透的工具
type NatThrough struct {
	Name                string
	OriginLocalPort     int
	Forward             netforward.NetForward
	Type                string
	STUN                string
	Remote              string
	KeepAlivePeriod     int
	KeepAliveIdel       int
	KeepAliveCount      int
	AutoUPnP            bool
	isOk                bool
	ExtUrl              string
	localipport         string
	keepaliveConn       net.Conn
	HealthCheckInterval int
	stopFn              context.CancelFunc
	stopCtx             context.Context
}

func (n *NatThrough) Close() {
	n.stopFn()
	n.Forward.Close()
}

func (c *NatThrough) Run() error {
	c.isOk = false
	c.stopCtx, c.stopFn = context.WithCancel(context.Background())
	c.OriginLocalPort = c.Forward.LocalPort
	if c.Forward.LocalPort == 0 {
		listener, err := net.Listen(c.Type, c.Forward.LocalAddr+":0")
		if err != nil {
			return fmt.Errorf("Failed to listen on %s: %v", c.Forward.LocalAddr, err)
		}
		if c.Type == "tcp" {
			c.Forward.LocalPort = listener.Addr().(*net.TCPAddr).Port
		} else {
			c.Forward.LocalPort = listener.Addr().(*net.UDPAddr).Port
		}
		listener.Close()
	}
	if c.Type == "tcp" {
		c.Forward.EnableTCP = true
		c.Forward.EnableUDP = false
	} else {
		c.Forward.EnableTCP = false
		c.Forward.EnableUDP = true
	}
	starlog.Infof("Local Port: %d\n", c.Forward.LocalPort)
	starlog.Infof("Keepalive To: %s\n", c.Remote)
	starlog.Infof("Forward To: %s\n", c.Forward.RemoteURI)

	innerIp, extIp, err := c.GetIPPortFromSTUN(c.Type, c.Forward.LocalAddr, c.Forward.LocalPort, c.STUN)
	if err != nil {
		return fmt.Errorf("Failed to get external IP and port: %v", err)
	}
	starlog.Infof("Internal Addr: %s \n", innerIp.String())
	starlog.Infof("External Addr: %s \n", extIp.String())
	getIP := func(ip net.Addr) string {
		switch ip.(type) {
		case *net.TCPAddr:
			return ip.(*net.TCPAddr).IP.String()
		case *net.UDPAddr:
			return ip.(*net.UDPAddr).IP.String()
		default:
			return ""
		}
	}
	getPort := func(ip net.Addr) int {
		switch ip.(type) {
		case *net.TCPAddr:
			return ip.(*net.TCPAddr).Port
		case *net.UDPAddr:
			return ip.(*net.UDPAddr).Port
		default:
			return 0
		}
	}
	go func() {
		if err := c.KeepAlive(c.Forward.LocalAddr, c.Forward.LocalPort); err != nil {
			starlog.Errorf("Failed to run forward: %v\n", err)
			c.Forward.Close()
			c.stopFn()
		}
	}()
	innerIp, extIp, err = c.GetIPPortFromSTUN(c.Type, c.Forward.LocalAddr, c.Forward.LocalPort, c.STUN)
	if err != nil {
		return fmt.Errorf("Failed to get external IP and port: %v", err)
	}
	starlog.Infof("Retest Internal Addr: %s \n", innerIp.String())
	starlog.Infof("Retest External Addr: %s \n", extIp.String())
	if c.AutoUPnP {
		go c.HandleUPnP(getIP(innerIp), uint16(getPort(extIp)))
	}
	err = c.Forward.Run()
	if err != nil {
		return fmt.Errorf("Failed to run forward: %v", err)
	}
	c.isOk = true
	c.localipport = fmt.Sprintf("%s:%d", getIP(innerIp), getPort(innerIp))
	c.ExtUrl = fmt.Sprintf("%s:%d", getIP(extIp), getPort(extIp))
	return nil
}

func (c *NatThrough) HealthCheck() {
	count := 0
	if c.HealthCheckInterval == 0 {
		c.HealthCheckInterval = 30
	}
	for {
		select {
		case <-c.stopCtx.Done():
			return
		case <-time.After(time.Second * time.Duration(c.HealthCheckInterval)):
		}
		if c.Type == "udp" {
			continue
		}
		conn, err := net.DialTimeout("tcp", c.ExtUrl, time.Second*2)
		if err != nil {
			starlog.Warningf("Health Check Fail: %v\n", err)
			count++
		} else {
			starlog.Infof("Health Check Ok\n")
			conn.(*net.TCPConn).SetLinger(0)
			conn.Close()
		}
		if count >= 3 {
			count = 0
			starlog.Errorf("Failed to connect to remote, close connection retrying\n")
			c.stopFn()
			c.keepaliveConn.Close()
			c.Forward.Close()
			forward := netforward.NetForward{
				LocalAddr:       c.Forward.LocalAddr,
				LocalPort:       c.OriginLocalPort,
				RemoteURI:       c.Forward.RemoteURI,
				KeepAlivePeriod: c.KeepAlivePeriod,
				KeepAliveIdel:   c.KeepAliveIdel,
				KeepAliveCount:  c.KeepAliveCount,
				UsingKeepAlive:  true,
			}
			time.Sleep(time.Second * 22)
			c.Forward = forward
			c.Run()
		}
	}
}

func (c *NatThrough) KeepAlive(localAddr string, localPort int) error {
	for {
		select {
		case <-c.stopCtx.Done():
			return nil
		default:
		}
		if c.Type == "tcp" {
			dialer := net.Dialer{
				Control:   netforward.ControlSetReUseAddr,
				LocalAddr: &net.TCPAddr{IP: net.ParseIP(localAddr), Port: localPort},
			}
			conn, err := dialer.Dial("tcp", c.Remote)
			if err != nil {
				starlog.Errorf("Failed to dial remote: %v\n", err)
				time.Sleep(time.Second * 5)
				continue
			}
			c.keepaliveConn = conn
			conn.(*net.TCPConn).SetLinger(0)
			netforward.SetTcpInfo(conn.(*net.TCPConn), true, c.KeepAliveIdel, c.KeepAlivePeriod, c.KeepAliveCount, 0)
			starlog.Infof("Keepalive local:%s remote: %s\n", conn.LocalAddr().String(), conn.RemoteAddr().String())
			go func() {
				for {
					str := fmt.Sprintf("HEAD /keep-alive HTTP/1.1\r\n"+
						"Host: %s\r\n"+
						"User-Agent: curl/8.0.0 (B612)\r\n"+
						"Accept: */*\r\n"+
						"Connection: keep-alive\r\n\r\n", strings.Split(c.Remote, ":")[0])
					//fmt.Println(str)
					if _, err = conn.Write([]byte(str)); err != nil {
						fmt.Println(err)
					}
					time.Sleep(time.Second * 20)
				}
			}()
			for {
				_, err := conn.Read(make([]byte, 4096))
				if err != nil {
					starlog.Warningf("Failed to keepalive remote: %v\n", err)
					conn.Close()
					break
				}
			}
		} else if c.Type == "udp" {
			rmtUdpAddr, err := net.ResolveUDPAddr("udp", c.Remote)
			if err != nil {
				return err
			}
			conn, err := net.DialUDP("udp", &net.UDPAddr{IP: net.ParseIP(localAddr), Port: localPort}, rmtUdpAddr)
			if err != nil {
				starlog.Errorf("Failed to dial remote: %v\n", err)
				time.Sleep(time.Second * 5)
				continue
			}
			c.keepaliveConn = conn
			for {
				_, err = conn.Write([]byte("b612 tcp nat through"))
				if err != nil {
					conn.Close()
					starlog.Warningf("Failed to keepalive remote: %v\n", err)
					break
				}
				time.Sleep(time.Second * 30)
			}
		}
	}
}

func (c *NatThrough) HandleUPnP(localaddr string, extPort uint16) {
	for {
		select {
		case <-c.stopCtx.Done():
			return
		default:
		}
		client, err := c.FoundUsableUPnP()
		if err != nil {
			starlog.Errorf("Failed to find UPnP device: %v\n", err)
			time.Sleep(time.Second * 20)
			continue
		}
		starlog.Infof("Found UPnP device!\n")
		_, _, _, _, _, err = client.GetSpecificPortMappingEntry("", uint16(c.Forward.LocalPort), "TCP")
		if err == nil {
			starlog.Infof("Port mapping Ok\n")
			time.Sleep(time.Second * 20)
			continue
		}
		err = client.AddPortMapping("", uint16(c.Forward.LocalPort), strings.ToUpper(c.Type), uint16(c.Forward.LocalPort), localaddr, true, "B612 TCP Nat PassThrough", 75)
		if err != nil {
			starlog.Errorf("Failed to add port mapping: %v\n", err)
			time.Sleep(time.Second * 20)
			continue
		}
		starlog.Infof("Port mapping added:externalPort %d,localAddr %s,localPort %d\n", extPort, localaddr, c.Forward.LocalPort)
		time.Sleep(time.Second * 20)
	}
}

func (c *NatThrough) GetIPPortFromSTUN(netType string, localip string, localPort int, stunServer string) (net.Addr, net.Addr, error) {
	// 替换为你的 TURN 服务器地址
	stunAddr, err := net.ResolveUDPAddr("udp", stunServer)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to resolve STUN server address: %v", err)
	}
	var conn net.Conn
	if netType == "tcp" {
		dialer := net.Dialer{
			Control:   netforward.ControlSetReUseAddr,
			LocalAddr: &net.TCPAddr{IP: net.ParseIP(localip), Port: localPort},
		}
		conn, err = dialer.Dial("tcp", stunAddr.String())
		if err != nil {
			return nil, nil, err
		}
		conn.(*net.TCPConn).SetLinger(0)
	}
	if netType == "udp" {
		conn, err = net.DialUDP(netType, &net.UDPAddr{IP: net.ParseIP(localip), Port: localPort}, stunAddr)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to connect to STUN server: %v", err)
		}
	}
	defer conn.Close()

	innerAddr := conn.LocalAddr()

	// Create STUN request
	transactionID := make([]byte, 12)
	rand.Read(transactionID)
	stunRequest := make([]byte, 20)
	binary.BigEndian.PutUint16(stunRequest[0:], 0x0001)   // Message Type: Binding Request
	binary.BigEndian.PutUint16(stunRequest[2:], 0x0000)   // Message Length
	copy(stunRequest[4:], []byte{0x21, 0x12, 0xa4, 0x42}) // Magic Cookie
	copy(stunRequest[8:], transactionID)                  // Transaction ID

	_, err = conn.Write(stunRequest)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to send STUN request: %v", err)
	}

	buf := make([]byte, 1500)
	conn.SetReadDeadline(time.Now().Add(3 * time.Second))
	n, err := conn.Read(buf)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to receive STUN response: %v", err)
	}

	// Parse STUN response
	if n < 20 {
		return nil, nil, fmt.Errorf("invalid STUN response")
	}

	payload := buf[20:n]
	var ip uint32
	var port uint16
	for len(payload) > 0 {
		attrType := binary.BigEndian.Uint16(payload[0:])
		attrLen := binary.BigEndian.Uint16(payload[2:])
		if len(payload) < int(4+attrLen) {
			return nil, nil, fmt.Errorf("invalid STUN attribute length")
		}

		if attrType == 0x0001 || attrType == 0x0020 {
			port = binary.BigEndian.Uint16(payload[6:])
			ip = binary.BigEndian.Uint32(payload[8:])
			if attrType == 0x0020 {
				port ^= 0x2112
				ip ^= 0x2112a442
			}
			break
		}
		payload = payload[4+attrLen:]
	}

	if ip == 0 || port == 0 {
		return nil, nil, fmt.Errorf("invalid STUN response")
	}

	outerAddr := &net.UDPAddr{
		IP:   net.IPv4(byte(ip>>24), byte(ip>>16), byte(ip>>8), byte(ip)),
		Port: int(port),
	}

	return innerAddr, outerAddr, nil
}

func (c *NatThrough) GetMyOutIP() string {
	tmp, err := net.Dial("udp", "8.8.8.8:53")
	if err != nil {
		return ""
	}
	return tmp.LocalAddr().(*net.UDPAddr).IP.String()
}
func (c *NatThrough) FoundUsableUPnP() (RouterClient, error) {
	wg := sync.WaitGroup{}
	found := false
	result := make(chan RouterClient, 3)
	defer close(result)
	wg.Add(3)
	go func() {
		defer wg.Done()
		clients, errors, err := internetgateway2.NewWANIPConnection2Clients()
		if err != nil {
			return
		}
		if len(errors) > 0 {
			return
		}
		if len(clients) == 0 {
			return
		}
		starlog.Infof("Found WANIPConnection2 clients:%s\n", clients[0].Location.String())
		found = true

		result <- clients[0]
	}()
	go func() {
		defer wg.Done()
		clients, errors, err := internetgateway2.NewWANIPConnection1Clients()
		if err != nil {
			return
		}
		if len(errors) > 0 {
			return
		}
		if len(clients) == 0 {
			return
		}
		starlog.Infof("Found WANIPConnection1 clients:%s\n", clients[0].Location.String())
		found = true
		result <- clients[0]
	}()
	go func() {
		defer wg.Done()
		clients, errors, err := internetgateway2.NewWANPPPConnection1Clients()
		if err != nil {
			return
		}
		if len(errors) > 0 {
			return
		}
		if len(clients) == 0 {
			return
		}
		starlog.Infof("Found WANPPPConnection1 clients:%s\n", clients[0].Location.String())
		found = true
		result <- clients[0]
	}()
	wg.Wait()
	if found {
		return <-result, nil
	}
	return nil, fmt.Errorf("no UPnP devices discovered")
}

type RouterClient interface {
	AddPortMapping(
		NewRemoteHost string,
		NewExternalPort uint16,
		NewProtocol string,
		NewInternalPort uint16,
		NewInternalClient string,
		NewEnabled bool,
		NewPortMappingDescription string,
		NewLeaseDuration uint32,
	) (err error)

	GetExternalIPAddress() (
		NewExternalIPAddress string,
		err error,
	)

	DeletePortMapping(NewRemoteHost string, NewExternalPort uint16, NewProtocol string) (err error)
	GetSpecificPortMappingEntry(NewRemoteHost string, NewExternalPort uint16, NewProtocol string) (NewInternalPort uint16, NewInternalClient string, NewEnabled bool, NewPortMappingDescription string, NewLeaseDuration uint32, err error)
}