portrd option to choose b/w admin, tunnel or both servers to run

.air.toml

@@ -4,7 +4,7 @@ tmp_dir = "tmp"
 
 [build]
 args_bin = []
-bin = "./tmp/main start -c configs/server.yaml"
+bin = "./tmp/main -c configs/server.yaml start all"
 cmd = "go build -o ./tmp/main cmd/portrd/main.go"
 delay = 1000
 exclude_dir = [

cmd/portrd/main.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"log"
 	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/amalshaji/portr/internal/server/admin"
 	"github.com/amalshaji/portr/internal/server/admin/service"
@@ -19,39 +21,61 @@ import (
 
 const VERSION = "0.0.1-beta"
 
+type ServiceToRun int
+
+const (
+	ADMIN_SERVER  ServiceToRun = iota + 1
+	TUNNEL_SERVER              = iota + 1
+	ALL_SERVERS                = iota + 1
+)
+
 func main() {
 	app := &cli.App{
 		Name:    "portrd",
 		Usage:   "portr server",
 		Version: VERSION,
+		Flags: []cli.Flag{
+			&cli.StringFlag{
+				Name:    "config",
+				Aliases: []string{"c"},
+				Usage:   "config file",
+				Value:   "config.yaml",
+			},
+		},
 		Commands: []*cli.Command{
 			{
 				Name:  "start",
-				Usage: "Start the portr server",
-				Flags: []cli.Flag{
-					&cli.StringFlag{
-						Name:    "config",
-						Aliases: []string{"c"},
-						Usage:   "config file",
-						Value:   "config.yaml",
+				Usage: "Specify the server to start",
+				Subcommands: []*cli.Command{
+					{
+						Name:  "admin",
+						Usage: "Start the admin server",
+						Action: func(c *cli.Context) error {
+							start(c.String("config"), ADMIN_SERVER)
+							return nil
+						},
+					},
+					{
+						Name:  "tunnel",
+						Usage: "Start the tunnel server",
+						Action: func(c *cli.Context) error {
+							start(c.String("config"), TUNNEL_SERVER)
+							return nil
+						},
+					},
+					{
+						Name:  "all",
+						Usage: "Start both admin and tunnel servers",
+						Action: func(c *cli.Context) error {
+							start(c.String("config"), ALL_SERVERS)
+							return nil
+						},
 					},
-				},
-				Action: func(c *cli.Context) error {
-					start(c.String("config"))
-					return nil
 				},
 			},
 			{
 				Name:  "migrate",
 				Usage: "Run database migrations",
-				Flags: []cli.Flag{
-					&cli.StringFlag{
-						Name:    "config",
-						Aliases: []string{"c"},
-						Usage:   "config file",
-						Value:   "config.yaml",
-					},
-				},
 				Action: func(c *cli.Context) error {
 					migrate(c.String("config"))
 					return nil
@@ -65,7 +89,7 @@ func main() {
 	}
 }
 
-func start(configFilePath string) {
+func start(configFilePath string, toRun ServiceToRun) {
 	config, err := config.Load(configFilePath)
 	if err != nil {
 		log.Fatal(err)
@@ -91,10 +115,26 @@ func start(configFilePath string) {
 	adminServer := admin.New(config, service)
 	cron := cron.New(_db, config)
 
-	go proxyServer.Start()
-	go sshServer.Start()
-	go cron.Start()
-	adminServer.Start()
+	if toRun == TUNNEL_SERVER || toRun == ALL_SERVERS {
+		go proxyServer.Start()
+		defer proxyServer.Shutdown(context.TODO())
+
+		go sshServer.Start()
+		defer sshServer.Shutdown(context.TODO())
+	}
+
+	if toRun == ADMIN_SERVER || toRun == ALL_SERVERS {
+		go adminServer.Start()
+		defer adminServer.Shutdown()
+
+		go cron.Start()
+		defer cron.Shutdown()
+	}
+
+	done := make(chan os.Signal, 1)
+	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
+
+	<-done
 }
 
 func migrate(configFilePath string) {

docker/Dockerfile

@@ -30,5 +30,5 @@ COPY --from=frontend-builder /app/dist/ ./internal/server/admin/web/dist/
 VOLUME [ "/app/configs" ]
 VOLUME [ "/app/data" ]
 
-ENTRYPOINT ["./portrd", "start"]
-CMD [ "--config", "/app/config.yaml" ]
+ENTRYPOINT ["./portrd", "--config", "/app/config.yaml"]
+

docker/docker-compose.dev.yaml

@@ -3,6 +3,7 @@ services:
     build:
       context: ..
       dockerfile: docker/Dockerfile
+    command: ["start", "all"]
     volumes:
       - ../configs/server.yaml:/app/config.yaml
       - ../data:/app/data/

docker/docker-compose.yaml

@@ -13,6 +13,7 @@ services:
       # use an image
       context: .
       dockerfile: docker/Dockerfile
+    command: ["start", "all"]
     volumes:
       - ./server.yaml:/app/config.yaml
       - ./data:/app/data/

internal/server/admin/admin.go

@@ -4,11 +4,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"log"
 	"log/slog"
 	"net/http"
-	"os"
-	"os/signal"
-	"syscall"
 	"time"
 
 	"github.com/amalshaji/portr/internal/server/admin/handler"
@@ -141,20 +139,18 @@ func New(config *config.Config, service *service.Service) *AdminServer {
 func (s *AdminServer) Start() {
 	s.log.Info("starting admin server", "port", s.config.ListenAddress())
 
-	done := make(chan os.Signal, 1)
-	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
-
-	go func() {
-		if err := s.app.Listen(s.config.ListenAddress()); err != nil && !errors.Is(err, http.ErrServerClosed) {
-			s.log.Error("failed to start admin server", "error", err)
-			done <- nil
-		}
-	}()
+	if err := s.app.Listen(s.config.ListenAddress()); err != nil && !errors.Is(err, http.ErrServerClosed) {
+		log.Fatalf("failed to start admin server: %v", err)
+	}
+}
 
-	<-done
+func (s *AdminServer) Shutdown() {
 	s.log.Info("stopping admin server")
+
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+
 	defer func() { cancel() }()
+
 	if err := s.app.ShutdownWithContext(ctx); err != nil {
 		s.log.Error("failed to stop proxy server", "error", err)
 	}

internal/server/cron/cron.go

@@ -4,8 +4,6 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
-	"os"
-	"os/signal"
 	"time"
 
 	"github.com/amalshaji/portr/internal/server/config"
@@ -14,9 +12,10 @@ import (
 )
 
 type Cron struct {
-	db     *db.Db
-	logger *slog.Logger
-	config *config.Config
+	db         *db.Db
+	logger     *slog.Logger
+	config     *config.Config
+	cancelFunc context.CancelFunc
 }
 
 func New(db *db.Db, config *config.Config) *Cron {
@@ -25,6 +24,7 @@ func New(db *db.Db, config *config.Config) *Cron {
 
 func (c *Cron) Start() {
 	ctx, cancel := context.WithCancel(context.Background())
+	c.cancelFunc = cancel
 
 	c.logger.Info(fmt.Sprintf("Starting %d cron jobs", len(crons)))
 	for _, job := range crons {
@@ -42,11 +42,8 @@ func (c *Cron) Start() {
 			}
 		}(job)
 	}
+}
 
-	sigCh := make(chan os.Signal, 1)
-	signal.Notify(sigCh, os.Interrupt)
-
-	<-sigCh
-
-	cancel()
+func (c *Cron) Shutdown() {
+	c.cancelFunc()
 }

internal/server/proxy/proxy.go

@@ -4,14 +4,12 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"log"
 	"log/slog"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
-	"os"
-	"os/signal"
 	"sync"
-	"syscall"
 	"time"
 
 	"github.com/amalshaji/portr/internal/server/config"
@@ -111,21 +109,20 @@ func (p *Proxy) handleRequest(w http.ResponseWriter, r *http.Request) {
 func (p *Proxy) Start() {
 	p.log.Info("starting proxy server", "port", p.GetServerAddr())
 
-	done := make(chan os.Signal, 1)
-	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
-
 	http.HandleFunc("/", p.handleRequest)
-	go func() {
-		if err := p.server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
-			p.log.Error("failed to start proxy server", "error", err)
-			done <- nil
-		}
-	}()
-
-	<-done
+
+	if err := p.server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+		log.Fatalf("failed to start proxy server: %v", err)
+	}
+}
+
+func (p *Proxy) Shutdown(_ context.Context) {
 	p.log.Info("stopping proxy server")
+
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+
 	defer func() { cancel() }()
+
 	if err := p.server.Shutdown(ctx); err != nil {
 		p.log.Error("failed to stop proxy server", "error", err)
 	}

internal/server/ssh/sshd.go

@@ -7,8 +7,6 @@ import (
 	"log"
 	"log/slog"
 	"os"
-	"os/signal"
-	"syscall"
 	"time"
 
 	"github.com/amalshaji/portr/internal/constants"
@@ -24,6 +22,7 @@ type SshServer struct {
 	log     *slog.Logger
 	proxy   *proxy.Proxy
 	service *service.Service
+	server  *ssh.Server
 }
 
 func New(config *config.SshConfig, proxy *proxy.Proxy, service *service.Service) *SshServer {
@@ -117,23 +116,23 @@ func (s *SshServer) Start() {
 		},
 	}
 
-	done := make(chan os.Signal, 1)
-	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
+	s.server = &server
 
 	s.log.Info("starting SSH server", "port", s.GetServerAddr())
 
-	go func() {
-		if err := server.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
-			s.log.Error("failed to start SSH server", "error", err)
-			done <- nil
-		}
-	}()
+	if err := server.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
+		log.Fatalf("failed to start SSH server: %v", err)
+	}
+}
 
-	<-done
+func (s *SshServer) Shutdown(_ context.Context) {
 	s.log.Info("stopping SSH server")
+
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+
 	defer func() { cancel() }()
-	if err := server.Shutdown(ctx); err != nil {
+
+	if err := s.server.Shutdown(ctx); err != nil {
 		s.log.Error("failed to stop SSH server", "error", err)
 	}
 }