golang: do not crash if dialing the VNC server fails

golang/websockify.go

@@ -42,9 +42,18 @@ var upgrader = websocket.Upgrader{
 
 func forwardTcp(wsConn *websocket.Conn, conn net.Conn) {
 	var tcpBuffer [1024]byte
-	defer wsConn.Close()
-	defer conn.Close()
+	defer func() {
+		if conn != nil {
+			conn.Close()
+		}
+		if wsConn != nil {
+			wsConn.Close()
+		}
+	}()
 	for {
+		if (conn == nil) || (wsConn == nil) {
+			return
+		}
 		n, err := conn.Read(tcpBuffer[0:])
 		if err != nil {
 			log.Printf("%s: reading from TCP failed: %s", time.Now().Format(time.Stamp), err)
@@ -58,15 +67,22 @@ func forwardTcp(wsConn *websocket.Conn, conn net.Conn) {
 }
 
 func forwardWeb(wsConn *websocket.Conn, conn net.Conn) {
-	defer wsConn.Close()
-	defer conn.Close()
+	defer func() {
+		if err := recover(); err != nil {
+			log.Printf("%s: reading from WS failed: %s", time.Now().Format(time.Stamp), err)
+		}
+		if conn != nil {
+			conn.Close()
+		}
+		if wsConn != nil {
+			wsConn.Close()
+		}
+	}()
 	for {
-		defer func() {
-			if err := recover(); err != nil {
-				log.Printf("%s: reading from WS failed: %s", time.Now().Format(time.Stamp), err)
-				return
-			}
-		}()
+		if (conn == nil) || (wsConn == nil) {
+			return
+		}
+
 		_, buffer, err := wsConn.ReadMessage()
 		if err == nil {
 			if _, err := conn.Write(buffer); err != nil {
@@ -82,10 +98,14 @@ func serveWs(w http.ResponseWriter, r *http.Request) {
 		log.Printf("%s: failed to upgrade to WS: %s", time.Now().Format(time.Stamp), err)
 		return
 	}
+
 	vnc, err := net.Dial("tcp", *targetAddr)
+	if err != nil {
+		log.Printf("%s: failed to bind to the VNC Server: %s", time.Now().Format(time.Stamp), err)
+	}
+
 	go forwardTcp(ws, vnc)
 	go forwardWeb(ws, vnc)
-
 }
 
 func main() {