more e2e tests (#1845)

This commit is contained in:
fatedier
2020-09-07 14:57:23 +08:00
committed by GitHub
parent 268afb3438
commit c9fe23eb10
21 changed files with 617 additions and 51 deletions

View File

@ -479,12 +479,25 @@ func (pxy *UDPProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
// close resources releated with old workConn
pxy.Close()
var rwc io.ReadWriteCloser = conn
var err error
if pxy.limiter != nil {
rwc := frpIo.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
rwc = frpIo.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
return conn.Close()
})
conn = frpNet.WrapReadWriteCloserToConn(rwc, conn)
}
if pxy.cfg.UseEncryption {
rwc, err = frpIo.WithEncryption(rwc, []byte(pxy.clientCfg.Token))
if err != nil {
conn.Close()
xl.Error("create encryption stream error: %v", err)
return
}
}
if pxy.cfg.UseCompression {
rwc = frpIo.WithCompression(rwc)
}
conn = frpNet.WrapReadWriteCloserToConn(rwc, conn)
pxy.mu.Lock()
pxy.workConn = conn
@ -579,12 +592,25 @@ func (pxy *SUDPProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
xl := pxy.xl
xl.Info("incoming a new work connection for sudp proxy, %s", conn.RemoteAddr().String())
var rwc io.ReadWriteCloser = conn
var err error
if pxy.limiter != nil {
rwc := frpIo.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
rwc = frpIo.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
return conn.Close()
})
conn = frpNet.WrapReadWriteCloserToConn(rwc, conn)
}
if pxy.cfg.UseEncryption {
rwc, err = frpIo.WithEncryption(rwc, []byte(pxy.clientCfg.Token))
if err != nil {
conn.Close()
xl.Error("create encryption stream error: %v", err)
return
}
}
if pxy.cfg.UseCompression {
rwc = frpIo.WithCompression(rwc)
}
conn = frpNet.WrapReadWriteCloserToConn(rwc, conn)
workConn := conn
readCh := make(chan *msg.UDPPacket, 1024)

View File

@ -488,8 +488,9 @@ func (sv *SUDPVisitor) worker(workConn net.Conn) {
xl.Info("sudp worker is closed")
}
func (sv *SUDPVisitor) getNewVisitorConn() (visitorConn net.Conn, err error) {
visitorConn, err = sv.ctl.connectServer()
func (sv *SUDPVisitor) getNewVisitorConn() (net.Conn, error) {
xl := xlog.FromContextSafe(sv.ctx)
visitorConn, err := sv.ctl.connectServer()
if err != nil {
return nil, fmt.Errorf("frpc connect frps error: %v", err)
}
@ -518,7 +519,20 @@ func (sv *SUDPVisitor) getNewVisitorConn() (visitorConn net.Conn, err error) {
if newVisitorConnRespMsg.Error != "" {
return nil, fmt.Errorf("start new visitor connection error: %s", newVisitorConnRespMsg.Error)
}
return
var remote io.ReadWriteCloser
remote = visitorConn
if sv.cfg.UseEncryption {
remote, err = frpIo.WithEncryption(remote, []byte(sv.cfg.Sk))
if err != nil {
xl.Error("create encryption stream error: %v", err)
return nil, err
}
}
if sv.cfg.UseCompression {
remote = frpIo.WithCompression(remote)
}
return frpNet.WrapReadWriteCloserToConn(remote, visitorConn), nil
}
func (sv *SUDPVisitor) Close() {