frps: allow to use custom status code for not found page/vhost

pkg/config/server.go

@@ -121,6 +121,9 @@ type ServerCommonConf struct {
 	// value is "", a default page will be displayed. By default, this value is
 	// "".
 	Custom404Page string `ini:"custom_404_page" json:"custom_404_page"`
+	// Custom404StatusCode specifies a status code for response when 404 page
+	// is being sent. By default, this value is 404.
+	Custom404StatusCode int `ini:"custom_404_status_code" json:"custom_404_status_code"`
 
 	// AllowPorts specifies a set of ports that clients are able to proxy to.
 	// If the length of this value is 0, all ports are allowed. By default,
@@ -203,6 +206,7 @@ func GetDefaultServerConf() ServerCommonConf {
 		HeartbeatTimeout:       90,
 		UserConnTimeout:        10,
 		Custom404Page:          "",
+		Custom404StatusCode:    404,
 		HTTPPlugins:            make(map[string]plugin.HTTPPluginOptions),
 		UDPPacketSize:          1500,
 	}

pkg/config/server_test.go

@@ -75,6 +75,7 @@ func Test_LoadServerCommonConf(t *testing.T) {
 				subdomain_host = frps.com
 				tcp_mux
 				udp_packet_size = 1509
+				custom_404_status_code = 400
 				[plugin.user-manager]
 				addr = 127.0.0.1:9009
 				path = /handler
@@ -125,6 +126,7 @@ func Test_LoadServerCommonConf(t *testing.T) {
 				DetailedErrorsToClient: true,
 				HeartbeatTimeout:       99,
 				UserConnTimeout:        9,
+				Custom404StatusCode:    400,
 				AllowPorts: map[int]struct{}{
 					10: struct{}{},
 					11: struct{}{},
@@ -189,6 +191,7 @@ func Test_LoadServerCommonConf(t *testing.T) {
 				LogMaxDays:             3,
 				DetailedErrorsToClient: true,
 				TCPMux:                 true,
+				Custom404StatusCode:    404,
 				AllowPorts:             make(map[int]struct{}),
 				MaxPoolCount:           5,
 				HeartbeatTimeout:       90,

pkg/util/vhost/http.go

@@ -20,6 +20,7 @@ import (
 	"encoding/base64"
 	"errors"
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"net/http"
@@ -90,8 +91,9 @@ func NewHTTPReverseProxy(option HTTPReverseProxyOptions, vhostRouter *Routers) *
 		ErrorLog:   log.New(newWrapLogger(), "", 0),
 		ErrorHandler: func(rw http.ResponseWriter, req *http.Request, err error) {
 			frpLog.Warn("do http proxy request error: %v", err)
-			rw.WriteHeader(http.StatusNotFound)
+			res := notFoundResponse()
-			rw.Write(getNotFoundPageContent())
+			rw.WriteHeader(res.StatusCode)
+			io.Copy(rw, res.Body)
 		},
 	}
 	rp.proxy = proxy

pkg/util/vhost/resource.go

@@ -24,7 +24,8 @@ import (
 )
 
 var (
-	NotFoundPagePath = ""
+	NotFoundPagePath   = ""
+	NotFoundStatusCode = 404
 )
 
 const (
@@ -74,8 +75,7 @@ func notFoundResponse() *http.Response {
 	header.Set("Content-Type", "text/html")
 
 	res := &http.Response{
-		Status:     "Not Found",
+		StatusCode: NotFoundStatusCode,
-		StatusCode: 404,
 		Proto:      "HTTP/1.0",
 		ProtoMajor: 1,
 		ProtoMinor: 0,
@@ -89,7 +89,6 @@ func noAuthResponse() *http.Response {
 	header := make(map[string][]string)
 	header["WWW-Authenticate"] = []string{`Basic realm="Restricted"`}
 	res := &http.Response{
-		Status:     "401 Not authorized",
 		StatusCode: 401,
 		Proto:      "HTTP/1.1",
 		ProtoMajor: 1,

server/service.go

@@ -163,6 +163,9 @@ func NewService(cfg config.ServerCommonConf) (svr *Service, err error) {
 	// Init 404 not found page
 	vhost.NotFoundPagePath = cfg.Custom404Page
 
+	// Init not found status code
+	vhost.NotFoundStatusCode = cfg.Custom404StatusCode
+
 	var (
 		httpMuxOn  bool
 		httpsMuxOn bool