simplify setup process

cmd/add_agent.go

@@ -0,0 +1,65 @@
+package main
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"os"
+	"time"
+
+	"github.com/yusing/go-proxy/agent/pkg/certs"
+	"github.com/yusing/go-proxy/internal/logging"
+)
+
+func AddAgent(args []string) {
+	if len(args) != 1 {
+		logging.Fatal().Msgf("invalid arguments: %v, expect host", args)
+	}
+	host := args[0]
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	req, err := http.NewRequestWithContext(ctx, "GET", "http://"+host, nil)
+	if err != nil {
+		logging.Fatal().Err(err).Msg("failed to create request")
+	}
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		logging.Fatal().Err(err).Msg("failed to send request")
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		logging.Fatal().Int("status", resp.StatusCode).Msg("failed to add agent")
+	}
+
+	zip, err := io.ReadAll(resp.Body)
+	if err != nil {
+		logging.Fatal().Err(err).Msg("failed to read response body")
+	}
+
+	f, err := os.OpenFile(certs.AgentCertsFilename(host), os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
+	if err != nil {
+		logging.Fatal().Err(err).Msg("failed to create client certs file")
+	}
+	defer f.Close()
+
+	if _, err := f.Write(zip); err != nil {
+		logging.Fatal().Err(err).Msg("failed to save client certs")
+	}
+
+	logging.Info().Msgf("agent %s added, certs saved to %s", host, certs.AgentCertsFilename(host))
+
+	req, err = http.NewRequestWithContext(ctx, "GET", "http://"+host+"/done", nil)
+	if err != nil {
+		logging.Fatal().Err(err).Msg("failed to create done request")
+	}
+
+	resp, err = http.DefaultClient.Do(req)
+	if err != nil {
+		logging.Fatal().Err(err).Msg("failed to send done request")
+	}
+	defer resp.Body.Close()
+}

cmd/main.go

@@ -2,11 +2,9 @@ package main
 
 import (
 	"encoding/json"
-	"io"
 	"log"
 	"os"
 
-	"github.com/rs/zerolog"
 	"github.com/yusing/go-proxy/internal"
 	"github.com/yusing/go-proxy/internal/api/v1/auth"
 	"github.com/yusing/go-proxy/internal/api/v1/favicon"
@@ -16,7 +14,6 @@ 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/utils"
@@ -25,14 +22,6 @@ import (
 
 var rawLogger = log.New(os.Stdout, "", 0)
 
-func init() {
-	var out io.Writer = os.Stderr
-	if common.EnableLogStreaming {
-		out = zerolog.MultiLevelWriter(out, memlogger.GetMemLogger())
-	}
-	logging.InitLogger(out)
-}
-
 func main() {
 	initProfiling()
 	args := pkg.GetArgs(common.MainServerCommandValidator{})
@@ -41,8 +30,8 @@ func main() {
 	case common.CommandSetup:
 		Setup()
 		return
-	case common.CommandNewAgent:
-		NewAgent(args.Args)
+	case common.CommandAddAgent:
+		AddAgent(args.Args)
 		return
 	case common.CommandReload:
 		if err := query.ReloadServer(); err != nil {

cmd/new_agent.go

@@ -1,46 +0,0 @@
-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))
-}