more better CNPeer.Send()

this fixes a race condition where if 2 goroutines try to send a packet at the same time, they could end up being
malformed due to the 2 separate calls to peer.conn.Write().

instead of writing the packet size to peer.conn.Write() directly, we make space in buf for the packet size,
and patch it in place. this lets us get away with only having 1 call to peer.conn.Write() which will ensure that
the full packet is written properly and be goroutine safe :3
This commit is contained in:
CPunch 2023-08-23 17:38:10 -05:00
parent 7ebe80d6e3
commit 670d4a514c

View File

@ -51,6 +51,9 @@ func (peer *CNPeer) Send(typeID uint32, data ...interface{}) error {
buf := pool.Get() buf := pool.Get()
defer pool.Put(buf) defer pool.Put(buf)
// allocate space for packet size
buf.Write(make([]byte, 4))
// body start // body start
pkt := NewPacket(buf) pkt := NewPacket(buf)
@ -66,25 +69,22 @@ func (peer *CNPeer) Send(typeID uint32, data ...interface{}) error {
} }
} }
// prepend the packet size
binary.LittleEndian.PutUint32(buf.Bytes()[:4], uint32(buf.Len()-4))
// encrypt body // encrypt body
switch peer.whichKey { switch peer.whichKey {
case USE_E: case USE_E:
EncryptData(buf.Bytes(), peer.E_key) EncryptData(buf.Bytes()[4:], peer.E_key)
case USE_FE: case USE_FE:
EncryptData(buf.Bytes(), peer.FE_key) EncryptData(buf.Bytes()[4:], peer.FE_key)
} }
// write packet size // send full packet
if err := binary.Write(peer.conn, binary.LittleEndian, uint32(buf.Len())); err != nil { log.Printf("Sending %#v, sizeof: %d, buffer: %v", data, buf.Len(), buf.Bytes())
return err
}
// write packet body
log.Printf("Sending %#v, sizeof: %d", data, buf.Len())
if _, err := peer.conn.Write(buf.Bytes()); err != nil { if _, err := peer.conn.Write(buf.Bytes()); err != nil {
return fmt.Errorf("[FATAL] failed to write packet body! %v", err) return fmt.Errorf("failed to write packet body! %v", err)
} }
return nil return nil
} }