gitee.com/aurawing/surguard-go@v0.3.1-0.20240409071558-96509a61ecf3/device/postconfig.go

package device

import (
	"encoding/hex"
	"errors"
	"fmt"
	"net"
	"net/netip"
	"os"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/jackpal/gateway"
)

const (
	ENV_SG_GROUP_NAME          = "SG_GROUP_NAME"
	ENV_SG_CONFIG_ROOT         = "SG_CONFIG_ROOT"
	ENV_SG_KMS_URL             = "SG_KMS_URL"
	ENV_SG_ZK_URL              = "SG_ZK_URL"
	ENV_SG_DEFAULT_IFACE       = "SG_DEFAULT_IFACE"
	ENV_SG_LISTEN_PORT         = "SG_LISTEN_PORT"
	ENV_SG_KEEPALIVE_INTERVAL  = "SG_KEEPALIVE_INTERVAL"
	ENV_SG_MONITOR_IP_INTERVAL = "SG_MONITOR_IP_INTERVAL"
	ENV_SG_PRIVATEKEY          = "SG_PRIVATEKEY"
	//ENV_SG_IF_BIND_INTERFACE   = "SG_IF_BIND_INTERFACE"
)

const DEFAULT_GROUP_NAME = "SURGUARD"
const DEFAULT_CONFIG_ROOT = "/etc/surguard"
const DEFAULT_LISTEN_PORT = 50001
const DEFAULT_KEEPALIVE_INTERVAL = 30
const DEFAULT_MONITOR_IP_INTERVAL = 5

var groupName string
var zkCli *ZkDiscovery
var interfaceIP string
var interfaceIPArr [][4]byte
var sgIPArr [][4]byte
var interfaceIndex byte
var finishCh chan chan bool
var keepaliveInterval int
var ifBindInterface bool

func IfBindInterface() bool {
	return ifBindInterface
}

func (device *Device) PostConfig() {
	//parse environment variables
	groupNameStr := os.Getenv(ENV_SG_GROUP_NAME)
	if groupNameStr == "" {
		groupName = DEFAULT_GROUP_NAME
	} else {
		groupName = groupNameStr
	}
	configRoot := os.Getenv(ENV_SG_CONFIG_ROOT)
	if configRoot == "" {
		configRoot = DEFAULT_CONFIG_ROOT
	}
	kmsURL := os.Getenv(ENV_SG_KMS_URL)
	if kmsURL == "" {
		device.log.Errorf("PostConfig: KMS URL is empty")
		panic(errors.New("invalid KMS URL"))
	}
	zkUrl := os.Getenv(ENV_SG_ZK_URL)
	if zkUrl == "" {
		device.log.Errorf("PostConfig: zookeeper URL is invalid: %s", zkUrl)
		panic(errors.New("invalid zookeeper URL"))
	}
	keepaliveIntervalStr := os.Getenv(ENV_SG_KEEPALIVE_INTERVAL)
	if keepaliveIntervalStr != "" {
		var err error
		keepaliveInterval, err = strconv.Atoi(keepaliveIntervalStr)
		if err != nil {
			device.log.Errorf("PostConfig: parsing keepalive interval failed: %s", err)
			keepaliveInterval = DEFAULT_KEEPALIVE_INTERVAL
		}
	} else {
		keepaliveInterval = DEFAULT_KEEPALIVE_INTERVAL
	}

	monitorIPInterval := DEFAULT_MONITOR_IP_INTERVAL
	monitorIPIntervalStr := os.Getenv(ENV_SG_MONITOR_IP_INTERVAL)
	if monitorIPIntervalStr != "" {
		var err error
		monitorIPInterval, err = strconv.Atoi(monitorIPIntervalStr)
		if err != nil {
			device.log.Errorf("PostConfig: parse monitor IP interval failed: %s", err)
			monitorIPInterval = DEFAULT_MONITOR_IP_INTERVAL
		}
	}

	// ifBindInterfaceStr := os.Getenv(ENV_SG_IF_BIND_INTERFACE)
	// if ifBindInterfaceStr == "" {
	// 	ifBindInterface = false
	// } else {
	// 	var err error
	// 	ifBindInterface, err = strconv.ParseBool(ifBindInterfaceStr)
	// 	if err != nil {
	// 		device.log.Errorf("PostConfig: ifBindInterface is invalid: %s -> %s", ifBindInterfaceStr, err)
	// 		panic(err)
	// 	}
	// }

	var err error
	finishCh = make(chan chan bool)

	// ipaddrs, err := device.getIP()
	// if err != nil {
	// 	device.log.Errorf("PostConfig: failed to get IP address: %s", err)
	// 	panic(err)
	// }
	// interfaceIP = strings.Join(ipaddrs, ",")
	// interfaceIPArr = make([][4]byte, 0)
	// for _, ipaddr := range ipaddrs {
	// 	tmp := net.ParseIP(ipaddr)
	// 	interfaceIPArr[0] = tmp[12]
	// 	interfaceIPArr[1] = tmp[13]
	// 	interfaceIPArr[2] = tmp[14]
	// 	interfaceIPArr[3] = tmp[15]
	// }

	err = device.configTunDevice()
	if err != nil {
		device.log.Errorf("PostConfig: config TUN device failed: %s", err)
		panic(err)
	}
	sk, err := device.loginKMS(kmsURL, configRoot)
	if err != nil {
		device.log.Errorf("PostConfig: login KMS failed: %s", err)
		panic(err)
	}
	err = device.setSKandPort(sk)
	if err != nil {
		device.log.Errorf("PostConfig: set private key and listen port failed: %s", err)
		panic(err)
	}
	err = device.initRules()
	if err != nil {
		device.log.Errorf("PostConfig: initialize rules failed: %s", err)
		panic(err)
	}
	// interfaceIPStr, err := getIP()
	// if err != nil {
	// 	panic(err)
	// }
	// interfaceIP = interfaceIPStr

	zkCli, err = CreateZkDiscovery(zkUrl, device, fmt.Sprintf("/%s", groupName))
	if err != nil {
		device.log.Errorf("PostConfig: create zookeeper client failed: %s", err)
		panic(err)
	}
	device.monitorIP(true)
	go func() {
		var ch chan bool
	OUTER:
		for {
			select {
			case ch = <-finishCh:
				break OUTER
			case <-time.After(time.Second * time.Duration(monitorIPInterval)):
				device.monitorIP(false)
			}
		}
		ch <- true
	}()
}

func (device *Device) monitorIP(init bool) {
	zkCli.Lock()
	pathExist, err := zkCli.ExistPeer(device.staticIdentity.publicKey)
	if err != nil {
		device.log.Errorf("PostConfig: error when get path from zk: %s", err)
	}
	if !pathExist && interfaceIP != "" && !init {
		ifAddrs := strings.Split(interfaceIP, ",")
		rpAddrs := make([]string, 0)
		for _, rpAddr := range ifAddrs {
			rpAddrs = append(rpAddrs, fmt.Sprintf("%s:%d", rpAddr, device.net.port))
		}
		zkCli.AddPeer(device.staticIdentity.publicKey, strings.Join(rpAddrs, ","))
	}
	zkCli.Unlock()
	//TODO: 检查zookeeper上的IP地址是否为空
	data := ""
	data, _ = zkCli.GetData(device.staticIdentity.publicKey)
	ipaddrs, err := device.getIP()
	if err != nil {
		device.log.Errorf("PostConfig: error when get interface IP: %s", err)
	} else {
		sort.Strings(ipaddrs)
		if data == "" || strings.Join(ipaddrs, ",") != interfaceIP || (ifBindInterface && init) {
			zkCli.Lock()
			if data == "" || !(interfaceIP == "" || (ifBindInterface && init)) {
				zkCli.RemovePeer(device.staticIdentity.publicKey)
			}
			reportAddrs := make([]string, 0)
			for _, reportAddr := range ipaddrs {
				reportAddrs = append(reportAddrs, fmt.Sprintf("%s:%d", reportAddr, device.net.port))
			}
			zkCli.AddPeer(device.staticIdentity.publicKey, strings.Join(reportAddrs, ","))
			interfaceIP = strings.Join(ipaddrs, ",")
			// tmp := net.ParseIP(ipaddrs[0])
			// interfaceIPArr[0] = tmp[12]
			// interfaceIPArr[1] = tmp[13]
			// interfaceIPArr[2] = tmp[14]
			// interfaceIPArr[3] = tmp[15]
			if ifBindInterface {
				oldsgLen := len(sgIPArr)
				interfaceIPArr = make([][4]byte, 0)
				sgIPArr = make([][4]byte, 0)

				for i, ipaddr := range ipaddrs {
					tmp := net.ParseIP(ipaddr)
					tmp2 := [4]byte{tmp[12], tmp[13], tmp[14], tmp[15]}
					interfaceIPArr = append(interfaceIPArr, tmp2)
					sgIPArr = append(sgIPArr, [4]byte{169, 254, interfaceIndex, byte(i + 1)})
				}
				if oldsgLen < len(sgIPArr) {
					for i := oldsgLen + 1; i <= len(sgIPArr); i++ {
						staticip := fmt.Sprintf("169.254.%d.%d", interfaceIndex, i)
						device.runCmd(false, "netsh", "interface", "ipv4", "add", "address", "name=\""+deviceName+"\"", staticip, "255.255.255.0")
						device.log.Verbosef("configTunDevice: add address %s/24 for device %s\n", staticip, deviceName)
					}
				} else if oldsgLen > len(sgIPArr) {
					for i := oldsgLen; i > len(sgIPArr); i-- {
						staticip := fmt.Sprintf("169.254.%d.%d", interfaceIndex, i)
						device.runCmd(false, "netsh", "interface", "ipv4", "delete", "address", "name=\""+deviceName+"\"", staticip)
						device.log.Verbosef("configTunDevice: delete address %s/24 for device %s\n", staticip, deviceName)
					}
				}
			}

			zkCli.Unlock()
			if ifBindInterface {
				device.net.Lock()
				device.net.ipv4Addr = ipaddrs[0]
				device.net.Unlock()
				if err := device.BindUpdate(); err != nil {
					device.log.Errorf("PostConfig: failed to set ip address %s: %s", ipaddrs[0], err)
				}
			}
		}
	}
}

// login to KMS and get private key
func (device *Device) loginKMS(kmsURL, configRoot string) (NoisePrivateKey, error) {
	skstr := os.Getenv(ENV_SG_PRIVATEKEY)
	if skstr != "" {
		return device.parsePrivateKey(skstr)
	}
	yidRoot := fmt.Sprintf("%s/%s", configRoot, groupName)
	err := os.MkdirAll(yidRoot, os.ModePerm)
	if err != nil {
		device.log.Errorf("PostConfig: create config path %s failed: %s", yidRoot, err)
		return NoisePrivateKey{}, err
	}
	yidPath := fmt.Sprintf("%s/%s", yidRoot, "kms.conf")

	//login and get private key
	var vui, yid, yentity string
	ystr, err := os.ReadFile(yidPath)
	if err != nil {
		//create entity
		vuin, yidn, yentityn, err := device.createEntity(kmsURL)
		if err != nil {
			device.log.Errorf("PostConfig: create entity failed: %s", err)
			return NoisePrivateKey{}, err
		}
		err = os.WriteFile(yidPath, []byte(fmt.Sprintf("%s|%s|%s", vuin, yidn, yentityn)), 0666)
		if err != nil {
			device.log.Errorf("PostConfig: write YID and entity to file failed: %s", err)
			return NoisePrivateKey{}, err
		}
		vui = vuin
		yid = yidn
		yentity = yentityn
	} else {
		ystrs := strings.Split(string(ystr), "|")
		if len(ystrs) != 3 {
			err := fmt.Errorf("format of config file is corrupt: %s", yidPath)
			device.log.Errorf("PostConfig: %s", err)
			return NoisePrivateKey{}, err
		}
		vui = strings.TrimSpace(ystrs[0])
		yid = strings.TrimSpace(ystrs[1])
		yentity = strings.TrimSpace(ystrs[2])
	}
	err = device.loginEntity(kmsURL, vui, yid)
	if err != nil {
		device.log.Errorf("PostConfig: login entity failed: %s", err)
		return NoisePrivateKey{}, err
	}
	yentityid, err := strconv.Atoi(yentity)
	if err != nil {
		device.log.Errorf("PostConfig: parse entity to number failed: %s", err)
		return NoisePrivateKey{}, err
	}
	skstr, err = device.getPrivateKey(kmsURL, yentityid)
	if err != nil {
		device.log.Errorf("PostConfig: get private key from KMS failed: %s", err)
		return NoisePrivateKey{}, err
	}
	return device.parsePrivateKey(skstr)
}

func (device *Device) getIP() ([]string, error) {
	interfaceName := os.Getenv(ENV_SG_DEFAULT_IFACE)
	if interfaceName == "" {
		ip, err := gateway.DiscoverInterface()
		if err != nil {
			device.log.Errorf("PostConfig: error when discovery interface IP: %s", err)
			return nil, err
		}
		if ip.To4() != nil {
			return []string{ip.String()}, nil
		} else {
			err := errors.New("no valid IPv4 address")
			device.log.Errorf("PostConfig: error when convert interface IP: %s", err)
			return nil, err
		}
	}

	iface, err := net.InterfaceByName(interfaceName)
	if err != nil {
		device.log.Errorf("PostConfig: error when get interface IP by name: %s", err)
		return nil, err
	}

	addrs, err := iface.Addrs()
	if err != nil {
		device.log.Errorf("PostConfig: error when get IP address of interface: %s", err)
		return nil, err
	}
	addrstrs := make([]string, 0)
	for _, addr := range addrs {
		ipNet, ok := addr.(*net.IPNet)
		if ok && !ipNet.IP.IsLoopback() {
			if ipNet.IP.To4() != nil {
				addrstrs = append(addrstrs, ipNet.IP.String())
				//return ipNet.IP.String(), nil
			}
		}
	}
	//TODO：增加指定IP地址功能
	if len(addrstrs) == 0 {
		err = errors.New("no valid IPv4 address")
		device.log.Errorf("PostConfig: error when get IP address of interface: %s", err)
		return nil, err
	}
	sort.Strings(addrstrs)
	return addrstrs, nil
}

// set up private key and UDP listen port
func (device *Device) setSKandPort(sk NoisePrivateKey) error {
	listenPort := DEFAULT_LISTEN_PORT
	portStr := os.Getenv(ENV_SG_LISTEN_PORT)
	var err error
	if portStr != "" {
		listenPort, err = strconv.Atoi(portStr)
		if err != nil {
			device.log.Errorf("PostConfig: failed to parse listen port %s: %s", portStr, err)
		}
	}
	//ischanged := false
	device.ipcMutex.Lock()
	//oldPK := device.staticIdentity.publicKey
	defer device.ipcMutex.Unlock()
	if !device.staticIdentity.privateKey.Equals(sk) {
		device.SetPrivateKey(sk)
	}
	// ipv4 := ""
	// ipv4s, err := device.getIP()
	// if err != nil {
	// 	device.log.Errorf("PostConfig: failed to get IP address: %s", err)
	// } else if len(ipv4s) > 0 {
	// 	ipv4 = ipv4s[0]
	// }
	var ipv4 string
	if ifBindInterface && len(interfaceIPArr) > 0 {
		ipv4 = fmt.Sprintf("%d.%d.%d.%d", interfaceIPArr[0][0], interfaceIPArr[0][1], interfaceIPArr[0][2], interfaceIPArr[0][3])
	} else {
		ipv4s, err := device.getIP()
		if err != nil {
			device.log.Errorf("PostConfig: failed to get IP address: %s", err)
		} else if len(ipv4s) > 0 {
			sort.Strings(ipv4s)
			ipv4 = ipv4s[0]
		}
	}
	if device.net.port != uint16(listenPort) {
		device.net.Lock()
		if ifBindInterface && ipv4 != "" {
			device.net.ipv4Addr = ipv4
		}
		device.net.port = uint16(listenPort)
		device.net.Unlock()
		if err := device.BindUpdate(); err != nil {
			device.log.Errorf("PostConfig: failed to set listen_port %d: %s", listenPort, err)
			return err
		}
		//ischanged = true
	}
	// if ischanged {
	// 	//TODO: update zk
	// 	zkCli.Lock()
	// 	zkCli.RemovePeer(oldPK)
	// 	zkCli.AddPeer(device.staticIdentity.publicKey)
	// 	zkCli.Unlock()
	// }
	return nil
}

func (device *Device) AddPeer(pk NoisePublicKey, endPoints string) error {
	endPointArr := strings.Split(endPoints, ",")
	if len(endPointArr) == 0 {
		device.log.Verbosef("PostConfig: no endpoints under public key %s", hex.EncodeToString(pk[:]))
		return errors.New("no endpoints")
	}
	device.ipcMutex.Lock()
	defer device.ipcMutex.Unlock()
	peer := device.LookupPeer(pk)
	if peer != nil {
		//err := errors.New("peer exists")
		device.log.Verbosef("PostConfig: peer %v exists", peer)
		return nil
	}
	peer, err := device.NewPeer(pk)
	if err != nil {
		device.log.Errorf("PostConfig: create peer failed: %s", err)
		return err
	}
	device.log.Verbosef("%v - PostConfig: Created", peer)
	endpoint, err := device.net.bind.ParseEndpoint(endPointArr[0])
	if err != nil {
		device.log.Errorf("%v - PostConfig: parse endpoint failed: %s", peer, err)
		return err
	}
	peer.endpoint.Lock()
	peer.endpoint.val = endpoint
	peer.endpoint.Unlock()
	device.log.Verbosef("%v - PostConfig: Updating endpoint %s", peer, endPointArr[0])

	device.allowedips.RemoveByPeer(peer)
	for _, endPoint := range endPointArr {
		allowedIP := fmt.Sprintf("%s/32", strings.Split(endPoint, ":")[0])
		prefix, err := netip.ParsePrefix(allowedIP)
		if err != nil {
			device.log.Errorf("%v - PostConfig: parse allowedIP failed: %s", peer, err)
			return err
		}
		device.allowedips.Insert(prefix, peer)
		device.log.Verbosef("%v - PostConfig: Adding allowedip: %s", peer, allowedIP)
	}

	old := peer.persistentKeepaliveInterval.Swap(uint32(keepaliveInterval))
	device.log.Verbosef("%v - PostConfig: Updating persistent keepalive interval %d", peer, keepaliveInterval)

	ipcPeer := new(ipcSetPeer)
	ipcPeer.Peer = peer
	ipcPeer.dummy = false
	ipcPeer.created = true
	ipcPeer.pkaOn = old == 0 && keepaliveInterval != 0

	if ipcPeer.created {
		ipcPeer.endpoint.disableRoaming = ipcPeer.device.net.brokenRoaming && ipcPeer.endpoint.val != nil
	}
	if ipcPeer.device.isUp() {
		ipcPeer.Start()
		if ipcPeer.pkaOn {
			ipcPeer.SendKeepalive()
		}
		ipcPeer.SendStagedPackets()
	}
	return nil
}

func (device *Device) DeletePeer(pk NoisePublicKey) {
	device.ipcMutex.Lock()
	defer device.ipcMutex.Unlock()
	peer := device.LookupPeer(pk)
	if peer != nil {
		device.RemovePeer(peer.handshake.remoteStatic)
	}
}

func (device *Device) ClearConfig() {
	ch := make(chan bool)
	finishCh <- ch
	<-ch
	zkCli.Close()
	device.clearConfigOSSpecific()
}

func (device *Device) IterPeerEndpoint(f func(string)) {
	device.peers.RLock()
	defer device.peers.RUnlock()
	for _, peer := range device.peers.keyMap {
		dstIP := peer.endpoint.val.DstIP().String()
		f(dstIP)
	}
}