From 670d4a514cb13aae2d918e650efc6c69964cd82b Mon Sep 17 00:00:00 2001 From: CPunch Date: Wed, 23 Aug 2023 17:38:10 -0500 Subject: [PATCH] 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 --- core/protocol/cnpeer.go | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/core/protocol/cnpeer.go b/core/protocol/cnpeer.go index 83d00f0..e75e2ff 100644 --- a/core/protocol/cnpeer.go +++ b/core/protocol/cnpeer.go @@ -51,6 +51,9 @@ func (peer *CNPeer) Send(typeID uint32, data ...interface{}) error { buf := pool.Get() defer pool.Put(buf) + // allocate space for packet size + buf.Write(make([]byte, 4)) + // body start 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 switch peer.whichKey { case USE_E: - EncryptData(buf.Bytes(), peer.E_key) + EncryptData(buf.Bytes()[4:], peer.E_key) case USE_FE: - EncryptData(buf.Bytes(), peer.FE_key) + EncryptData(buf.Bytes()[4:], peer.FE_key) } - // write packet size - if err := binary.Write(peer.conn, binary.LittleEndian, uint32(buf.Len())); err != nil { - return err - } - - // write packet body - log.Printf("Sending %#v, sizeof: %d", data, buf.Len()) + // send full packet + log.Printf("Sending %#v, sizeof: %d, buffer: %v", data, buf.Len(), buf.Bytes()) 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 }