return 504 instead of 404 for proxy type http request timeout (#4151)

pkg/util/log/log.go

@@ -15,11 +15,20 @@
 package log
 
 import (
+	"bytes"
 	"os"
 
 	"github.com/fatedier/golib/log"
 )
 
+var (
+	TraceLevel = log.TraceLevel
+	DebugLevel = log.DebugLevel
+	InfoLevel  = log.InfoLevel
+	WarnLevel  = log.WarnLevel
+	ErrorLevel = log.ErrorLevel
+)
+
 var Logger *log.Logger
 
 func init() {
@@ -77,3 +86,24 @@ func Debugf(format string, v ...interface{}) {
 func Tracef(format string, v ...interface{}) {
 	Logger.Tracef(format, v...)
 }
+
+func Logf(level log.Level, offset int, format string, v ...interface{}) {
+	Logger.Logf(level, offset, format, v...)
+}
+
+type WriteLogger struct {
+	level  log.Level
+	offset int
+}
+
+func NewWriteLogger(level log.Level, offset int) *WriteLogger {
+	return &WriteLogger{
+		level:  level,
+		offset: offset,
+	}
+}
+
+func (w *WriteLogger) Write(p []byte) (n int, err error) {
+	Logger.Log(w.level, w.offset, string(bytes.TrimRight(p, "\n")))
+	return len(p), nil
+}

pkg/util/net/dial.go

@@ -5,11 +5,11 @@ import (
 	"net"
 	"net/url"
 
-	libdial "github.com/fatedier/golib/net/dial"
+	libnet "github.com/fatedier/golib/net"
 	"golang.org/x/net/websocket"
 )
 
-func DialHookCustomTLSHeadByte(enableTLS bool, disableCustomTLSHeadByte bool) libdial.AfterHookFunc {
+func DialHookCustomTLSHeadByte(enableTLS bool, disableCustomTLSHeadByte bool) libnet.AfterHookFunc {
 	return func(ctx context.Context, c net.Conn, addr string) (context.Context, net.Conn, error) {
 		if enableTLS && !disableCustomTLSHeadByte {
 			_, err := c.Write([]byte{byte(FRPTLSHeadByte)})
@@ -21,7 +21,7 @@ func DialHookCustomTLSHeadByte(enableTLS bool, disableCustomTLSHeadByte bool) li
 	}
 }
 
-func DialHookWebsocket(protocol string, host string) libdial.AfterHookFunc {
+func DialHookWebsocket(protocol string, host string) libnet.AfterHookFunc {
 	return func(ctx context.Context, c net.Conn, addr string) (context.Context, net.Conn, error) {
 		if protocol != "wss" {
 			protocol = "ws"

pkg/util/net/websocket.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"net"
 	"net/http"
+	"time"
 
 	"golang.org/x/net/websocket"
 )
@@ -39,8 +40,9 @@ func NewWebsocketListener(ln net.Listener) (wl *WebsocketListener) {
 	}))
 
 	wl.server = &http.Server{
-		Addr:    ln.Addr().String(),
-		Handler: muxer,
+		Addr:              ln.Addr().String(),
+		Handler:           muxer,
+		ReadHeaderTimeout: 60 * time.Second,
 	}
 
 	go func() {

pkg/util/vhost/http.go

@@ -15,7 +15,6 @@
 package vhost
 
 import (
-	"bytes"
 	"context"
 	"encoding/base64"
 	"errors"
@@ -116,10 +115,16 @@ func NewHTTPReverseProxy(option HTTPReverseProxyOptions, vhostRouter *Routers) *
 				return nil, nil
 			},
 		},
-		BufferPool: newWrapPool(),
-		ErrorLog:   stdlog.New(newWrapLogger(), "", 0),
+		BufferPool: pool.NewBuffer(32 * 1024),
+		ErrorLog:   stdlog.New(log.NewWriteLogger(log.WarnLevel, 2), "", 0),
 		ErrorHandler: func(rw http.ResponseWriter, req *http.Request, err error) {
-			log.Warnf("do http proxy request [host: %s] error: %v", req.Host, err)
+			log.Logf(log.WarnLevel, 1, "do http proxy request [host: %s] error: %v", req.Host, err)
+			if err != nil {
+				if e, ok := err.(net.Error); ok && e.Timeout() {
+					rw.WriteHeader(http.StatusGatewayTimeout)
+					return
+				}
+			}
 			rw.WriteHeader(http.StatusNotFound)
 			_, _ = rw.Write(getNotFoundPageContent())
 		},
@@ -322,20 +327,3 @@ func (rp *HTTPReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request)
 		rp.proxy.ServeHTTP(rw, newreq)
 	}
 }
-
-type wrapPool struct{}
-
-func newWrapPool() *wrapPool { return &wrapPool{} }
-
-func (p *wrapPool) Get() []byte { return pool.GetBuf(32 * 1024) }
-
-func (p *wrapPool) Put(buf []byte) { pool.PutBuf(buf) }
-
-type wrapLogger struct{}
-
-func newWrapLogger() *wrapLogger { return &wrapLogger{} }
-
-func (l *wrapLogger) Write(p []byte) (n int, err error) {
-	log.Warnf("%s", string(bytes.TrimRight(p, "\n")))
-	return len(p), nil
-}

pkg/util/vhost/https_test.go

@@ -12,7 +12,7 @@ import (
 func TestGetHTTPSHostname(t *testing.T) {
 	require := require.New(t)
 
-	l, err := net.Listen("tcp", ":")
+	l, err := net.Listen("tcp", "127.0.0.1:")
 	require.NoError(err)
 	defer l.Close()