From 7bc5a01ba7dff89f0344976e78145dd481b114f4 Mon Sep 17 00:00:00 2001 From: bagas Date: Thu, 18 Dec 2025 18:30:49 +0700 Subject: [PATCH] feat: add pprof for debuging --- README.md | 48 ++++++++++++++++++++ internal/port/port.go | 2 +- main.go | 24 +++++++++- server/http.go | 4 +- server/https.go | 2 +- server/server.go | 4 +- session/handler.go | 6 +-- session/interaction/interaction.go | 12 ++--- utils/utils.go | 71 +++++++++++++++++++++++++++--- 9 files changed, 150 insertions(+), 23 deletions(-) diff --git a/README.md b/README.md index 7a765ef..9e8bf54 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,54 @@ A lightweight SSH-based tunnel server written in Go that enables secure TCP and - Go 1.18 or higher - Valid domain name for subdomain routing +## Environment Variables + +The following environment variables can be configured in the `.env` file: + +| Variable | Description | Default | Required | +|----------|-------------|---------|----------| +| `DOMAIN` | Domain name for subdomain routing | `localhost` | No | +| `PORT` | SSH server port | `2200` | No | +| `TLS_ENABLED` | Enable TLS/HTTPS | `false` | No | +| `TLS_REDIRECT` | Redirect HTTP to HTTPS | `false` | No | +| `CERT_LOC` | Path to TLS certificate | `certs/cert.pem` | No | +| `KEY_LOC` | Path to TLS private key | `certs/privkey.pem` | No | +| `SSH_PRIVATE_KEY` | Path to SSH private key (auto-generated if missing) | `certs/id_rsa` | No | +| `CORS_LIST` | Comma-separated list of allowed CORS origins | - | No | +| `ALLOWED_PORTS` | Port range for TCP tunnels (e.g., 40000-41000) | `40000-41000` | No | +| `PPROF_ENABLED` | Enable pprof profiling server | `false` | No | +| `PPROF_PORT` | Port for pprof server | `6060` | No | + +**Note:** All environment variables now use UPPERCASE naming. The application includes sensible defaults for all variables, so you can run it without a `.env` file for basic functionality. + +### SSH Key Auto-Generation + +If the SSH private key specified in `SSH_PRIVATE_KEY` doesn't exist, the application will automatically generate a new 4096-bit RSA key pair at the specified location. This makes it easier to get started without manually creating SSH keys. + +### Profiling with pprof + +To enable profiling for performance analysis: + +1. Set `PPROF_ENABLED=true` in your `.env` file +2. Optionally set `PPROF_PORT` to your desired port (default: 6060) +3. Access profiling data at `http://localhost:6060/debug/pprof/` + +Common pprof endpoints: +- `/debug/pprof/` - Index page with available profiles +- `/debug/pprof/heap` - Memory allocation profile +- `/debug/pprof/goroutine` - Stack traces of all current goroutines +- `/debug/pprof/profile` - CPU profile (30-second sample by default) +- `/debug/pprof/trace` - Execution trace + +Example usage with `go tool pprof`: +```bash +# Analyze CPU profile +go tool pprof http://localhost:6060/debug/pprof/profile?seconds=30 + +# Analyze memory heap +go tool pprof http://localhost:6060/debug/pprof/heap +``` + ## Contributing Contributions are welcome! diff --git a/internal/port/port.go b/internal/port/port.go index 1f05eab..6512f40 100644 --- a/internal/port/port.go +++ b/internal/port/port.go @@ -21,7 +21,7 @@ var Manager = PortManager{ } func init() { - rawRange := utils.Getenv("ALLOWED_PORTS") + rawRange := utils.Getenv("ALLOWED_PORTS", "40000-41000") splitRange := strings.Split(rawRange, "-") if len(splitRange) != 2 { Manager.AddPortRange(30000, 31000) diff --git a/main.go b/main.go index 3dd3811..65fb3ca 100644 --- a/main.go +++ b/main.go @@ -1,7 +1,10 @@ package main import ( + "fmt" "log" + "net/http" + _ "net/http/pprof" "os" "tunnel_pls/server" "tunnel_pls/utils" @@ -12,13 +15,30 @@ import ( func main() { log.SetOutput(os.Stdout) log.SetFlags(log.LstdFlags | log.Lshortfile) - + + pprofEnabled := utils.Getenv("PPROF_ENABLED", "false") + if pprofEnabled == "true" { + pprofPort := utils.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: "SSH-2.0-TunnlPls-1.0", } - privateBytes, err := os.ReadFile(utils.Getenv("ssh_private_key")) + sshKeyPath := utils.Getenv("SSH_PRIVATE_KEY", "certs/id_rsa") + if err := utils.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) } diff --git a/server/http.go b/server/http.go index 9eafb7d..34fad0e 100644 --- a/server/http.go +++ b/server/http.go @@ -194,7 +194,7 @@ func NewHTTPServer() error { if err != nil { return errors.New("Error listening: " + err.Error()) } - if utils.Getenv("tls_enabled") == "true" && utils.Getenv("tls_redirect") == "true" { + if utils.Getenv("TLS_ENABLED", "false") == "true" && utils.Getenv("TLS_REDIRECT", "false") == "true" { redirectTLS = true } go func() { @@ -246,7 +246,7 @@ func Handler(conn net.Conn) { if redirectTLS { _, err = conn.Write([]byte("HTTP/1.1 301 Moved Permanently\r\n" + - fmt.Sprintf("Location: https://%s.%s/\r\n", slug, utils.Getenv("domain")) + + fmt.Sprintf("Location: https://%s.%s/\r\n", slug, utils.Getenv("DOMAIN", "localhost")) + "Content-Length: 0\r\n" + "Connection: close\r\n" + "\r\n")) diff --git a/server/https.go b/server/https.go index 4e23d17..f65a0ba 100644 --- a/server/https.go +++ b/server/https.go @@ -13,7 +13,7 @@ import ( ) func NewHTTPSServer() error { - cert, err := tls.LoadX509KeyPair(utils.Getenv("cert_loc"), utils.Getenv("key_loc")) + cert, err := tls.LoadX509KeyPair(utils.Getenv("CERT_LOC", "certs/cert.pem"), utils.Getenv("KEY_LOC", "certs/privkey.pem")) if err != nil { return err } diff --git a/server/server.go b/server/server.go index 8051a02..e81ac62 100644 --- a/server/server.go +++ b/server/server.go @@ -17,12 +17,12 @@ type Server struct { } func NewServer(config *ssh.ServerConfig) *Server { - listener, err := net.Listen("tcp", fmt.Sprintf(":%s", utils.Getenv("port"))) + listener, err := net.Listen("tcp", fmt.Sprintf(":%s", utils.Getenv("PORT", "2200"))) if err != nil { log.Fatalf("failed to listen on port 2200: %v", err) return nil } - if utils.Getenv("tls_enabled") == "true" { + if utils.Getenv("TLS_ENABLED", "false") == "true" { go func() { err = NewHTTPSServer() if err != nil { diff --git a/session/handler.go b/session/handler.go index 79ee46c..962aa65 100644 --- a/session/handler.go +++ b/session/handler.go @@ -179,9 +179,9 @@ func (s *SSHSession) HandleHTTPForward(req *ssh.Request, portToBind uint16) { } log.Printf("HTTP forwarding approved on port: %d", portToBind) - domain := utils.Getenv("domain") + domain := utils.Getenv("DOMAIN", "localhost") protocol := "http" - if utils.Getenv("tls_enabled") == "true" { + if utils.Getenv("TLS_ENABLED", "false") == "true" { protocol = "https" } @@ -261,7 +261,7 @@ func (s *SSHSession) HandleTCPForward(req *ssh.Request, addr string, portToBind s.Forwarder.SetForwardedPort(portToBind) s.Interaction.SendMessage("\033[H\033[2J") s.Interaction.ShowWelcomeMessage() - s.Interaction.SendMessage(fmt.Sprintf("Forwarding your traffic to tcp://%s:%d \r\n", utils.Getenv("domain"), s.Forwarder.GetForwardedPort())) + s.Interaction.SendMessage(fmt.Sprintf("Forwarding your traffic to tcp://%s:%d \r\n", utils.Getenv("DOMAIN", "localhost"), s.Forwarder.GetForwardedPort())) s.Lifecycle.SetStatus(types.RUNNING) go s.Forwarder.AcceptTCPConnections() s.Interaction.HandleUserInput() diff --git a/session/interaction/interaction.go b/session/interaction/interaction.go index 17020a8..d8cf4c7 100644 --- a/session/interaction/interaction.go +++ b/session/interaction/interaction.go @@ -208,7 +208,7 @@ func (i *Interaction) appendToSlug(char byte) { } func (i *Interaction) refreshSlugDisplay() { - domain := utils.Getenv("domain") + domain := utils.Getenv("DOMAIN", "localhost") i.SendMessage(clearToLineEnd) i.SendMessage("➤ " + i.EditSlug + "." + domain) } @@ -238,7 +238,7 @@ func (i *Interaction) updateSlug() { return } - domain := utils.Getenv("domain") + domain := utils.Getenv("DOMAIN", "localhost") i.SendMessage("\r\n\r\n✅ SUBDOMAIN UPDATED ✅\r\n\r\n") i.SendMessage("Your new address is: " + newSlug + "." + domain + "\r\n\r\n") i.SendMessage("Press any key to continue...\r\n") @@ -340,16 +340,16 @@ func (i *Interaction) handleSlugCommand() { i.SendMessage(clearScreen) i.DisplaySlugEditor() - domain := utils.Getenv("domain") + domain := utils.Getenv("DOMAIN", "localhost") i.SendMessage("➤ " + i.EditSlug + "." + domain) } func (i *Interaction) ShowForwardingMessage() { - domain := utils.Getenv("domain") + domain := utils.Getenv("DOMAIN", "localhost") if i.Forwarder.GetTunnelType() == types.HTTP { protocol := "http" - if utils.Getenv("tls_enabled") == "true" { + if utils.Getenv("TLS_ENABLED", "false") == "true" { protocol = "https" } i.SendMessage(fmt.Sprintf("Forwarding your traffic to %s://%s.%s \r\n", protocol, i.SlugManager.Get(), domain)) @@ -384,7 +384,7 @@ func (i *Interaction) ShowWelcomeMessage() { } func (i *Interaction) DisplaySlugEditor() { - domain := utils.Getenv("domain") + domain := utils.Getenv("DOMAIN", "localhost") fullDomain := i.SlugManager.Get() + "." + domain contentLine := " ║ Current: " + fullDomain diff --git a/utils/utils.go b/utils/utils.go index d5d05da..2518627 100644 --- a/utils/utils.go +++ b/utils/utils.go @@ -1,14 +1,20 @@ package utils import ( + "crypto/rand" + "crypto/rsa" + "crypto/x509" + "encoding/pem" "log" - "math/rand" + mathrand "math/rand" "os" + "path/filepath" "strings" "sync" "time" "github.com/joho/godotenv" + "golang.org/x/crypto/ssh" ) type Env struct { @@ -24,7 +30,7 @@ func init() { func GenerateRandomString(length int) string { const charset = "abcdefghijklmnopqrstuvwxyz" - seededRand := rand.New(rand.NewSource(time.Now().UnixNano() + int64(rand.Intn(9999)))) + seededRand := mathrand.New(mathrand.NewSource(time.Now().UnixNano() + int64(mathrand.Intn(9999)))) var result strings.Builder for i := 0; i < length; i++ { randomIndex := seededRand.Intn(len(charset)) @@ -33,7 +39,7 @@ func GenerateRandomString(length int) string { return result.String() } -func Getenv(key string) string { +func Getenv(key, defaultValue string) string { env.mu.Lock() defer env.mu.Unlock() if val, ok := env.value[key]; ok { @@ -48,11 +54,64 @@ func Getenv(key string) string { } val := os.Getenv(key) - env.value[key] = val - if val == "" { - panic("Asking for env: " + key + " but got nothing, please set your environment first") + val = defaultValue } + env.value[key] = val return val } + +func GenerateSSHKeyIfNotExist(keyPath string) error { + if _, err := os.Stat(keyPath); err == nil { + log.Printf("SSH key already exists at %s", keyPath) + return nil + } + + log.Printf("SSH key not found at %s, generating new key pair...", keyPath) + + privateKey, err := rsa.GenerateKey(rand.Reader, 4096) + if err != nil { + return err + } + + privateKeyPEM := &pem.Block{ + Type: "RSA PRIVATE KEY", + Bytes: x509.MarshalPKCS1PrivateKey(privateKey), + } + + dir := filepath.Dir(keyPath) + if err := os.MkdirAll(dir, 0700); err != nil { + return err + } + + privateKeyFile, err := os.OpenFile(keyPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600) + if err != nil { + return err + } + defer privateKeyFile.Close() + + if err := pem.Encode(privateKeyFile, privateKeyPEM); err != nil { + return err + } + + publicKey, err := ssh.NewPublicKey(&privateKey.PublicKey) + if err != nil { + return err + } + + pubKeyPath := keyPath + ".pub" + pubKeyFile, err := os.OpenFile(pubKeyPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644) + if err != nil { + return err + } + defer pubKeyFile.Close() + + _, err = pubKeyFile.Write(ssh.MarshalAuthorizedKey(publicKey)) + if err != nil { + return err + } + + log.Printf("SSH key pair generated successfully at %s and %s", keyPath, pubKeyPath) + return nil +}