Add -d command-line flag for pre-defining variables

Compare the -D flag in Clang

Author: mcy

cmd/ascii2der/main.go

@@ -20,12 +20,70 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"strings"
 )
 
+// pairs conforms to flag.Value. Each time Set() is called, it collects another
+// k=v pair into itself.
+type pairs map[string]string
+
+func (p pairs) String() string {
+	return ""
+}
+
+func (p pairs) Set(pair string) error {
+	if pair == "" || p == nil {
+		return nil
+	}
+
+	split := strings.SplitN(pair, "=", 2)
+	if len(split) != 2 {
+		return fmt.Errorf("missing \"=\": %q", pair)
+	}
+
+	p[split[0]] = split[1]
+	return nil
+}
+
+var defines = make(map[string]string)
+var fileDefines = make(map[string]string)
+
+func init() {
+	flag.Var(pairs(defines), "d",
+		`pair of the form a=b; define("a", "b") is inserted at the start of the input`+
+			"\nmay occur multiple times")
+	flag.Var(pairs(fileDefines), "df",
+		`like -d, except the second value is interpreted as a binary file to read`+
+			"\nmay occur multiple times")
+}
+
 var inPath = flag.String("i", "", "input file to use (defaults to stdin)")
 var outPath = flag.String("o", "", "output file to use (defaults to stdout)")
 var pemType = flag.String("pem", "", "if provided, format the output as a PEM block with this type")
 
+func readAll(path string) []byte {
+	var file *os.File
+	if path == "" {
+		file = os.Stdin
+	} else {
+		var err error
+		file, err = os.Open(path)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Error opening %s: %s\n", path, err)
+			os.Exit(1)
+		}
+		defer file.Close()
+	}
+
+	buf, err := ioutil.ReadAll(file)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error reading %s: %s\n", path, err)
+		os.Exit(1)
+	}
+
+	return buf
+}
+
 func main() {
 	flag.Parse()
 
@@ -35,24 +93,25 @@ func main() {
 		os.Exit(1)
 	}
 
-	inFile := os.Stdin
-	if *inPath != "" {
-		var err error
-		inFile, err = os.Open(*inPath)
-		if err != nil {
-			fmt.Fprintf(os.Stderr, "Error opening %s: %s\n", *inPath, err)
+	inBytes := readAll(*inPath)
+
+	vars := make(map[string][]byte)
+	for k, v := range defines {
+		if _, ok := vars[k]; ok {
+			fmt.Fprintf(os.Stderr, "Error: tried to define %q with flags twice\n", k)
 			os.Exit(1)
 		}
-		defer inFile.Close()
+		vars[k] = []byte(v)
 	}
-
-	inBytes, err := ioutil.ReadAll(inFile)
-	if err != nil {
-		fmt.Fprintf(os.Stderr, "Error reading input: %s\n", err)
-		os.Exit(1)
+	for k, v := range fileDefines {
+		if _, ok := vars[k]; ok {
+			fmt.Fprintf(os.Stderr, "Error: tried to define %q with flags twice\n", k)
+			os.Exit(1)
+		}
+		vars[k] = readAll(v)
 	}
 
-	outBytes, err := asciiToDER(string(inBytes))
+	outBytes, err := asciiToDER(string(inBytes), vars)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "Syntax error: %s\n", err)
 		os.Exit(1)

cmd/ascii2der/scanner.go

@@ -594,8 +594,9 @@ func asciiToDERImpl(scanner *scanner, left *token) ([]byte, []token, error) {
 	}
 }
 
-func asciiToDER(input string) ([]byte, error) {
+func asciiToDER(input string, vars map[string][]byte) ([]byte, error) {
 	scanner := newScanner(input)
+	scanner.vars = vars
 	bytes, _, err := asciiToDERImpl(scanner, nil)
 	return bytes, err
 }

cmd/ascii2der/scanner_test.go

@@ -426,7 +426,7 @@ var asciiToDERTests = []struct {
 	{"1...2", nil, false},
 	{"true", nil, false},
 	{"false", nil, false},
-	// Mismatched curlies.
+	// Mismatched brackets.
 	{"{", nil, false},
 	{"}", nil, false},
 	{"(", nil, false},
@@ -472,7 +472,7 @@ var asciiToDERTests = []struct {
 
 func TestASCIIToDER(t *testing.T) {
 	for i, tt := range asciiToDERTests {
-		out, err := asciiToDER(tt.in)
+		out, err := asciiToDER(tt.in, make(map[string][]byte))
 		ok := err == nil
 		if !tt.ok {
 			if ok {