From 6c9fa3f342a1e0172b154f3ea192112d2dbfb075 Mon Sep 17 00:00:00 2001 From: JackDoan Date: Fri, 16 Jan 2026 12:26:01 -0600 Subject: [PATCH] use bufio.Scanner --- cert/ca_pool.go | 63 +++++++++++++++++++-------------------- pki_hup_benchmark_test.go | 3 +- util/pem.go | 52 ++++++++++++++++++++++++++++++++ 3 files changed, 84 insertions(+), 34 deletions(-) create mode 100644 util/pem.go diff --git a/cert/ca_pool.go b/cert/ca_pool.go index 5439b06a..0d5f98e0 100644 --- a/cert/ca_pool.go +++ b/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) - if n > 0 { - 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) { + ready := scanner.Scan() + if !ready { 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 { - return nil, ErrInvalidPEMBlock + err = pool.AddCA(c) + if errors.Is(err, ErrExpired) { + expired = true + continue + } else if err != nil { + return nil, err + } } if expired { diff --git a/pki_hup_benchmark_test.go b/pki_hup_benchmark_test.go index 3a201070..39f648ff 100644 --- a/pki_hup_benchmark_test.go +++ b/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) } diff --git a/util/pem.go b/util/pem.go new file mode 100644 index 00000000..300ab62f --- /dev/null +++ b/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 +}