implement godoxy-agent

cmd/main.go

@@ -5,13 +5,9 @@ import (
 	"io"
 	"log"
 	"os"
-	"os/signal"
-	"syscall"
-	"time"
 
 	"github.com/rs/zerolog"
 	"github.com/yusing/go-proxy/internal"
-	v1 "github.com/yusing/go-proxy/internal/api/v1"
 	"github.com/yusing/go-proxy/internal/api/v1/auth"
 	"github.com/yusing/go-proxy/internal/api/v1/favicon"
 	"github.com/yusing/go-proxy/internal/api/v1/query"
@@ -20,9 +16,10 @@ import (
 	E "github.com/yusing/go-proxy/internal/error"
 	"github.com/yusing/go-proxy/internal/homepage"
 	"github.com/yusing/go-proxy/internal/logging"
+	"github.com/yusing/go-proxy/internal/logging/memlogger"
 	"github.com/yusing/go-proxy/internal/net/http/middleware"
 	"github.com/yusing/go-proxy/internal/route/routes/routequery"
-	"github.com/yusing/go-proxy/internal/task"
+	"github.com/yusing/go-proxy/internal/utils"
 	"github.com/yusing/go-proxy/pkg"
 )
 
@@ -31,19 +28,21 @@ var rawLogger = log.New(os.Stdout, "", 0)
 func init() {
 	var out io.Writer = os.Stderr
 	if common.EnableLogStreaming {
-		out = zerolog.MultiLevelWriter(out, v1.GetMemLogger())
+		out = zerolog.MultiLevelWriter(out, memlogger.GetMemLogger())
 	}
 	logging.InitLogger(out)
-	// logging.AddHook(v1.GetMemLogger())
 }
 
 func main() {
 	initProfiling()
-	args := common.GetArgs()
+	args := pkg.GetArgs(common.MainServerCommandValidator{})
 
 	switch args.Command {
 	case common.CommandSetup:
-		internal.Setup()
+		Setup()
+		return
+	case common.CommandNewAgent:
+		NewAgent(args.Args)
 		return
 	case common.CommandReload:
 		if err := query.ReloadServer(); err != nil {
@@ -141,17 +140,7 @@ func main() {
 
 	config.WatchChanges()
 
-	sig := make(chan os.Signal, 1)
-	signal.Notify(sig, syscall.SIGINT)
-	signal.Notify(sig, syscall.SIGTERM)
-	signal.Notify(sig, syscall.SIGHUP)
-
-	// wait for signal
-	<-sig
-
-	// gracefully shutdown
-	logging.Info().Msg("shutting down")
-	_ = task.GracefulShutdown(time.Second * time.Duration(cfg.Value().TimeoutShutdown))
+	utils.WaitExit(cfg.Value().TimeoutShutdown)
 }
 
 func prepareDirectory(dir string) {

cmd/new_agent.go

@@ -0,0 +1,46 @@
+package main
+
+import (
+	"encoding/base64"
+	"log"
+	"net"
+	"os"
+
+	"github.com/yusing/go-proxy/agent/pkg/certs"
+)
+
+func NewAgent(args []string) {
+	if len(args) != 2 {
+		log.Fatalf("invalid arguments: %v", args)
+	}
+	host := args[0]
+	certDataBase64 := args[1]
+
+	ip, _, err := net.SplitHostPort(host)
+	if err != nil {
+		log.Fatalf("invalid host: %v", err)
+	}
+
+	_, err = net.ResolveIPAddr("ip", ip)
+	if err != nil {
+		log.Fatalf("invalid host: %v", err)
+	}
+
+	certData, err := base64.StdEncoding.DecodeString(certDataBase64)
+	if err != nil {
+		log.Fatalf("invalid cert data: %v", err)
+	}
+
+	f, err := os.OpenFile(certs.AgentCertsFilename(host), os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
+	if err != nil {
+		log.Fatalf("failed to create file: %v", err)
+	}
+	defer f.Close()
+
+	_, err = f.Write(certData)
+	if err != nil {
+		log.Fatalf("failed to write cert data: %v", err)
+	}
+
+	log.Printf("agent cert created: %s", certs.AgentCertsFilename(host))
+}

cmd/setup.go

@@ -0,0 +1,127 @@
+package main
+
+import (
+	"io"
+	"log"
+	"net/http"
+	"net/url"
+	"os"
+	"path"
+
+	"github.com/yusing/go-proxy/internal/common"
+)
+
+var (
+	branch          = common.GetEnvString("BRANCH", "v0.9")
+	baseURL         = "https://github.com/yusing/go-proxy/raw/" + branch
+	requiredConfigs = []Config{
+		{common.ConfigBasePath, true, false, ""},
+		{common.DotEnvPath, false, true, common.DotEnvExamplePath},
+		{common.ComposeFileName, false, true, common.ComposeExampleFileName},
+		{path.Join(common.ConfigBasePath, common.ConfigFileName), false, true, common.ConfigExampleFileName},
+	}
+)
+
+type Config struct {
+	Pathname         string
+	IsDir            bool
+	NeedDownload     bool
+	DownloadFileName string
+}
+
+func Setup() {
+	log.Println("setting up go-proxy")
+	log.Println("branch:", branch)
+
+	if err := os.Chdir("/setup"); err != nil {
+		log.Fatalf("failed: %s\n", err)
+	}
+
+	for _, config := range requiredConfigs {
+		config.setup()
+	}
+
+	log.Println("setup finished")
+}
+
+func (c *Config) setup() {
+	if c.IsDir {
+		mkdir(c.Pathname)
+		return
+	}
+	if !c.NeedDownload {
+		touch(c.Pathname)
+		return
+	}
+
+	fetch(c.DownloadFileName, c.Pathname)
+}
+
+func hasFileOrDir(path string) bool {
+	_, err := os.Stat(path)
+	return err == nil
+}
+
+func mkdir(pathname string) {
+	_, err := os.Stat(pathname)
+	if err != nil && os.IsNotExist(err) {
+		log.Printf("creating directory %q\n", pathname)
+		err := os.MkdirAll(pathname, 0o755)
+		if err != nil {
+			log.Fatalf("failed: %s\n", err)
+		}
+		return
+	}
+	if err != nil {
+		log.Fatalf("failed: %s\n", err)
+	}
+}
+
+func touch(pathname string) {
+	if hasFileOrDir(pathname) {
+		return
+	}
+	log.Printf("creating file %q\n", pathname)
+	_, err := os.Create(pathname)
+	if err != nil {
+		log.Fatalf("failed: %s\n", err)
+	}
+}
+
+func fetch(remoteFilename string, outFileName string) {
+	if hasFileOrDir(outFileName) {
+		if remoteFilename == outFileName {
+			log.Printf("%q already exists, not overwriting\n", outFileName)
+			return
+		}
+		log.Printf("%q already exists, downloading to %q\n", outFileName, remoteFilename)
+		outFileName = remoteFilename
+	}
+	log.Printf("downloading %q to %q\n", remoteFilename, outFileName)
+
+	url, err := url.JoinPath(baseURL, remoteFilename)
+	if err != nil {
+		log.Fatalf("unexpected error: %s\n", err)
+	}
+
+	resp, err := http.Get(url)
+	if err != nil {
+		log.Fatalf("http request failed: %s\n", err)
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		resp.Body.Close()
+		log.Fatalf("error reading response body: %s\n", err)
+	}
+
+	err = os.WriteFile(outFileName, body, 0o644)
+	if err != nil {
+		resp.Body.Close()
+		log.Fatalf("failed to write to file: %s\n", err)
+	}
+
+	log.Print("done")
+
+	resp.Body.Close()
+}