Merge pull request #53 from antoniomika/debug_panics

Work on fixing panics

.github/workflows/docker.yml

@@ -10,21 +10,10 @@ jobs:
     - name: Lint the codebase
       run: |
         docker run --rm -v $(pwd):/app -w /app golangci/golangci-lint:v1.21.0 golangci-lint run -E goimports
-    - name: Register multiarch for Docker
-      run: |
-        docker run --privileged linuxkit/binfmt:v0.7
-    - name: Install buildx
-      run: |
-        mkdir -p ~/.docker/cli-plugins
-        wget https://github.com/docker/buildx/releases/download/v0.3.0/buildx-v0.3.0.linux-amd64 \
-          -O ~/.docker/cli-plugins/docker-buildx
-        chmod a+x ~/.docker/cli-plugins/docker-buildx
-    - name: Setup buildkit
-      run: |
-        docker buildx create --name builder
-        docker buildx use builder
-        docker buildx inspect --bootstrap
-        docker buildx ls
+    - name: Set up Docker Buildx
+      uses: crazy-max/ghaction-docker-buildx@v1
+      with:
+        version: latest
     - name: Login to GitHub Docker Registry
       env:
         DOCKER_USERNAME: ${{ secrets.GITHUB_DOCKER_USERNAME }}

channels.go

@@ -24,22 +24,19 @@ func handleSession(newChannel ssh.NewChannel, sshConn *SSHConnection, state *Sta
 		log.Println("Handling session for connection:", connection)
 	}
 
+	writeToSession(connection, aurora.BgRed("Press Ctrl-C to close the session.").String()+"\r\n")
+
 	go func() {
 		for {
 			select {
 			case c := <-sshConn.Messages:
-				_, err := connection.Write(append([]byte(c), []byte{'\r', '\n'}...))
-				if err != nil && *debug {
-					log.Println("Error trying to write message to socket:", err)
-				}
+				writeToSession(connection, c)
 			case <-sshConn.Close:
 				return
 			}
 		}
 	}()
 
-	sshConn.Messages <- aurora.BgRed("Press Ctrl-C to close the session.").String()
-
 	go func() {
 		for {
 			data := make([]byte, 4096)
@@ -75,7 +72,7 @@ func handleSession(newChannel ssh.NewChannel, sshConn *SSHConnection, state *Sta
 				if strings.HasPrefix(payloadString, proxyProtoPrefix) && *proxyProtoEnabled {
 					sshConn.ProxyProto = getProxyProtoVersion(strings.TrimPrefix(payloadString, proxyProtoPrefix))
 					if sshConn.ProxyProto != 0 {
-						sendMessage(sshConn, fmt.Sprintf("Proxy protocol enabled for TCP connections. Using protocol version %d", int(sshConn.ProxyProto)))
+						sendMessage(sshConn, fmt.Sprintf("Proxy protocol enabled for TCP connections. Using protocol version %d", int(sshConn.ProxyProto)), true)
 					}
 				}
 			default:
@@ -129,6 +126,13 @@ func handleAlias(newChannel ssh.NewChannel, sshConn *SSHConnection, state *State
 	sshConn.CleanUp(state)
 }
 
+func writeToSession(connection ssh.Channel, c string) {
+	_, err := connection.Write(append([]byte(c), []byte{'\r', '\n'}...))
+	if err != nil && *debug {
+		log.Println("Error trying to write message to socket:", err)
+	}
+}
+
 func getProxyProtoVersion(proxyProtoUserVersion string) byte {
 	if *proxyProtoVersion != "userdefined" {
 		proxyProtoUserVersion = *proxyProtoVersion

http.go

@@ -69,7 +69,7 @@ func startHTTPHandler(state *State) {
 			loc, ok := state.HTTPListeners.Load(hostname)
 			if ok {
 				proxyHolder := loc.(*ProxyHolder)
-				sendMessage(proxyHolder.SSHConn, strings.TrimSpace(logLine))
+				sendMessage(proxyHolder.SSHConn, strings.TrimSpace(logLine), true)
 			}
 		}
 

main.go

@@ -271,7 +271,7 @@ func main() {
 						})
 
 						if count == 0 {
-							sendMessage(holderConn, "No forwarding requests sent. Closing connection.")
+							sendMessage(holderConn, "No forwarding requests sent. Closing connection.", true)
 							time.Sleep(1 * time.Millisecond)
 							holderConn.CleanUp(state)
 						}
@@ -287,7 +287,6 @@ func main() {
 // CleanUp closes all allocated resources and cleans them up
 func (s *SSHConnection) CleanUp(state *State) {
 	close(s.Close)
-	close(s.Messages)
 	s.SSHConn.Close()
 	state.SSHConnections.Delete(s.SSHConn.RemoteAddr())
 	log.Println("Closed SSH connection for:", s.SSHConn.RemoteAddr(), "user:", s.SSHConn.User())

requests.go

@@ -108,7 +108,7 @@ func handleRemoteForward(newRequest *ssh.Request, sshConn *SSHConnection, state
 		connType = "https"
 	}
 
-	requestMessages := fmt.Sprintf("\nStarting SSH Fowarding service for %s. Forwarded connections can be accessed via the following methods:\r\n", aurora.Sprintf(aurora.Green("%s:%s"), connType, stringPort))
+	requestMessages := fmt.Sprintf("Starting SSH Fowarding service for %s. Forwarded connections can be accessed via the following methods:\r\n", aurora.Sprintf(aurora.Green("%s:%s"), connType, stringPort))
 
 	if stringPort == "80" || stringPort == "443" {
 		scheme := "http"
@@ -142,7 +142,7 @@ func handleRemoteForward(newRequest *ssh.Request, sshConn *SSHConnection, state
 				httpsPortString = fmt.Sprintf(":%d", httpsPort)
 			}
 
-			requestMessages += fmt.Sprintf("%s: https://%s%s", aurora.BgBlue("HTTPS"), host, httpsPortString)
+			requestMessages += fmt.Sprintf("%s: https://%s%s\r\n", aurora.BgBlue("HTTPS"), host, httpsPortString)
 			log.Printf("%s forwarding started: https://%s%s -> %s for client: %s\n", aurora.BgBlue("HTTPS"), host, httpPortString, chanListener.Addr().String(), sshConn.SSHConn.RemoteAddr().String())
 		}
 	} else {
@@ -152,15 +152,15 @@ func handleRemoteForward(newRequest *ssh.Request, sshConn *SSHConnection, state
 			state.TCPListeners.Store(validAlias, chanListener.Addr().String())
 			defer state.TCPListeners.Delete(validAlias)
 
-			requestMessages += fmt.Sprintf("%s: %s", aurora.BgBlue("TCP Alias"), validAlias)
+			requestMessages += fmt.Sprintf("%s: %s\r\n", aurora.BgBlue("TCP Alias"), validAlias)
 			log.Printf("%s forwarding started: %s -> %s for client: %s\n", aurora.BgBlue("TCP Alias"), validAlias, chanListener.Addr().String(), sshConn.SSHConn.RemoteAddr().String())
 		} else {
-			requestMessages += fmt.Sprintf("%s: %s:%d", aurora.BgBlue("TCP"), *rootDomain, chanListener.Addr().(*net.TCPAddr).Port)
+			requestMessages += fmt.Sprintf("%s: %s:%d\r\n", aurora.BgBlue("TCP"), *rootDomain, chanListener.Addr().(*net.TCPAddr).Port)
 			log.Printf("%s forwarding started: %s:%d -> %s for client: %s\n", aurora.BgBlue("TCP"), *rootDomain, chanListener.Addr().(*net.TCPAddr).Port, chanListener.Addr().String(), sshConn.SSHConn.RemoteAddr().String())
 		}
 	}
 
-	sendMessage(sshConn, requestMessages)
+	sendMessage(sshConn, requestMessages, false)
 
 	go func() {
 		<-sshConn.Close
@@ -180,7 +180,7 @@ func handleRemoteForward(newRequest *ssh.Request, sshConn *SSHConnection, state
 			log.Println(logLine)
 
 			if *logToClient {
-				sendMessage(sshConn, logLine)
+				sendMessage(sshConn, logLine, true)
 			}
 		}
 
@@ -193,7 +193,7 @@ func handleRemoteForward(newRequest *ssh.Request, sshConn *SSHConnection, state
 
 		newChan, newReqs, err := sshConn.SSHConn.OpenChannel("forwarded-tcpip", ssh.Marshal(resp))
 		if err != nil {
-			sendMessage(sshConn, err.Error())
+			sendMessage(sshConn, err.Error(), true)
 			cl.Close()
 			continue
 		}

utils.go

@@ -20,6 +20,7 @@ import (
 	"time"
 
 	"github.com/fsnotify/fsnotify"
+	"github.com/logrusorgru/aurora"
 	"golang.org/x/crypto/ssh"
 )
 
@@ -302,7 +303,7 @@ func getOpenHost(addr string, state *State, sshConn *SSHConnection) string {
 		}
 		reportUnavailable := func(unavailable bool) {
 			if first && unavailable {
-				sendMessage(sshConn, "This subdomain is unavaible. Assigning a random subdomain.")
+				sendMessage(sshConn, aurora.Sprintf("The subdomain %s is unavaible. Assigning a random subdomain.", aurora.Red(host)), true)
 			}
 		}
 
@@ -337,7 +338,7 @@ func getOpenAlias(addr string, port string, state *State, sshConn *SSHConnection
 		}
 		reportUnavailable := func(unavailable bool) {
 			if first && unavailable {
-				sendMessage(sshConn, "This alias is unavaible. Assigning a random alias.")
+				sendMessage(sshConn, aurora.Sprintf("The alias %s is unavaible. Assigning a random alias.", aurora.Red(alias)), true)
 			}
 		}
 
@@ -391,7 +392,12 @@ func RandStringBytesMaskImprSrc(n int) string {
 	return string(b)
 }
 
-func sendMessage(sshConn *SSHConnection, message string) {
+func sendMessage(sshConn *SSHConnection, message string, block bool) {
+	if block {
+		sshConn.Messages <- message
+		return
+	}
+
 	for i := 0; i < 2; {
 		select {
 		case <-sshConn.Close: