Skip to content

Instantly share code, notes, and snippets.

@anitgandhi
Last active May 19, 2020 22:22
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save anitgandhi/58b0618512fdb3caa89e86c8a6a536ab to your computer and use it in GitHub Desktop.
Save anitgandhi/58b0618512fdb3caa89e86c8a6a536ab to your computer and use it in GitHub Desktop.
A drop-in replacement for the Go pem package that can also error if there's an invalid block - https://github.com/golang/go/issues/34069
// Package pemerr is a modified version of the stdlib pem package that adds an additional
// Decode function to support error handling.
// It is otherwise a drop-in replacement.
package pemerr
import (
"bytes"
"encoding/base64"
"encoding/pem"
"errors"
"io"
)
// A Block is an alias for pem.Block
type Block = pem.Block
var (
ErrInvalidBlock = errors.New("pem: invalid block")
)
// getLine results the first \r\n or \n delineated line from the given byte
// array. The line does not include trailing whitespace or the trailing new
// line bytes. The remainder of the byte array (also not including the new line
// bytes) is also returned and this will always be smaller than the original
// argument.
func getLine(data []byte) (line, rest []byte) {
i := bytes.IndexByte(data, '\n')
var j int
if i < 0 {
i = len(data)
j = i
} else {
j = i + 1
if i > 0 && data[i-1] == '\r' {
i--
}
}
return bytes.TrimRight(data[0:i], " \t"), data[j:]
}
// removeSpacesAndTabs returns a copy of its input with all spaces and tabs
// removed, if there were any. Otherwise, the input is returned unchanged.
//
// The base64 decoder already skips newline characters, so we don't need to
// filter them out here.
func removeSpacesAndTabs(data []byte) []byte {
if !bytes.ContainsAny(data, " \t") {
// Fast path; most base64 data within PEM contains newlines, but
// no spaces nor tabs. Skip the extra alloc and work.
return data
}
result := make([]byte, len(data))
n := 0
for _, b := range data {
if b == ' ' || b == '\t' {
continue
}
result[n] = b
n++
}
return result[0:n]
}
var pemStart = []byte("\n-----BEGIN ")
var pemEnd = []byte("\n-----END ")
var pemEndOfLine = []byte("-----")
// Decode will find the next PEM formatted block (certificate, private key
// etc) in the input. It returns that block and the remainder of the input. If
// no PEM data is found, p is nil and the whole of the input is returned in
// rest.
func Decode(data []byte) (p *Block, rest []byte) {
return decodeWithErrorHandler(data, decodeError)
}
// DecodeStrict behaves like Decode but will return an error upon encountering
// the first invalid PEM block, if any.
func DecodeStrict(data []byte) (p *Block, rest []byte, err error) {
p, rest = decodeWithErrorHandler(data, func(_, _ []byte) (*pem.Block, []byte) {
err = ErrInvalidBlock
return nil, nil
})
return
}
// decodeWithERrorHandler behaves like Decode except it will execute the callback f
// when a likely invalid PEM block is found, instead of Decode's default behavior of
// recursing through the remainder of the input.
// This provides error handling capabilities
func decodeWithErrorHandler(data []byte, f func([]byte, []byte) (*Block, []byte)) (p *Block, rest []byte) {
// pemStart begins with a newline. However, at the very beginning of
// the byte array, we'll accept the start string without it.
rest = data
if bytes.HasPrefix(data, pemStart[1:]) {
rest = rest[len(pemStart)-1 : len(data)]
} else if i := bytes.Index(data, pemStart); i >= 0 {
rest = rest[i+len(pemStart) : len(data)]
} else {
return nil, data
}
typeLine, rest := getLine(rest)
if !bytes.HasSuffix(typeLine, pemEndOfLine) {
return f(data, rest)
}
typeLine = typeLine[0 : len(typeLine)-len(pemEndOfLine)]
p = &Block{
Headers: make(map[string]string),
Type: string(typeLine),
}
for {
// This loop terminates because getLine's second result is
// always smaller than its argument.
if len(rest) == 0 {
return nil, data
}
line, next := getLine(rest)
i := bytes.IndexByte(line, ':')
if i == -1 {
break
}
// TODO(agl): need to cope with values that spread across lines.
key, val := line[:i], line[i+1:]
key = bytes.TrimSpace(key)
val = bytes.TrimSpace(val)
p.Headers[string(key)] = string(val)
rest = next
}
var endIndex, endTrailerIndex int
// If there were no headers, the END line might occur
// immediately, without a leading newline.
if len(p.Headers) == 0 && bytes.HasPrefix(rest, pemEnd[1:]) {
endIndex = 0
endTrailerIndex = len(pemEnd) - 1
} else {
endIndex = bytes.Index(rest, pemEnd)
endTrailerIndex = endIndex + len(pemEnd)
}
if endIndex < 0 {
return f(data, rest)
}
// After the "-----" of the ending line, there should be the same type
// and then a final five dashes.
endTrailer := rest[endTrailerIndex:]
endTrailerLen := len(typeLine) + len(pemEndOfLine)
if len(endTrailer) < endTrailerLen {
return f(data, rest)
}
restOfEndLine := endTrailer[endTrailerLen:]
endTrailer = endTrailer[:endTrailerLen]
if !bytes.HasPrefix(endTrailer, typeLine) ||
!bytes.HasSuffix(endTrailer, pemEndOfLine) {
return f(data, rest)
}
// The line must end with only whitespace.
if s, _ := getLine(restOfEndLine); len(s) != 0 {
return f(data, rest)
}
base64Data := removeSpacesAndTabs(rest[:endIndex])
p.Bytes = make([]byte, base64.StdEncoding.DecodedLen(len(base64Data)))
n, err := base64.StdEncoding.Decode(p.Bytes, base64Data)
if err != nil {
return f(data, rest)
}
p.Bytes = p.Bytes[:n]
// the -1 is because we might have only matched pemEnd without the
// leading newline if the PEM block was empty.
_, rest = getLine(rest[endIndex+len(pemEnd)-1:])
return
}
func decodeError(data, rest []byte) (*Block, []byte) {
// If we get here then we have rejected a likely looking, but
// ultimately invalid PEM block. We need to start over from a new
// position. We have consumed the preamble line and will have consumed
// any lines which could be header lines. However, a valid preamble
// line is not a valid header line, therefore we cannot have consumed
// the preamble line for the any subsequent block. Thus, we will always
// find any valid block, no matter what bytes precede it.
//
// For example, if the input is
//
// -----BEGIN MALFORMED BLOCK-----
// junk that may look like header lines
// or data lines, but no END line
//
// -----BEGIN ACTUAL BLOCK-----
// realdata
// -----END ACTUAL BLOCK-----
//
// we've failed to parse using the first BEGIN line
// and now will try again, using the second BEGIN line.
p, rest := Decode(rest)
if p == nil {
rest = data
}
return p, rest
}
// Encode is a warpper for pem.Encode
func Encode(out io.Writer, b *Block) error {
return pem.Encode(out, b)
}
// EncodeToMemory is a wrapper for pem.EncodeToMemory
func EncodeToMemory(b *Block) []byte {
return pem.EncodeToMemory(b)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment