Refactor Mihomo WebSocket handling and enhance tests
- Updated the WebSocket upgrade detection logic in `isWebSocketUpgrade` to improve header handling and added a fallback for the "Sec-WebSocket-Key" header. - Refactored the WebSocket proxy logic to use a new `newMihomoWSTunnel` function, enhancing the connection handling process. - Introduced comprehensive test cases in `TestIsWebSocketUpgrade` to validate various WebSocket upgrade scenarios, ensuring robust functionality. - Improved error handling and request building for WebSocket upgrades, ensuring secure and efficient communication.
This commit is contained in:
Denozordec
+63
-65
@@ -16,10 +16,8 @@ import (
|
||||
|
||||
// NewMihomoForward proxies /api/{alias}/mihomo/... to Mihomo external-controller.
|
||||
//
|
||||
// REST и обычные GET/POST идут тем же путём, что и NewAliasForward (http.NewRequest + RoundTrip):
|
||||
// httputil.ReverseProxy для всего Mihomo часто давал 400 на upstream при том, что curl к контроллеру работал.
|
||||
// Только WebSocket (traffic, memory, …) остаётся на ReverseProxy + Rewrite.
|
||||
// Чеклист при повторении проблемы: docs/GATEWAY_RUN.md#mihomo-debug-400
|
||||
// REST requests use the same path as NewAliasForward (http.NewRequest + RoundTrip).
|
||||
// WebSocket (traffic, memory, …) uses raw TCP tunnel: hijack + dial + handshake + bidirectional copy.
|
||||
func NewMihomoForward(
|
||||
target *url.URL,
|
||||
stripPrefix string,
|
||||
@@ -28,7 +26,7 @@ func NewMihomoForward(
|
||||
errHandler func(http.ResponseWriter, *http.Request, error),
|
||||
) http.Handler {
|
||||
httpH := newAliasForward(target, stripPrefix, "", auth, rt, errHandler, false)
|
||||
wsH := newMihomoWebSocketReverseProxy(target, stripPrefix, auth, rt, errHandler)
|
||||
wsH := newMihomoWSTunnel(target, stripPrefix, auth, errHandler)
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if isWebSocketUpgrade(r) {
|
||||
wsH.ServeHTTP(w, r)
|
||||
@@ -42,17 +40,14 @@ func isWebSocketUpgrade(r *http.Request) bool {
|
||||
if r == nil {
|
||||
return false
|
||||
}
|
||||
// Be tolerant to proxy/header quirks:
|
||||
// - RFC path: Connection: upgrade + Upgrade: websocket
|
||||
// - Fallback: Sec-WebSocket-Key presence strongly indicates WS handshake.
|
||||
if headerHasToken(r.Header, "Upgrade", "websocket") &&
|
||||
headerHasToken(r.Header, "Connection", "upgrade") {
|
||||
if headerContainsToken(r.Header, "Upgrade", "websocket") &&
|
||||
headerContainsToken(r.Header, "Connection", "upgrade") {
|
||||
return true
|
||||
}
|
||||
return strings.TrimSpace(r.Header.Get("Sec-WebSocket-Key")) != ""
|
||||
}
|
||||
|
||||
func headerHasToken(h http.Header, key, token string) bool {
|
||||
func headerContainsToken(h http.Header, key, token string) bool {
|
||||
for _, v := range h.Values(key) {
|
||||
for _, part := range strings.Split(v, ",") {
|
||||
if strings.EqualFold(strings.TrimSpace(part), token) {
|
||||
@@ -63,11 +58,12 @@ func headerHasToken(h http.Header, key, token string) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func newMihomoWebSocketReverseProxy(
|
||||
// newMihomoWSTunnel creates a handler that tunnels WebSocket connections to Mihomo
|
||||
// by hijacking the client TCP connection and building the upstream request from scratch.
|
||||
func newMihomoWSTunnel(
|
||||
target *url.URL,
|
||||
stripPrefix string,
|
||||
auth string,
|
||||
rt http.RoundTripper,
|
||||
errHandler func(http.ResponseWriter, *http.Request, error),
|
||||
) http.Handler {
|
||||
if errHandler == nil {
|
||||
@@ -88,46 +84,36 @@ func newMihomoWebSocketReverseProxy(
|
||||
|
||||
hj, ok := w.(http.Hijacker)
|
||||
if !ok {
|
||||
errHandler(w, r, fmt.Errorf("websocket hijack unsupported"))
|
||||
errHandler(w, r, fmt.Errorf("websocket: hijack not supported by ResponseWriter"))
|
||||
return
|
||||
}
|
||||
clientConn, clientRW, err := hj.Hijack()
|
||||
clientConn, clientBuf, err := hj.Hijack()
|
||||
if err != nil {
|
||||
errHandler(w, r, fmt.Errorf("hijack client conn: %w", err))
|
||||
errHandler(w, r, fmt.Errorf("websocket: hijack failed: %w", err))
|
||||
return
|
||||
}
|
||||
defer clientConn.Close()
|
||||
|
||||
upConn, err := dialWebSocketUpstream(dest)
|
||||
upConn, err := dialUpstream(dest)
|
||||
if err != nil {
|
||||
_ = writeRawHTTPError(clientConn, http.StatusBadGateway, "bad gateway")
|
||||
_ = rawHTTPError(clientConn, http.StatusBadGateway, "upstream connect failed")
|
||||
return
|
||||
}
|
||||
defer upConn.Close()
|
||||
|
||||
outReq := r.Clone(r.Context())
|
||||
outReq.URL = dest
|
||||
outReq.Host = dest.Host
|
||||
outReq.RequestURI = ""
|
||||
outReq.Proto = "HTTP/1.1"
|
||||
outReq.ProtoMajor = 1
|
||||
outReq.ProtoMinor = 1
|
||||
outReq.Header = cloneHeader(r.Header)
|
||||
outReq.Header.Del("Authorization")
|
||||
if auth != "" {
|
||||
outReq.Header.Set("Authorization", auth)
|
||||
}
|
||||
if err := outReq.Write(upConn); err != nil {
|
||||
_ = writeRawHTTPError(clientConn, http.StatusBadGateway, "bad gateway")
|
||||
reqBytes := buildWSUpgradeRequest(r, dest, auth)
|
||||
if _, err := upConn.Write(reqBytes); err != nil {
|
||||
_ = rawHTTPError(clientConn, http.StatusBadGateway, "upstream write failed")
|
||||
return
|
||||
}
|
||||
|
||||
upBr := bufio.NewReader(upConn)
|
||||
resp, err := http.ReadResponse(upBr, outReq)
|
||||
upBuf := bufio.NewReader(upConn)
|
||||
resp, err := http.ReadResponse(upBuf, nil)
|
||||
if err != nil {
|
||||
_ = writeRawHTTPError(clientConn, http.StatusBadGateway, "bad gateway")
|
||||
_ = rawHTTPError(clientConn, http.StatusBadGateway, "upstream response read failed")
|
||||
return
|
||||
}
|
||||
|
||||
if err := resp.Write(clientConn); err != nil {
|
||||
return
|
||||
}
|
||||
@@ -135,18 +121,48 @@ func newMihomoWebSocketReverseProxy(
|
||||
return
|
||||
}
|
||||
|
||||
clientSrc := io.MultiReader(clientRW, clientConn)
|
||||
upSrc := io.MultiReader(upBr, upConn)
|
||||
errCh := make(chan error, 2)
|
||||
go proxyCopy(errCh, upConn, clientSrc)
|
||||
go proxyCopy(errCh, clientConn, upSrc)
|
||||
<-errCh
|
||||
done := make(chan struct{}, 2)
|
||||
go func() { _, _ = io.Copy(upConn, clientBuf); done <- struct{}{} }()
|
||||
go func() { _, _ = io.Copy(clientConn, upBuf); done <- struct{}{} }()
|
||||
<-done
|
||||
})
|
||||
}
|
||||
|
||||
func dialWebSocketUpstream(dest *url.URL) (net.Conn, error) {
|
||||
// buildWSUpgradeRequest constructs a raw HTTP/1.1 WebSocket upgrade request
|
||||
// with only the headers required by RFC 6455 + Authorization for Mihomo.
|
||||
// This avoids any extra headers that Go's Request.Write may add.
|
||||
func buildWSUpgradeRequest(orig *http.Request, dest *url.URL, auth string) []byte {
|
||||
reqURI := dest.RequestURI()
|
||||
if reqURI == "" {
|
||||
reqURI = "/"
|
||||
}
|
||||
var b strings.Builder
|
||||
fmt.Fprintf(&b, "GET %s HTTP/1.1\r\n", reqURI)
|
||||
fmt.Fprintf(&b, "Host: %s\r\n", dest.Host)
|
||||
b.WriteString("Connection: Upgrade\r\n")
|
||||
b.WriteString("Upgrade: websocket\r\n")
|
||||
if v := orig.Header.Get("Sec-WebSocket-Version"); v != "" {
|
||||
fmt.Fprintf(&b, "Sec-WebSocket-Version: %s\r\n", v)
|
||||
}
|
||||
if v := orig.Header.Get("Sec-WebSocket-Key"); v != "" {
|
||||
fmt.Fprintf(&b, "Sec-WebSocket-Key: %s\r\n", v)
|
||||
}
|
||||
if v := orig.Header.Get("Sec-WebSocket-Protocol"); v != "" {
|
||||
fmt.Fprintf(&b, "Sec-WebSocket-Protocol: %s\r\n", v)
|
||||
}
|
||||
if v := orig.Header.Get("Sec-WebSocket-Extensions"); v != "" {
|
||||
fmt.Fprintf(&b, "Sec-WebSocket-Extensions: %s\r\n", v)
|
||||
}
|
||||
if auth != "" {
|
||||
fmt.Fprintf(&b, "Authorization: %s\r\n", auth)
|
||||
}
|
||||
b.WriteString("\r\n")
|
||||
return []byte(b.String())
|
||||
}
|
||||
|
||||
func dialUpstream(dest *url.URL) (net.Conn, error) {
|
||||
if dest == nil {
|
||||
return nil, fmt.Errorf("nil destination")
|
||||
return nil, errors.New("nil destination URL")
|
||||
}
|
||||
hostPort := dest.Host
|
||||
if !strings.Contains(hostPort, ":") {
|
||||
@@ -166,7 +182,7 @@ func dialWebSocketUpstream(dest *url.URL) (net.Conn, error) {
|
||||
}
|
||||
}
|
||||
|
||||
func writeRawHTTPError(conn net.Conn, code int, text string) error {
|
||||
func rawHTTPError(conn net.Conn, code int, text string) error {
|
||||
reason := http.StatusText(code)
|
||||
if reason == "" {
|
||||
reason = "Error"
|
||||
@@ -174,30 +190,12 @@ func writeRawHTTPError(conn net.Conn, code int, text string) error {
|
||||
if text == "" {
|
||||
text = reason
|
||||
}
|
||||
_, err := fmt.Fprintf(conn, "HTTP/1.1 %d %s\r\nContent-Type: text/plain; charset=utf-8\r\nContent-Length: %d\r\nConnection: close\r\n\r\n%s", code, reason, len(text), text)
|
||||
_, err := fmt.Fprintf(conn,
|
||||
"HTTP/1.1 %d %s\r\nContent-Type: text/plain; charset=utf-8\r\nContent-Length: %d\r\nConnection: close\r\n\r\n%s",
|
||||
code, reason, len(text), text)
|
||||
return err
|
||||
}
|
||||
|
||||
func proxyCopy(errCh chan<- error, dst io.Writer, src io.Reader) {
|
||||
_, err := io.Copy(dst, src)
|
||||
if err != nil && !isNetClosed(err) {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
errCh <- nil
|
||||
}
|
||||
|
||||
func isNetClosed(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
if errors.Is(err, net.ErrClosed) {
|
||||
return true
|
||||
}
|
||||
msg := err.Error()
|
||||
return strings.Contains(msg, "use of closed network connection")
|
||||
}
|
||||
|
||||
// MihomoMetaJSON returns a JSON body for GET .../mihomo/meta (display URL without credentials).
|
||||
func MihomoMetaJSON(target *url.URL) []byte {
|
||||
display := strings.TrimSuffix(target.String(), "/")
|
||||
|
||||
Reference in New Issue
Block a user