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	"github.com/golang/protobuf/proto"
9	netcontext "golang.org/x/net/context"
10)
11
12type CallOverrideFunc func(ctx netcontext.Context, service, method string, in, out proto.Message) error
13
14var callOverrideKey = "holds []CallOverrideFunc"
15
16func WithCallOverride(ctx netcontext.Context, f CallOverrideFunc) netcontext.Context {
17	// We avoid appending to any existing call override
18	// so we don't risk overwriting a popped stack below.
19	var cofs []CallOverrideFunc
20	if uf, ok := ctx.Value(&callOverrideKey).([]CallOverrideFunc); ok {
21		cofs = append(cofs, uf...)
22	}
23	cofs = append(cofs, f)
24	return netcontext.WithValue(ctx, &callOverrideKey, cofs)
25}
26
27func callOverrideFromContext(ctx netcontext.Context) (CallOverrideFunc, netcontext.Context, bool) {
28	cofs, _ := ctx.Value(&callOverrideKey).([]CallOverrideFunc)
29	if len(cofs) == 0 {
30		return nil, nil, false
31	}
32	// We found a list of overrides; grab the last, and reconstitute a
33	// context that will hide it.
34	f := cofs[len(cofs)-1]
35	ctx = netcontext.WithValue(ctx, &callOverrideKey, cofs[:len(cofs)-1])
36	return f, ctx, true
37}
38
39type logOverrideFunc func(level int64, format string, args ...interface{})
40
41var logOverrideKey = "holds a logOverrideFunc"
42
43func WithLogOverride(ctx netcontext.Context, f logOverrideFunc) netcontext.Context {
44	return netcontext.WithValue(ctx, &logOverrideKey, f)
45}
46
47var appIDOverrideKey = "holds a string, being the full app ID"
48
49func WithAppIDOverride(ctx netcontext.Context, appID string) netcontext.Context {
50	return netcontext.WithValue(ctx, &appIDOverrideKey, appID)
51}
52
53var namespaceKey = "holds the namespace string"
54
55func withNamespace(ctx netcontext.Context, ns string) netcontext.Context {
56	return netcontext.WithValue(ctx, &namespaceKey, ns)
57}
58
59func NamespaceFromContext(ctx netcontext.Context) string {
60	// If there's no namespace, return the empty string.
61	ns, _ := ctx.Value(&namespaceKey).(string)
62	return ns
63}
64
65// FullyQualifiedAppID returns the fully-qualified application ID.
66// This may contain a partition prefix (e.g. "s~" for High Replication apps),
67// or a domain prefix (e.g. "example.com:").
68func FullyQualifiedAppID(ctx netcontext.Context) string {
69	if id, ok := ctx.Value(&appIDOverrideKey).(string); ok {
70		return id
71	}
72	return fullyQualifiedAppID(ctx)
73}
74
75func Logf(ctx netcontext.Context, level int64, format string, args ...interface{}) {
76	if f, ok := ctx.Value(&logOverrideKey).(logOverrideFunc); ok {
77		f(level, format, args...)
78		return
79	}
80	logf(fromContext(ctx), level, format, args...)
81}
82
83// NamespacedContext wraps a Context to support namespaces.
84func NamespacedContext(ctx netcontext.Context, namespace string) netcontext.Context {
85	return withNamespace(ctx, namespace)
86}
87