session/interaction/constants.go

@@ -0,0 +1,27 @@
+package interaction
+
+const (
+	backspaceChar    = 8
+	deleteChar       = 127
+	enterChar        = 13
+	escapeChar       = 27
+	ctrlC            = 3
+	forwardSlash     = '/'
+	minPrintableChar = 32
+	maxPrintableChar = 126
+
+	minSlugLength = 3
+	maxSlugLength = 20
+
+	clearScreen    = "\033[H\033[2J"
+	clearLine      = "\033[K"
+	clearToLineEnd = "\r\033[K"
+	backspaceSeq   = "\b \b"
+
+	minBoxWidth  = 50
+	paddingRight = 4
+)
+
+var forbiddenSlugs = []string{
+	"ping",
+}

session/interaction/interaction.go

@@ -14,10 +14,6 @@ import (
 	"golang.org/x/crypto/ssh"
 )
 
-var forbiddenSlug = []string{
-	"ping",
-}
-
 type Lifecycle interface {
 	Close() error
 }
@@ -70,228 +66,290 @@ func (i *Interaction) SetChannel(channel ssh.Channel) {
 }
 
 func (i *Interaction) SendMessage(message string) {
-	if i.channel != nil {
-		_, err := i.channel.Write([]byte(message))
-		if err != nil && err != io.EOF {
-			log.Printf("Error writing to channel: %v", err)
-			return
-		}
+	if i.channel == nil {
+		log.Printf("channel is nil")
 	}
+
+	_, err := i.channel.Write([]byte(message))
+	if err != nil && err != io.EOF {
+		log.Printf("error writing to channel: %s", err)
+	}
+	return
 }
 
 func (i *Interaction) HandleUserInput() {
 	buf := make([]byte, 1)
 	i.InteractiveMode = false
+
 	for {
 		n, err := i.channel.Read(buf)
 		if err != nil {
-			if err != io.EOF {
-				log.Printf("Error reading from client: %s", err)
-			}
+			i.handleReadError(err)
 			break
 		}
 
 		if n > 0 {
-			char := buf[0]
-			if i.InteractiveMode {
-				if i.InteractionType == types.Slug {
-					i.HandleSlugEditMode(char)
-				} else if i.InteractionType == types.Drop {
-					i.HandleDropMode(char)
-				}
-				continue
-			}
-
-			if i.pendingExit {
-				if char != 3 {
-					i.pendingExit = false
-					i.SendMessage("Operation canceled.\r\n")
-				}
-			}
-
-			if char == 3 {
-				if i.pendingExit {
-					i.SendMessage("Closing connection...\r\n")
-					err = i.Lifecycle.Close()
-					if err != nil {
-						log.Printf("failed to close session: %v", err)
-						return
-					}
-					return
-				}
-				i.SendMessage("Please press Ctrl+C again to disconnect.\r\n")
-				i.pendingExit = true
-			}
-
-			i.SendMessage(string(buf[:n]))
-
-			if char == 8 || char == 127 {
-				if i.InputLength > 0 {
-					i.SendMessage("\b \b")
-				}
-				if i.CommandBuffer.Len() > 0 {
-					i.CommandBuffer.Truncate(i.CommandBuffer.Len() - 1)
-				}
-				continue
-			}
-
-			i.InputLength += n
-
-			if char == '/' {
-				i.CommandBuffer.Reset()
-				i.CommandBuffer.WriteByte(char)
-				continue
-			}
-
-			if i.CommandBuffer.Len() > 0 {
-				if char == 13 {
-					i.SendMessage("\033[K")
-					i.HandleCommand(i.CommandBuffer.String())
-					continue
-				}
-				i.CommandBuffer.WriteByte(char)
-			}
-
-			if char == 13 {
-				i.SendMessage("\033[K")
-			}
+			i.processCharacter(buf[0])
 		}
 	}
 }
 
+func (i *Interaction) handleReadError(err error) {
+	if err != io.EOF {
+		log.Printf("Error reading from client: %s", err)
+	}
+}
+
+func (i *Interaction) processCharacter(char byte) {
+	if i.InteractiveMode {
+		i.handleInteractiveMode(char)
+		return
+	}
+
+	if i.handleExitSequence(char) {
+		return
+	}
+
+	i.SendMessage(string(char))
+	i.handleNonInteractiveInput(char)
+}
+
+func (i *Interaction) handleInteractiveMode(char byte) {
+	switch i.InteractionType {
+	case types.Slug:
+		i.HandleSlugEditMode(char)
+	case types.Drop:
+		i.HandleDropMode(char)
+	}
+}
+
+func (i *Interaction) handleExitSequence(char byte) bool {
+	if char == ctrlC {
+		if i.pendingExit {
+			i.SendMessage("Closing connection...\r\n")
+			if err := i.Lifecycle.Close(); err != nil {
+				log.Printf("failed to close session: %v", err)
+			}
+			return true
+		}
+		i.SendMessage("Please press Ctrl+C again to disconnect.\r\n")
+		i.pendingExit = true
+		return true
+	}
+
+	if i.pendingExit && char != ctrlC {
+		i.pendingExit = false
+		i.SendMessage("Operation canceled.\r\n")
+	}
+
+	return false
+}
+
+func (i *Interaction) handleNonInteractiveInput(char byte) {
+	switch {
+	case char == backspaceChar || char == deleteChar:
+		i.handleBackspace()
+	case char == forwardSlash:
+		i.handleCommandStart()
+	case i.CommandBuffer.Len() > 0:
+		i.handleCommandInput(char)
+	case char == enterChar:
+		i.SendMessage(clearLine)
+	default:
+		i.InputLength++
+	}
+}
+
+func (i *Interaction) handleBackspace() {
+	if i.InputLength > 0 {
+		i.SendMessage(backspaceSeq)
+	}
+	if i.CommandBuffer.Len() > 0 {
+		i.CommandBuffer.Truncate(i.CommandBuffer.Len() - 1)
+	}
+}
+
+func (i *Interaction) handleCommandStart() {
+	i.CommandBuffer.Reset()
+	i.CommandBuffer.WriteByte(forwardSlash)
+}
+
+func (i *Interaction) handleCommandInput(char byte) {
+	if char == enterChar {
+		i.SendMessage(clearLine)
+		i.HandleCommand(i.CommandBuffer.String())
+		return
+	}
+	i.CommandBuffer.WriteByte(char)
+	i.InputLength++
+}
+
 func (i *Interaction) HandleSlugEditMode(char byte) {
-	if char == 13 {
+	switch {
+	case char == enterChar:
 		i.HandleSlugSave()
-	} else if char == 27 || char == 3 {
+	case char == escapeChar || char == ctrlC:
 		i.HandleSlugCancel()
-	} else if char == 8 || char == 127 {
-		if len(i.EditSlug) > 0 {
-			i.EditSlug = (i.EditSlug)[:len(i.EditSlug)-1]
-			i.SendMessage("\r\033[K")
-			i.SendMessage("➤ " + i.EditSlug + "." + utils.Getenv("domain"))
-		}
-	} else if char >= 32 && char <= 126 {
-		if (char >= 'a' && char <= 'z') || (char >= '0' && char <= '9') || char == '-' {
-			i.EditSlug += string(char)
-			i.SendMessage("\r\033[K")
-			i.SendMessage("➤ " + i.EditSlug + "." + utils.Getenv("domain"))
-		}
+	case char == backspaceChar || char == deleteChar:
+		i.handleSlugBackspace()
+	case char >= minPrintableChar && char <= maxPrintableChar:
+		i.appendToSlug(char)
 	}
 }
 
+func (i *Interaction) handleSlugBackspace() {
+	if len(i.EditSlug) > 0 {
+		i.EditSlug = i.EditSlug[:len(i.EditSlug)-1]
+		i.refreshSlugDisplay()
+	}
+}
+
+func (i *Interaction) appendToSlug(char byte) {
+	if isValidSlugChar(char) {
+		i.EditSlug += string(char)
+		i.refreshSlugDisplay()
+	}
+}
+
+func (i *Interaction) refreshSlugDisplay() {
+	domain := utils.Getenv("domain")
+	i.SendMessage(clearToLineEnd)
+	i.SendMessage("➤ " + i.EditSlug + "." + domain)
+}
+
 func (i *Interaction) HandleSlugSave() {
-	isValid := isValidSlug(i.EditSlug)
+	i.SendMessage(clearScreen)
 
-	i.SendMessage("\033[H\033[2J")
-	if isValid {
-		oldSlug := i.SlugManager.Get()
-		newSlug := i.EditSlug
-
-		if !i.updateClientSlug(oldSlug, newSlug) {
-			i.HandleSlugUpdateError()
-			return
-		}
-
-		i.SendMessage("\r\n\r\n✅ SUBDOMAIN UPDATED ✅\r\n\r\n")
-		i.SendMessage("Your new address is: " + newSlug + "." + utils.Getenv("domain") + "\r\n\r\n")
-		i.SendMessage("Press any key to continue...\r\n")
-	} else if isForbiddenSlug(i.EditSlug) {
-		i.SendMessage("\r\n\r\n❌ FORBIDDEN SUBDOMAIN ❌\r\n\r\n")
-		i.SendMessage("This subdomain is not allowed.\r\n")
-		i.SendMessage("Please try a different subdomain.\r\n\r\n")
-		i.SendMessage("Press any key to continue...\r\n")
-	} else {
-		i.SendMessage("\r\n\r\n❌ INVALID SUBDOMAIN ❌\r\n\r\n")
-		i.SendMessage("Use only lowercase letters, numbers, and hyphens.\r\n")
-		i.SendMessage("Length must be 3-20 characters and cannot start or end with a hyphen.\r\n\r\n")
-		i.SendMessage("Press any key to continue...\r\n")
+	switch {
+	case isForbiddenSlug(i.EditSlug):
+		i.showForbiddenSlugMessage()
+	case !isValidSlug(i.EditSlug):
+		i.showInvalidSlugMessage()
+	default:
+		i.updateSlug()
 	}
 
 	i.WaitForKeyPress()
-	i.SendMessage("\033[H\033[2J")
-	i.ShowWelcomeMessage()
+	i.returnToMainScreen()
+}
+
+func (i *Interaction) updateSlug() {
+	oldSlug := i.SlugManager.Get()
+	newSlug := i.EditSlug
+
+	if !i.updateClientSlug(oldSlug, newSlug) {
+		i.HandleSlugUpdateError()
+		return
+	}
 
 	domain := utils.Getenv("domain")
-	protocol := "http"
-	if utils.Getenv("tls_enabled") == "true" {
-		protocol = "https"
-	}
-	i.SendMessage(fmt.Sprintf("Forwarding your traffic to %s://%s.%s \r\n", protocol, i.SlugManager.Get(), domain))
+	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")
+}
 
+func (i *Interaction) showForbiddenSlugMessage() {
+	i.SendMessage("\r\n\r\n❌ FORBIDDEN SUBDOMAIN ❌\r\n\r\n")
+	i.SendMessage("This subdomain is not allowed.\r\n")
+	i.SendMessage("Please try a different subdomain.\r\n\r\n")
+	i.SendMessage("Press any key to continue...\r\n")
+}
+
+func (i *Interaction) showInvalidSlugMessage() {
+	i.SendMessage("\r\n\r\n❌ INVALID SUBDOMAIN ❌\r\n\r\n")
+	i.SendMessage("Use only lowercase letters, numbers, and hyphens.\r\n")
+	i.SendMessage(fmt.Sprintf("Length must be %d-%d characters and cannot start or end with a hyphen.\r\n\r\n", minSlugLength, maxSlugLength))
+	i.SendMessage("Press any key to continue...\r\n")
+}
+
+func (i *Interaction) returnToMainScreen() {
+	i.SendMessage(clearScreen)
+	i.ShowWelcomeMessage()
+	i.ShowForwardingMessage()
 	i.InteractiveMode = false
 	i.CommandBuffer.Reset()
 }
 
 func (i *Interaction) HandleSlugCancel() {
 	i.InteractiveMode = false
-	i.SendMessage("\033[H\033[2J")
-	i.SendMessage("\r\n\r\n⚠️ SUBDOMAIN EDIT CANCELLED ⚠️\r\n\r\n")
-	i.SendMessage("Press any key to continue...\r\n")
-
-	i.WaitForKeyPress()
-
-	i.SendMessage("\033[H\033[2J")
-	i.ShowWelcomeMessage()
-	i.ShowForwardingMessage()
-	i.CommandBuffer.Reset()
+	i.showMessageAndWait("\r\n\r\n⚠️ SUBDOMAIN EDIT CANCELLED ⚠️\r\n\r\n")
 }
 
 func (i *Interaction) HandleSlugUpdateError() {
 	i.SendMessage("\r\n\r\n❌ SERVER ERROR ❌\r\n\r\n")
 	i.SendMessage("Failed to update subdomain. You will be disconnected in 5 seconds.\r\n\r\n")
 
-	for iter := 5; iter > 0; iter-- {
-		i.SendMessage(fmt.Sprintf("Disconnecting in %d...\r\n", iter))
+	for countdown := 5; countdown > 0; countdown-- {
+		i.SendMessage(fmt.Sprintf("Disconnecting in %d...\r\n", countdown))
 		time.Sleep(1 * time.Second)
 	}
-	err := i.Lifecycle.Close()
-	if err != nil {
+
+	if err := i.Lifecycle.Close(); err != nil {
 		log.Printf("failed to close session: %v", err)
-		return
 	}
 }
 
 func (i *Interaction) HandleCommand(command string) {
-	switch command {
-	case "/bye":
-		i.SendMessage("Closing connection...\r\n")
-		err := i.Lifecycle.Close()
-		if err != nil {
-			log.Printf("failed to close session: %v", err)
-			return
-		}
-		return
-	case "/help":
-		i.SendMessage("\r\nAvailable commands: /bye, /help, /clear, /slug\r\n")
-	case "/clear":
-		i.SendMessage("\033[H\033[2J")
-		i.ShowWelcomeMessage()
-		i.ShowForwardingMessage()
-	case "/slug":
-		if i.Forwarder.GetTunnelType() != types.HTTP {
-			i.SendMessage(fmt.Sprintf("\r\n%s tunnels cannot have custom subdomains", i.Forwarder.GetTunnelType()))
-		} else {
-			i.InteractiveMode = true
-			i.InteractionType = types.Slug
-			i.EditSlug = i.SlugManager.Get()
-			i.SendMessage("\033[H\033[2J")
-			i.DisplaySlugEditor()
-			i.SendMessage("➤ " + i.EditSlug + "." + utils.Getenv("domain"))
-		}
-	case "/drop":
-		i.InteractiveMode = true
-		i.InteractionType = types.Drop
-		i.SendMessage("\033[H\033[2J")
-		i.ShowDropMessage()
-	default:
+	handlers := map[string]func(){
+		"/bye":   i.handleByeCommand,
+		"/help":  i.handleHelpCommand,
+		"/clear": i.handleClearCommand,
+		"/slug":  i.handleSlugCommand,
+		"/drop":  i.handleDropCommand,
+	}
+
+	if handler, exists := handlers[command]; exists {
+		handler()
+	} else {
 		i.SendMessage("Unknown command\r\n")
 	}
 
 	i.CommandBuffer.Reset()
 }
 
+func (i *Interaction) handleByeCommand() {
+	i.SendMessage("Closing connection...\r\n")
+	if err := i.Lifecycle.Close(); err != nil {
+		log.Printf("failed to close session: %v", err)
+	}
+}
+
+func (i *Interaction) handleHelpCommand() {
+	i.SendMessage("\r\nAvailable commands: /bye, /help, /clear, /slug, /drop\r\n")
+}
+
+func (i *Interaction) handleClearCommand() {
+	i.SendMessage(clearScreen)
+	i.ShowWelcomeMessage()
+	i.ShowForwardingMessage()
+}
+
+func (i *Interaction) handleSlugCommand() {
+	if i.Forwarder.GetTunnelType() != types.HTTP {
+		i.SendMessage(fmt.Sprintf("\r\n%s tunnels cannot have custom subdomains\r\n", i.Forwarder.GetTunnelType()))
+		return
+	}
+
+	i.InteractiveMode = true
+	i.InteractionType = types.Slug
+	i.EditSlug = i.SlugManager.Get()
+	i.SendMessage(clearScreen)
+	i.DisplaySlugEditor()
+
+	domain := utils.Getenv("domain")
+	i.SendMessage("➤ " + i.EditSlug + "." + domain)
+}
+
+func (i *Interaction) handleDropCommand() {
+	i.InteractiveMode = true
+	i.InteractionType = types.Drop
+	i.SendMessage(clearScreen)
+	i.ShowDropMessage()
+}
+
 func (i *Interaction) ShowForwardingMessage() {
 	domain := utils.Getenv("domain")
+
 	if i.Forwarder.GetTunnelType() == types.HTTP {
 		protocol := "http"
 		if utils.Getenv("tls_enabled") == "true" {
@@ -304,39 +362,47 @@ func (i *Interaction) ShowForwardingMessage() {
 }
 
 func (i *Interaction) HandleDropMode(char byte) {
-	if char == 13 || char == 121 || char == 89 {
-		count := i.Forwarder.DropAllForwarder()
-		i.SendMessage("\033[H\033[2J")
-		i.SendMessage(fmt.Sprintf("Dropped %d forwarders\r\n", count))
-		i.SendMessage("Press any key to continue...\r\n")
-		i.InteractiveMode = false
-		i.InteractionType = ""
-		i.WaitForKeyPress()
-		i.SendMessage("\033[H\033[2J")
-		i.ShowWelcomeMessage()
-		i.ShowForwardingMessage()
-	} else if char == 27 || char == 110 || char == 78 || char == 3 {
-		i.SendMessage("\033[H\033[2J")
-		i.SendMessage(fmt.Sprintf("Dropping canceled.\r\n"))
-		i.SendMessage("Press any key to continue...\r\n")
-		i.InteractiveMode = false
-		i.InteractionType = ""
-		i.WaitForKeyPress()
-		i.SendMessage("\033[H\033[2J")
-		i.ShowWelcomeMessage()
-		i.ShowForwardingMessage()
+	switch {
+	case char == enterChar || char == 'y' || char == 'Y':
+		i.executeDropAll()
+	case char == escapeChar || char == 'n' || char == 'N' || char == ctrlC:
+		i.cancelDrop()
 	}
 }
 
+func (i *Interaction) executeDropAll() {
+	count := i.Forwarder.DropAllForwarder()
+	message := fmt.Sprintf("Dropped %d forwarders\r\n", count)
+	i.showMessageAndWait(message)
+}
+
+func (i *Interaction) cancelDrop() {
+	i.showMessageAndWait("Dropping canceled.\r\n")
+}
+
+func (i *Interaction) showMessageAndWait(message string) {
+	i.SendMessage(clearScreen)
+	i.SendMessage(message)
+	i.SendMessage("Press any key to continue...\r\n")
+
+	i.InteractiveMode = false
+	i.InteractionType = ""
+	i.WaitForKeyPress()
+
+	i.SendMessage(clearScreen)
+	i.ShowWelcomeMessage()
+	i.ShowForwardingMessage()
+}
+
 func (i *Interaction) ShowDropMessage() {
-	const paddingRight = 4
-
 	confirmText := fmt.Sprintf("  ║  Drop ALL %d active connections?", i.Forwarder.GetForwarderCount())
-	boxWidth := len(confirmText) + paddingRight + 1
-	if boxWidth < 50 {
-		boxWidth = 50
-	}
+	boxWidth := calculateBoxWidth(confirmText)
 
+	box := buildDropConfirmationBox(boxWidth, confirmText)
+	i.SendMessage("\r\n" + box + "\r\n\r\n")
+}
+
+func buildDropConfirmationBox(boxWidth int, confirmText string) string {
 	topBorder := "  ╔" + strings.Repeat("═", boxWidth-4) + "╗\r\n"
 	title := centerText("DROP CONFIRMATION", boxWidth-4)
 	header := "  ║" + title + "║\r\n"
@@ -350,18 +416,7 @@ func (i *Interaction) ShowDropMessage() {
 
 	bottomBorder := "  ╚" + strings.Repeat("═", boxWidth-4) + "╝\r\n"
 
-	asciiArt := topBorder +
-		header +
-		midBorder +
-		emptyLine +
-		confirmLine +
-		emptyLine +
-		controlLine +
-		emptyLine +
-		bottomBorder
-
-	i.SendMessage("\r\n" + asciiArt)
-	i.SendMessage("\r\n\r\n")
+	return topBorder + header + midBorder + emptyLine + confirmLine + emptyLine + controlLine + emptyLine + bottomBorder
 }
 
 func (i *Interaction) ShowWelcomeMessage() {
@@ -394,14 +449,14 @@ func (i *Interaction) DisplaySlugEditor() {
 	domain := utils.Getenv("domain")
 	fullDomain := i.SlugManager.Get() + "." + domain
 
-	const paddingRight = 4
-
 	contentLine := "  ║  Current:  " + fullDomain
-	boxWidth := len(contentLine) + paddingRight + 1
-	if boxWidth < 50 {
-		boxWidth = 50
-	}
+	boxWidth := calculateBoxWidth(contentLine)
 
+	box := buildSlugEditorBox(boxWidth, fullDomain)
+	i.SendMessage("\r\n\r\n" + box + "\r\n\r\n")
+}
+
+func buildSlugEditorBox(boxWidth int, fullDomain string) string {
 	topBorder := "  ╔" + strings.Repeat("═", boxWidth-4) + "╗\r\n"
 	title := centerText("SUBDOMAIN EDITOR", boxWidth-4)
 	header := "  ║" + title + "║\r\n"
@@ -414,18 +469,7 @@ func (i *Interaction) DisplaySlugEditor() {
 	saveCancel := "  ║  [Enter] Save  |  [Esc] Cancel" + strings.Repeat(" ", boxWidth-35) + "║\r\n"
 	bottomBorder := "  ╚" + strings.Repeat("═", boxWidth-4) + "╝\r\n"
 
-	i.SendMessage("\r\n\r\n")
-	i.SendMessage(topBorder)
-	i.SendMessage(header)
-	i.SendMessage(midBorder)
-	i.SendMessage(emptyLine)
-	i.SendMessage(currentLine)
-	i.SendMessage(emptyLine)
-	i.SendMessage(emptyLine)
-	i.SendMessage(midBorder)
-	i.SendMessage(saveCancel)
-	i.SendMessage(bottomBorder)
-	i.SendMessage("\r\n\r\n")
+	return topBorder + header + midBorder + emptyLine + currentLine + emptyLine + emptyLine + midBorder + saveCancel + bottomBorder
 }
 
 func (i *Interaction) SetSlugModificator(modificator func(oldSlug, newSlug string) bool) {
@@ -442,6 +486,14 @@ func (i *Interaction) WaitForKeyPress() {
 	}
 }
 
+func calculateBoxWidth(contentLine string) int {
+	boxWidth := len(contentLine) + paddingRight + 1
+	if boxWidth < minBoxWidth {
+		boxWidth = minBoxWidth
+	}
+	return boxWidth
+}
+
 func centerText(text string, width int) string {
 	padding := (width - len(text)) / 2
 	if padding < 0 {
@@ -451,7 +503,7 @@ func centerText(text string, width int) string {
 }
 
 func isValidSlug(slug string) bool {
-	if len(slug) < 3 || len(slug) > 20 {
+	if len(slug) < minSlugLength || len(slug) > maxSlugLength {
 		return false
 	}
 
@@ -460,7 +512,7 @@ func isValidSlug(slug string) bool {
 	}
 
 	for _, c := range slug {
-		if !((c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '-') {
+		if !isValidSlugChar(byte(c)) {
 			return false
 		}
 	}
@@ -468,8 +520,12 @@ func isValidSlug(slug string) bool {
 	return true
 }
 
+func isValidSlugChar(c byte) bool {
+	return (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '-'
+}
+
 func isForbiddenSlug(slug string) bool {
-	for _, s := range forbiddenSlug {
+	for _, s := range forbiddenSlugs {
 		if slug == s {
 			return true
 		}