use bufio.Scanner

cert/ca_pool.go

@@ -1,6 +1,7 @@
 package cert
 
 import (
+	"bufio"
 	"bytes"
 	"encoding/pem"
 	"errors"
@@ -9,6 +10,8 @@ import (
 	"net/netip"
 	"slices"
 	"time"
+
+	"github.com/slackhq/nebula/util"
 )
 
 type CAPool struct {
@@ -38,48 +41,42 @@ func NewCAPoolFromPEM(caPEMs []byte) (*CAPool, error) {
 // The reader must contain a PEM-encoded set of nebula certificates.
 func NewCAPoolFromPEMReader(r io.Reader) (*CAPool, error) {
 	pool := NewCAPool()
-	buf := make([]byte, 0, 64*1024)
+
-	tmp := make([]byte, 32*1024)
 	var expired bool
 
+	scanner := bufio.NewScanner(r)
+	scanner.Split(util.SplitPEM)
+
 	for {
-		n, err := r.Read(tmp)
+		ready := scanner.Scan()
-		if n > 0 {
+		if !ready {
-			buf = append(buf, tmp[:n]...)
-
-			for {
-				var block *pem.Block
-				block, buf = pem.Decode(buf)
-				if block == nil {
-					break
-				}
-
-				c, err := unmarshalCertificateBlock(block)
-				if err != nil {
-					return nil, err
-				}
-
-				err = pool.AddCA(c)
-				if errors.Is(err, ErrExpired) {
-					expired = true
-					continue
-				}
-				if err != nil {
-					return nil, err
-				}
-			}
-		}
-
-		if errors.Is(err, io.EOF) {
 			break
 		}
+		pemBytes := scanner.Bytes()
+		if scanner.Err() != nil {
+			return nil, scanner.Err()
+		}
+
+		block, rest := pem.Decode(pemBytes)
+		if len(bytes.TrimSpace(rest)) > 0 {
+			return nil, ErrInvalidPEMBlock
+		}
+		if block == nil {
+			break
+		}
+
+		c, err := unmarshalCertificateBlock(block)
 		if err != nil {
 			return nil, err
 		}
-	}
 
-	if len(bytes.TrimSpace(buf)) > 0 {
+		err = pool.AddCA(c)
-		return nil, ErrInvalidPEMBlock
+		if errors.Is(err, ErrExpired) {
+			expired = true
+			continue
+		} else if err != nil {
+			return nil, err
+		}
 	}
 
 	if expired {

pki_hup_benchmark_test.go

@@ -72,6 +72,7 @@ func buildCABundle(b *testing.B, count int) (cert.Certificate, []byte, []byte) {
 	)
 
 	buf := bytes.NewBuffer(pem)
+	buf.Write([]byte("\n# a comment!\n"))
 
 	for i := 1; i < count; i++ {
 		_, _, _, extraPEM := cert_test.NewTestCaCert(
@@ -83,7 +84,7 @@ func buildCABundle(b *testing.B, count int) (cert.Certificate, []byte, []byte) {
 			nil,
 			nil,
 		)
-
+		buf.Write([]byte("\n# a comment!\n"))
 		buf.Write(extraPEM)
 	}
 

util/pem.go

@@ -0,0 +1,52 @@
+package util
+
+import (
+	"bufio"
+	"bytes"
+)
+
+// SplitPEM is a split function for bufio.Scanner that returns each PEM block.
+func SplitPEM(data []byte, atEOF bool) (advance int, token []byte, err error) {
+	// Look for the start of a PEM block
+	start := bytes.Index(data, []byte("-----BEGIN "))
+	if start == -1 {
+		if atEOF && len(data) > 0 {
+			// No PEM block found, skip remaining data
+			return len(data), nil, nil
+		}
+		// Request more data
+		return 0, nil, nil
+	}
+
+	// Look for the end marker
+	endMarkerStart := bytes.Index(data[start:], []byte("-----END "))
+	if endMarkerStart == -1 {
+		if atEOF {
+			// Incomplete PEM block at EOF
+			return 0, nil, bufio.ErrFinalToken
+		}
+		// Need more data to find the end
+		return 0, nil, nil
+	}
+
+	// Find the actual end of the END line (after the newline)
+	endMarkerStart += start
+	endLineEnd := bytes.IndexByte(data[endMarkerStart:], '\n')
+	if endLineEnd == -1 {
+		if atEOF {
+			// END marker without newline at EOF - take it anyway
+			endLineEnd = len(data) - endMarkerStart
+		} else {
+			// Need more data
+			return 0, nil, nil
+		}
+	}
+
+	end := endMarkerStart + endLineEnd + 1
+
+	// Extract the PEM block
+	pemBlock := data[start:end]
+
+	// Return the valid PEM block
+	return end, pemBlock, nil
+}