hoist PEM reading function

Signed-off-by: eternal-flame-AD <[email protected]>

Author: eternal-flame-AD

v2/shim_v1.go

@@ -10,7 +10,6 @@ import (
 	"encoding/json"
 	"encoding/pem"
 	"errors"
-	"io"
 	"log"
 	"math"
 	"net/http"
@@ -120,35 +119,28 @@ func NewCompatV1Rpc(compatV1 *CompatV1, cliArgs []string) (*CompatV1Shim, error)
 		return nil, err
 	}
 
-	var certBytes []byte
 	var certificateChain []tls.Certificate
-	for {
-		var buf [2048]byte
-		n, err := cliFlags.KexRespFile.Read(buf[:])
-		if err != nil {
-			if err == io.EOF {
-				break
-			}
-			return nil, err
-		}
-		certBytes = append(certBytes, buf[:n]...)
 
-		for block, rest := pem.Decode(certBytes); block != nil; block, rest = pem.Decode(rest) {
-			if block.Type == "CERTIFICATE" {
-				parsedCert, err := x509.ParseCertificate(block.Bytes)
-				if err != nil {
-					return nil, err
-				}
-				// Server signs with IsCA=false, so we can add all of them to the root CA pool without
-				// trusting things we shouldn't.
-				rootCAs.AddCert(parsedCert)
-				certificateChain = append(certificateChain, tls.Certificate{
-					Certificate: [][]byte{block.Bytes},
-					Leaf:        parsedCert,
-				})
+	if err := transport.IteratePEMFile(cliFlags.KexRespFile, func(block *pem.Block) (continueIterate bool, err error) {
+		if block.Type == "CERTIFICATE" {
+			parsedCert, err := x509.ParseCertificate(block.Bytes)
+			if err != nil {
+				return false, err
 			}
-			certBytes = rest
+			rootCAs.AddCert(parsedCert)
+			certificateChain = append(certificateChain, tls.Certificate{
+				Certificate: [][]byte{block.Bytes},
+				Leaf:        parsedCert,
+			})
+			return true, nil
 		}
+		return true, nil
+	}); err != nil {
+		return nil, err
+	}
+
+	if len(certificateChain) == 0 {
+		return nil, errors.New("no certificate chain found in kex response file")
 	}
 
 	certificateChain[0].PrivateKey = priv

v2/transport/pem_file.go

@@ -0,0 +1,32 @@
+package transport
+
+import (
+	"encoding/pem"
+	"io"
+)
+
+func IteratePEMFile(r io.Reader, callback func(block *pem.Block) (continueIterate bool, err error)) error {
+	var bufferBytes []byte
+	for {
+		var buf [2048]byte
+		n, err := r.Read(buf[:])
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+			return err
+		}
+		bufferBytes = append(bufferBytes, buf[:n]...)
+
+		for block, rest := pem.Decode(bufferBytes); block != nil; block, rest = pem.Decode(rest) {
+			continueIterate, err := callback(block)
+			if err != nil {
+				return err
+			}
+			if !continueIterate {
+				return nil
+			}
+		}
+	}
+	return nil
+}

v2/transport/transport_auth.go

@@ -8,6 +8,7 @@ import (
 	"crypto/x509/pkix"
 	"encoding/hex"
 	"encoding/pem"
+	"errors"
 	"fmt"
 	"io"
 	"slices"
@@ -102,27 +103,25 @@ func (s *EphemeralTLSClient) SignPluginCSR(moduleName string, csr *x509.Certific
 
 func (s *EphemeralTLSClient) Kex(req io.Reader, resp io.Writer) error {
 	var csr *x509.CertificateRequest
-	var csrBytes []byte
-	for csr == nil {
-		var buf [2048]byte
-		n, err := req.Read(buf[:])
-		if err != nil {
-			return err
-		}
-		csrBytes = append(csrBytes, buf[:n]...)
-		block, _ := pem.Decode(csrBytes)
-		if block == nil {
-			continue
-		}
 
+	if err := IteratePEMFile(req, func(block *pem.Block) (continueIterate bool, err error) {
 		if block.Type == "CERTIFICATE REQUEST" {
-			csrParsed, err := x509.ParseCertificateRequest(block.Bytes)
+			csr, err = x509.ParseCertificateRequest(block.Bytes)
 			if err != nil {
-				return err
+				return false, err
 			}
-			csr = csrParsed
+
+			return false, nil
 		}
+		return true, nil
+	}); err != nil {
+		return err
+	}
+
+	if csr == nil {
+		return errors.New("no certificate request found in kex request file")
 	}
+
 	dnsName := csr.Subject.CommonName
 	certBytes, err := s.SignCSR(dnsName, csr)
 	if err != nil {
@@ -165,6 +164,9 @@ func NewEphemeralTLSClient() (*EphemeralTLSClient, error) {
 		IsCA: true,
 	}
 	caCertBytes, err := x509.CreateCertificate(rand.Reader, caCertTemplate, caCertTemplate, caPub, caPriv)
+	if err != nil {
+		return nil, err
+	}
 	caCert, err := x509.ParseCertificate(caCertBytes)
 	if err != nil {
 		return nil, err
@@ -203,7 +205,6 @@ func NewEphemeralTLSClient() (*EphemeralTLSClient, error) {
 			},
 			{
 				Certificate: [][]byte{caCertBytes},
-				PrivateKey:  caPriv,
 			},
 		},
 		RootCAs: certPool,