Repository: google/go-flow-levee
Pull request #85 (2020-09-08 13:27:04)

Refactor config code related to calls and functions
The tests permitting this refactoring were added in #69 and #75. 

- [x] Tests pass
- [x] Appropriate changes to README are included in PR


Patch:
diff --git a/internal/pkg/config/config.go b/internal/pkg/config/config.go
index 1a290fe..a0dba33 100644
--- a/internal/pkg/config/config.go
+++ b/internal/pkg/config/config.go
@@ -39,18 +39,16 @@ func init() {
 // config contains matchers and analysis scope information
 type Config struct {
 	Sources    []sourceMatcher
-	Sinks      []callMatcher
-	Sanitizers []callMatcher
+	Sinks      []funcMatcher
+	Sanitizers []funcMatcher
 }
 
-func (c Config) IsSink(call *ssa.Call) bool {
-	for _, p := range c.Sinks {
-		if p.Match(call) {
-			return true
-		}
+func (c Config) IsSinkCall(call *ssa.Call) bool {
+	callee := call.Call.StaticCallee()
+	if callee == nil {
+		return false
 	}
-
-	return false
+	return c.IsSinkFunction(callee)
 }
 
 func (c Config) IsSinkFunction(f *ssa.Function) bool {
@@ -60,16 +58,8 @@ func (c Config) IsSinkFunction(f *ssa.Function) bool {
 		return false
 	}
 
-	var recvName string
-	if recv := f.Signature.Recv(); recv != nil {
-		recvName = unqualifiedName(recv)
-	}
-
 	for _, p := range c.Sinks {
-		if !p.PackageRE.MatchString(f.Pkg.Pkg.Path()) || !p.MethodRE.MatchString(f.Name()) {
-			continue
-		}
-		if p.ReceiverRE.MatchString(recvName) {
+		if p.Match(f) {
 			return true
 		}
 	}
@@ -77,12 +67,16 @@ func (c Config) IsSinkFunction(f *ssa.Function) bool {
 }
 
 func (c Config) IsSanitizer(call *ssa.Call) bool {
+	callee := call.Call.StaticCallee()
+	if callee == nil {
+		return false
+	}
+
 	for _, p := range c.Sanitizers {
-		if p.Match(call) {
+		if p.Match(callee) {
 			return true
 		}
 	}
-
 	return false
 }
 
@@ -151,36 +145,30 @@ func (s sourceMatcher) match(n *types.Named) bool {
 	return s.PackageRE.MatchString(n.Obj().Pkg().Path()) && s.TypeRE.MatchString(n.Obj().Name())
 }
 
-type callMatcher struct {
+type funcMatcher struct {
 	PackageRE  regexp.Regexp
 	ReceiverRE regexp.Regexp
 	MethodRE   regexp.Regexp
 }
 
-func (r callMatcher) matchPackage(p *types.Package) bool {
-	return r.PackageRE.MatchString(p.Path())
+func (fm funcMatcher) matchPackage(p *types.Package) bool {
+	return fm.PackageRE.MatchString(p.Path())
 }
 
 // Match matches methods based on package, method, and receiver regexp.
 // To explicitly match a method with no receiver (i.e., a top-level function),
 // provide the ReceiverRE regexp `^$`.
-func (r callMatcher) Match(c *ssa.Call) bool {
-	callee := c.Call.StaticCallee()
-	if callee == nil || callee.Pkg == nil {
+func (fm funcMatcher) Match(f *ssa.Function) bool {
+	if !fm.matchPackage(f.Pkg.Pkg) || !fm.MethodRE.MatchString(f.Name()) {
 		return false
 	}
 
-	if !r.matchPackage(callee.Pkg.Pkg) || !r.MethodRE.MatchString(callee.Name()) {
-		return false
-	}
-
-	recv := c.Call.Signature().Recv()
+	recv := f.Signature.Recv()
 	var recvName string
 	if recv != nil {
 		recvName = unqualifiedName(recv)
 	}
-
-	return r.ReceiverRE.MatchString(recvName)
+	return fm.ReceiverRE.MatchString(recvName)
 }
 
 func unqualifiedName(v *types.Var) string {
diff --git a/internal/pkg/debug/graph/graph.go b/internal/pkg/debug/graph/graph.go
index 5b86ecd..35edfb1 100644
--- a/internal/pkg/debug/graph/graph.go
+++ b/internal/pkg/debug/graph/graph.go
@@ -88,8 +88,7 @@ func (g *FuncGraph) visit(b *ssa.BasicBlock) {
 }
 
 func (g *FuncGraph) visitOperands(n ssa.Node) {
-	var operands []*ssa.Value
-	operands = n.Operands(operands)
+	operands := n.Operands(nil)
 	if operands == nil {
 		return
 	}
diff --git a/internal/pkg/fieldpropagator/analyzer.go b/internal/pkg/fieldpropagator/analyzer.go
index 0806fbb..6812283 100644
--- a/internal/pkg/fieldpropagator/analyzer.go
+++ b/internal/pkg/fieldpropagator/analyzer.go
@@ -28,7 +28,13 @@ import (
 )
 
 // ResultType is a map telling whether an object (really a function) is a field propagator.
-type ResultType = FieldPropagators
+type ResultType map[types.Object]bool
+
+// IsFieldPropagator determines whether a call is a field propagator.
+func (r ResultType) IsFieldPropagator(c *ssa.Call) bool {
+	cf, ok := c.Call.Value.(*ssa.Function)
+	return ok && r[cf.Object()]
+}
 
 type isFieldPropagator struct{}
 
@@ -76,7 +82,7 @@ func run(pass *analysis.Pass) (interface{}, error) {
 	for _, f := range pass.AllObjectFacts() {
 		isFieldPropagator[f.Object] = true
 	}
-	return FieldPropagators(isFieldPropagator), nil
+	return ResultType(isFieldPropagator), nil
 }
 
 func analyzeBlocks(pass *analysis.Pass, conf *config.Config, tf fieldtags.ResultType, meth *ssa.Function) {
@@ -118,11 +124,3 @@ func fieldAddr(x ssa.Value) (*ssa.FieldAddr, bool) {
 	}
 	return nil, false
 }
-
-type FieldPropagators map[types.Object]bool
-
-// Contains determines whether a call is a field propagator, i.e. whether it is contained by the FieldPropagators map.
-func (f FieldPropagators) Contains(c *ssa.Call) bool {
-	cf, ok := c.Call.Value.(*ssa.Function)
-	return ok && f[cf.Object()]
-}
diff --git a/internal/pkg/levee/levee.go b/internal/pkg/levee/levee.go
index 4eb300f..6be9bc0 100644
--- a/internal/pkg/levee/levee.go
+++ b/internal/pkg/levee/levee.go
@@ -58,10 +58,10 @@ func run(pass *analysis.Pass) (interface{}, error) {
 					continue
 				}
 				switch {
-				case fieldPropagators.Contains(v):
+				case fieldPropagators.IsFieldPropagator(v):
 					sources = append(sources, source.New(v, conf))
 
-				case conf.IsSink(v):
+				case conf.IsSinkCall(v):
 					for _, s := range sources {
 						if s.HasPathTo(instr.(ssa.Node)) && !s.IsSanitizedAt(v) {
 							report(pass, s, v)
diff --git a/internal/pkg/source/source.go b/internal/pkg/source/source.go
index 3c9618d..74708f3 100644
--- a/internal/pkg/source/source.go
+++ b/internal/pkg/source/source.go
@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"go/token"
 	"go/types"
+	"log"
 	"strings"
 
 	"github.com/google/go-flow-levee/internal/pkg/sanitizer"
@@ -39,12 +40,13 @@ type classifier interface {
 // its referrers.
 // Source.sanitized notes sanitizer calls that sanitize this Source
 type Source struct {
-	node            ssa.Node
-	marked          map[ssa.Node]bool
-	preOrder        []ssa.Node
-	sanitizers      []*sanitizer.Sanitizer
-	config          classifier
-	maxInstrReached map[*ssa.BasicBlock]int
+	node             ssa.Node
+	marked           map[ssa.Node]bool
+	preOrder         []ssa.Node
+	sanitizers       []*sanitizer.Sanitizer
+	config           classifier
+	maxInstrReached  map[*ssa.BasicBlock]int
+	lastBlockVisited *ssa.BasicBlock
 }
 
 // Node returns the underlying ssa.Node of the Source.
@@ -72,22 +74,20 @@ func (a *Source) dfs(n ssa.Node) {
 	a.marked[n.(ssa.Node)] = true
 
 	if instr, ok := n.(ssa.Instruction); ok {
-		a.recordIndex(instr)
+		a.record(instr)
 	}
 
-	if n.Referrers() != nil {
-		a.visitReferrers(n.Referrers())
-	}
+	a.visitReferrers(n)
 
-	var operands []*ssa.Value
-	operands = n.Operands(operands)
+	operands := n.Operands(nil)
 	if operands != nil {
 		a.visitOperands(n, operands)
 	}
 }
 
-func (a *Source) recordIndex(target ssa.Instruction) {
+func (a *Source) record(target ssa.Instruction) {
 	b := target.Block()
+	a.lastBlockVisited = b
 	i, ok := indexInBlock(target)
 	if !ok {
 		return
@@ -97,41 +97,91 @@ func (a *Source) recordIndex(target ssa.Instruction) {
 	}
 }
 
-func (a *Source) visitReferrers(referrers *[]ssa.Instruction) {
-	for _, r := range *referrers {
+func (a *Source) visitReferrers(n ssa.Node) {
+	referrers := a.referrersToVisit(n)
+
+	for _, r := range referrers {
 		if a.marked[r.(ssa.Node)] {
 			continue
 		}
+
 		switch v := r.(type) {
 		case *ssa.Call:
-			// This is to avoid attaching calls where the source is the receiver, ex:
-			// core.Sinkf("Source id: %v", wrapper.Source.GetID())
-			if recv := v.Call.Signature().Recv(); recv != nil && a.config.IsSource(utils.Dereference(recv.Type())) {
+			if a.config.IsSanitizer(v) {
+				a.sanitizers = append(a.sanitizers, &sanitizer.Sanitizer{Call: v})
+			}
+
+		case *ssa.FieldAddr:
+			if !a.config.IsSourceFieldAddr(v) {
 				continue
 			}
+		}
 
-			if a.config.IsSanitizer(v) {
-				a.sanitizers = append(a.sanitizers, &sanitizer.Sanitizer{Call: v})
+		a.dfs(r.(ssa.Node))
+	}
+}
+
+// referrersToVisit produces a filtered list of Referrers for an ssa.Node.
+// Specifically, we avoid referrers that:
+// - Are in a block that is not reachable from the current instruction
+// - Are calls to a Source method
+// - Are calls that occur earlier in the same block as the value being referred
+func (a *Source) referrersToVisit(n ssa.Node) (referrers []ssa.Instruction) {
+	if n.Referrers() == nil {
+		return
+	}
+	for _, r := range *n.Referrers() {
+		// If the referrer is in a different block from the one we last visited,
+		// and it can't be reached from the block we are visiting, then stop visiting.
+		if rb := r.Block(); a.lastBlockVisited != nil &&
+			rb != a.lastBlockVisited &&
+			!a.canReach(a.lastBlockVisited, rb) {
+			continue
+		}
+
+		if c, ok := r.(*ssa.Call); ok {
+			// This is to avoid attaching calls where the source is the receiver, ex:
+			// core.Sinkf("Source id: %v", wrapper.Source.GetID())
+			if recv := c.Call.Signature().Recv(); recv != nil && a.config.IsSource(utils.Dereference(recv.Type())) {
+				continue
 			}
 
 			// If this call's index is lower than the highest in its block,
 			// then this call is "in the past" and we should stop traversing.
 			i, ok := indexInBlock(r)
 			if !ok {
-				break
+				continue
 			}
-
 			if i < a.maxInstrReached[r.Block()] {
 				continue
 			}
+		}
+		referrers = append(referrers, r)
+	}
+	return referrers
+}
 
-		case *ssa.FieldAddr:
-			if !a.config.IsSourceFieldAddr(v) {
+func (a *Source) canReach(start *ssa.BasicBlock, dest *ssa.BasicBlock) bool {
+	if start.Dominates(dest) {
+		return true
+	}
+
+	stack := stack([]*ssa.BasicBlock{start})
+	seen := map[*ssa.BasicBlock]bool{start: true}
+	for len(stack) > 0 {
+		current := stack.pop()
+		if current == dest {
+			return true
+		}
+		for _, s := range current.Succs {
+			if seen[s] {
 				continue
 			}
+			seen[s] = true
+			stack.push(s)
 		}
-		a.dfs(r.(ssa.Node))
 	}
+	return false
 }
 
 func (a *Source) visitOperands(n ssa.Node, operands []*ssa.Value) {
@@ -338,3 +388,18 @@ func indexInBlock(target ssa.Instruction) (int, bool) {
 	// i.e. an instruction does not appear within its parent block
 	return 0, false
 }
+
+type stack []*ssa.BasicBlock
+
+func (s *stack) pop() *ssa.BasicBlock {
+	if len(*s) == 0 {
+		log.Println("tried to pop from empty stack")
+	}
+	popped := (*s)[len(*s)-1]
+	*s = (*s)[:len(*s)-1]
+	return popped
+}
+
+func (s *stack) push(b *ssa.BasicBlock) {
+	*s = append(*s, b)
+}


Test patch:
diff --git a/internal/pkg/config/config_test.go b/internal/pkg/config/config_test.go
index 3da738b..bc51c5b 100644
--- a/internal/pkg/config/config_test.go
+++ b/internal/pkg/config/config_test.go
@@ -45,7 +45,7 @@ func runTest(pass *analysis.Pass) (interface{}, error) {
 		}
 		for _, b := range f.Blocks {
 			for _, i := range b.Instrs {
-				if c, ok := i.(*ssa.Call); ok && conf.IsSink(c) {
+				if c, ok := i.(*ssa.Call); ok && conf.IsSinkCall(c) {
 					pass.Reportf(i.Pos(), "sink call")
 				}
 			}
diff --git a/internal/pkg/levee/testdata/src/example.com/core/sanitize.go b/internal/pkg/levee/testdata/src/example.com/core/sanitize.go
index 6b24593..9474586 100644
--- a/internal/pkg/levee/testdata/src/example.com/core/sanitize.go
+++ b/internal/pkg/levee/testdata/src/example.com/core/sanitize.go
@@ -21,3 +21,7 @@ func Sanitize(args ...interface{}) []interface{} {
 func SanitizeSource(s Source) Source {
 	return Source{ID: s.ID}
 }
+
+func SanitizePtr(s *Source) {
+	s.Data = "<redacted>"
+}
diff --git a/internal/pkg/levee/testdata/src/example.com/tests/callorder/multiblock.go b/internal/pkg/levee/testdata/src/example.com/tests/callorder/multiblock.go
new file mode 100644
index 0000000..aacad5e
--- /dev/null
+++ b/internal/pkg/levee/testdata/src/example.com/tests/callorder/multiblock.go
@@ -0,0 +1,120 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package callorder
+
+import (
+	"fmt"
+	"io"
+
+	"example.com/core"
+)
+
+func TestSinkInIfBeforeTaint(s core.Source, w io.Writer) {
+	if true {
+		core.Sink(w)
+	}
+	fmt.Fprintf(w, "%v", s)
+}
+
+func TestTaintInIfBeforeSink(s core.Source, w io.Writer) {
+	if true {
+		fmt.Fprintf(w, "%v", s)
+	}
+	core.Sink(w) // want "a source has reached a sink"
+}
+
+func TestSinkAndTaintInDifferentIfBranches(s core.Source, w io.Writer) {
+	if true {
+		fmt.Fprintf(w, "%v", s)
+	} else {
+		core.Sink(w)
+	}
+}
+
+func TestSinkInIfBeforeTaintInIf(s core.Source, w io.Writer) {
+	if true {
+		core.Sink(w)
+	}
+	if true {
+		fmt.Fprintf(w, "%v", s)
+	}
+}
+
+func TestTaintInIfBeforeSinkInIf(s core.Source, w io.Writer) {
+	if true {
+		fmt.Fprintf(w, "%v", s)
+	}
+	if true {
+		core.Sink(w) // want "a source has reached a sink"
+	}
+}
+
+func TestSinkBeforeTaintInSameIfBlock(s core.Source, w io.Writer) {
+	if true {
+		core.Sink(w)
+		fmt.Fprintf(w, "%v", s)
+	}
+}
+
+func TestTaintBeforeSinkInSameIfBlock(s core.Source, w io.Writer) {
+	if true {
+		fmt.Fprintf(w, "%v", s)
+		core.Sink(w) // want "a source has reached a sink"
+	}
+}
+
+func TestSinkInNestedIfBeforeTaint(s core.Source, w io.Writer) {
+	if true {
+		if true {
+			core.Sink(w)
+		}
+	}
+	fmt.Fprintf(w, "%v", s)
+}
+
+func TestTaintInNestedIfBeforeSink(s core.Source, w io.Writer) {
+	if true {
+		if true {
+			fmt.Fprintf(w, "%v", s)
+			core.Sink(w) // want "a source has reached a sink"
+		}
+		core.Sink(w) // want "a source has reached a sink"
+	}
+	core.Sink(w) // want "a source has reached a sink"
+}
+
+func TestSinkAndTaintInSeparateSwitchCases(s core.Source, w io.Writer) {
+	switch "true" {
+	case "true":
+		core.Sink(w)
+	case "false":
+		fmt.Fprintf(w, "%v", s)
+	}
+}
+
+func TestSinkAfterTaintInSwitch(s core.Source, w io.Writer) {
+	switch "true" {
+	case "true":
+		fmt.Fprintf(w, "%v", s)
+	}
+	core.Sink(w) // want "a source has reached a sink"
+}
+
+func TestSinkAfterTaintInFor(sources []core.Source, w io.Writer) {
+	for i := 0; i < len(sources); i++ {
+		fmt.Fprintf(w, "%v", sources[i])
+	}
+	core.Sink(w) // want "a source has reached a sink"
+}
diff --git a/internal/pkg/levee/testdata/src/example.com/tests/callorder/tests.go b/internal/pkg/levee/testdata/src/example.com/tests/callorder/singleblock.go
similarity index 100%
rename from internal/pkg/levee/testdata/src/example.com/tests/callorder/tests.go
rename to internal/pkg/levee/testdata/src/example.com/tests/callorder/singleblock.go
diff --git a/internal/pkg/levee/testdata/src/example.com/tests/dominance/tests.go b/internal/pkg/levee/testdata/src/example.com/tests/dominance/tests.go
index cf06d62..b0e53be 100644
--- a/internal/pkg/levee/testdata/src/example.com/tests/dominance/tests.go
+++ b/internal/pkg/levee/testdata/src/example.com/tests/dominance/tests.go
@@ -43,3 +43,8 @@ func TestNotGuaranteedSanitization(c *core.Source) {
 	}
 	core.Sinkf("Sometimes sanitized: %v", p) // want "a source has reached a sink"
 }
+
+func TestSanitizationByPointer(c core.Source) {
+	core.SanitizePtr(&c)
+	core.Sink(c)
+}