extern/netbird
1
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2025-06-21 18:22:37 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix buffer handling

Browse Source
This commit is contained in:
Zoltan Papp 2024-05-17 23:29:47 +02:00
parent 49dfbc82d9
commit 6ae7a790f2
1 changed files with 12 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
relay/client/client.go
View File

@ -149,13 +149,13 @@ func (c *Client) readLoop() {
c.msgPool.Put(buf) c.msgPool.Put(buf)
continue continue
case messages.MsgTypeTransport: case messages.MsgTypeTransport:
channelId, payload, err := messages.UnmarshalTransportMsg(buf[:n]) channelId, err := messages.UnmarshalTransportID(buf[:n])
if err != nil { if err != nil {
log.Errorf("failed to parse transport message: %v", err) log.Errorf("failed to parse transport message: %v", err)
c.msgPool.Put(buf) c.msgPool.Put(buf)
continue continue
} }
c.handleTransport(channelId, payload) c.handleTransport(channelId, buf[:n])
} }
} }
@ -185,16 +185,15 @@ func (c *Client) handleBindResponse(channelId uint16, peerId string) {
bindSuccessChan <- conn bindSuccessChan <- conn
} }
func (c *Client) handleTransport(channelId uint16, payload []byte) { func (c *Client) handleTransport(channelId uint16, msg []byte) {
container, ok := c.channels[channelId] container, ok := c.channels[channelId]
if !ok { if !ok {
log.Errorf("c.channels: %v", c.peerID)
log.Errorf("unexpected transport message for channel: %d", channelId) log.Errorf("unexpected transport message for channel: %d", channelId)
return return
} }
select { select {
case container.messages <- payload: case container.messages <- msg:
default: default:
log.Errorf("dropping message for channel: %d", channelId) log.Errorf("dropping message for channel: %d", channelId)
} }
@ -211,13 +210,19 @@ func (c *Client) writeTo(channelID uint16, payload []byte) (int, error) {
func (c *Client) generateConnReaderFN(messageBufferChan chan []byte) func(b []byte) (n int, err error) { func (c *Client) generateConnReaderFN(messageBufferChan chan []byte) func(b []byte) (n int, err error) {
return func(b []byte) (n int, err error) { return func(b []byte) (n int, err error) {
defer c.msgPool.Put(b)
select { select {
case msg, ok := <-messageBufferChan: case msg, ok := <-messageBufferChan:
if !ok { if !ok {
return 0, io.EOF return 0, io.EOF
} }
n = copy(b, msg)
c.msgPool.Put(msg) payload, err := messages.UnmarshalTransportPayload(msg)
if err != nil {
return 0, err
}
n = copy(b, payload)
} }
return n, nil return n, nil
} }