Ver código fonte

support encryption and compression in new http reverser proxy

fatedier 7 anos atrás
pai
commit
46358d466d

+ 1 - 1
models/plugin/http_proxy.go

@@ -111,7 +111,7 @@ func (hp *HttpProxy) Handle(conn io.ReadWriteCloser) {
 	if realConn, ok := conn.(frpNet.Conn); ok {
 		wrapConn = realConn
 	} else {
-		wrapConn = frpNet.WrapReadWriteCloserToConn(conn)
+		wrapConn = frpNet.WrapReadWriteCloserToConn(conn, realConn)
 	}
 
 	sc, rd := frpNet.NewShareConn(wrapConn)

+ 1 - 1
models/plugin/socks5.go

@@ -50,7 +50,7 @@ func (sp *Socks5Plugin) Handle(conn io.ReadWriteCloser) {
 	if realConn, ok := conn.(frpNet.Conn); ok {
 		wrapConn = realConn
 	} else {
-		wrapConn = frpNet.WrapReadWriteCloserToConn(conn)
+		wrapConn = frpNet.WrapReadWriteCloserToConn(conn, realConn)
 	}
 
 	sp.Server.ServeConn(wrapConn)

+ 1 - 1
server/manager.go

@@ -146,7 +146,7 @@ func (vm *VisitorManager) NewConn(name string, conn frpNet.Conn, timestamp int64
 		if useCompression {
 			rwc = frpIo.WithCompression(rwc)
 		}
-		err = l.PutConn(frpNet.WrapReadWriteCloserToConn(rwc))
+		err = l.PutConn(frpNet.WrapReadWriteCloserToConn(rwc, conn))
 	} else {
 		err = fmt.Errorf("custom listener for [%s] doesn't exist", name)
 		return

+ 24 - 2
server/proxy.go

@@ -208,7 +208,7 @@ func (pxy *HttpProxy) Run() (err error) {
 		routeConfig.Domain = domain
 		for _, location := range locations {
 			routeConfig.Location = location
-			err := pxy.ctl.svr.httpReverseProxy.Register(routeConfig.Domain, routeConfig.Location, routeConfig.RewriteHost, pxy.GetWorkConnFromPool)
+			err := pxy.ctl.svr.httpReverseProxy.Register(routeConfig.Domain, routeConfig.Location, routeConfig.RewriteHost, pxy.GetRealConn)
 			if err != nil {
 				return err
 			}
@@ -225,7 +225,7 @@ func (pxy *HttpProxy) Run() (err error) {
 		routeConfig.Domain = pxy.cfg.SubDomain + "." + config.ServerCommonCfg.SubDomainHost
 		for _, location := range locations {
 			routeConfig.Location = location
-			err := pxy.ctl.svr.httpReverseProxy.Register(routeConfig.Domain, routeConfig.Location, routeConfig.RewriteHost, pxy.GetWorkConnFromPool)
+			err := pxy.ctl.svr.httpReverseProxy.Register(routeConfig.Domain, routeConfig.Location, routeConfig.RewriteHost, pxy.GetRealConn)
 			if err != nil {
 				return err
 			}
@@ -244,6 +244,28 @@ func (pxy *HttpProxy) GetConf() config.ProxyConf {
 	return pxy.cfg
 }
 
+func (pxy *HttpProxy) GetRealConn() (workConn frpNet.Conn, err error) {
+	tmpConn, errRet := pxy.GetWorkConnFromPool()
+	if errRet != nil {
+		err = errRet
+		return
+	}
+
+	var rwc io.ReadWriteCloser = tmpConn
+	if pxy.cfg.UseEncryption {
+		rwc, err = frpIo.WithEncryption(rwc, []byte(config.ServerCommonCfg.PrivilegeToken))
+		if err != nil {
+			pxy.Error("create encryption stream error: %v", err)
+			return
+		}
+	}
+	if pxy.cfg.UseCompression {
+		rwc = frpIo.WithCompression(rwc)
+	}
+	workConn = frpNet.WrapReadWriteCloserToConn(rwc, tmpConn)
+	return
+}
+
 func (pxy *HttpProxy) Close() {
 	pxy.BaseProxy.Close()
 	for _, closeFn := range pxy.closeFuncs {

+ 19 - 1
utils/net/conn.go

@@ -49,32 +49,50 @@ func WrapConn(c net.Conn) Conn {
 type WrapReadWriteCloserConn struct {
 	io.ReadWriteCloser
 	log.Logger
+
+	underConn net.Conn
 }
 
-func WrapReadWriteCloserToConn(rwc io.ReadWriteCloser) Conn {
+func WrapReadWriteCloserToConn(rwc io.ReadWriteCloser, underConn net.Conn) Conn {
 	return &WrapReadWriteCloserConn{
 		ReadWriteCloser: rwc,
 		Logger:          log.NewPrefixLogger(""),
+		underConn:       underConn,
 	}
 }
 
 func (conn *WrapReadWriteCloserConn) LocalAddr() net.Addr {
+	if conn.underConn != nil {
+		return conn.underConn.LocalAddr()
+	}
 	return (*net.TCPAddr)(nil)
 }
 
 func (conn *WrapReadWriteCloserConn) RemoteAddr() net.Addr {
+	if conn.underConn != nil {
+		return conn.underConn.RemoteAddr()
+	}
 	return (*net.TCPAddr)(nil)
 }
 
 func (conn *WrapReadWriteCloserConn) SetDeadline(t time.Time) error {
+	if conn.underConn != nil {
+		return conn.underConn.SetDeadline(t)
+	}
 	return &net.OpError{Op: "set", Net: "wrap", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
 }
 
 func (conn *WrapReadWriteCloserConn) SetReadDeadline(t time.Time) error {
+	if conn.underConn != nil {
+		return conn.underConn.SetReadDeadline(t)
+	}
 	return &net.OpError{Op: "set", Net: "wrap", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
 }
 
 func (conn *WrapReadWriteCloserConn) SetWriteDeadline(t time.Time) error {
+	if conn.underConn != nil {
+		return conn.underConn.SetWriteDeadline(t)
+	}
 	return &net.OpError{Op: "set", Net: "wrap", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
 }
 

+ 3 - 2
utils/vhost/newhttp.go

@@ -15,6 +15,7 @@
 package vhost
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"log"
@@ -74,8 +75,8 @@ func NewHttpReverseProxy() *HttpReverseProxy {
 			host = rp.GetRealHost(host, url)
 			if host != "" {
 				req.Host = host
-				req.URL.Host = req.Host
 			}
+			req.URL.Host = req.Host
 		},
 		Transport: &http.Transport{
 			ResponseHeaderTimeout: responseHeaderTimeout,
@@ -172,6 +173,6 @@ type wrapLogger struct{}
 func newWrapLogger() *wrapLogger { return &wrapLogger{} }
 
 func (l *wrapLogger) Write(p []byte) (n int, err error) {
-	frpLog.Warn("%s", string(p))
+	frpLog.Warn("%s", string(bytes.TrimRight(p, "\n")))
 	return len(p), nil
 }