1// Copyright 2015 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 internal
6
7import (
8	"errors"
9	"os"
10
11	"github.com/golang/protobuf/proto"
12	netcontext "golang.org/x/net/context"
13)
14
15var errNotAppEngineContext = errors.New("not an App Engine context")
16
17type CallOverrideFunc func(ctx netcontext.Context, service, method string, in, out proto.Message) error
18
19var callOverrideKey = "holds []CallOverrideFunc"
20
21func WithCallOverride(ctx netcontext.Context, f CallOverrideFunc) netcontext.Context {
22	// We avoid appending to any existing call override
23	// so we don't risk overwriting a popped stack below.
24	var cofs []CallOverrideFunc
25	if uf, ok := ctx.Value(&callOverrideKey).([]CallOverrideFunc); ok {
26		cofs = append(cofs, uf...)
27	}
28	cofs = append(cofs, f)
29	return netcontext.WithValue(ctx, &callOverrideKey, cofs)
30}
31
32func callOverrideFromContext(ctx netcontext.Context) (CallOverrideFunc, netcontext.Context, bool) {
33	cofs, _ := ctx.Value(&callOverrideKey).([]CallOverrideFunc)
34	if len(cofs) == 0 {
35		return nil, nil, false
36	}
37	// We found a list of overrides; grab the last, and reconstitute a
38	// context that will hide it.
39	f := cofs[len(cofs)-1]
40	ctx = netcontext.WithValue(ctx, &callOverrideKey, cofs[:len(cofs)-1])
41	return f, ctx, true
42}
43
44type logOverrideFunc func(level int64, format string, args ...interface{})
45
46var logOverrideKey = "holds a logOverrideFunc"
47
48func WithLogOverride(ctx netcontext.Context, f logOverrideFunc) netcontext.Context {
49	return netcontext.WithValue(ctx, &logOverrideKey, f)
50}
51
52var appIDOverrideKey = "holds a string, being the full app ID"
53
54func WithAppIDOverride(ctx netcontext.Context, appID string) netcontext.Context {
55	return netcontext.WithValue(ctx, &appIDOverrideKey, appID)
56}
57
58var namespaceKey = "holds the namespace string"
59
60func withNamespace(ctx netcontext.Context, ns string) netcontext.Context {
61	return netcontext.WithValue(ctx, &namespaceKey, ns)
62}
63
64func NamespaceFromContext(ctx netcontext.Context) string {
65	// If there's no namespace, return the empty string.
66	ns, _ := ctx.Value(&namespaceKey).(string)
67	return ns
68}
69
70// FullyQualifiedAppID returns the fully-qualified application ID.
71// This may contain a partition prefix (e.g. "s~" for High Replication apps),
72// or a domain prefix (e.g. "example.com:").
73func FullyQualifiedAppID(ctx netcontext.Context) string {
74	if id, ok := ctx.Value(&appIDOverrideKey).(string); ok {
75		return id
76	}
77	return fullyQualifiedAppID(ctx)
78}
79
80func Logf(ctx netcontext.Context, level int64, format string, args ...interface{}) {
81	if f, ok := ctx.Value(&logOverrideKey).(logOverrideFunc); ok {
82		f(level, format, args...)
83		return
84	}
85	c := fromContext(ctx)
86	if c == nil {
87		panic(errNotAppEngineContext)
88	}
89	logf(c, level, format, args...)
90}
91
92// NamespacedContext wraps a Context to support namespaces.
93func NamespacedContext(ctx netcontext.Context, namespace string) netcontext.Context {
94	return withNamespace(ctx, namespace)
95}
96
97// SetTestEnv sets the env variables for testing background ticket in Flex.
98func SetTestEnv() func() {
99	var environ = []struct {
100		key, value string
101	}{
102		{"GAE_LONG_APP_ID", "my-app-id"},
103		{"GAE_MINOR_VERSION", "067924799508853122"},
104		{"GAE_MODULE_INSTANCE", "0"},
105		{"GAE_MODULE_NAME", "default"},
106		{"GAE_MODULE_VERSION", "20150612t184001"},
107	}
108
109	for _, v := range environ {
110		old := os.Getenv(v.key)
111		os.Setenv(v.key, v.value)
112		v.value = old
113	}
114	return func() { // Restore old environment after the test completes.
115		for _, v := range environ {
116			if v.value == "" {
117				os.Unsetenv(v.key)
118				continue
119			}
120			os.Setenv(v.key, v.value)
121		}
122	}
123}
124