Merge pull request #4 from fatedier/master

update
max0y · May 31, 2019 · 04978bb · 04978bb
2 parents a75bafd + 95444ea
commit 04978bb
Show file tree

Hide file tree

Showing 14 changed files with 86 additions and 44 deletions.
diff --git a/client/proxy/proxy.go b/client/proxy/proxy.go
@@ -503,10 +503,43 @@ func HandleTcpWorkConnection(localInfo *config.LocalSvrConf, proxyPlugin plugin.
 		remote = frpIo.WithCompression(remote)
 	}
 
+	// check if we need to send proxy protocol info
+	var extraInfo []byte
+	if baseInfo.ProxyProtocolVersion != "" {
+		if m.SrcAddr != "" && m.SrcPort != 0 {
+			if m.DstAddr == "" {
+				m.DstAddr = "127.0.0.1"
+			}
+			h := &pp.Header{
+				Command:            pp.PROXY,
+				SourceAddress:      net.ParseIP(m.SrcAddr),
+				SourcePort:         m.SrcPort,
+				DestinationAddress: net.ParseIP(m.DstAddr),
+				DestinationPort:    m.DstPort,
+			}
+
+			if h.SourceAddress.To16() == nil {
+				h.TransportProtocol = pp.TCPv4
+			} else {
+				h.TransportProtocol = pp.TCPv6
+			}
+
+			if baseInfo.ProxyProtocolVersion == "v1" {
+				h.Version = 1
+			} else if baseInfo.ProxyProtocolVersion == "v2" {
+				h.Version = 2
+			}
+
+			buf := bytes.NewBuffer(nil)
+			h.WriteTo(buf)
+			extraInfo = buf.Bytes()
+		}
+	}
+
 	if proxyPlugin != nil {
 		// if plugin is set, let plugin handle connections first
 		workConn.Debug("handle by plugin: %s", proxyPlugin.Name())
-		proxyPlugin.Handle(remote, workConn)
+		proxyPlugin.Handle(remote, workConn, extraInfo)
 		workConn.Debug("handle by plugin finished")
 		return
 	} else {
@@ -520,34 +553,8 @@ func HandleTcpWorkConnection(localInfo *config.LocalSvrConf, proxyPlugin plugin.
 		workConn.Debug("join connections, localConn(l[%s] r[%s]) workConn(l[%s] r[%s])", localConn.LocalAddr().String(),
 			localConn.RemoteAddr().String(), workConn.LocalAddr().String(), workConn.RemoteAddr().String())
 
-		// check if we need to send proxy protocol info
-		if baseInfo.ProxyProtocolVersion != "" {
-			if m.SrcAddr != "" && m.SrcPort != 0 {
-				if m.DstAddr == "" {
-					m.DstAddr = "127.0.0.1"
-				}
-				h := &pp.Header{
-					Command:            pp.PROXY,
-					SourceAddress:      net.ParseIP(m.SrcAddr),
-					SourcePort:         m.SrcPort,
-					DestinationAddress: net.ParseIP(m.DstAddr),
-					DestinationPort:    m.DstPort,
-				}
-
-				if h.SourceAddress.To16() == nil {
-					h.TransportProtocol = pp.TCPv4
-				} else {
-					h.TransportProtocol = pp.TCPv6
-				}
-
-				if baseInfo.ProxyProtocolVersion == "v1" {
-					h.Version = 1
-				} else if baseInfo.ProxyProtocolVersion == "v2" {
-					h.Version = 2
-				}
-
-				h.WriteTo(localConn)
-			}
+		if len(extraInfo) > 0 {
+			localConn.Write(extraInfo)
 		}
 
 		frpIo.Join(localConn, remote)

diff --git a/conf/frps_full.ini b/conf/frps_full.ini
@@ -65,3 +65,6 @@ subdomain_host = frps.com
 
 # if tcp stream multiplexing is used, default is true
 tcp_mux = true
+
+# custom 404 page for HTTP requests
+# custom_404_page = /path/to/404.html
diff --git a/models/config/server_common.go b/models/config/server_common.go
@@ -69,6 +69,7 @@ type ServerCommonConf struct {
 	Token         string `json:"token"`
 	SubDomainHost string `json:"subdomain_host"`
 	TcpMux        bool   `json:"tcp_mux"`
+	Custom404Page string `json:"custom_404_page"`
 
 	AllowPorts        map[int]struct{}
 	MaxPoolCount      int64 `json:"max_pool_count"`
@@ -104,6 +105,7 @@ func GetDefaultServerConf() *ServerCommonConf {
 		MaxPortsPerClient: 0,
 		HeartBeatTimeout:  90,
 		UserConnTimeout:   10,
+		Custom404Page:     "",
 	}
 }
 
@@ -293,6 +295,10 @@ func UnmarshalServerConfFromIni(defaultCfg *ServerCommonConf, content string) (c
 		cfg.TcpMux = true
 	}
 
+	if tmpStr, ok = conf.Get("common", "custom_404_page"); ok {
+		cfg.Custom404Page = tmpStr
+	}
+
 	if tmpStr, ok = conf.Get("common", "heartbeat_timeout"); ok {
 		v, errRet := strconv.ParseInt(tmpStr, 10, 64)
 		if errRet != nil {

diff --git a/models/plugin/http_proxy.go b/models/plugin/http_proxy.go
@@ -64,7 +64,7 @@ func (hp *HttpProxy) Name() string {
 	return PluginHttpProxy
 }
 
-func (hp *HttpProxy) Handle(conn io.ReadWriteCloser, realConn frpNet.Conn) {
+func (hp *HttpProxy) Handle(conn io.ReadWriteCloser, realConn frpNet.Conn, extraBufToLocal []byte) {
 	wrapConn := frpNet.WrapReadWriteCloserToConn(conn, realConn)
 
 	sc, rd := gnet.NewSharedConn(wrapConn)

diff --git a/models/plugin/https2http.go b/models/plugin/https2http.go
@@ -100,16 +100,11 @@ func (p *HTTPS2HTTPPlugin) genTLSConfig() (*tls.Config, error) {
 	return config, nil
 }
 
-func (p *HTTPS2HTTPPlugin) Handle(conn io.ReadWriteCloser, realConn frpNet.Conn) {
+func (p *HTTPS2HTTPPlugin) Handle(conn io.ReadWriteCloser, realConn frpNet.Conn, extraBufToLocal []byte) {
 	wrapConn := frpNet.WrapReadWriteCloserToConn(conn, realConn)
 	p.l.PutConn(wrapConn)
 }
 
-func (p *HTTPS2HTTPPlugin) handleRequest(w http.ResponseWriter, r *http.Request) {
-	w.Write([]byte("hello"))
-	return
-}
-
 func (p *HTTPS2HTTPPlugin) Name() string {
 	return PluginHTTPS2HTTP
 }

diff --git a/models/plugin/plugin.go b/models/plugin/plugin.go
@@ -46,7 +46,7 @@ func Create(name string, params map[string]string) (p Plugin, err error) {
 
 type Plugin interface {
 	Name() string
-	Handle(conn io.ReadWriteCloser, realConn frpNet.Conn)
+	Handle(conn io.ReadWriteCloser, realConn frpNet.Conn, extraBufToLocal []byte)
 	Close() error
 }
 

diff --git a/models/plugin/socks5.go b/models/plugin/socks5.go
@@ -53,7 +53,7 @@ func NewSocks5Plugin(params map[string]string) (p Plugin, err error) {
 	return
 }
 
-func (sp *Socks5Plugin) Handle(conn io.ReadWriteCloser, realConn frpNet.Conn) {
+func (sp *Socks5Plugin) Handle(conn io.ReadWriteCloser, realConn frpNet.Conn, extraBufToLocal []byte) {
 	defer conn.Close()
 	wrapConn := frpNet.WrapReadWriteCloserToConn(conn, realConn)
 	sp.Server.ServeConn(wrapConn)

diff --git a/models/plugin/static_file.go b/models/plugin/static_file.go
@@ -72,7 +72,7 @@ func NewStaticFilePlugin(params map[string]string) (Plugin, error) {
 	return sp, nil
 }
 
-func (sp *StaticFilePlugin) Handle(conn io.ReadWriteCloser, realConn frpNet.Conn) {
+func (sp *StaticFilePlugin) Handle(conn io.ReadWriteCloser, realConn frpNet.Conn, extraBufToLocal []byte) {
 	wrapConn := frpNet.WrapReadWriteCloserToConn(conn, realConn)
 	sp.l.PutConn(wrapConn)
 }

diff --git a/models/plugin/unix_domain_socket.go b/models/plugin/unix_domain_socket.go
@@ -53,11 +53,14 @@ func NewUnixDomainSocketPlugin(params map[string]string) (p Plugin, err error) {
 	return
 }
 
-func (uds *UnixDomainSocketPlugin) Handle(conn io.ReadWriteCloser, realConn frpNet.Conn) {
+func (uds *UnixDomainSocketPlugin) Handle(conn io.ReadWriteCloser, realConn frpNet.Conn, extraBufToLocal []byte) {
 	localConn, err := net.DialUnix("unix", nil, uds.UnixAddr)
 	if err != nil {
 		return
 	}
+	if len(extraBufToLocal) > 0 {
+		localConn.Write(extraBufToLocal)
+	}
 
 	frpIo.Join(localConn, conn)
 }

diff --git a/server/dashboard_api.go b/server/dashboard_api.go
@@ -279,6 +279,7 @@ func (svr *Service) getProxyStatsByTypeAndName(proxyType string, proxyName strin
 		proxyInfo.CurConns = ps.CurConns
 		proxyInfo.LastStartTime = ps.LastStartTime
 		proxyInfo.LastCloseTime = ps.LastCloseTime
+		code = 200
 	}
 
 	return

diff --git a/server/service.go b/server/service.go
@@ -108,6 +108,9 @@ func NewService() (svr *Service, err error) {
 		return
 	}
 
+	// Init 404 not found page
+	vhost.NotFoundPagePath = cfg.Custom404Page
+
 	var (
 		httpMuxOn  bool
 		httpsMuxOn bool

diff --git a/utils/version/version.go b/utils/version/version.go
@@ -19,7 +19,7 @@ import (
 	"strings"
 )
 
-var version string = "0.26.0"
+var version string = "0.27.0"
 
 func Full() string {
 	return version

diff --git a/utils/vhost/resource.go b/utils/vhost/resource.go
@@ -15,13 +15,18 @@
 package vhost
 
 import (
+	"bytes"
 	"io/ioutil"
 	"net/http"
-	"strings"
 
+	frpLog "github.com/fatedier/frp/utils/log"
 	"github.com/fatedier/frp/utils/version"
 )
 
+var (
+	NotFoundPagePath = ""
+)
+
 const (
 	NotFound = `<!DOCTYPE html>
 <html>
@@ -46,18 +51,36 @@ Please try again later.</p>
 `
 )
 
+func getNotFoundPageContent() []byte {
+	var (
+		buf []byte
+		err error
+	)
+	if NotFoundPagePath != "" {
+		buf, err = ioutil.ReadFile(NotFoundPagePath)
+		if err != nil {
+			frpLog.Warn("read custom 404 page error: %v", err)
+			buf = []byte(NotFound)
+		}
+	} else {
+		buf = []byte(NotFound)
+	}
+	return buf
+}
+
 func notFoundResponse() *http.Response {
 	header := make(http.Header)
 	header.Set("server", "frp/"+version.Full())
 	header.Set("Content-Type", "text/html")
+
 	res := &http.Response{
 		Status:     "Not Found",
 		StatusCode: 404,
 		Proto:      "HTTP/1.0",
 		ProtoMajor: 1,
 		ProtoMinor: 0,
 		Header:     header,
-		Body:       ioutil.NopCloser(strings.NewReader(NotFound)),
+		Body:       ioutil.NopCloser(bytes.NewReader(getNotFoundPageContent())),
 	}
 	return res
 }

diff --git a/utils/vhost/reverseproxy.go b/utils/vhost/reverseproxy.go
@@ -254,7 +254,8 @@ func (p *ReverseProxy) serveHTTP(rw http.ResponseWriter, req *http.Request) {
 	if err != nil {
 		p.logf("http: proxy error: %v", err)
 		rw.WriteHeader(http.StatusNotFound)
-		rw.Write([]byte(NotFound))
+
+		rw.Write(getNotFoundPageContent())
 		return
 	}