Last active
May 19, 2020 22:22
-
-
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
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
// 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