simplify and tidy up the string obfuscation code

Mainly removing unnecessary indentation and newlines, but also other
minor things like making error handling a bit more consistent.

crypto.go

@@ -9,26 +9,29 @@ import (
 // If math/rand.Seed() is not called, the generator behaves as if seeded by rand.Seed(1),
 // so the generator is deterministic.
 
-// genAesKey generates a 128bit AES Key
+// genAesKey generates a 128-bit AES Key.
 func genAesKey() []byte {
 	return genRandBytes(16)
 }
 
-// genAesKey generates a 128bit nonce
+// genAesKey generates a 128-bit nonce.
 func genNonce() []byte {
 	return genRandBytes(12)
 }
 
-// genRandBytes return a random []byte with the length of size
+// genRandBytes return a random []byte with the length of size.
 func genRandBytes(size int) []byte {
 	buffer := make([]byte, size)
 	rand.Read(buffer) // error is always nil so save to ignore
 	return buffer
 }
 
-// encAes encrypt data with AesKey in AES gcm mode
+// encAES encrypt data with key in AES GCM mode.
-func encAes(data []byte, AesKey []byte) ([]byte, error) {
+func encAES(data, key []byte) ([]byte, error) {
-	block, _ := aes.NewCipher(AesKey)
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
 	nonce := genNonce()
 
 	aesgcm, err := cipher.NewGCM(block)

main.go

@@ -192,14 +192,14 @@ func mainErr(args []string) error {
 			}
 
 			if info, err := os.Stat(flagDebugDir); os.IsNotExist(err) {
-				err := os.MkdirAll(flagDebugDir, 0755)
+				err := os.MkdirAll(flagDebugDir, 0o755)
 				if err != nil {
 					return err
 				}
 			} else if err != nil {
-				return fmt.Errorf("Debugdir error: %v", err)
+				return fmt.Errorf("debugdir error: %v", err)
 			} else if !info.IsDir() {
-				return fmt.Errorf("Debugdir exists, but is a file not a directory")
+				return fmt.Errorf("debugdir exists, but is a file not a directory")
 			}
 		}
 
@@ -371,7 +371,7 @@ func transformCompile(args []string) ([]string, error) {
 	if envGarbleDebugDir != "" {
 		osPkgPath := filepath.FromSlash(pkgPath)
 		pkgDebugDir = filepath.Join(envGarbleDebugDir, osPkgPath)
-		err = os.MkdirAll(pkgDebugDir, 0755)
+		err = os.MkdirAll(pkgDebugDir, 0o755)
 		if err != nil {
 			return nil, err
 		}

main_test.go

@@ -54,7 +54,7 @@ func TestScripts(t *testing.T) {
 				"GONOSUMDB=*",
 			)
 			bindir := filepath.Join(env.WorkDir, ".bin")
-			if err := os.Mkdir(bindir, 0777); err != nil {
+			if err := os.Mkdir(bindir, 0o777); err != nil {
 				return err
 			}
 			binfile := filepath.Join(bindir, "garble")

strings.go

@@ -2,9 +2,9 @@ package main
 
 import (
 	"encoding/hex"
+	"fmt"
 	"go/ast"
 	"go/token"
-	"log"
 	"strconv"
 	"strings"
 
@@ -12,7 +12,6 @@ import (
 )
 
 func obfuscateLiterals(files []*ast.File) []*ast.File {
-
 	pre := func(cursor *astutil.Cursor) bool {
 		t, ok := cursor.Node().(*ast.GenDecl)
 		if !ok {
@@ -24,7 +23,6 @@ func obfuscateLiterals(files []*ast.File) []*ast.File {
 		if t.Tok == token.CONST {
 			t.Tok = token.VAR
 		}
-
 		return true
 	}
 
@@ -33,9 +31,7 @@ func obfuscateLiterals(files []*ast.File) []*ast.File {
 		fset       = token.NewFileSet()
 		addedToPkg bool // we only want to inject the code and imports once
 	)
-
 	post := func(cursor *astutil.Cursor) bool {
-
 		switch x := cursor.Node().(type) {
 		case *ast.File:
 			if !addedToPkg {
@@ -43,14 +39,8 @@ func obfuscateLiterals(files []*ast.File) []*ast.File {
 				x.Decls = append(x.Decls, keyStmt(key))
 
 				if x.Imports == nil {
-					var newDecls = []ast.Decl{
+					newDecls := []ast.Decl{cryptoAesImportSpec}
-						cryptoAesImportSpec,
+					newDecls = append(newDecls, x.Decls...)
-					}
-
-					for _, decl := range x.Decls {
-						newDecls = append(newDecls, decl)
-					}
-
 					x.Decls = newDecls
 				} else {
 					astutil.AddImport(fset, x, "crypto/aes")
@@ -58,7 +48,6 @@ func obfuscateLiterals(files []*ast.File) []*ast.File {
 				}
 
 				addedToPkg = true
-
 				return true
 			}
 		case *ast.BasicLit:
@@ -71,15 +60,12 @@ func obfuscateLiterals(files []*ast.File) []*ast.File {
 
 			value, err := strconv.Unquote(x.Value)
 			if err != nil {
-				log.Fatalln("[Fatal]: Could not unqote string", err)
+				panic(fmt.Sprintf("cannot unquote string: %v", err))
-				return false
 			}
 
-			ciphertext, err := encAes([]byte(value), key)
+			ciphertext, err := encAES([]byte(value), key)
 			if err != nil {
-
+				panic(fmt.Sprintf("cannot encrypt string: %v", err))
-				log.Fatalln("[Fatal]: Could not encrypt string:", err)
-				return false
 			}
 
 			cursor.Replace(ciphertextStmt(ciphertext))
@@ -88,14 +74,14 @@ func obfuscateLiterals(files []*ast.File) []*ast.File {
 		return true
 	}
 
-	for _, file := range files {
+	for i := range files {
-		file = astutil.Apply(file, pre, post).(*ast.File)
+		files[i] = astutil.Apply(files[i], pre, post).(*ast.File)
 	}
 
 	return files
 }
 
-// ast definitions for injection
+// AST definitions for injection
 var (
 	aesCipherStmt = &ast.AssignStmt{
 		Lhs: []ast.Expr{
@@ -103,17 +89,13 @@ var (
 			&ast.Ident{Name: "err"},
 		},
 		Tok: token.DEFINE,
-		Rhs: []ast.Expr{
+		Rhs: []ast.Expr{&ast.CallExpr{
-			&ast.CallExpr{
+			Fun: &ast.SelectorExpr{
-				Fun: &ast.SelectorExpr{
+				X:   &ast.Ident{Name: "aes"},
-					X:   &ast.Ident{Name: "aes"},
+				Sel: &ast.Ident{Name: "NewCipher"},
-					Sel: &ast.Ident{Name: "NewCipher"},
-				},
-				Args: []ast.Expr{
-					&ast.Ident{Name: "garbleKey"},
-				},
 			},
-		},
+			Args: []ast.Expr{&ast.Ident{Name: "garbleKey"}},
+		}},
 	}
 
 	aesGcmCipherStmt = &ast.AssignStmt{
@@ -122,17 +104,13 @@ var (
 			&ast.Ident{Name: "err"},
 		},
 		Tok: token.DEFINE,
-		Rhs: []ast.Expr{
+		Rhs: []ast.Expr{&ast.CallExpr{
-			&ast.CallExpr{
+			Fun: &ast.SelectorExpr{
-				Fun: &ast.SelectorExpr{
+				X:   &ast.Ident{Name: "cipher"},
-					X:   &ast.Ident{Name: "cipher"},
+				Sel: &ast.Ident{Name: "NewGCM"},
-					Sel: &ast.Ident{Name: "NewGCM"},
-				},
-				Args: []ast.Expr{
-					&ast.Ident{Name: "block"},
-				},
 			},
-		},
+			Args: []ast.Expr{&ast.Ident{Name: "block"}},
+		}},
 	}
 
 	plaintextStmt = &ast.AssignStmt{
@@ -141,44 +119,38 @@ var (
 			&ast.Ident{Name: "err"},
 		},
 		Tok: token.DEFINE,
-		Rhs: []ast.Expr{
+		Rhs: []ast.Expr{&ast.CallExpr{
-			&ast.CallExpr{
+			Fun: &ast.SelectorExpr{
-				Fun: &ast.SelectorExpr{
+				X:   &ast.Ident{Name: "aesgcm"},
-					X:   &ast.Ident{Name: "aesgcm"},
+				Sel: &ast.Ident{Name: "Open"},
-					Sel: &ast.Ident{Name: "Open"},
+			},
-				},
+			Args: []ast.Expr{
-				Args: []ast.Expr{
+				&ast.Ident{Name: "nil"},
-					&ast.Ident{Name: "nil"},
+				&ast.SliceExpr{
-					&ast.SliceExpr{
+					X: &ast.Ident{Name: "ciphertext"},
-						X: &ast.Ident{Name: "ciphertext"},
+					High: &ast.BasicLit{
-						High: &ast.BasicLit{
+						Kind:  token.INT,
-							Kind:  token.INT,
+						Value: "12",
-							Value: "12",
-						},
 					},
-					&ast.SliceExpr{
+				},
-						X: &ast.Ident{Name: "ciphertext"},
+				&ast.SliceExpr{
-						Low: &ast.BasicLit{
+					X: &ast.Ident{Name: "ciphertext"},
-							Kind:  token.INT,
+					Low: &ast.BasicLit{
-							Value: "12",
+						Kind:  token.INT,
-						},
+						Value: "12",
 					},
-					&ast.Ident{Name: "nil"},
 				},
+				&ast.Ident{Name: "nil"},
 			},
-		},
+		}},
 	}
 
-	returnStmt = &ast.ReturnStmt{
+	returnStmt = &ast.ReturnStmt{Results: []ast.Expr{
-		Results: []ast.Expr{
+		&ast.CallExpr{
-			&ast.CallExpr{
+			Fun:  &ast.Ident{Name: "string"},
-				Fun: &ast.Ident{Name: "string"},
+			Args: []ast.Expr{&ast.Ident{Name: "plaintext"}},
-				Args: []ast.Expr{
-					&ast.Ident{Name: "plaintext"},
-				},
-			},
 		},
-	}
+	}}
 )
 
 func decErrStmt() *ast.IfStmt {
@@ -188,137 +160,102 @@ func decErrStmt() *ast.IfStmt {
 			Op: token.NEQ,
 			Y:  &ast.Ident{Name: "nil"},
 		},
-		Body: &ast.BlockStmt{
+		Body: &ast.BlockStmt{List: []ast.Stmt{
-			List: []ast.Stmt{
+			&ast.ExprStmt{X: &ast.CallExpr{
-				&ast.ExprStmt{
+				Fun: &ast.Ident{Name: "panic"},
-					X: &ast.CallExpr{
+				Args: []ast.Expr{&ast.BinaryExpr{
-						Fun: &ast.Ident{Name: "panic"},
+					X: &ast.BasicLit{
-						Args: []ast.Expr{
+						Kind:  token.STRING,
-							&ast.BinaryExpr{
+						Value: `"garble: literal couldn't be decrypted: "`,
-								X: &ast.BasicLit{
-									Kind:  token.STRING,
-									Value: `"[garble] Literal couldn't be decrypted: "`,
-								},
-								Op: token.ADD,
-								Y: &ast.CallExpr{
-									Fun: &ast.SelectorExpr{
-										X:   &ast.Ident{Name: "err"},
-										Sel: &ast.Ident{Name: "Error"},
-									},
-								},
-							},
-						},
 					},
-				},
+					Op: token.ADD,
-			},
+					Y: &ast.CallExpr{Fun: &ast.SelectorExpr{
-		},
+						X:   &ast.Ident{Name: "err"},
+						Sel: &ast.Ident{Name: "Error"},
+					}},
+				}},
+			}},
+		}},
 	}
 }
 
 var funcStmt = &ast.FuncDecl{
 	Name: &ast.Ident{Name: "garbleDecrypt"},
 	Type: &ast.FuncType{
-		Params: &ast.FieldList{
+		Params: &ast.FieldList{List: []*ast.Field{{
-			List: []*ast.Field{
+			Names: []*ast.Ident{{Name: "ciphertext"}},
-				{
+			Type: &ast.ArrayType{
-					Names: []*ast.Ident{{Name: "ciphertext"}},
+				Elt: &ast.Ident{Name: "byte"},
-					Type: &ast.ArrayType{
-						Elt: &ast.Ident{Name: "byte"},
-					},
-				},
-			},
-		},
-		Results: &ast.FieldList{
-			List: []*ast.Field{
-				{
-					Type: &ast.Ident{Name: "string"},
-				},
 			},
-		},
+		}}},
-	},
+		Results: &ast.FieldList{List: []*ast.Field{{
-	Body: &ast.BlockStmt{
+			Type: &ast.Ident{Name: "string"},
-		List: []ast.Stmt{
+		}}},
-			aesCipherStmt,
-			decErrStmt(),
-			aesGcmCipherStmt,
-			decErrStmt(),
-			plaintextStmt,
-			decErrStmt(),
-			returnStmt,
-		},
 	},
+	Body: &ast.BlockStmt{List: []ast.Stmt{
+		aesCipherStmt,
+		decErrStmt(),
+		aesGcmCipherStmt,
+		decErrStmt(),
+		plaintextStmt,
+		decErrStmt(),
+		returnStmt,
+	}},
 }
 
 func ciphertextStmt(ciphertext []byte) *ast.CallExpr {
-	ciphertextLit := byteToByteLit(ciphertext)
+	ciphertextLit := dataAsByteSlice(ciphertext)
 
 	return &ast.CallExpr{
-		Fun: &ast.Ident{Name: "garbleDecrypt"},
+		Fun:  &ast.Ident{Name: "garbleDecrypt"},
-		Args: []ast.Expr{
+		Args: []ast.Expr{ciphertextLit},
-			ciphertextLit,
-		},
 	}
 }
 
-func byteToByteLit(buffer []byte) *ast.CallExpr {
+// dataAsByteSlice turns a byte slice like []byte{1, 2, 3} into an AST
-	hexstr := hex.EncodeToString(buffer)
+// expression which encodes it, such as []byte("\x01\x02\x03").
-
+func dataAsByteSlice(data []byte) *ast.CallExpr {
 	var b strings.Builder
 
-	b.WriteString(`"`)
+	b.WriteByte('"')
+	hexstr := hex.EncodeToString(data)
 	for i := 0; i < len(hexstr); i += 2 {
 		b.WriteString("\\x" + hexstr[i:i+2])
 	}
-
+	b.WriteByte('"')
-	b.WriteString(`"`)
 
 	return &ast.CallExpr{
 		Fun: &ast.ArrayType{
 			Elt: &ast.Ident{Name: "byte"},
 		},
-		Args: []ast.Expr{
+		Args: []ast.Expr{&ast.BasicLit{
-			&ast.BasicLit{
+			Kind:  token.STRING,
-				Kind:  token.STRING,
+			Value: b.String(),
-				Value: b.String(),
+		}},
-			},
-		},
 	}
 }
 
-func keyStmt(key []byte) (decl *ast.GenDecl) {
+func keyStmt(key []byte) *ast.GenDecl {
-	keyLit := byteToByteLit(key)
+	keyLit := dataAsByteSlice(key)
-
+	return &ast.GenDecl{
-	decl = &ast.GenDecl{
 		Tok: token.VAR,
-		Specs: []ast.Spec{
+		Specs: []ast.Spec{&ast.ValueSpec{
-			&ast.ValueSpec{
+			Names:  []*ast.Ident{{Name: "garbleKey"}},
-				Names: []*ast.Ident{
+			Values: []ast.Expr{keyLit},
-					{Name: "garbleKey"},
+		}},
-				},
-				Values: []ast.Expr{
-					keyLit,
-				},
-			},
-		},
 	}
-
-	return
 }
 
 var cryptoAesImportSpec = &ast.GenDecl{
 	Tok: token.IMPORT,
 	Specs: []ast.Spec{
-		&ast.ImportSpec{
+		&ast.ImportSpec{Path: &ast.BasicLit{
-			Path: &ast.BasicLit{
+			Kind:  token.STRING,
-				Kind:  token.STRING,
+			Value: `"crypto/aes"`,
-				Value: `"crypto/aes"`,
+		}},
-			},
+		&ast.ImportSpec{Path: &ast.BasicLit{
-		},
+			Kind:  token.STRING,
-		&ast.ImportSpec{
+			Value: `"crypto/cipher"`,
-			Path: &ast.BasicLit{
+		}},
-				Kind:  token.STRING,
-				Value: `"crypto/cipher"`,
-			},
-		},
 	},
 }