suyono/netbounce
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor the error handling and optimize the heap allocation

Browse Source
This commit is contained in:
Suyono
2025-05-04 13:32:31 +10:00
parent 4abc00f07c
commit 0c1a9fb7e7
6 changed files with 242 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
backend_udp.go
View File

@@ -1,7 +1,6 @@
package netbounce
import (
"bytes"
"context"
"errors"
"fmt"
@@ -26,13 +25,13 @@ type backendUDP struct {
relMtx *sync.Mutex
client net.PacketConn
cfg abstract.UDPConnectionConfig
bufPool sync.Pool
bufPool *BufPool
msgChan chan udpMessage
}
type udpMessage struct {
addr string
buf *bytes.Buffer
buf []byte
}
func (r udpRel) Network() string {
@@ -43,18 +42,14 @@ func (r udpRel) String() string {
return r.clientAddr
}
func initUDP(wg *sync.WaitGroup, ctx context.Context, cfg abstract.UDPConnectionConfig, client net.PacketConn) *backendUDP {
func initUDP(wg *sync.WaitGroup, ctx context.Context, bp *BufPool, cfg abstract.UDPConnectionConfig, client net.PacketConn) *backendUDP {
backend := &backendUDP{
relations: make(map[string]udpRel),
relMtx: new(sync.Mutex),
cfg: cfg,
client: client,
msgChan: make(chan udpMessage),
bufPool: sync.Pool{
New: func() any {
return new(bytes.Buffer)
},
},
bufPool: bp,
}
go func(wg *sync.WaitGroup, ctx context.Context, cfg abstract.UDPConnectionConfig, backend *backendUDP) {
@@ -104,7 +99,6 @@ func (b *backendUDP) createRelSend(clientAddr string, buf []byte) (udpRel, error
var (
udpAddr *net.UDPAddr
udpConn *net.UDPConn
n int
err error
)
@@ -120,7 +114,8 @@ func (b *backendUDP) createRelSend(clientAddr string, buf []byte) (udpRel, error
return rel, fmt.Errorf("create udp relation and send message: dial udp: %w", err)
}
if n, err = udpConn.WriteTo(buf, b.cfg.BackendAddr()); err != nil && n == 0 {
if _, err = udpConn.WriteTo(buf, b.cfg.BackendAddr()); err != nil {
//TODO: I think I need to fix this. This error handling is odd.
_ = udpConn.Close()
return rel, fmt.Errorf("create udp relation and send message: write udp: %w", err)
}
@@ -156,14 +151,12 @@ func (b *backendUDP) handle(wg *sync.WaitGroup, ctx context.Context, msg udpMess
)
if rel, ok = b.findRel(msg.addr); !ok {
if rel, err = b.createRelSend(msg.addr, msg.buf.Bytes()); err != nil {
msg.buf.Reset()
if rel, err = b.createRelSend(msg.addr, msg.buf); err != nil {
b.bufPool.Put(msg.buf)
log.Error().Err(err).Str(DIRECTION, CLIENT_TO_BACKEND).Msg("establish relation with udp backend")
return
}
msg.buf.Reset()
b.bufPool.Put(msg.buf)
rel.ctx, rel.ctxCancel = context.WithCancel(ctx)
@@ -174,14 +167,12 @@ func (b *backendUDP) handle(wg *sync.WaitGroup, ctx context.Context, msg udpMess
return
}
if err = b.relSend(rel, msg.buf.Bytes()); err != nil {
msg.buf.Reset()
if err = b.relSend(rel, msg.buf); err != nil {
b.bufPool.Put(msg.buf)
log.Error().Err(err).Msg("handle: send for existing relation")
return
}
msg.buf.Reset()
b.bufPool.Put(msg.buf)
rel.expiry = time.Now().Add(b.cfg.Timeout())
@@ -190,18 +181,13 @@ func (b *backendUDP) handle(wg *sync.WaitGroup, ctx context.Context, msg udpMess
func (b *backendUDP) Send(ctx context.Context, addr string, p []byte) error {
var (
n int
err error
n int
)
buf := b.bufPool.Get().(*bytes.Buffer)
if n, err = buf.Write(p); err != nil {
buf.Reset()
b.bufPool.Put(buf)
return fmt.Errorf("send udp message to handler: %w", err)
}
buf := b.bufPool.Get()
n = copy(buf, p)
if len(p) != n {
buf.Reset()
b.bufPool.Put(buf)
return fmt.Errorf("send udp message to handler: failed to write complete message")
}
@@ -209,7 +195,7 @@ func (b *backendUDP) Send(ctx context.Context, addr string, p []byte) error {
select {
case b.msgChan <- udpMessage{
addr: addr,
buf: buf,
buf: buf[:n],
}:
case <-ctx.Done():
}