Add callgraphutil.WriteCSV (#29)

Author: picatz

callgraphutil/csv.go

@@ -0,0 +1,118 @@
+package callgraphutil
+
+import (
+	"encoding/csv"
+	"fmt"
+	"io"
+	"runtime"
+	"strings"
+
+	"golang.org/x/tools/go/callgraph"
+)
+
+// WriteCSV writes the given callgraph.Graph to the given io.Writer in CSV
+// format. This format can be used to generate a visual representation of the
+// call graph using many different tools.
+func WriteCSV(w io.Writer, g *callgraph.Graph) error {
+	cw := csv.NewWriter(w)
+	cw.Comma = ','
+	defer cw.Flush()
+
+	// Write header.
+	if err := cw.Write([]string{
+		"source_pkg",
+		"source_pkg_go_version",
+		"source_pkg_origin",
+		"source_func",
+		"source_func_name",
+		"source_func_signature",
+		"target_pkg",
+		"target_pkg_go_version",
+		"target_pkg_origin",
+		"target_func",
+		"target_func_name",
+		"target_func_signature",
+	}); err != nil {
+		return fmt.Errorf("failed to write header: %w", err)
+	}
+
+	// Write edges.
+	for _, n := range g.Nodes {
+		source, err := getNodeInfo(n)
+		if err != nil {
+			return fmt.Errorf("failed to get node info: %w", err)
+		}
+
+		for _, e := range n.Out {
+			target, err := getNodeInfo(e.Callee)
+			if err != nil {
+				return fmt.Errorf("failed to get node info: %w", err)
+			}
+
+			record := []string{}
+			record = append(record, source.CSV()...)
+			record = append(record, target.CSV()...)
+
+			// Write edge.
+			if err := cw.Write(record); err != nil {
+				return fmt.Errorf("failed to write edge: %w", err)
+			}
+		}
+	}
+
+	return nil
+}
+
+// nodeInfo is a struct that contains information about a callgraph.Node used
+// to generate CSV output.
+type nodeInfo struct {
+	pkgPath          string
+	pkgGoVersion     string
+	pkgOrigin        string
+	pkgFunc          string
+	pkgFuncName      string
+	pkgFuncSignature string
+}
+
+// CSV returns single record for the node.
+func (n *nodeInfo) CSV() []string {
+	return []string{
+		n.pkgPath,
+		n.pkgGoVersion,
+		n.pkgOrigin,
+		n.pkgFunc,
+		n.pkgFuncName,
+		n.pkgFuncSignature,
+	}
+}
+
+// getNodeInfo returns a nodeInfo struct for the given callgraph.Node.
+func getNodeInfo(n *callgraph.Node) (*nodeInfo, error) {
+	info := &nodeInfo{
+		pkgPath:          "unknown",
+		pkgGoVersion:     runtime.Version(),
+		pkgOrigin:        "unknown",
+		pkgFunc:          n.Func.String(),
+		pkgFuncName:      n.Func.Name(),
+		pkgFuncSignature: n.Func.Signature.String(),
+	}
+
+	if n.Func.Pkg != nil {
+		info.pkgPath = n.Func.Pkg.Pkg.Path()
+
+		if goVersion := n.Func.Pkg.Pkg.GoVersion(); goVersion != "" {
+			info.pkgGoVersion = goVersion
+		}
+	}
+
+	if strings.Contains(info.pkgPath, ".") {
+		info.pkgOrigin = strings.Split(info.pkgPath, "/")[0]
+	} else {
+		// If the package path doesn't contain a dot, then it's
+		// probably a standard library package? This is a pattern
+		// I've used and seen elsewhere.
+		info.pkgOrigin = "stdlib"
+	}
+
+	return info, nil
+}

callgraphutil/csv_test.go

@@ -0,0 +1,48 @@
+package callgraphutil_test
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/picatz/taint/callgraphutil"
+)
+
+func TestWriteCSV(t *testing.T) {
+	var (
+		ownerName = "picatz"
+		repoName  = "taint"
+	)
+
+	repo, _, err := cloneGitHubRepository(context.Background(), ownerName, repoName)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	pkgs, err := loadPackages(context.Background(), repo, "./...")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	mainFn, srcFns, err := loadSSA(context.Background(), pkgs)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	cg, err := loadCallGraph(context.Background(), mainFn, srcFns)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	fh, err := os.Create(fmt.Sprintf("%s.csv", repoName))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer fh.Close()
+
+	err = callgraphutil.WriteCSV(fh, cg)
+	if err != nil {
+		t.Fatal(err)
+	}
+}

callgraphutil/dot_test.go

@@ -8,7 +8,6 @@ import (
 	"go/parser"
 	"go/token"
 	"os"
-	"path/filepath"
 	"testing"
 
 	"github.com/go-git/go-git/v5"
@@ -24,25 +23,9 @@ func cloneGitHubRepository(ctx context.Context, ownerName, repoName string) (str
 	ownerAndRepo := ownerName + "/" + repoName
 
 	// Get the directory path.
-	dir := filepath.Join(os.TempDir(), "taint", "github", ownerAndRepo)
-
-	// Check if the directory exists.
-	_, err := os.Stat(dir)
-	if err == nil {
-		// If the directory exists, we'll assume it's a valid repository,
-		// and return the directory. Open the directory to
-		repo, err := git.PlainOpen(dir)
-		if err != nil {
-			return dir, "", fmt.Errorf("%w", err)
-		}
-
-		// Get the repository's HEAD.
-		head, err := repo.Head()
-		if err != nil {
-			return dir, "", fmt.Errorf("%w", err)
-		}
-
-		return dir, head.Hash().String(), nil
+	dir, err := os.MkdirTemp(os.TempDir(), fmt.Sprintf("callgraphutil_csv-%s-%s", ownerName, repoName))
+	if err != nil {
+		return "", "", fmt.Errorf("failed to create temp dir: %w", err)
 	}
 
 	// Clone the repository.