1// Copyright 2016 Google Inc. All rights reserved.
2// Use of this source code is governed by the Apache 2.0
3// license that can be found in the LICENSE file.
4
5package main
6
7import (
8	"go/ast"
9	"path"
10	"strconv"
11	"strings"
12)
13
14const (
15	ctxPackage = "golang.org/x/net/context"
16
17	newPackageBase = "google.golang.org/"
18	stutterPackage = false
19)
20
21func init() {
22	register(fix{
23		"ae",
24		"2016-04-15",
25		aeFn,
26		`Update old App Engine APIs to new App Engine APIs`,
27	})
28}
29
30// logMethod is the set of methods on appengine.Context used for logging.
31var logMethod = map[string]bool{
32	"Debugf":    true,
33	"Infof":     true,
34	"Warningf":  true,
35	"Errorf":    true,
36	"Criticalf": true,
37}
38
39// mapPackage turns "appengine" into "google.golang.org/appengine", etc.
40func mapPackage(s string) string {
41	if stutterPackage {
42		s += "/" + path.Base(s)
43	}
44	return newPackageBase + s
45}
46
47func aeFn(f *ast.File) bool {
48	// During the walk, we track the last thing seen that looks like
49	// an appengine.Context, and reset it once the walk leaves a func.
50	var lastContext *ast.Ident
51
52	fixed := false
53
54	// Update imports.
55	mainImp := "appengine"
56	for _, imp := range f.Imports {
57		pth, _ := strconv.Unquote(imp.Path.Value)
58		if pth == "appengine" || strings.HasPrefix(pth, "appengine/") {
59			newPth := mapPackage(pth)
60			imp.Path.Value = strconv.Quote(newPth)
61			fixed = true
62
63			if pth == "appengine" {
64				mainImp = newPth
65			}
66		}
67	}
68
69	// Update any API changes.
70	walk(f, func(n interface{}) {
71		if ft, ok := n.(*ast.FuncType); ok && ft.Params != nil {
72			// See if this func has an `appengine.Context arg`.
73			// If so, remember its identifier.
74			for _, param := range ft.Params.List {
75				if !isPkgDot(param.Type, "appengine", "Context") {
76					continue
77				}
78				if len(param.Names) == 1 {
79					lastContext = param.Names[0]
80					break
81				}
82			}
83			return
84		}
85
86		if as, ok := n.(*ast.AssignStmt); ok {
87			if len(as.Lhs) == 1 && len(as.Rhs) == 1 {
88				// If this node is an assignment from an appengine.NewContext invocation,
89				// remember the identifier on the LHS.
90				if isCall(as.Rhs[0], "appengine", "NewContext") {
91					if ident, ok := as.Lhs[0].(*ast.Ident); ok {
92						lastContext = ident
93						return
94					}
95				}
96				// x (=|:=) appengine.Timeout(y, z)
97				//   should become
98				// x, _ (=|:=) context.WithTimeout(y, z)
99				if isCall(as.Rhs[0], "appengine", "Timeout") {
100					addImport(f, ctxPackage)
101					as.Lhs = append(as.Lhs, ast.NewIdent("_"))
102					// isCall already did the type checking.
103					sel := as.Rhs[0].(*ast.CallExpr).Fun.(*ast.SelectorExpr)
104					sel.X = ast.NewIdent("context")
105					sel.Sel = ast.NewIdent("WithTimeout")
106					fixed = true
107					return
108				}
109			}
110			return
111		}
112
113		// If this node is a FuncDecl, we've finished the function, so reset lastContext.
114		if _, ok := n.(*ast.FuncDecl); ok {
115			lastContext = nil
116			return
117		}
118
119		if call, ok := n.(*ast.CallExpr); ok {
120			if isPkgDot(call.Fun, "appengine", "Datacenter") && len(call.Args) == 0 {
121				insertContext(f, call, lastContext)
122				fixed = true
123				return
124			}
125			if isPkgDot(call.Fun, "taskqueue", "QueueStats") && len(call.Args) == 3 {
126				call.Args = call.Args[:2] // drop last arg
127				fixed = true
128				return
129			}
130
131			sel, ok := call.Fun.(*ast.SelectorExpr)
132			if !ok {
133				return
134			}
135			if lastContext != nil && refersTo(sel.X, lastContext) && logMethod[sel.Sel.Name] {
136				// c.Errorf(...)
137				//   should become
138				// log.Errorf(c, ...)
139				addImport(f, mapPackage("appengine/log"))
140				sel.X = &ast.Ident{ // ast.NewIdent doesn't preserve the position.
141					NamePos: sel.X.Pos(),
142					Name:    "log",
143				}
144				insertContext(f, call, lastContext)
145				fixed = true
146				return
147			}
148		}
149	})
150
151	// Change any `appengine.Context` to `context.Context`.
152	// Do this in a separate walk because the previous walk
153	// wants to identify "appengine.Context".
154	walk(f, func(n interface{}) {
155		expr, ok := n.(ast.Expr)
156		if ok && isPkgDot(expr, "appengine", "Context") {
157			addImport(f, ctxPackage)
158			// isPkgDot did the type checking.
159			n.(*ast.SelectorExpr).X.(*ast.Ident).Name = "context"
160			fixed = true
161			return
162		}
163	})
164
165	// The changes above might remove the need to import "appengine".
166	// Check if it's used, and drop it if it isn't.
167	if fixed && !usesImport(f, mainImp) {
168		deleteImport(f, mainImp)
169	}
170
171	return fixed
172}
173
174// ctx may be nil.
175func insertContext(f *ast.File, call *ast.CallExpr, ctx *ast.Ident) {
176	if ctx == nil {
177		// context is unknown, so use a plain "ctx".
178		ctx = ast.NewIdent("ctx")
179	} else {
180		// Create a fresh *ast.Ident so we drop the position information.
181		ctx = ast.NewIdent(ctx.Name)
182	}
183
184	call.Args = append([]ast.Expr{ctx}, call.Args...)
185}
186