waf-platform/EdgeHttpDNS/internal/nodes/httpdns_node.go

package nodes

import (
	"errors"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"runtime"
	"sync"
	"time"

	teaconst "github.com/TeaOSLab/EdgeHttpDNS/internal/const"
	"github.com/TeaOSLab/EdgeHttpDNS/internal/utils"
	"github.com/iwind/TeaGo/maps"
	"github.com/iwind/gosock/pkg/gosock"
)

var DaemonIsOn = false
var DaemonPid = 0

type HTTPDNSNode struct {
	sock *gosock.Sock

	quitOnce sync.Once
	quitCh   chan struct{}
}

func NewHTTPDNSNode() *HTTPDNSNode {
	return &HTTPDNSNode{
		sock:   gosock.NewTmpSock(teaconst.ProcessName),
		quitCh: make(chan struct{}),
	}
}

func (n *HTTPDNSNode) Run() {
	_, ok := os.LookupEnv("EdgeDaemon")
	if ok {
		DaemonIsOn = true
		DaemonPid = os.Getppid()
	}

	err := n.listenSock()
	if err != nil {
		log.Println("[HTTPDNS_NODE]" + err.Error())
		return
	}

	go func() {
		defer func() {
			if r := recover(); r != nil {
				log.Println(fmt.Sprintf("[HTTPDNS_NODE]goroutine panic: %v", r))
			}
		}()
		n.start()
	}()
	select {}
}

func (n *HTTPDNSNode) Daemon() {
	path := os.TempDir() + "/" + teaconst.ProcessName + ".sock"
	for {
		conn, err := net.DialTimeout("unix", path, 1*time.Second)
		if err != nil {
			exe, exeErr := os.Executable()
			if exeErr != nil {
				log.Println("[DAEMON]", exeErr)
				time.Sleep(1 * time.Second)
				continue
			}

			cmd := exec.Command(exe)
			cmd.Env = append(os.Environ(), "EdgeBackground=on", "EdgeDaemon=on")
			if runtime.GOOS != "windows" {
				cmd.Stdout = os.Stdout
				cmd.Stderr = os.Stderr
			}
			startErr := cmd.Start()
			if startErr != nil {
				log.Println("[DAEMON]", startErr)
				time.Sleep(1 * time.Second)
				continue
			}
			_ = cmd.Wait()
			time.Sleep(5 * time.Second)
			continue
		}

		_ = conn.Close()
		time.Sleep(5 * time.Second)
	}
}

func (n *HTTPDNSNode) InstallSystemService() error {
	exe, err := os.Executable()
	if err != nil {
		return err
	}

	manager := utils.NewServiceManager(teaconst.SystemdServiceName, teaconst.ProductName)
	return manager.Install(exe, []string{})
}

func (n *HTTPDNSNode) listenSock() error {
	if runtime.GOOS == "windows" {
		return nil
	}

	if n.sock.IsListening() {
		reply, err := n.sock.Send(&gosock.Command{Code: "pid"})
		if err == nil {
			return errors.New("the process is already running, pid: " + maps.NewMap(reply.Params).GetString("pid"))
		}
		return errors.New("the process is already running")
	}

	go func() {
		defer func() {
			if r := recover(); r != nil {
				log.Println(fmt.Sprintf("[HTTPDNS_NODE]sock goroutine panic: %v", r))
			}
		}()

		n.sock.OnCommand(func(cmd *gosock.Command) {
			switch cmd.Code {
			case "pid":
				_ = cmd.Reply(&gosock.Command{
					Code: "pid",
					Params: map[string]interface{}{
						"pid": os.Getpid(),
					},
				})
			case "info":
				exePath, _ := os.Executable()
				_ = cmd.Reply(&gosock.Command{
					Code: "info",
					Params: map[string]interface{}{
						"pid":     os.Getpid(),
						"version": teaconst.Version,
						"path":    exePath,
					},
				})
			case "stop":
				_ = cmd.ReplyOk()
				n.stop()
				time.Sleep(100 * time.Millisecond)
				os.Exit(0)
			}
		})

		err := n.sock.Listen()
		if err != nil {
			log.Println("[HTTPDNS_NODE][sock]", err.Error())
		}
	}()

	return nil
}

func (n *HTTPDNSNode) start() {
	log.Println("[HTTPDNS_NODE]started")

	snapshotManager := NewSnapshotManager(n.quitCh)
	statusManager := NewStatusManager(n.quitCh)
	taskManager := NewTaskManager(n.quitCh, snapshotManager)
	resolveServer := NewResolveServer(n.quitCh, snapshotManager)

	go func() {
		defer func() {
			if r := recover(); r != nil {
				log.Println(fmt.Sprintf("[HTTPDNS_NODE]snapshot goroutine panic: %v", r))
			}
		}()
		snapshotManager.Start()
	}()
	go func() {
		defer func() {
			if r := recover(); r != nil {
				log.Println(fmt.Sprintf("[HTTPDNS_NODE]status goroutine panic: %v", r))
			}
		}()
		statusManager.Start()
	}()
	go func() {
		defer func() {
			if r := recover(); r != nil {
				log.Println(fmt.Sprintf("[HTTPDNS_NODE]task goroutine panic: %v", r))
			}
		}()
		taskManager.Start()
	}()
	go func() {
		defer func() {
			if r := recover(); r != nil {
				log.Println(fmt.Sprintf("[HTTPDNS_NODE]resolve goroutine panic: %v", r))
			}
		}()
		resolveServer.Start()
	}()

	go func() {
		defer func() {
			if r := recover(); r != nil {
				log.Println(fmt.Sprintf("[HTTPDNS_NODE]upgrade goroutine panic: %v", r))
			}
		}()
		NewUpgradeManager().Loop()
	}()
}

func (n *HTTPDNSNode) stop() {
	n.quitOnce.Do(func() {
		close(n.quitCh)
		_ = n.sock.Close()
	})
}