feat: add pprof for debuging

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!
 

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)

main.go

@@ -1,7 +1,10 @@
 package main
 
 import (
+	"fmt"
 	"log"
+	"net/http"
+	_ "net/http/pprof"
 	"os"
 	"tunnel_pls/server"
 	"tunnel_pls/utils"
@@ -13,12 +16,29 @@ 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)
 	}

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"))

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
 	}

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 {

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()

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

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
+}