142 lines
3.5 KiB
Go
142 lines
3.5 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
_ "net/http/pprof"
|
|
"os"
|
|
"os/signal"
|
|
"strings"
|
|
"syscall"
|
|
"time"
|
|
"tunnel_pls/internal/config"
|
|
"tunnel_pls/internal/grpc/client"
|
|
"tunnel_pls/internal/key"
|
|
"tunnel_pls/server"
|
|
"tunnel_pls/session"
|
|
"tunnel_pls/version"
|
|
|
|
"golang.org/x/crypto/ssh"
|
|
)
|
|
|
|
func main() {
|
|
if len(os.Args) > 1 && (os.Args[1] == "--version" || os.Args[1] == "-v") {
|
|
fmt.Println(version.GetVersion())
|
|
os.Exit(0)
|
|
}
|
|
|
|
log.SetOutput(os.Stdout)
|
|
log.SetFlags(log.LstdFlags | log.Lshortfile)
|
|
|
|
log.Printf("Starting %s", version.GetVersion())
|
|
|
|
mode := strings.ToLower(config.Getenv("MODE", "standalone"))
|
|
isNodeMode := mode == "node"
|
|
|
|
pprofEnabled := config.Getenv("PPROF_ENABLED", "false")
|
|
if pprofEnabled == "true" {
|
|
pprofPort := config.Getenv("PPROF_PORT", "6060")
|
|
go func() {
|
|
pprofAddr := fmt.Sprintf("localhost:%s", pprofPort)
|
|
log.Printf("Starting pprof server on http://%s/debug/pprof/", pprofAddr)
|
|
if err := http.ListenAndServe(pprofAddr, nil); err != nil {
|
|
log.Printf("pprof server error: %v", err)
|
|
}
|
|
}()
|
|
}
|
|
|
|
sshConfig := &ssh.ServerConfig{
|
|
NoClientAuth: true,
|
|
ServerVersion: fmt.Sprintf("SSH-2.0-TunnlPls-%s", version.GetShortVersion()),
|
|
}
|
|
|
|
sshKeyPath := "certs/ssh/id_rsa"
|
|
if err := key.GenerateSSHKeyIfNotExist(sshKeyPath); err != nil {
|
|
log.Fatalf("Failed to generate SSH key: %s", err)
|
|
}
|
|
|
|
privateBytes, err := os.ReadFile(sshKeyPath)
|
|
if err != nil {
|
|
log.Fatalf("Failed to load private key: %s", err)
|
|
}
|
|
|
|
private, err := ssh.ParsePrivateKey(privateBytes)
|
|
if err != nil {
|
|
log.Fatalf("Failed to parse private key: %s", err)
|
|
}
|
|
|
|
sshConfig.AddHostKey(private)
|
|
sessionRegistry := session.NewRegistry()
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
errChan := make(chan error, 2)
|
|
shutdownChan := make(chan os.Signal, 1)
|
|
signal.Notify(shutdownChan, os.Interrupt, syscall.SIGTERM)
|
|
|
|
var grpcClient *client.Client
|
|
if isNodeMode {
|
|
grpcHost := config.Getenv("GRPC_ADDRESS", "localhost")
|
|
grpcPort := config.Getenv("GRPC_PORT", "8080")
|
|
grpcAddr := fmt.Sprintf("%s:%s", grpcHost, grpcPort)
|
|
nodeToken := config.Getenv("NODE_TOKEN", "")
|
|
if nodeToken == "" {
|
|
log.Fatalf("NODE_TOKEN is required in node mode")
|
|
}
|
|
|
|
c, err := client.New(&client.GrpcConfig{
|
|
Address: grpcAddr,
|
|
UseTLS: false,
|
|
InsecureSkipVerify: false,
|
|
Timeout: 10 * time.Second,
|
|
KeepAlive: true,
|
|
MaxRetries: 3,
|
|
}, sessionRegistry)
|
|
if err != nil {
|
|
log.Fatalf("failed to create grpc client: %v", err)
|
|
}
|
|
grpcClient = c
|
|
|
|
healthCtx, healthCancel := context.WithTimeout(ctx, 5*time.Second)
|
|
if err := grpcClient.CheckServerHealth(healthCtx); err != nil {
|
|
healthCancel()
|
|
log.Fatalf("gRPC health check failed: %v", err)
|
|
}
|
|
healthCancel()
|
|
|
|
go func() {
|
|
identity := config.Getenv("DOMAIN", "localhost")
|
|
if err := grpcClient.SubscribeEvents(ctx, identity, nodeToken); err != nil {
|
|
errChan <- fmt.Errorf("failed to subscribe to events: %w", err)
|
|
}
|
|
}()
|
|
}
|
|
|
|
go func() {
|
|
app, err := server.NewServer(sshConfig, sessionRegistry, grpcClient)
|
|
if err != nil {
|
|
errChan <- fmt.Errorf("failed to start server: %s", err)
|
|
return
|
|
}
|
|
app.Start()
|
|
}()
|
|
|
|
select {
|
|
case err := <-errChan:
|
|
log.Printf("error happen : %s", err)
|
|
case sig := <-shutdownChan:
|
|
log.Printf("received signal %s, shutting down", sig)
|
|
}
|
|
|
|
cancel()
|
|
|
|
if grpcClient != nil {
|
|
if err := grpcClient.Close(); err != nil {
|
|
log.Printf("failed to close grpc conn : %s", err)
|
|
}
|
|
}
|
|
}
|