WIP: backend udp
This commit is contained in:
168
backend_udp.go
Normal file
168
backend_udp.go
Normal file
@@ -0,0 +1,168 @@
|
||||
package netbounce
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.suyono.dev/suyono/netbounce/abstract"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
type udpRel struct {
|
||||
backend net.PacketConn
|
||||
clientAddr string
|
||||
expiry time.Time
|
||||
}
|
||||
|
||||
type backendUDP struct {
|
||||
relations map[string]udpRel
|
||||
relMtx *sync.Mutex
|
||||
client net.PacketConn
|
||||
cfg abstract.UDPConnectionConfig
|
||||
bufPool sync.Pool
|
||||
msgChan chan udpMessage
|
||||
}
|
||||
|
||||
type udpMessage struct {
|
||||
addr string
|
||||
buf *bytes.Buffer
|
||||
}
|
||||
|
||||
func initUDP(wg *sync.WaitGroup, ctx context.Context, 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)
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
go func(wg *sync.WaitGroup, ctx context.Context, cfg abstract.UDPConnectionConfig, backend *backendUDP) {
|
||||
defer wg.Done()
|
||||
|
||||
incoming := backend.msgChan
|
||||
readIncoming:
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
break readIncoming
|
||||
case m := <-incoming:
|
||||
backend.handle(m)
|
||||
}
|
||||
}
|
||||
}(wg, ctx, cfg, backend)
|
||||
|
||||
return backend
|
||||
}
|
||||
|
||||
func (b *backendUDP) findRel(clientAddr string) (udpRel, bool) {
|
||||
b.relMtx.Lock()
|
||||
defer b.relMtx.Unlock()
|
||||
|
||||
rel, ok := b.relations[clientAddr]
|
||||
if ok && rel.expiry.Before(time.Now()) {
|
||||
delete(b.relations, clientAddr)
|
||||
return rel, false
|
||||
}
|
||||
|
||||
return rel, ok
|
||||
}
|
||||
|
||||
func (b *backendUDP) addRel(clientAddr string, rel udpRel) {
|
||||
b.relMtx.Lock()
|
||||
defer b.relMtx.Unlock()
|
||||
|
||||
b.relations[clientAddr] = rel
|
||||
}
|
||||
|
||||
func (b *backendUDP) createRelSend(clientAddr string, buf []byte) (udpRel, error) {
|
||||
var (
|
||||
udpAddr *net.UDPAddr
|
||||
udpConn *net.UDPConn
|
||||
n int
|
||||
err error
|
||||
)
|
||||
|
||||
rel := udpRel{
|
||||
clientAddr: clientAddr,
|
||||
}
|
||||
|
||||
if udpAddr, err = net.ResolveUDPAddr("udp", b.cfg.Backend()); err != nil {
|
||||
return rel, fmt.Errorf("create udp relation and send message: resolve udp addr: %w", err)
|
||||
}
|
||||
|
||||
if udpConn, err = net.DialUDP("udp", nil, udpAddr); err != nil {
|
||||
return rel, fmt.Errorf("create udp relation and send message: dial udp: %w", err)
|
||||
}
|
||||
|
||||
if n, err = udpConn.Write(buf); err != nil && n == 0 {
|
||||
_ = udpConn.Close()
|
||||
return rel, fmt.Errorf("create udp relation and send message: write udp: %w", err)
|
||||
}
|
||||
|
||||
rel.backend = udpConn
|
||||
rel.expiry = time.Now().Add(b.cfg.Timeout())
|
||||
|
||||
return rel, nil
|
||||
}
|
||||
|
||||
func (b *backendUDP) handle(msg udpMessage) {
|
||||
var (
|
||||
rel udpRel
|
||||
ok bool
|
||||
err error
|
||||
)
|
||||
|
||||
if rel, ok = b.findRel(msg.addr); !ok {
|
||||
if rel, err = b.createRelSend(msg.addr, msg.buf.Bytes()); err != nil {
|
||||
log.Error().Err(err).Str(DIRECTION, CLIENT_TO_BACKEND).Msg("establish relation with udp backend")
|
||||
}
|
||||
|
||||
b.addRel(msg.addr, rel)
|
||||
return
|
||||
}
|
||||
|
||||
_ = rel
|
||||
// if rel.expiry.Before(time.Now()) {
|
||||
// //TODO: handle expiry
|
||||
// }
|
||||
|
||||
}
|
||||
|
||||
func (b *backendUDP) Send(ctx context.Context, addr string, p []byte) error {
|
||||
var (
|
||||
n int
|
||||
err error
|
||||
)
|
||||
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)
|
||||
}
|
||||
|
||||
if len(p) != n {
|
||||
buf.Reset()
|
||||
b.bufPool.Put(buf)
|
||||
return fmt.Errorf("send udp message to handler: failed to write complete message")
|
||||
}
|
||||
|
||||
select {
|
||||
case b.msgChan <- udpMessage{
|
||||
addr: addr,
|
||||
buf: buf,
|
||||
}:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user