Skip to content

Commit

Permalink
perf: handle specified tcp relay directly
Browse files Browse the repository at this point in the history
  • Loading branch information
vyloy committed Apr 9, 2024
1 parent 1e213cc commit dec8241
Show file tree
Hide file tree
Showing 4 changed files with 57 additions and 28 deletions.
1 change: 0 additions & 1 deletion bin/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
* limitations under the License.
*/

use std::env;
use std::path::PathBuf;

use clap::Parser;
Expand Down
56 changes: 30 additions & 26 deletions libcs/client/conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -315,7 +315,7 @@ func (c *conn) readLoop(connID uint) {
if rErr != nil {
err = wErr
if !errors.Is(rErr, net.ErrClosed) {
c.Logger.Warn().Err(rErr).Msg("failed to read data in processData")
c.Logger.Warn().Err(rErr).Msg("failed to read data in processServiceData")
}
return
}
Expand All @@ -327,7 +327,7 @@ func (c *conn) readLoop(connID uint) {
}
if wErr != nil {
if !errors.Is(wErr, net.ErrClosed) {
c.Logger.Warn().Err(wErr).Msg("failed to write data in processData")
c.Logger.Warn().Err(wErr).Msg("failed to write data in processServiceData")
}
continue
}
Expand Down Expand Up @@ -387,22 +387,24 @@ func (c *conn) dial(s *service) (task *httpTask, err error) {
}

func (c *conn) processServiceData(connID uint, taskID uint32, s *service, r *bufio.LimitedReader) (readErr, writeErr error) {
var peekBytes []byte
peekBytes, readErr = r.Peek(2)
if readErr != nil {
return
}
// first 2 bytes of p2p sdp request is "XP"(0x5850)
isP2P := (uint16(peekBytes[1]) | uint16(peekBytes[0])<<8) == 0x5850
if isP2P {
if len(c.stuns) < 1 {
respAndClose(taskID, c, [][]byte{
[]byte("HTTP/1.1 403 Forbidden\r\nConnection: Closed\r\n\r\n"),
})
if r.N > 0 {
var peekBytes []byte
peekBytes, readErr = r.Peek(2)
if readErr != nil {
return
}
// first 2 bytes of p2p sdp request is "XP"(0x5850)
isP2P := (uint16(peekBytes[1]) | uint16(peekBytes[0])<<8) == 0x5850
if isP2P {
if len(c.stuns) < 1 {
respAndClose(taskID, c, [][]byte{
[]byte("HTTP/1.1 403 Forbidden\r\nConnection: Closed\r\n\r\n"),
})
return
}
c.processP2P(taskID, r)
return
}
c.processP2P(taskID, r)
return
}

var task *httpTask
Expand All @@ -429,18 +431,20 @@ func (c *conn) processServiceData(connID uint, taskID uint32, s *service, r *buf
c.tasksRWMtx.Unlock()
go task.process(connID, taskID, c)

_, err := r.WriteTo(task)
if err != nil {
switch e := err.(type) {
case *net.OpError:
switch e.Op {
case "write":
if r.N > 0 {
_, err := r.WriteTo(task)
if err != nil {
switch e := err.(type) {
case *net.OpError:
switch e.Op {
case "write":
writeErr = err
}
case *bufio.WriteErr:
writeErr = err
default:
readErr = err
}
case *bufio.WriteErr:
writeErr = err
default:
readErr = err
}
}
if task.service.LocalTimeout.Duration > 0 {
Expand Down
2 changes: 1 addition & 1 deletion libcs/server/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -448,7 +448,7 @@ func (c *client) openSpecifiedTCPPort(serviceIndex uint16, l *tcpListener, tcpPo
}()
tunnel.Logger.Info().Uint16("serviceIndex", serviceIndex).Uint16("tcpPort", tcpPort).Msg("tcp forward start")
conn.serviceIndex = serviceIndex
conn.handle(func() {
conn.handleTCP(func() {
err = c.process(conn)
if err != nil {
conn.Logger.Error().Err(err).Msg("tcp handle")
Expand Down
26 changes: 26 additions & 0 deletions libcs/server/conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,32 @@ func (c *conn) handle(handleFunc func()) {
handleFunc()
}

func (c *conn) handleTCP(handleFunc func()) {
startTime := time.Now()
reader := pool.GetReader(c.Conn)
c.Reader = reader
defer func() {
c.Close()
pool.PutReader(reader)
endTime := time.Now()
if !predef.Debug {
if e := recover(); e != nil {
c.Logger.Error().Msgf("recovered panic: %#v\n%s", e, debug.Stack())
}
}
c.Logger.Info().Dur("cost", endTime.Sub(startTime)).Msg("closed")
}()
if c.server.config.Timeout.Duration > 0 {
dl := startTime.Add(c.server.config.Timeout.Duration)
err := c.SetReadDeadline(dl)
if err != nil {
c.Logger.Debug().Err(err).Msg("handle set deadline failed")
return
}
}
handleFunc()
}

func (c *conn) handleProbe(r *bufio.Reader) {
op, err := r.ReadByte()
if err != nil {
Expand Down

0 comments on commit dec8241

Please sign in to comment.