feat: add dedicated WebSocket service for subdomain ping

server/http.go

@@ -7,11 +7,57 @@ import (
 	"fmt"
 	"log"
 	"net"
+	"net/http"
 	"strings"
 	"tunnel_pls/session"
 	"tunnel_pls/utils"
+
+	"github.com/gorilla/websocket"
 )
 
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+}
+
+type connResponseWriter struct {
+	conn   net.Conn
+	header http.Header
+	wrote  bool
+}
+
+func (w *connResponseWriter) Header() http.Header {
+	if w.header == nil {
+		w.header = make(http.Header)
+	}
+	return w.header
+}
+
+func (w *connResponseWriter) WriteHeader(statusCode int) {
+	if w.wrote {
+		return
+	}
+	w.wrote = true
+	fmt.Fprintf(w.conn, "HTTP/1.1 %d %s\r\n", statusCode, http.StatusText(statusCode))
+	w.header.Write(w.conn)
+	fmt.Fprint(w.conn, "\r\n")
+}
+
+func (w *connResponseWriter) Write(b []byte) (int, error) {
+	if !w.wrote {
+		w.WriteHeader(http.StatusOK)
+	}
+	return w.conn.Write(b)
+}
+
+func (w *connResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	rw := bufio.NewReadWriter(
+		bufio.NewReader(w.conn),
+		bufio.NewWriter(w.conn),
+	)
+	return w.conn, rw, nil
+}
+
 var redirectTLS = false
 
 func NewHTTPServer() error {
@@ -72,6 +118,40 @@ func Handler(conn net.Conn) {
 		return
 	}
 
+	if slug == "ping" {
+		req, err := http.ReadRequest(reader)
+		if err != nil {
+			log.Println("failed to parse HTTP request:", err)
+			return
+		}
+		rw := &connResponseWriter{conn: conn}
+
+		wsConn, err := upgrader.Upgrade(rw, req, nil)
+		if err != nil {
+			if !strings.Contains(err.Error(), "the client is not using the websocket protocol") {
+				log.Println("Upgrade failed:", err)
+			}
+			err := conn.Close()
+			if err != nil {
+				log.Println("failed to close connection:", err)
+				return
+			}
+			return
+		}
+
+		err = wsConn.WriteMessage(websocket.TextMessage, []byte("pong"))
+		if err != nil {
+			log.Println("failed to write pong:", err)
+			return
+		}
+		err = wsConn.Close()
+		if err != nil {
+			log.Println("websocket close failed :", err)
+			return
+		}
+		return
+	}
+
 	sshSession, ok := session.Clients[slug]
 	if !ok {
 		conn.Write([]byte("HTTP/1.1 400 Bad Request\r\n\r\n"))

server/https.go

@@ -6,9 +6,12 @@ import (
 	"errors"
 	"log"
 	"net"
+	"net/http"
 	"strings"
 	"tunnel_pls/session"
 	"tunnel_pls/utils"
+
+	"github.com/gorilla/websocket"
 )
 
 func NewHTTPSServer() error {
@@ -62,6 +65,40 @@ func HandlerTLS(conn net.Conn) {
 	}
 	slug := host[0]
 
+	if slug == "ping" {
+		req, err := http.ReadRequest(reader)
+		if err != nil {
+			log.Println("failed to parse HTTP request:", err)
+			return
+		}
+		rw := &connResponseWriter{conn: conn}
+
+		wsConn, err := upgrader.Upgrade(rw, req, nil)
+		if err != nil {
+			if !strings.Contains(err.Error(), "the client is not using the websocket protocol") {
+				log.Println("Upgrade failed:", err)
+			}
+			err := conn.Close()
+			if err != nil {
+				log.Println("failed to close connection:", err)
+				return
+			}
+			return
+		}
+
+		err = wsConn.WriteMessage(websocket.TextMessage, []byte("pong"))
+		if err != nil {
+			log.Println("failed to write pong:", err)
+			return
+		}
+		err = wsConn.Close()
+		if err != nil {
+			log.Println("websocket close failed :", err)
+			return
+		}
+		return
+	}
+
 	sshSession, ok := session.Clients[slug]
 	if !ok {
 		conn.Write([]byte("HTTP/1.1 400 Bad Request\r\n\r\n"))