fix: avoid accidental OOM due to bogus allocation

Author: devgianlu

ap/ap.go

@@ -425,7 +425,7 @@ func (ap *Accesspoint) performKeyExchange() ([]byte, error) {
 
 	// receive APResponseMessage message
 	var apResponse pb.APResponseMessage
-	if err := readMessage(cc, &apResponse); err != nil {
+	if err := readMessage(cc, -1, &apResponse); err != nil {
 		return nil, fmt.Errorf("failed reading APResponseMessage message: %w", err)
 	}
 
@@ -502,7 +502,7 @@ func (ap *Accesspoint) authenticate(ctx context.Context, credentials *pb.LoginCr
 	var challengeResp pb.APResponseMessage
 	if peekBytes, err := ap.encConn.peekUnencrypted(9); err != nil {
 		return fmt.Errorf("failed peeking unencrypted bytes: %w", err)
-	} else if err = readMessage(bytes.NewReader(peekBytes), &challengeResp); err == nil {
+	} else if err = readMessage(bytes.NewReader(peekBytes), 9, &challengeResp); err == nil {
 		return &AccesspointLoginError{Message: challengeResp.LoginFailed}
 	}
 

ap/conn.go

@@ -37,13 +37,18 @@ func writeMessage(w io.Writer, withHello bool, m proto.Message) error {
 	return nil
 }
 
-func readMessage(r io.Reader, m proto.Message) error {
+func readMessage(r io.Reader, maxLength int, m proto.Message) error {
 	// read length
 	var length uint32
 	if err := binary.Read(r, binary.BigEndian, &length); err != nil {
 		return fmt.Errorf("failed reading message length: %w", err)
 	}
 
+	// check length to avoid a mega allocation
+	if maxLength > 0 && length > uint32(maxLength) {
+		return fmt.Errorf("message too long: %d", length)
+	}
+
 	// read message
 	data := make([]byte, length-4)
 	if _, err := io.ReadFull(r, data); err != nil {