146 lines
3.8 KiB
Go
146 lines
3.8 KiB
Go
package main
|
|
|
|
import (
|
|
"log"
|
|
"net"
|
|
"os"
|
|
"os/signal"
|
|
"strconv"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/nixc/reverse-ssh-traefik/internal/buildinfo"
|
|
"github.com/nixc/reverse-ssh-traefik/internal/server"
|
|
)
|
|
|
|
func envOr(key, fallback string) string {
|
|
if v := os.Getenv(key); v != "" {
|
|
return v
|
|
}
|
|
return fallback
|
|
}
|
|
|
|
func envRequired(key string) string {
|
|
v := os.Getenv(key)
|
|
if v == "" {
|
|
log.Fatalf("Required environment variable %s is not set", key)
|
|
}
|
|
return v
|
|
}
|
|
|
|
func envInt(key string, fallback int) int {
|
|
v := os.Getenv(key)
|
|
if v == "" {
|
|
return fallback
|
|
}
|
|
n, err := strconv.Atoi(v)
|
|
if err != nil {
|
|
log.Printf("WARN: invalid %s=%q, using default %d", key, v, fallback)
|
|
return fallback
|
|
}
|
|
return n
|
|
}
|
|
|
|
func envBool(key string, fallback bool) bool {
|
|
v := os.Getenv(key)
|
|
if v == "" {
|
|
return fallback
|
|
}
|
|
b, err := strconv.ParseBool(v)
|
|
if err != nil {
|
|
log.Printf("WARN: invalid %s=%q, using default %t", key, v, fallback)
|
|
return fallback
|
|
}
|
|
return b
|
|
}
|
|
|
|
func main() {
|
|
log.SetFlags(log.LstdFlags | log.Lshortfile)
|
|
|
|
if len(os.Args) > 1 && os.Args[1] == "healthcheck" {
|
|
runHealthcheck()
|
|
return
|
|
}
|
|
|
|
log.Printf("tunnel-server starting commit=%s", buildinfo.Commit)
|
|
|
|
// SSH server config (for accepting tunnel clients).
|
|
sshPort := envOr("SSH_PORT", "2222")
|
|
hostKeyPath := envOr("SSH_HOST_KEY", "/keys/host_key")
|
|
authKeysPath := envOr("AUTHORIZED_KEYS", "/keys/authorized_keys")
|
|
portStart := envInt("PORT_RANGE_START", 10000)
|
|
portEnd := envInt("PORT_RANGE_END", 10100)
|
|
|
|
// Swarm manager SSH config (for updating service labels).
|
|
traefikHost := envRequired("TRAEFIK_SSH_HOST")
|
|
traefikUser := envOr("TRAEFIK_SSH_USER", "root")
|
|
traefikKey := envRequired("TRAEFIK_SSH_KEY")
|
|
serviceName := envOr("SWARM_SERVICE_NAME", "better-argo-tunnels_tunnel-server")
|
|
entrypoint := envOr("TRAEFIK_ENTRYPOINT", "websecure")
|
|
certResolver := envOr("TRAEFIK_CERT_RESOLVER", "letsencryptresolver")
|
|
|
|
// Load the SSH key for connecting to the Swarm manager.
|
|
traefikSigner, err := server.LoadSigner(traefikKey)
|
|
if err != nil {
|
|
log.Fatalf("Failed to load Traefik SSH key: %v", err)
|
|
}
|
|
log.Printf("Loaded Swarm manager SSH key")
|
|
|
|
// Initialize port pool.
|
|
pool := server.NewPortPool(portStart, portEnd)
|
|
log.Printf("Port pool: %d-%d (%d ports)", portStart, portEnd, portEnd-portStart+1)
|
|
|
|
// Initialize label manager (Swarm service update via SSH).
|
|
labels, err := server.NewLabelManager(
|
|
traefikHost, traefikUser, traefikSigner,
|
|
serviceName, entrypoint, certResolver,
|
|
)
|
|
if err != nil {
|
|
log.Fatalf("Failed to init label manager: %v", err)
|
|
}
|
|
defer labels.Close()
|
|
|
|
if envBool("PURGE_STALE_LABELS_ON_START", false) {
|
|
log.Printf("PURGE_STALE_LABELS_ON_START enabled; removing all tunnel labels")
|
|
if err := labels.PurgeAllTunnelLabels(); err != nil {
|
|
log.Printf("WARN: failed to purge stale labels: %v", err)
|
|
}
|
|
} else if err := labels.ReconcileExistingTunnelLabels(); err != nil {
|
|
log.Printf("WARN: failed to reconcile existing tunnel labels: %v", err)
|
|
}
|
|
|
|
// Initialize SSH server for tunnel clients.
|
|
sshSrv, err := server.NewSSHServer(hostKeyPath, authKeysPath, pool, labels)
|
|
if err != nil {
|
|
log.Fatalf("Failed to init SSH server: %v", err)
|
|
}
|
|
|
|
// Handle graceful shutdown.
|
|
sigCh := make(chan os.Signal, 1)
|
|
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
|
|
go func() {
|
|
sig := <-sigCh
|
|
log.Printf("Received signal %s, shutting down", sig)
|
|
os.Exit(0)
|
|
}()
|
|
|
|
// Start SSH server.
|
|
addr := "0.0.0.0:" + sshPort
|
|
if err := sshSrv.ListenAndServe(addr); err != nil {
|
|
log.Fatalf("SSH server error: %v", err)
|
|
}
|
|
}
|
|
|
|
func runHealthcheck() {
|
|
sshPort := envOr("SSH_PORT", "2222")
|
|
addr := envOr("HEALTHCHECK_ADDR", "127.0.0.1:"+sshPort)
|
|
timeout := time.Duration(envInt("HEALTHCHECK_TIMEOUT_SECONDS", 2)) * time.Second
|
|
|
|
conn, err := net.DialTimeout("tcp", addr, timeout)
|
|
if err != nil {
|
|
log.Printf("healthcheck failed for %s: %v", addr, err)
|
|
os.Exit(1)
|
|
}
|
|
conn.Close()
|
|
}
|