| // Copyright 2016 The LUCI Authors. |
| // |
| // 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 |
| // |
| // http://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 lucictx implements a Go client for the protocol defined here: |
| // |
| // https://github.com/luci/luci-py/blob/master/client/LUCI_CONTEXT.md |
| // |
| // It differs from the python client in a couple ways: |
| // - The initial LUCI_CONTEXT value is captured once at application start. |
| // - Writes are cached into the golang context.Context, not a global variable. |
| // - The LUCI_CONTEXT environment variable is not changed automatically when |
| // using the Set function. To pass the new context on to a child process, |
| // you must use the Export function to dump the current context state to |
| // disk and call exported.SetInCmd(cmd) to configure new command's |
| // environment. |
| package lucictx |
| |
| import ( |
| "bytes" |
| "context" |
| "encoding/json" |
| "fmt" |
| "io" |
| "io/ioutil" |
| "os" |
| "reflect" |
| "sync" |
| |
| "github.com/golang/protobuf/jsonpb" |
| protoV1 "github.com/golang/protobuf/proto" |
| "google.golang.org/protobuf/encoding/protojson" |
| "google.golang.org/protobuf/proto" |
| |
| "go.chromium.org/luci/common/errors" |
| ) |
| |
| // EnvKey is the environment variable key for the LUCI_CONTEXT file. |
| const EnvKey = "LUCI_CONTEXT" |
| |
| // lctx is wrapper around top-level JSON dict of a LUCI_CONTEXT file. |
| // |
| // Note that we must use '*json.RawMessage' as dict value type since only |
| // pointer *json.RawMessage type implements json.Marshaler and json.Unmarshaler |
| // interfaces. Without '*' the JSON library treats json.RawMessage as []byte and |
| // marshals it as base64 blob. |
| type lctx struct { |
| sections map[string]*json.RawMessage // readonly! lives outside the lock |
| |
| lock sync.Mutex |
| path string // non-empty if exists as file on disk |
| refs int // number of open references to the dropped file |
| } |
| |
| func alloc(size int) *lctx { |
| return &lctx{sections: make(map[string]*json.RawMessage, size)} |
| } |
| |
| func (l *lctx) clone() *lctx { |
| ret := alloc(len(l.sections)) |
| for k, v := range l.sections { |
| ret.sections[k] = v |
| } |
| return ret |
| } |
| |
| var lctxKey = "Holds the current lctx" |
| |
| // This is the LUCI_CONTEXT loaded from the environment once when the process |
| // starts. |
| var externalContext = extractFromEnv(os.Stderr) |
| |
| func extractFromEnv(out io.Writer) *lctx { |
| path := os.Getenv(EnvKey) |
| if path == "" { |
| return &lctx{} |
| } |
| f, err := os.Open(path) |
| if err != nil { |
| fmt.Fprintf(out, "Could not open LUCI_CONTEXT file %q: %s\n", path, err) |
| return &lctx{} |
| } |
| defer f.Close() |
| |
| dec := json.NewDecoder(f) |
| dec.UseNumber() |
| tmp := map[string]interface{}{} |
| if err := dec.Decode(&tmp); err != nil { |
| fmt.Fprintf(out, "Could not decode LUCI_CONTEXT file %q: %s\n", path, err) |
| return &lctx{} |
| } |
| |
| ret := alloc(len(tmp)) |
| for k, v := range tmp { |
| if reflect.TypeOf(v).Kind() != reflect.Map { |
| fmt.Fprintf(out, "Could not re-encode LUCI_CONTEXT file %q, section %q: Not a map.\n", path, k) |
| continue |
| } |
| item, err := json.Marshal(v) |
| if err != nil { |
| fmt.Fprintf(out, "Could not marshal LUCI_CONTEXT %v: %s\n", v, err) |
| return &lctx{} |
| } |
| |
| // This section just came from json.Unmarshal, so we know that json.Marshal |
| // will work on it. |
| raw := json.RawMessage(item) |
| ret.sections[k] = &raw |
| } |
| |
| ret.path = path // reuse existing external file in Export() |
| ret.refs = 1 // never decremented, ensuring we don't delete the external file |
| return ret |
| } |
| |
| // Note: it never returns nil. |
| func getCurrent(ctx context.Context) *lctx { |
| if l := ctx.Value(&lctxKey); l != nil { |
| return l.(*lctx) |
| } |
| return externalContext |
| } |
| |
| // Get retrieves the current section from the current LUCI_CONTEXT, and |
| // deserializes it into out. Out may be any target for json.Unmarshal. If the |
| // section exists, it deserializes it into the provided out object. If not, then |
| // out is unmodified. |
| func Get(ctx context.Context, section string, out proto.Message) error { |
| _, err := Lookup(ctx, section, out) |
| return err |
| } |
| |
| // Lookup retrieves the current section from the current LUCI_CONTEXT, and |
| // deserializes it into out. Out may be any target for json.Unmarshal. It |
| // returns a deserialization error (if any), and a boolean indicating if the |
| // section was actually found. |
| func Lookup(ctx context.Context, section string, out proto.Message) (bool, error) { |
| data, _ := getCurrent(ctx).sections[section] |
| if data == nil { |
| return false, nil |
| } |
| unmarshaler := &jsonpb.Unmarshaler{ |
| AllowUnknownFields: true, |
| } |
| if err := unmarshaler.Unmarshal(bytes.NewReader(*data), protoV1.MessageV1(out)); err != nil { |
| return true, errors.Annotate(err, "failed to unmarshal json: %s", string(*data)).Err() |
| } |
| return true, nil |
| } |
| |
| // Set writes the json serialization of `in` as the given section into the |
| // LUCI_CONTEXT, returning the new ctx object containing it. This ctx can be |
| // passed to Export to serialize it to disk. |
| // |
| // If in is nil, it will clear that section of the LUCI_CONTEXT. |
| // |
| // The returned context is always safe to use, even if this returns an error. |
| func Set(ctx context.Context, section string, in proto.Message) context.Context { |
| var data json.RawMessage |
| if in != nil && !reflect.ValueOf(in).IsNil() { |
| buf, err := protojson.Marshal(in) |
| if err != nil { |
| panic(err) // Only errors could be from writing to buf. |
| } |
| data = buf |
| } |
| cur := getCurrent(ctx) |
| if _, alreadyHas := cur.sections[section]; data == nil && !alreadyHas { |
| // Removing a section which is already missing is a no-op |
| return ctx |
| } |
| newLctx := cur.clone() |
| if data == nil { |
| delete(newLctx.sections, section) |
| } else { |
| newLctx.sections[section] = &data |
| } |
| return context.WithValue(ctx, &lctxKey, newLctx) |
| } |
| |
| // Export takes the current LUCI_CONTEXT information from ctx, writes it to |
| // a file in os.TempDir and returns a wrapping Exported object. This exported |
| // value must then be installed into the environment of any subcommands (see |
| // the methods on Exported). |
| // |
| // It is required that the caller of this function invoke Close() on the |
| // returned Exported object, or they will leak temporary files. |
| // |
| // Internally this function reuses existing files, when possible, so if you |
| // anticipate calling a lot of subcommands with exported LUCI_CONTEXT, you can |
| // export it in advance (thus grabbing a reference to the exported file). Then |
| // subsequent Export() calls with this context will be extremely cheap, since |
| // they will just reuse the existing file. Don't forget to release it with |
| // Close() when done. |
| func Export(ctx context.Context) (Exported, error) { |
| return getCurrent(ctx).export("") |
| } |
| |
| // ExportInto is like Export, except it places the temporary file into the given |
| // directory. |
| // |
| // Exports done via this method are not reused: each individual ExportInto call |
| // produces a new temporary file. |
| func ExportInto(ctx context.Context, dir string) (Exported, error) { |
| return getCurrent(ctx).export(dir) |
| } |
| |
| func (l *lctx) export(dir string) (Exported, error) { |
| if len(l.sections) == 0 { |
| return &nullExport{}, nil |
| } |
| |
| if dir != "" { |
| path, err := dropToDisk(l.sections, dir) |
| if err != nil { |
| return nil, err |
| } |
| return &liveExport{ |
| path: path, |
| closer: func() { removeFromDisk(path) }, |
| }, nil |
| } |
| |
| l.lock.Lock() |
| defer l.lock.Unlock() |
| |
| if l.refs == 0 { |
| if l.path != "" { |
| panic("lctx.path is supposed to be empty here") |
| } |
| path, err := dropToDisk(l.sections, "") |
| if err != nil { |
| return nil, err |
| } |
| l.path = path |
| } |
| |
| l.refs++ |
| return &liveExport{ |
| path: l.path, |
| closer: func() { |
| l.lock.Lock() |
| defer l.lock.Unlock() |
| if l.refs == 0 { |
| panic("lctx.refs can't be zero here") |
| } |
| l.refs-- |
| if l.refs == 0 { |
| removeFromDisk(l.path) |
| l.path = "" |
| } |
| }, |
| }, nil |
| } |
| |
| func dropToDisk(sections map[string]*json.RawMessage, dir string) (string, error) { |
| // Note: this makes a file in 0600 mode. This is what we want, the context |
| // may have secrets. |
| f, err := ioutil.TempFile(dir, "luci_context.") |
| if err != nil { |
| return "", errors.Annotate(err, "creating luci_context file").Err() |
| } |
| |
| err = json.NewEncoder(f).Encode(sections) |
| if clErr := f.Close(); err == nil { |
| err = clErr |
| } |
| if err != nil { |
| removeFromDisk(f.Name()) |
| return "", errors.Annotate(err, "writing luci_context").Err() |
| } |
| |
| return f.Name(), nil |
| } |
| |
| func removeFromDisk(path string) { |
| if err := os.Remove(path); err != nil && !os.IsNotExist(err) { |
| fmt.Fprintf(os.Stderr, "Could not remove LUCI_CONTEXT file %q: %s\n", path, err) |
| } |
| } |