vendor/go.chromium.org/luci/lucictx/lucictx.go - infra/infra - Git at Google

 // 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]any{}
 	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)
 	}
 }
	// 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]any{}
	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)
	}
	}