add X11 socket/port probe

auth.go

@@ -2,15 +2,16 @@ package main
 
 import (
 	"fmt"
+	"net"
 	"os"
 	"os/exec"
 	"path"
 	"strconv"
 	"strings"
-	"net"
+	"x11proxy/x11probe"
 )
 
-func resolveDisplay(display string) (ConnType, string) {
+func resolveDisplay(display string) (x11probe.ConnType, string) {
 	var dispNum int
 	var err error
 
@@ -26,17 +27,17 @@ func resolveDisplay(display string) (ConnType, string) {
 
 	unixPath := path.Join("/tmp/.X11-unix", fmt.Sprintf("X%d", dispNum))
 	if _, err := os.Stat(unixPath); err == nil {
-        return Unix, unixPath
+		return x11probe.Unix, unixPath
 	}
 
 	tcpAddr := fmt.Sprintf("127.0.0.1:%d", 6000+dispNum)
 	conn, err := net.Dial("tcp", tcpAddr)
 	if err == nil {
 		conn.Close()
-        return TCP, tcpAddr
+		return x11probe.TCP, tcpAddr
 	}
 
-    return Unix, unixPath
+	return x11probe.Unix, unixPath
 }
 
 func getXAuthCookie(display string) ([]byte, error) {
@@ -68,7 +69,6 @@ func getXAuthCookie(display string) ([]byte, error) {
 	return nil, fmt.Errorf("no matching cookie found for display %s", display)
 }
 
-
 func parseHexCookie(hexStr string) []byte {
 	var cookie []byte
 	for i := 0; i < len(hexStr); i += 2 {
@@ -113,4 +113,3 @@ func PatchAuth(data []byte, cookie []byte) []byte {
 
 	return patched
 }
-

main.go

@@ -6,6 +6,8 @@ import (
 	"log"
 	"os"
 	"strings"
+	"time"
+	"x11proxy/x11probe"
 )
 
 const version = "0.1"
@@ -40,14 +42,38 @@ func main() {
 	connType, target := resolveDisplay(display)
 	fmt.Printf("Proxying to %s (%s)\n", target, connTypeString(connType))
 
-    err := StartProxy(*overrideSocket, target, connType, display)
+	timeout := 10 * time.Second
+	status, err := x11probe.ProbeX11Socket(connType, target, timeout)
 	if err != nil {
-        log.Fatalf("Proxy error: %v", err)
+		log.Fatalf("Connection probe error: %v", err)
+	}
+
+	switch status {
+	case x11probe.X11Dead:
+		log.Fatalf("Target connection is dead or unreachable")
+	case x11probe.X11AuthRequired:
+		fmt.Println("Authentication required for the target connection.")
+	case x11probe.X11Live:
+		fmt.Println("Connection to target is live and valid.")
+	}
+
+	switch status {
+	case x11probe.X11Live:
+		log.Println("X11 socket is live and accepting connections")
+	case x11probe.X11AuthRequired:
+		log.Println("X11 socket is live but requires authentication")
+	case x11probe.X11Dead:
+		log.Fatalf("Target connection is dead or unreachable")
+	}
+
+	err2 := StartProxy(*overrideSocket, target, connType, display)
+	if err2 != nil {
+		log.Fatalf("Proxy error: %v", err2)
 	}
 }
 
-func connTypeString(t ConnType) string {
-    if t == Unix {
+func connTypeString(t x11probe.ConnType) string {
+	if t == x11probe.Unix {
 		return "Unix socket"
 	}
 	return "TCP"

proxy.go

@@ -8,16 +8,10 @@ import (
 	"log"
 	"net"
 	"os"
+	"x11proxy/x11probe"
 )
 
-type ConnType int
-
-const (
-    Unix ConnType = iota
-    TCP
-)
-
-func StartProxy(proxyPath, target string, connType ConnType, display string) error {
+func StartProxy(proxyPath, target string, connType x11probe.ConnType, display string) error {
 	os.Remove(proxyPath)
 	listener, err := net.Listen("unix", proxyPath)
 	if err != nil {
@@ -113,12 +107,11 @@ func binaryOrder(isLittle bool) binary.ByteOrder {
 	return binary.BigEndian
 }
 
-
-func handleConnection(client net.Conn, connType ConnType, target string, display string) {
+func handleConnection(client net.Conn, connType x11probe.ConnType, target string, display string) {
 	var serverConn net.Conn
 	var err error
 
-    if connType == Unix {
+	if connType == x11probe.Unix {
 		serverConn, err = net.Dial("unix", target)
 	} else {
 		serverConn, err = net.Dial("tcp", target)

x11probe/probe.go

@@ -0,0 +1,125 @@
+package x11probe
+
+import (
+	"fmt"
+	"io"
+	"log"
+	"net"
+	"strings"
+	"time"
+)
+
+type ConnType int
+
+const (
+	Unix ConnType = iota
+	TCP
+)
+
+// X11Status represents the result of probing an X11 socket.
+type X11Status int
+
+const (
+	X11Dead X11Status = iota
+	X11Live
+	X11AuthRequired
+)
+
+func ProbeX11Socket(connType ConnType, target string, timeout time.Duration) (X11Status, error) {
+	var (
+		conn net.Conn
+		err  error
+	)
+
+	switch connType {
+	case Unix:
+		conn, err = net.DialTimeout("unix", target, timeout)
+		if err != nil {
+			return X11Dead, fmt.Errorf("unix dial failed: %w", err)
+		}
+	case TCP:
+		addr := parseTCPAddress(target)
+		conn, err = net.DialTimeout("tcp", addr, timeout)
+		if err != nil {
+			return X11Dead, fmt.Errorf("TCP dial failed: %w", err)
+		}
+	}
+
+	defer conn.Close()
+	log.Printf("Successfully connected to %s", target)
+
+	if err := performHandshake(conn); err != nil {
+		return X11Dead, err
+	}
+
+	status, err := readResponse(conn, timeout)
+	if err != nil {
+		return X11Dead, err
+	}
+
+	return status, nil
+}
+
+func parseTCPAddress(target string) string {
+	parts := strings.SplitN(target, ":", 2)
+
+	ipPart := parts[0]
+	portPart := "0"
+
+	if len(parts) == 2 && portPart != "" {
+		portPart = parts[1]
+	}
+
+	return fmt.Sprintf("%s:%s", ipPart, portPart)
+}
+
+func performHandshake(conn net.Conn) error {
+	// Minimal valid handshake with correct padding
+	handshake := []byte{
+		'l', 0, // Byte order
+		0, 11, 0, 0, // Protocol major/minor
+		0, 0, // Auth name length
+		0, 0, // Auth data length
+		0, 0, 0, 0, // Padding
+	}
+
+	if _, err := conn.Write(handshake); err != nil {
+		return fmt.Errorf("write failed: %w", err)
+	}
+	return nil
+}
+
+func readResponse(conn net.Conn, timeout time.Duration) (X11Status, error) {
+	buf := make([]byte, 1)
+
+	deadline := time.Now().Add(timeout)
+	if err := conn.SetReadDeadline(deadline); err != nil {
+		return X11Dead, fmt.Errorf("SetReadDeadline failed: %w", err)
+	}
+
+	n, err := conn.Read(buf)
+	if err != nil {
+		if err == io.EOF {
+			return X11AuthRequired, nil
+		}
+		return X11Dead, fmt.Errorf("read failed: %w", err)
+	}
+
+	if n == 0 {
+		return X11Dead, fmt.Errorf("read returned 0 bytes")
+	}
+
+	switch buf[0] {
+	case 0x01:
+		return X11Live, nil
+	case 0x02:
+		return X11AuthRequired, nil
+	case 0x00:
+		fmt.Printf("X11 handshake rejected (value %x)\n", buf[0])
+		return X11AuthRequired, nil
+	default:
+		fmt.Printf("unexpected value %x\n", buf[0])
+		return X11Dead, nil
+	}
+
+}