Improve logging and connection handling

Dockerfile

@@ -7,4 +7,5 @@ RUN go build -ldflags="-w -s" -o /go/bin/sshway
 
 FROM alpine:3.18
 COPY --from=builder /go/bin/sshway /bin/sshway
-CMD ["/bin/sshway"]
+EXPOSE 2222
\ No newline at end of file
+CMD ["/bin/sshway"]

main.go

@@ -67,17 +67,20 @@ func main() {
 	fmt.Printf("Listening on %s\n", localAddr)
 	defer server.Close()
 
+	client_counter := 0
+
 	for {
 		connection, err := server.Accept()
 		if err != nil {
 			log.Fatalf("Error accepting local client: %v\n", err)
 		}
-		fmt.Printf("client connected %s\n", connection.RemoteAddr().String())
-		go handleClient(connection)
+		log.Printf("%6d New connection from client: %s\n", client_counter, connection.RemoteAddr().String())
+		go handleClient(client_counter, connection)
+		client_counter++
 	}
 }
 
-func handleClient(localSock net.Conn) {
+func handleClient(client_counter int, localSock net.Conn) {
 	defer localSock.Close()
 
 	config := &ssh.ClientConfig{
@@ -92,21 +95,24 @@ func handleClient(localSock net.Conn) {
 
 	conn, err := ssh.Dial("tcp", sshHost, config)
 	if err != nil {
-		log.Printf("SSH connection failed: %v\n", err)
+		log.Printf("%6d SSH connection failed: %v\n", client_counter, err)
 		return
 	}
-	defer conn.Close()
 
 	forwardClient, err := conn.Dial("tcp", remoteAddr)
 	if err != nil {
-		log.Printf("Port forward failed: %v\n", err)
+		log.Printf("%6d Port forward failed: %v\n", client_counter, err)
+		conn.Close()
 		return
 	}
 
 	go func() {
-		io.Copy(localSock, forwardClient)
+		io.Copy(forwardClient, localSock)
+		log.Printf("%6d Local client disconnected\n", client_counter)
+		conn.Close()
 	}()
 
-	io.Copy(forwardClient, localSock)
-	log.Printf("Disconnected: %v\n", localSock)
+	io.Copy(localSock, forwardClient)
+	log.Printf("%6d Remote socket closed for local client\n", client_counter)
+	conn.Close()
 }