Restricting what packages are allowed to import in Golang using ASTs
I’m working on refactoring a monolithic end-to-end test suite for my company. It’s grown organically over the past 2 years as many, many new engineers and teams have onboarded. As you may know, “grown organically” is a euphemism for “it’s a hot mess”.
The question I wanted to answer today was “how do I restrict what packages my golang code can import?”
The solution I whipped up is a “linter” that runs as part of CI and fails if my import rules aren’t met. But what does it look like for real?
After a bunch of refactoring, we have test directories roughly corresponding to individual teams, and a set of library packages with stable interfaces that are intended to be used by those teams for testing. Now we need to make sure that each teams’ tests depend _only_ on the stable library packages and not random helpers in the test packages of other teams.
Say this is our repo layout:
<br />
$ tree .<br />
.<br />
├── pkg<br />
└── test<br />
├── connect<br />
│ ├── connector_sink_test.go<br />
│ ├── ...<br />
│ └── utils.go<br />
├── kafka<br />
│ ├── byok<br />
│ │ ├── aws_test.go<br />
│ │ └── gcp_test.go<br />
│ ├── kafka_suite<br />
│ │ ├── ...<br />
│ │ └── storage_test.go<br />
│ └── test_helpers.go<br />
└── test_helpers.go<br />
With this layout, we’d expect that
-
connector_sink_test.gocan import fromtest/connect/utils.goortest/test_helpers.go -
connector_sink_test.gocan NOT import fromtest/kafka/test_helpers.go -
aws_test.gocan import from anything intest/kafka/byok,test/kafka/test_helpers.go, ortest/test_helpers.go -
aws_test.gocan NOT import from anything fromtest/kafka/kafka_suiteortest/connect
Roughly speaking, we can break this down into like 3 phases.
1. Parse all the golang packages in the test directories (recursively) into ASTs
2. For each golang file, define the list of allowed import directories
3. Iterate through the ASTs and evaluate what we import versus what we allow
Here’s the heart of the approach. We have a list of ast.Packages and we want to validate that they only import paths that are in the parent directory hierarchy, rooted at testDir. Since this uses local file paths for parsing ASTs, we also pass our goSrcDir for translating between import paths and disk locations (using $GOPATH/src).
func validateAllowedImports(pkgs []*ast.Package, testDir string, goSrcDir string) []error {
testImport := strings.TrimPrefix(testDir, goSrcDir)
<p>var errors []error<br />
for _, pkg := range pkgs {<br />
for filename, file := range pkg.Files {<br />
allowedDirs := directoriesBetween(filename, testDir)<br />
allowedImports := make(map[string]struct{})<br />
for _, path := range allowedDirs {<br />
importLine := strings.TrimPrefix(path, goSrcDir)<br />
allowedImports[importLine] = struct{}{}<br />
}<br />
for _, imp := range file.Imports {<br />
path := imp.Path.Value[1 : len(imp.Path.Value)-1] // trim " from start and end<br />
// if we're importing from another test directory<br />
if strings.HasPrefix(path, testImport) {<br />
// make sure that its one of our allowed imports from a parent package<br />
if _, allowed := allowedImports[path]; !allowed {<br />
filename := strings.TrimPrefix(filename, goSrcDir)<br />
errors = append(errors, &IllegalImportError{ImportPath: path, Filepath: filename})<br />
}<br />
}<br />
}<br />
}<br />
}</p>
<p>return errors<br />
}<br />
This should be pretty straightforward to follow. There’s only a couple weird things in here.
1. The import line from the AST includes quotes, like "github.com/codyaray/break" so we trim them off
2. We use a map[string]struct{} as a poor man’s “set” data structure in golang for easy “set contains” operations (the if _, allowed := allowedImports[path]; !allowed line).
3. The “magic” that must be in the directoriesBetween(filename, testDir) function
Ok, ok. There’s not really much magic here. We each start with the filename and walk up the file path directory by directory, until we hit the root directory. That set of directories is what we allow tests to import from.
func directoriesBetween(filename string, rootDir string) []string {
d := filepath.Dir(filename)
allowedImports := []string{d}
for d != rootDir {
d = filepath.Dir(d)
allowedImports = append(allowedImports, d)
}
return allowedImports
}
The rest is just boring boilerplate code, but for the sake of completeness, here it is in full:
<br />
package main</p>
<p>import (<br />
"fmt"<br />
"go/ast"<br />
"go/parser"<br />
"go/token"<br />
"os"<br />
"path/filepath"<br />
"strings"<br />
)</p>
<p>func main() {<br />
rootDir, err := os.Getwd()<br />
if err != nil {<br />
panic(err)<br />
}<br />
testDir := filepath.Join(rootDir, "./test")</p>
<p>dirs, err := listDirectoriesRecursive(testDir)<br />
if err != nil {<br />
panic(err)<br />
}</p>
<p>pkgs, err := parsePackages(dirs)<br />
if err != nil {<br />
panic(err)<br />
}</p>
<p>goPath := os.Getenv("GOPATH")<br />
if goPath == "" {<br />
panic(fmt.Errorf("$GOPATH unset"))<br />
}<br />
goSrcDir := fmt.Sprintf("%s/src/", goPath)</p>
<p>errs := validateAllowedImports(pkgs, testDir, goSrcDir)<br />
for _, err := range errs {<br />
fmt.Println(err)<br />
}<br />
if len(errs) > 0 {<br />
os.Exit(1)<br />
}<br />
}</p>
<p>func validateAllowedImports(pkgs []*ast.Package, testDir string, goSrcDir string) []error {<br />
testImport := strings.TrimPrefix(testDir, goSrcDir)</p>
<p>var errors []error<br />
for _, pkg := range pkgs {<br />
for filename, file := range pkg.Files {<br />
allowedDirs := directoriesBetween(filename, testDir)<br />
allowedImports := make(map[string]struct{})<br />
for _, path := range allowedDirs {<br />
importLine := strings.TrimPrefix(path, goSrcDir)<br />
allowedImports[importLine] = struct{}{}<br />
}<br />
for _, imp := range file.Imports {<br />
path := imp.Path.Value[1 : len(imp.Path.Value)-1] // trim " from start and end<br />
// if we're importing from another test directory<br />
if strings.HasPrefix(path, testImport) {<br />
// make sure that its one of our allowed imports from a parent package<br />
if _, allowed := allowedImports[path]; !allowed {<br />
filename := strings.TrimPrefix(filename, goSrcDir)<br />
errors = append(errors, &IllegalImportError{ImportPath: path, Filepath: filename})<br />
}<br />
}<br />
}<br />
}<br />
}</p>
<p>return errors<br />
}</p>
<p>type IllegalImportError struct {<br />
ImportPath string<br />
Filepath string<br />
}</p>
<p>func (e *IllegalImportError) Error() string {<br />
return fmt.Sprintf("Illegal import of %s from %s", e.ImportPath, e.Filepath)<br />
}</p>
<p>// Returns a list of the directories between filename and rootDir<br />
//<br />
// e.g., given filename=test/kafka/byok/aws_test.go and rootDir=test,<br />
// this will return test/kafka/byok, test/kafka, and test<br />
func directoriesBetween(filename string, rootDir string) []string {<br />
d := filepath.Dir(filename)<br />
allowedImports := []string{d}<br />
for d != rootDir {<br />
d = filepath.Dir(d)<br />
allowedImports = append(allowedImports, d)<br />
}<br />
return allowedImports<br />
}</p>
<p>func parsePackages(dirs []string) ([]*ast.Package, error) {<br />
var packages []*ast.Package<br />
fset := token.NewFileSet()<br />
for _, dir := range dirs {<br />
pkgs, err := parser.ParseDir(fset, dir, nil, 0)<br />
if err != nil {<br />
return nil, err<br />
}<br />
for _, pkg := range pkgs {<br />
packages = append(packages, pkg)<br />
}<br />
}<br />
return packages, nil<br />
}</p>
<p>func listDirectoriesRecursive(rootDir string) ([]string, error) {<br />
dirs := []string{rootDir}<br />
err := filepath.Walk(rootDir,<br />
func(path string, info os.FileInfo, err error) error {<br />
if err != nil {<br />
return err<br />
}<br />
if info.IsDir() {<br />
dirs = append(dirs, path)<br />
}<br />
return nil<br />
})<br />
if err != nil {<br />
return nil, err<br />
}<br />
return dirs, nil<br />
}<br />