Repository: google/go-flow-levee
Pull request #78 (2020-08-31 18:21:25)

Remove unnecessary alias
This mirrors the removal in #70.

- [x] Tests pass
- [x] Appropriate changes to README are included in PR


Patch:
diff --git a/go.mod b/go.mod
index 34f01d8..787f951 100644
--- a/go.mod
+++ b/go.mod
@@ -2,4 +2,7 @@ module github.com/google/go-flow-levee
 
 go 1.14
 
-require golang.org/x/tools v0.0.0-20200416214402-fc959738d646
+require (
+	github.com/google/go-cmp v0.5.2
+	golang.org/x/tools v0.0.0-20200416214402-fc959738d646
+)
diff --git a/go.sum b/go.sum
index f91eab3..4d08442 100644
--- a/go.sum
+++ b/go.sum
@@ -1,3 +1,5 @@
+github.com/google/go-cmp v0.5.2 h1:X2ev0eStA3AbceY54o37/0PQ/UWqKEiiO2dKL5OPaFM=
+github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
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 c5f25d7..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,10 +82,10 @@ 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.TaggedFields, meth *ssa.Function) {
+func analyzeBlocks(pass *analysis.Pass, conf *config.Config, tf fieldtags.ResultType, meth *ssa.Function) {
 	// Function does not return anything
 	if res := meth.Signature.Results(); res == nil || (*res).Len() == 0 {
 		return
@@ -97,13 +103,13 @@ func analyzeBlocks(pass *analysis.Pass, conf *config.Config, tf fieldtags.Tagged
 	}
 }
 
-func analyzeResults(pass *analysis.Pass, conf *config.Config, tf fieldtags.TaggedFields, meth *ssa.Function, results []ssa.Value) {
+func analyzeResults(pass *analysis.Pass, conf *config.Config, tf fieldtags.ResultType, meth *ssa.Function, results []ssa.Value) {
 	for _, r := range results {
 		fa, ok := fieldAddr(r)
 		if !ok {
 			continue
 		}
-		if conf.IsSourceFieldAddr(fa) || tf.IsSource(fa) {
+		if conf.IsSourceFieldAddr(fa) || tf.IsSourceFieldAddr(fa) {
 			pass.ExportObjectFact(meth.Object(), &isFieldPropagator{})
 		}
 	}
@@ -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/fieldtags/analyzer.go b/internal/pkg/fieldtags/analyzer.go
index 93bac7a..bc303a0 100644
--- a/internal/pkg/fieldtags/analyzer.go
+++ b/internal/pkg/fieldtags/analyzer.go
@@ -17,14 +17,11 @@
 package fieldtags
 
 import (
-	"fmt"
 	"go/ast"
 	"go/types"
 	"reflect"
 	"strings"
 
-	"github.com/google/go-flow-levee/internal/pkg/utils"
-
 	"golang.org/x/tools/go/analysis"
 	"golang.org/x/tools/go/analysis/passes/inspect"
 	"golang.org/x/tools/go/ast/inspector"
@@ -42,16 +39,9 @@ var patterns sourcePatterns = []keyValue{
 	{"levee", "source"},
 }
 
-// TaggedFields is named type representing a slice of TaggedFields.
-type TaggedFields []TaggedField
-
-// ResultType is a slice of TaggedFields.
-type ResultType = TaggedFields
-
-// A TaggedField contains the necessary information to identify a tagged struct field.
-type TaggedField struct {
-	pkg, typ, field string
-}
+// ResultType is a map from types.Object to bool.
+// It can be used to determine whether a field is a tagged Source field.
+type ResultType map[types.Object]bool
 
 var Analyzer = &analysis.Analyzer{
 	Name: "fieldtags",
@@ -65,7 +55,7 @@ var Analyzer = &analysis.Analyzer{
 
 func run(pass *analysis.Pass) (interface{}, error) {
 	inspectResult := pass.ResultOf[inspect.Analyzer].(*inspector.Inspector)
-	var taggedFields TaggedFields
+	taggedFields := map[types.Object]bool{}
 
 	nodeFilter := []ast.Node{
 		(*ast.TypeSpec)(nil),
@@ -82,17 +72,17 @@ func run(pass *analysis.Pass) (interface{}, error) {
 		}
 		for _, f := range (*s).Fields.List {
 			if patterns.isSource(f) && len(f.Names) > 0 {
-				tf := TaggedField{
-					pkg:   pass.Pkg.Name(),
-					typ:   t.Name.Name,
-					field: f.Names[0].Name,
+				fNames := make([]string, len(f.Names))
+				for i, ident := range f.Names {
+					fNames[i] = ident.Name
+					taggedFields[pass.TypesInfo.ObjectOf(ident)] = true
 				}
-				taggedFields = append(taggedFields, tf)
-				pass.Reportf(f.Pos(), "tagged field: %s", tf)
+				pass.Reportf(f.Pos(), "tagged field: %s", strings.Join(fNames, ", "))
 			}
 		}
 	})
-	return taggedFields, nil
+
+	return ResultType(taggedFields), nil
 }
 
 func (sp *sourcePatterns) isSource(field *ast.Field) bool {
@@ -150,32 +140,14 @@ func (sp *sourcePatterns) isSource(field *ast.Field) bool {
 	return false
 }
 
-// IsSource determines whether a FieldAddr is a source, that is whether it refers to a field previously identified as a source.
-func (t TaggedFields) IsSource(f *ssa.FieldAddr) bool {
-	n, ok := namedType(f)
-	if !ok {
-		return false
-	}
-	fName := utils.FieldName(f)
-	for _, tf := range t {
-		if tf.matches(n.String(), fName) {
-			return true
-		}
-	}
-	return false
-}
-
-func namedType(f *ssa.FieldAddr) (*types.Named, bool) {
-	structType := f.X.Type()
-	deref := utils.Dereference(structType)
-	named, ok := deref.(*types.Named)
-	return named, ok
-}
-
-func (t TaggedField) matches(typeName, fieldName string) bool {
-	return fmt.Sprintf("%s.%s", t.pkg, t.typ) == typeName && t.field == fieldName
+// IsSourceFieldAddr determines whether a ssa.FieldAddr is a source, that is whether it refers to a field previously identified as a source.
+func (r ResultType) IsSourceFieldAddr(fa *ssa.FieldAddr) bool {
+	// incantation plundered from the docstring for ssa.FieldAddr.Field
+	field := fa.X.Type().Underlying().(*types.Pointer).Elem().Underlying().(*types.Struct).Field(fa.Field)
+	return r.IsSource(field)
 }
 
-func (t TaggedField) String() string {
-	return fmt.Sprintf("%s.%s.%s", t.pkg, t.typ, t.field)
+// IsSource determines whether a types.Var is a source, that is whether it refers to a field previously identified as a source.
+func (r ResultType) IsSource(field *types.Var) bool {
+	return r[(types.Object)(field)]
 }
diff --git a/internal/pkg/levee/levee.go b/internal/pkg/levee/levee.go
index 4eb300f..262cf34 100644
--- a/internal/pkg/levee/levee.go
+++ b/internal/pkg/levee/levee.go
@@ -58,7 +58,7 @@ 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):
diff --git a/internal/pkg/source/source.go b/internal/pkg/source/source.go
index d33c7f9..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(operands)
+		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,49 +97,111 @@ 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(operands []*ssa.Value) {
+func (a *Source) visitOperands(n ssa.Node, operands []*ssa.Value) {
+	_, visitingFromExtract := n.(*ssa.Extract)
+
 	for _, o := range operands {
 		n, ok := (*o).(ssa.Node)
 		if !ok || a.marked[n] {
 			continue
 		}
+
+		// Do not visit a Call if the current node is an Extract.
+		// This is to avoid incorrectly tainting non-Source values that are
+		// returned from a call that has a Source among its return values,
+		// e.g. a call to a function with a signature like:
+		// func CreateSource() (core.Source, error)
+		if _, ok := (*o).(*ssa.Call); visitingFromExtract && ok {
+			continue
+		}
+
 		// An Alloc represents the allocation of space for a variable. If a Node is an Alloc,
 		// and the thing being allocated is not an array, then either:
 		// a) it is a Source value, in which case it will get its own traversal when sourcesFromBlocks
@@ -326,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)
+}
diff --git a/internal/pkg/sourcetype/analyzer.go b/internal/pkg/sourcetype/analyzer.go
index afb9fb7..9ec612f 100644
--- a/internal/pkg/sourcetype/analyzer.go
+++ b/internal/pkg/sourcetype/analyzer.go
@@ -24,6 +24,7 @@ import (
 	"sort"
 
 	"github.com/google/go-flow-levee/internal/pkg/config"
+	"github.com/google/go-flow-levee/internal/pkg/fieldtags"
 	"golang.org/x/tools/go/analysis"
 	"golang.org/x/tools/go/analysis/passes/buildssa"
 	"golang.org/x/tools/go/ssa"
@@ -78,7 +79,7 @@ var Analyzer = &analysis.Analyzer{
 	Doc:        "This analyzer identifies types.Types values which contain dataflow sources.",
 	Flags:      config.FlagSet,
 	Run:        run,
-	Requires:   []*analysis.Analyzer{buildssa.Analyzer},
+	Requires:   []*analysis.Analyzer{buildssa.Analyzer, fieldtags.Analyzer},
 	ResultType: reflect.TypeOf(new(sourceClassifier)),
 	FactTypes:  []analysis.Fact{new(typeDeclFact), new(fieldDeclFact)},
 }
@@ -86,6 +87,7 @@ var Analyzer = &analysis.Analyzer{
 var Report bool
 
 func run(pass *analysis.Pass) (interface{}, error) {
+	taggedFields := pass.ResultOf[fieldtags.Analyzer].(fieldtags.ResultType)
 	ssaInput := pass.ResultOf[buildssa.Analyzer].(*buildssa.SSA)
 	conf, err := config.ReadConfig()
 	if err != nil {
@@ -95,7 +97,7 @@ func run(pass *analysis.Pass) (interface{}, error) {
 	// Members contains all named entities
 	for _, mem := range ssaInput.Pkg.Members {
 		if ssaType, ok := mem.(*ssa.Type); ok && conf.IsSource(ssaType.Type()) {
-			exportSourceFacts(pass, ssaType, conf)
+			exportSourceFacts(pass, ssaType, conf, taggedFields)
 		}
 	}
 
@@ -107,11 +109,15 @@ func run(pass *analysis.Pass) (interface{}, error) {
 	return classifier, nil
 }
 
-func exportSourceFacts(pass *analysis.Pass, ssaType *ssa.Type, conf *config.Config) {
+func exportSourceFacts(pass *analysis.Pass, ssaType *ssa.Type, conf *config.Config, taggedFields fieldtags.ResultType) {
 	pass.ExportObjectFact(ssaType.Object(), &typeDeclFact{})
 	if under, ok := ssaType.Type().Underlying().(*types.Struct); ok {
 		for i := 0; i < under.NumFields(); i++ {
-			if fld := under.Field(i); conf.IsSourceField(ssaType.Type(), fld) && fld.Pkg() == pass.Pkg {
+			fld := under.Field(i)
+			if fld.Pkg() != pass.Pkg {
+				continue
+			}
+			if conf.IsSourceField(ssaType.Type(), fld) || taggedFields.IsSource(fld) {
 				pass.ExportObjectFact(fld, &fieldDeclFact{})
 			}
 		}


Test patch:
diff --git a/internal/pkg/fieldtags/analyzer_test.go b/internal/pkg/fieldtags/analyzer_test.go
index 40abfec..2885548 100644
--- a/internal/pkg/fieldtags/analyzer_test.go
+++ b/internal/pkg/fieldtags/analyzer_test.go
@@ -17,11 +17,33 @@ package fieldtags
 import (
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"golang.org/x/tools/go/analysis/analysistest"
 )
 
 func TestFieldTagsAnalysis(t *testing.T) {
 	testdata := analysistest.TestData()
 
-	analysistest.Run(t, testdata, Analyzer, "tests")
+	results := analysistest.Run(t, testdata, Analyzer, "tests")
+
+	if len(results) != 1 {
+		t.Fatalf("expected 1 result, got %d", len(results))
+	}
+
+	want := []string{
+		"password",
+		"creds",
+		"secret",
+		"another",
+	}
+
+	var got []string
+	for o := range results[0].Result.(ResultType) {
+		got = append(got, o.Name())
+	}
+
+	if diff := cmp.Diff(want, got, cmpopts.SortSlices(func(a, b string) bool { return a < b })); diff != "" {
+		t.Errorf("Tagged Fields diff (-want +got):\n%s", diff)
+	}
 }
diff --git a/internal/pkg/fieldtags/testdata/src/tests/test.go b/internal/pkg/fieldtags/testdata/src/tests/test.go
index f635d8e..9974403 100644
--- a/internal/pkg/fieldtags/testdata/src/tests/test.go
+++ b/internal/pkg/fieldtags/testdata/src/tests/test.go
@@ -15,9 +15,9 @@
 package fieldtags
 
 type Person struct {
-	password             string      `levee:"source"`               // want "tagged field: fieldtags.Person.password"
-	secret               string      `json:"secret" levee:"source"` // want "tagged field: fieldtags.Person.secret"
-	another              interface{} "levee:\"source\""             // want "tagged field: fieldtags.Person.another"
+	password, creds      string      `levee:"source"`               // want "tagged field: password, creds"
+	secret               string      `json:"secret" levee:"source"` // want "tagged field: secret"
+	another              interface{} "levee:\"source\""             // want "tagged field: another"
 	name                 string      `some_key:"non_secret"`
 	spaceAfterFinalQuote string      `key:"value" `
 	someNotTaggedField   int
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)
+}
diff --git a/internal/pkg/levee/testdata/src/example.com/tests/extracts/tests.go b/internal/pkg/levee/testdata/src/example.com/tests/extracts/tests.go
new file mode 100644
index 0000000..d457fb5
--- /dev/null
+++ b/internal/pkg/levee/testdata/src/example.com/tests/extracts/tests.go
@@ -0,0 +1,62 @@
+// 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 extracts
+
+import (
+	"example.com/core"
+)
+
+func CreateSource() (core.Source, error) {
+	return core.Source{}, nil
+}
+
+func CreateSourceFlipped() (error, core.Source) {
+	return nil, core.Source{}
+}
+
+func TakeSource(s core.Source) (string, int, interface{}) {
+	return "", 0, nil
+}
+
+func TestOnlySourceExtractIsTainted() {
+	s, err := CreateSource()
+	core.Sink(s) // want "a source has reached a sink"
+	core.Sink(err)
+}
+
+func TestOnlySourceExtractIsTaintedInstructionOrder() {
+	s, err := CreateSource()
+	core.Sink(err)
+	core.Sink(s) // want "a source has reached a sink"
+}
+
+func TestOnlySourceExtractIsTaintedFlipped() {
+	err, s := CreateSourceFlipped()
+	core.Sink(s) // want "a source has reached a sink"
+	core.Sink(err)
+}
+
+func TestOnlySourceExtractIsTaintedInstructionOrderFlipped() {
+	err, s := CreateSourceFlipped()
+	core.Sink(err)
+	core.Sink(s) // want "a source has reached a sink"
+}
+
+func TestExtractsFromCallWithSourceArgAreTainted(s core.Source) {
+	str, i, e := TakeSource(s)
+	core.Sink(str) // want "a source has reached a sink"
+	core.Sink(i)   // want "a source has reached a sink"
+	core.Sink(e)   // want "a source has reached a sink"
+}
diff --git a/internal/pkg/sourcetype/testdata/src/sourcetype/test.go b/internal/pkg/sourcetype/testdata/src/sourcetype/test.go
index 5189af5..b7dd560 100644
--- a/internal/pkg/sourcetype/testdata/src/sourcetype/test.go
+++ b/internal/pkg/sourcetype/testdata/src/sourcetype/test.go
@@ -15,8 +15,10 @@
 package sourcetype
 
 type Source struct { // want Source:"source type"
-	Data string // want Data:"source field"
-	ID   int
+	Data         string // want Data:"source field"
+	Tagged       string `levee:"source"` // want Tagged:"source field"
+	Name1, Name2 string `levee:"source"` // want Name1:"source field" Name2:"source field"
+	ID           int
 }
 
 type AliasStruct = Source // want AliasStruct:"source type"