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)