Selaa lähdekoodia

lib/discover, lib/protocol: Buffer allocation

greatroar 4 vuotta sitten
vanhempi
sitoutus
eafb40460d
2 muutettua tiedostoa jossa 7 lisäystä ja 8 poistoa
  1. 6 7
      lib/discover/local.go
  2. 1 1
      lib/protocol/hello.go

+ 6 - 7
lib/discover/local.go

@@ -114,19 +114,18 @@ func (c *localClient) announcementPkt(instanceID int64, msg []byte) ([]byte, boo
 		return msg, false
 	}
 
-	if cap(msg) >= 4 {
-		msg = msg[:4]
-	} else {
-		msg = make([]byte, 4)
-	}
-	binary.BigEndian.PutUint32(msg, Magic)
-
 	pkt := Announce{
 		ID:         c.myID,
 		Addresses:  addrs,
 		InstanceID: instanceID,
 	}
 	bs, _ := pkt.Marshal()
+
+	if pktLen := 4 + len(bs); cap(msg) < pktLen {
+		msg = make([]byte, 0, pktLen)
+	}
+	msg = msg[:4]
+	binary.BigEndian.PutUint32(msg, Magic)
 	msg = append(msg, bs...)
 
 	return msg, true

+ 1 - 1
lib/protocol/hello.go

@@ -90,7 +90,7 @@ func writeHello(c io.Writer, h HelloIntf) error {
 		panic("bug: attempting to serialize too large hello message")
 	}
 
-	header := make([]byte, 6)
+	header := make([]byte, 6, 6+len(msg))
 	binary.BigEndian.PutUint32(header[:4], h.Magic())
 	binary.BigEndian.PutUint16(header[4:], uint16(len(msg)))