<commit_before>package adaptor

import (
	"fmt"
	"reflect"

	"github.com/compose/transporter/pkg/pipe"
)

var (
	// Adaptors is a registry of adaptor types, constructors and configs
	Adaptors = make(Registry)
)

func init() {
	Register("mongo", "a mongodb adaptor that functions as both a source and a sink", NewMongodb, MongodbConfig{})
	Register("file", "an adaptor that reads / writes files", NewFile, FileConfig{})
	Register("elasticsearch", "an elasticsearch sink adaptor", NewElasticsearch, dbConfig{})
	Register("influx", "an InfluxDB sink adaptor", NewInfluxdb, dbConfig{})
	Register("transformer", "an adaptor that transforms documents using a javascript function", NewTransformer, TransformerConfig{})
}

// Register registers an adaptor (database adaptor) for use with Transporter
// The second argument, fn, is a constructor that returns an instance of the
// given adaptor, config is an instance of the adaptor's config struct
func Register(name, desc string, fn func(*pipe.Pipe, string, Config) (StopStartListener, error), config interface{}) {
	Adaptors[name] = RegistryEntry{
		Name:        name,
		Description: desc,
		Constructor: fn,
		Config:      config,
	}
}

// Registry maps the adaptor's name to the RegistryEntry
type Registry map[string]RegistryEntry

// RegistryEntry stores the adaptor constructor and configuration struct
type RegistryEntry struct {
	Name        string
	Description string
	Constructor func(*pipe.Pipe, string, Config) (StopStartListener, error)
	Config      interface{}
}

// About inspects the  RegistryEntry's Config object, and uses
// each field's tags as a docstring
func (r *RegistryEntry) About() string {
	doc := fmt.Sprintf("%s %s\n\n", r.Name, r.Description)
	t := reflect.TypeOf(r.Config)
	doc += fmt.Sprintf("%-15s %-10s %s\n", "name", "type", "description")
	for i := 0; i < t.NumField(); i++ {
		f := t.Field(i)
		doc += fmt.Sprintf("%-15s %-10s %s\n", f.Tag.Get("json"), f.Type, f.Tag.Get("doc"))
	}

	return doc
}
<commit_msg>enable rethinkdb adaptor<commit_after>package adaptor

import (
	"fmt"
	"reflect"

	"github.com/compose/transporter/pkg/pipe"
)

var (
	// Adaptors is a registry of adaptor types, constructors and configs
	Adaptors = make(Registry)
)

func init() {
	Register("mongo", "a mongodb adaptor that functions as both a source and a sink", NewMongodb, MongodbConfig{})
	Register("file", "an adaptor that reads / writes files", NewFile, FileConfig{})
	Register("elasticsearch", "an elasticsearch sink adaptor", NewElasticsearch, dbConfig{})
	Register("influx", "an InfluxDB sink adaptor", NewInfluxdb, dbConfig{})
	Register("transformer", "an adaptor that transforms documents using a javascript function", NewTransformer, TransformerConfig{})
	Register("rethinkdb", "a rethinkdb sink adaptor", NewRethinkdb, dbConfig{})
}

// Register registers an adaptor (database adaptor) for use with Transporter
// The second argument, fn, is a constructor that returns an instance of the
// given adaptor, config is an instance of the adaptor's config struct
func Register(name, desc string, fn func(*pipe.Pipe, string, Config) (StopStartListener, error), config interface{}) {
	Adaptors[name] = RegistryEntry{
		Name:        name,
		Description: desc,
		Constructor: fn,
		Config:      config,
	}
}

// Registry maps the adaptor's name to the RegistryEntry
type Registry map[string]RegistryEntry

// RegistryEntry stores the adaptor constructor and configuration struct
type RegistryEntry struct {
	Name        string
	Description string
	Constructor func(*pipe.Pipe, string, Config) (StopStartListener, error)
	Config      interface{}
}

// About inspects the  RegistryEntry's Config object, and uses
// each field's tags as a docstring
func (r *RegistryEntry) About() string {
	doc := fmt.Sprintf("%s %s\n\n", r.Name, r.Description)
	t := reflect.TypeOf(r.Config)
	doc += fmt.Sprintf("%-15s %-10s %s\n", "name", "type", "description")
	for i := 0; i < t.NumField(); i++ {
		f := t.Field(i)
		doc += fmt.Sprintf("%-15s %-10s %s\n", f.Tag.Get("json"), f.Type, f.Tag.Get("doc"))
	}

	return doc
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes 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 admission

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"

	"k8s.io/klog"
	"sigs.k8s.io/yaml"

	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/serializer"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apiserver/pkg/apis/apiserver"
	apiserverv1alpha1 "k8s.io/apiserver/pkg/apis/apiserver/v1alpha1"
)

func makeAbs(path, base string) (string, error) {
	if filepath.IsAbs(path) {
		return path, nil
	}
	if len(base) == 0 || base == "." {
		cwd, err := os.Getwd()
		if err != nil {
			return "", err
		}
		base = cwd
	}
	return filepath.Join(base, path), nil
}

// ReadAdmissionConfiguration reads the admission configuration at the specified path.
// It returns the loaded admission configuration if the input file aligns with the required syntax.
// If it does not align with the provided syntax, it returns a default configuration for the enumerated
// set of pluginNames whose config location references the specified configFilePath.
// It does this to preserve backward compatibility when admission control files were opaque.
// It returns an error if the file did not exist.
func ReadAdmissionConfiguration(pluginNames []string, configFilePath string, configScheme *runtime.Scheme) (ConfigProvider, error) {
	if configFilePath == "" {
		return configProvider{config: &apiserver.AdmissionConfiguration{}}, nil
	}
	// a file was provided, so we just read it.
	data, err := ioutil.ReadFile(configFilePath)
	if err != nil {
		return nil, fmt.Errorf("unable to read admission control configuration from %q [%v]", configFilePath, err)
	}
	codecs := serializer.NewCodecFactory(configScheme)
	decoder := codecs.UniversalDecoder()
	decodedObj, err := runtime.Decode(decoder, data)
	// we were able to decode the file successfully
	if err == nil {
		decodedConfig, ok := decodedObj.(*apiserver.AdmissionConfiguration)
		if !ok {
			return nil, fmt.Errorf("unexpected type: %T", decodedObj)
		}
		baseDir := path.Dir(configFilePath)
		for i := range decodedConfig.Plugins {
			if decodedConfig.Plugins[i].Path == "" {
				continue
			}
			// we update relative file paths to absolute paths
			absPath, err := makeAbs(decodedConfig.Plugins[i].Path, baseDir)
			if err != nil {
				return nil, err
			}
			decodedConfig.Plugins[i].Path = absPath
		}
		return configProvider{
			config: decodedConfig,
			scheme: configScheme,
		}, nil
	}
	// we got an error where the decode wasn't related to a missing type
	if !(runtime.IsMissingVersion(err) || runtime.IsMissingKind(err) || runtime.IsNotRegisteredError(err)) {
		return nil, err
	}

	// Only tolerate load errors if the file appears to be one of the two legacy plugin configs
	unstructuredData := map[string]interface{}{}
	if err2 := yaml.Unmarshal(data, &unstructuredData); err2 != nil {
		return nil, err
	}
	_, isLegacyImagePolicy := unstructuredData["imagePolicy"]
	_, isLegacyPodNodeSelector := unstructuredData["podNodeSelectorPluginConfig"]
	if !isLegacyImagePolicy && !isLegacyPodNodeSelector {
		return nil, err
	}

	// convert the legacy format to the new admission control format
	// in order to preserve backwards compatibility, we set plugins that
	// previously read input from a non-versioned file configuration to the
	// current input file.
	legacyPluginsWithUnversionedConfig := sets.NewString("ImagePolicyWebhook", "PodNodeSelector")
	externalConfig := &apiserverv1alpha1.AdmissionConfiguration{}
	for _, pluginName := range pluginNames {
		if legacyPluginsWithUnversionedConfig.Has(pluginName) {
			externalConfig.Plugins = append(externalConfig.Plugins,
				apiserverv1alpha1.AdmissionPluginConfiguration{
					Name: pluginName,
					Path: configFilePath})
		}
	}
	configScheme.Default(externalConfig)
	internalConfig := &apiserver.AdmissionConfiguration{}
	if err := configScheme.Convert(externalConfig, internalConfig, nil); err != nil {
		return nil, err
	}
	return configProvider{
		config: internalConfig,
		scheme: configScheme,
	}, nil
}

type configProvider struct {
	config *apiserver.AdmissionConfiguration
	scheme *runtime.Scheme
}

// GetAdmissionPluginConfigurationFor returns a reader that holds the admission plugin configuration.
func GetAdmissionPluginConfigurationFor(pluginCfg apiserver.AdmissionPluginConfiguration) (io.Reader, error) {
	// if there is a nest object, return it directly
	if pluginCfg.Configuration != nil {
		return bytes.NewBuffer(pluginCfg.Configuration.Raw), nil
	}
	// there is nothing nested, so we delegate to path
	if pluginCfg.Path != "" {
		content, err := ioutil.ReadFile(pluginCfg.Path)
		if err != nil {
			klog.Fatalf("Couldn't open admission plugin configuration %s: %#v", pluginCfg.Path, err)
			return nil, err
		}
		return bytes.NewBuffer(content), nil
	}
	// there is no special config at all
	return nil, nil
}

// ConfigFor returns a reader for the specified plugin.
// If no specific configuration is present, we return a nil reader.
func (p configProvider) ConfigFor(pluginName string) (io.Reader, error) {
	// there is no config, so there is no potential config
	if p.config == nil {
		return nil, nil
	}
	// look for matching plugin and get configuration
	for _, pluginCfg := range p.config.Plugins {
		if pluginName != pluginCfg.Name {
			continue
		}
		pluginConfig, err := GetAdmissionPluginConfigurationFor(pluginCfg)
		if err != nil {
			return nil, err
		}
		return pluginConfig, nil
	}
	// there is no registered config that matches on plugin name.
	return nil, nil
}
<commit_msg>remove unnecessary scheme from config struct<commit_after>/*
Copyright 2017 The Kubernetes 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 admission

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"

	"k8s.io/klog"
	"sigs.k8s.io/yaml"

	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/serializer"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apiserver/pkg/apis/apiserver"
	apiserverv1alpha1 "k8s.io/apiserver/pkg/apis/apiserver/v1alpha1"
)

func makeAbs(path, base string) (string, error) {
	if filepath.IsAbs(path) {
		return path, nil
	}
	if len(base) == 0 || base == "." {
		cwd, err := os.Getwd()
		if err != nil {
			return "", err
		}
		base = cwd
	}
	return filepath.Join(base, path), nil
}

// ReadAdmissionConfiguration reads the admission configuration at the specified path.
// It returns the loaded admission configuration if the input file aligns with the required syntax.
// If it does not align with the provided syntax, it returns a default configuration for the enumerated
// set of pluginNames whose config location references the specified configFilePath.
// It does this to preserve backward compatibility when admission control files were opaque.
// It returns an error if the file did not exist.
func ReadAdmissionConfiguration(pluginNames []string, configFilePath string, configScheme *runtime.Scheme) (ConfigProvider, error) {
	if configFilePath == "" {
		return configProvider{config: &apiserver.AdmissionConfiguration{}}, nil
	}
	// a file was provided, so we just read it.
	data, err := ioutil.ReadFile(configFilePath)
	if err != nil {
		return nil, fmt.Errorf("unable to read admission control configuration from %q [%v]", configFilePath, err)
	}
	codecs := serializer.NewCodecFactory(configScheme)
	decoder := codecs.UniversalDecoder()
	decodedObj, err := runtime.Decode(decoder, data)
	// we were able to decode the file successfully
	if err == nil {
		decodedConfig, ok := decodedObj.(*apiserver.AdmissionConfiguration)
		if !ok {
			return nil, fmt.Errorf("unexpected type: %T", decodedObj)
		}
		baseDir := path.Dir(configFilePath)
		for i := range decodedConfig.Plugins {
			if decodedConfig.Plugins[i].Path == "" {
				continue
			}
			// we update relative file paths to absolute paths
			absPath, err := makeAbs(decodedConfig.Plugins[i].Path, baseDir)
			if err != nil {
				return nil, err
			}
			decodedConfig.Plugins[i].Path = absPath
		}
		return configProvider{
			config: decodedConfig,
		}, nil
	}
	// we got an error where the decode wasn't related to a missing type
	if !(runtime.IsMissingVersion(err) || runtime.IsMissingKind(err) || runtime.IsNotRegisteredError(err)) {
		return nil, err
	}

	// Only tolerate load errors if the file appears to be one of the two legacy plugin configs
	unstructuredData := map[string]interface{}{}
	if err2 := yaml.Unmarshal(data, &unstructuredData); err2 != nil {
		return nil, err
	}
	_, isLegacyImagePolicy := unstructuredData["imagePolicy"]
	_, isLegacyPodNodeSelector := unstructuredData["podNodeSelectorPluginConfig"]
	if !isLegacyImagePolicy && !isLegacyPodNodeSelector {
		return nil, err
	}

	// convert the legacy format to the new admission control format
	// in order to preserve backwards compatibility, we set plugins that
	// previously read input from a non-versioned file configuration to the
	// current input file.
	legacyPluginsWithUnversionedConfig := sets.NewString("ImagePolicyWebhook", "PodNodeSelector")
	externalConfig := &apiserverv1alpha1.AdmissionConfiguration{}
	for _, pluginName := range pluginNames {
		if legacyPluginsWithUnversionedConfig.Has(pluginName) {
			externalConfig.Plugins = append(externalConfig.Plugins,
				apiserverv1alpha1.AdmissionPluginConfiguration{
					Name: pluginName,
					Path: configFilePath})
		}
	}
	configScheme.Default(externalConfig)
	internalConfig := &apiserver.AdmissionConfiguration{}
	if err := configScheme.Convert(externalConfig, internalConfig, nil); err != nil {
		return nil, err
	}
	return configProvider{
		config: internalConfig,
	}, nil
}

type configProvider struct {
	config *apiserver.AdmissionConfiguration
}

// GetAdmissionPluginConfigurationFor returns a reader that holds the admission plugin configuration.
func GetAdmissionPluginConfigurationFor(pluginCfg apiserver.AdmissionPluginConfiguration) (io.Reader, error) {
	// if there is a nest object, return it directly
	if pluginCfg.Configuration != nil {
		return bytes.NewBuffer(pluginCfg.Configuration.Raw), nil
	}
	// there is nothing nested, so we delegate to path
	if pluginCfg.Path != "" {
		content, err := ioutil.ReadFile(pluginCfg.Path)
		if err != nil {
			klog.Fatalf("Couldn't open admission plugin configuration %s: %#v", pluginCfg.Path, err)
			return nil, err
		}
		return bytes.NewBuffer(content), nil
	}
	// there is no special config at all
	return nil, nil
}

// ConfigFor returns a reader for the specified plugin.
// If no specific configuration is present, we return a nil reader.
func (p configProvider) ConfigFor(pluginName string) (io.Reader, error) {
	// there is no config, so there is no potential config
	if p.config == nil {
		return nil, nil
	}
	// look for matching plugin and get configuration
	for _, pluginCfg := range p.config.Plugins {
		if pluginName != pluginCfg.Name {
			continue
		}
		pluginConfig, err := GetAdmissionPluginConfigurationFor(pluginCfg)
		if err != nil {
			return nil, err
		}
		return pluginConfig, nil
	}
	// there is no registered config that matches on plugin name.
	return nil, nil
}
<|endoftext|>


<commit_before>package dai

import (
	r "github.com/dancannon/gorethink"
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/db/model"
	"github.com/materials-commons/mcstore/pkg/db/schema"
)

// rProjects implements the Projects interface for RethinkDB
type rProjects struct {
	session *r.Session
}

const OwnedProjects = true
const AllProjects = false

// NewRProjects creates a new instance of rProjects.
func NewRProjects(session *r.Session) rProjects {
	return rProjects{
		session: session,
	}
}

// ByID looks up a project by the given id.
func (p rProjects) ByID(id string) (*schema.Project, error) {
	var project schema.Project
	if err := model.Projects.Qs(p.session).ByID(id, &project); err != nil {
		return nil, err
	}
	return &project, nil
}

func (p rProjects) ByName(name string, owner string) (*schema.Project, error) {
	var projects []schema.Project
	rql := r.Table("projects").GetAllByIndex("owner", owner).
		Filter(r.Row.Field("name").Eq(name))

	if err := model.Projects.Qs(p.session).Rows(rql, &projects); err != nil {
		return nil, err
	}

	numProjects := len(projects)
	switch {
	case numProjects == 0:
		return nil, nil
	case numProjects > 1:
		app.Log.Critf("Projects table corrupted, there are multiple projects for user '%s' with name '%s'", owner, name)
		return &projects[0], nil
	default:
		return &projects[0], nil
	}
}

// ForUser
func (p rProjects) ForUser(user string, ownedOnly bool) ([]schema.Project, error) {
	if ownedOnly {
		return p.projectsOwnedByUser(user)
	} else {
		return p.projectsUserHasAccessTo(user)
	}
}

// projectsOwnerByUser returns all projects owned by the user.
func (p rProjects) projectsOwnedByUser(user string) ([]schema.Project, error) {
	var projects []schema.Project
	rql := r.Table("projects").GetAllByIndex("owner", user)
	if err := model.Projects.Qs(p.session).Rows(rql, &projects); err != nil {
		return nil, err
	}
	return projects, nil
}

// projectsUserHasAccessTo returns all projects that user can access, including
// projects the user owns.
func (p rProjects) projectsUserHasAccessTo(user string) ([]schema.Project, error) {
	var projects []schema.Project
	rql := r.Table("access").GetAllByIndex("user_id", user).
		EqJoin("project_id", r.Table("projects")).Zip()
	if err := model.Projects.Qs(p.session).Rows(rql, &projects); err != nil {
		return nil, err
	}
	return projects, nil
}

// Insert creates a new project for the given owner. It also creates the directory associated
// with the project.
func (p rProjects) Insert(project *schema.Project) (*schema.Project, error) {
	if project.DataDir != "" {
		return nil, app.ErrInvalid
	}

	var (
		newProject schema.Project
		newDir     *schema.Directory
		err        error
	)

	if err = model.Projects.Qs(p.session).Insert(project, &newProject); err != nil {
		return nil, app.ErrCreate
	}

	dir := schema.NewDirectory(project.Name, project.Owner, newProject.ID, "")
	rdirs := NewRDirs(p.session)

	if newDir, err = rdirs.Insert(&dir); err != nil {
		return nil, app.ErrCreate
	}

	newProject.DataDir = newDir.ID
	if err = model.Projects.Qs(p.session).Update(newProject.ID, &newProject); err != nil {
		return &newProject, err
	}

	err = p.AddDirectories(&newProject, newDir.ID)

	return &newProject, err
}

// AddDirectories adds new directories to the project.
func (p rProjects) AddDirectories(project *schema.Project, directoryIDs ...string) error {
	var rverror error

	// Add each directory to the project2datadir table. If there are any errors,
	// remember that we saw an error, but continue on.
	for _, dirID := range directoryIDs {
		p2d := schema.Project2DataDir{
			ProjectID: project.ID,
			DataDirID: dirID,
		}
		if err := model.Projects.Qs(p.session).InsertRaw("project2datadir", p2d, nil); err != nil {
			rverror = app.ErrCreate
		}
	}

	return rverror
}

// HasDirectory checks if the given directoryID is in the given project.
func (p rProjects) HasDirectory(projectID, dirID string) bool {
	rql := model.ProjectDirs.T().GetAllByIndex("datadir_id", dirID)
	var proj2dir []schema.Project2DataDir
	if err := model.ProjectDirs.Qs(p.session).Rows(rql, &proj2dir); err != nil {
		return false
	}

	// Look for matching projectID
	for _, entry := range proj2dir {
		if entry.ProjectID == projectID {
			return true
		}
	}

	return false
}

// Get the access list for this project.
func (p rProjects) AccessList(projectID string) ([]schema.Access, error) {
	rql := r.Table("access").Filter(r.Row.Field("project_id").Eq(projectID))
	var access []schema.Access
	if err := model.Access.Qs(p.session).Rows(rql, &access); err != nil {
		return nil, err
	}
	return access, nil
}
<commit_msg>Add an access entry for the owner when creating a project.<commit_after>package dai

import (
	r "github.com/dancannon/gorethink"
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/db/model"
	"github.com/materials-commons/mcstore/pkg/db/schema"
)

// rProjects implements the Projects interface for RethinkDB
type rProjects struct {
	session *r.Session
}

const OwnedProjects = true
const AllProjects = false

// NewRProjects creates a new instance of rProjects.
func NewRProjects(session *r.Session) rProjects {
	return rProjects{
		session: session,
	}
}

// ByID looks up a project by the given id.
func (p rProjects) ByID(id string) (*schema.Project, error) {
	var project schema.Project
	if err := model.Projects.Qs(p.session).ByID(id, &project); err != nil {
		return nil, err
	}
	return &project, nil
}

func (p rProjects) ByName(name string, owner string) (*schema.Project, error) {
	var projects []schema.Project
	rql := r.Table("projects").GetAllByIndex("owner", owner).
		Filter(r.Row.Field("name").Eq(name))

	if err := model.Projects.Qs(p.session).Rows(rql, &projects); err != nil {
		return nil, err
	}

	numProjects := len(projects)
	switch {
	case numProjects == 0:
		return nil, nil
	case numProjects > 1:
		app.Log.Critf("Projects table corrupted, there are multiple projects for user '%s' with name '%s'", owner, name)
		return &projects[0], nil
	default:
		return &projects[0], nil
	}
}

// ForUser
func (p rProjects) ForUser(user string, ownedOnly bool) ([]schema.Project, error) {
	if ownedOnly {
		return p.projectsOwnedByUser(user)
	} else {
		return p.projectsUserHasAccessTo(user)
	}
}

// projectsOwnerByUser returns all projects owned by the user.
func (p rProjects) projectsOwnedByUser(user string) ([]schema.Project, error) {
	var projects []schema.Project
	rql := r.Table("projects").GetAllByIndex("owner", user)
	if err := model.Projects.Qs(p.session).Rows(rql, &projects); err != nil {
		return nil, err
	}
	return projects, nil
}

// projectsUserHasAccessTo returns all projects that user can access, including
// projects the user owns.
func (p rProjects) projectsUserHasAccessTo(user string) ([]schema.Project, error) {
	var projects []schema.Project
	rql := r.Table("access").GetAllByIndex("user_id", user).
		EqJoin("project_id", r.Table("projects")).Zip()
	if err := model.Projects.Qs(p.session).Rows(rql, &projects); err != nil {
		return nil, err
	}
	return projects, nil
}

// Insert creates a new project for the given owner. It also creates the directory associated
// with the project.
func (p rProjects) Insert(project *schema.Project) (*schema.Project, error) {
	if project.DataDir != "" {
		return nil, app.ErrInvalid
	}

	var (
		newProject schema.Project
		newDir     *schema.Directory
		err        error
	)

	if err = model.Projects.Qs(p.session).Insert(project, &newProject); err != nil {
		return nil, app.ErrCreate
	}

	// Add owner to access list
	accessEntry := schema.NewAccess(newProject.ID, newProject.Name, newProject.Owner)
	if err := model.Access.Qs(p.session).Insert(&accessEntry, nil); err != nil {
		return nil, app.ErrCreate
	}

	dir := schema.NewDirectory(project.Name, project.Owner, newProject.ID, "")
	rdirs := NewRDirs(p.session)

	if newDir, err = rdirs.Insert(&dir); err != nil {
		return nil, app.ErrCreate
	}

	newProject.DataDir = newDir.ID
	if err = model.Projects.Qs(p.session).Update(newProject.ID, &newProject); err != nil {
		return &newProject, err
	}

	err = p.AddDirectories(&newProject, newDir.ID)

	return &newProject, err
}

// AddDirectories adds new directories to the project.
func (p rProjects) AddDirectories(project *schema.Project, directoryIDs ...string) error {
	var rverror error

	// Add each directory to the project2datadir table. If there are any errors,
	// remember that we saw an error, but continue on.
	for _, dirID := range directoryIDs {
		p2d := schema.Project2DataDir{
			ProjectID: project.ID,
			DataDirID: dirID,
		}
		if err := model.Projects.Qs(p.session).InsertRaw("project2datadir", p2d, nil); err != nil {
			rverror = app.ErrCreate
		}
	}

	return rverror
}

// HasDirectory checks if the given directoryID is in the given project.
func (p rProjects) HasDirectory(projectID, dirID string) bool {
	rql := model.ProjectDirs.T().GetAllByIndex("datadir_id", dirID)
	var proj2dir []schema.Project2DataDir
	if err := model.ProjectDirs.Qs(p.session).Rows(rql, &proj2dir); err != nil {
		return false
	}

	// Look for matching projectID
	for _, entry := range proj2dir {
		if entry.ProjectID == projectID {
			return true
		}
	}

	return false
}

// Get the access list for this project.
func (p rProjects) AccessList(projectID string) ([]schema.Access, error) {
	rql := r.Table("access").Filter(r.Row.Field("project_id").Eq(projectID))
	var access []schema.Access
	if err := model.Access.Qs(p.session).Rows(rql, &access); err != nil {
		return nil, err
	}
	return access, nil
}
<|endoftext|>


<commit_before>package distributor

import (
	"compress/gzip"
	"fmt"
	"math"
	"net/http"
	"strings"
	"time"

	"github.com/cortexproject/cortex/pkg/util"
	util_log "github.com/cortexproject/cortex/pkg/util/log"
	"github.com/dustin/go-humanize"
	gokit "github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/weaveworks/common/httpgrpc"
	"github.com/weaveworks/common/user"

	"github.com/grafana/loki/pkg/loghttp"
	"github.com/grafana/loki/pkg/logproto"
	"github.com/grafana/loki/pkg/logql/unmarshal"
	unmarshal_legacy "github.com/grafana/loki/pkg/logql/unmarshal/legacy"
	lokiutil "github.com/grafana/loki/pkg/util"
)

var (
	contentType = http.CanonicalHeaderKey("Content-Type")
	contentEnc  = http.CanonicalHeaderKey("Content-Encoding")

	bytesIngested = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "loki",
		Name:      "distributor_bytes_received_total",
		Help:      "The total number of uncompressed bytes received per tenant",
	}, []string{"tenant"})
	linesIngested = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "loki",
		Name:      "distributor_lines_received_total",
		Help:      "The total number of lines received per tenant",
	}, []string{"tenant"})
)

const applicationJSON = "application/json"

// PushHandler reads a snappy-compressed proto from the HTTP body.
func (d *Distributor) PushHandler(w http.ResponseWriter, r *http.Request) {
	logger := util_log.WithContext(r.Context(), util_log.Logger)
	userID, _ := user.ExtractOrgID(r.Context())
	req, err := ParseRequest(logger, userID, r)
	if err != nil {
		if d.tenantConfigs.LogPushRequest(userID) {
			level.Debug(logger).Log(
				"msg", "push request failed",
				"code", http.StatusBadRequest,
				"err", err,
			)
		}
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	if d.tenantConfigs.LogPushRequestStreams(userID) {
		var sb strings.Builder
		for _, s := range req.Streams {
			sb.WriteString(s.Labels)
		}
		level.Debug(logger).Log(
			"msg", "push request streams",
			"streams", sb.String(),
		)
	}

	_, err = d.Push(r.Context(), req)
	if err == nil {
		if d.tenantConfigs.LogPushRequest(userID) {
			level.Debug(logger).Log(
				"msg", "push request successful",
			)
		}
		w.WriteHeader(http.StatusNoContent)
		return
	}

	resp, ok := httpgrpc.HTTPResponseFromError(err)
	if ok {
		body := string(resp.Body)
		if d.tenantConfigs.LogPushRequest(userID) {
			level.Debug(logger).Log(
				"msg", "push request failed",
				"code", resp.Code,
				"err", body,
			)
		}
		http.Error(w, body, int(resp.Code))
	} else {
		if d.tenantConfigs.LogPushRequest(userID) {
			level.Debug(logger).Log(
				"msg", "push request failed",
				"code", http.StatusInternalServerError,
				"err", err.Error(),
			)
		}
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func ParseRequest(logger gokit.Logger, userID string, r *http.Request) (*logproto.PushRequest, error) {

	var body lokiutil.SizeReader

	contentEncoding := r.Header.Get(contentEnc)
	switch contentEncoding {
	case "":
		body = lokiutil.NewSizeReader(r.Body)
	case "snappy":
		// Snappy-decoding is done by `util.ParseProtoReader(..., util.RawSnappy)` below.
		// Pass on body bytes. Note: HTTP clients do not need to set this header,
		// but they sometimes do. See #3407.
		body = lokiutil.NewSizeReader(r.Body)
	case "gzip":
		gzipReader, err := gzip.NewReader(r.Body)
		if err != nil {
			return nil, err
		}
		defer gzipReader.Close()
		body = lokiutil.NewSizeReader(gzipReader)
	default:
		return nil, fmt.Errorf("Content-Encoding %q not supported", contentEncoding)
	}

	contentType := r.Header.Get(contentType)
	var req logproto.PushRequest

	defer func() {
		var (
			entriesSize      int64
			streamLabelsSize int64
			totalEntries     int64
		)

		mostRecentEntry := time.Unix(0, 0)

		for _, s := range req.Streams {
			streamLabelsSize += int64(len(s.Labels))
			for _, e := range s.Entries {
				totalEntries++
				entriesSize += int64(len(e.Line))
				if e.Timestamp.After(mostRecentEntry) {
					mostRecentEntry = e.Timestamp
				}
			}
		}

		// incrementing tenant metrics if we have a tenant.
		if totalEntries != 0 && userID != "" {
			bytesIngested.WithLabelValues(userID).Add(float64(entriesSize))
			linesIngested.WithLabelValues(userID).Add(float64(totalEntries))
		}

		level.Debug(logger).Log(
			"msg", "push request parsed",
			"path", r.URL.Path,
			"contentType", contentType,
			"contentEncoding", contentEncoding,
			"bodySize", humanize.Bytes(uint64(body.Size())),
			"streams", len(req.Streams),
			"entries", totalEntries,
			"streamLabelsSize", humanize.Bytes(uint64(streamLabelsSize)),
			"entriesSize", humanize.Bytes(uint64(entriesSize)),
			"totalSize", humanize.Bytes(uint64(entriesSize+streamLabelsSize)),
			"mostRecentLagMs", time.Since(mostRecentEntry).Milliseconds(),
		)
	}()

	switch contentType {
	case applicationJSON:
		var err error

		// todo once https://github.com/weaveworks/common/commit/73225442af7da93ec8f6a6e2f7c8aafaee3f8840 is in Loki.
		// We can try to pass the body as bytes.buffer instead to avoid reading into another buffer.
		if loghttp.GetVersion(r.RequestURI) == loghttp.VersionV1 {
			err = unmarshal.DecodePushRequest(body, &req)
		} else {
			err = unmarshal_legacy.DecodePushRequest(body, &req)
		}

		if err != nil {
			return nil, err
		}

	default:
		// When no content-type header is set or when it is set to
		// `application/x-protobuf`: expect snappy compression.
		if err := util.ParseProtoReader(r.Context(), body, int(r.ContentLength), math.MaxInt32, &req, util.RawSnappy); err != nil {
			return nil, err
		}
	}
	return &req, nil
}
<commit_msg>Loki: Distributor log message bodySize should always reflect the compressed size (#3572)<commit_after>package distributor

import (
	"compress/gzip"
	"fmt"
	"io"
	"math"
	"net/http"
	"strings"
	"time"

	"github.com/cortexproject/cortex/pkg/util"
	util_log "github.com/cortexproject/cortex/pkg/util/log"
	"github.com/dustin/go-humanize"
	gokit "github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/weaveworks/common/httpgrpc"
	"github.com/weaveworks/common/user"

	"github.com/grafana/loki/pkg/loghttp"
	"github.com/grafana/loki/pkg/logproto"
	"github.com/grafana/loki/pkg/logql/unmarshal"
	unmarshal_legacy "github.com/grafana/loki/pkg/logql/unmarshal/legacy"
	lokiutil "github.com/grafana/loki/pkg/util"
)

var (
	contentType = http.CanonicalHeaderKey("Content-Type")
	contentEnc  = http.CanonicalHeaderKey("Content-Encoding")

	bytesIngested = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "loki",
		Name:      "distributor_bytes_received_total",
		Help:      "The total number of uncompressed bytes received per tenant",
	}, []string{"tenant"})
	linesIngested = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "loki",
		Name:      "distributor_lines_received_total",
		Help:      "The total number of lines received per tenant",
	}, []string{"tenant"})
)

const applicationJSON = "application/json"

// PushHandler reads a snappy-compressed proto from the HTTP body.
func (d *Distributor) PushHandler(w http.ResponseWriter, r *http.Request) {
	logger := util_log.WithContext(r.Context(), util_log.Logger)
	userID, _ := user.ExtractOrgID(r.Context())
	req, err := ParseRequest(logger, userID, r)
	if err != nil {
		if d.tenantConfigs.LogPushRequest(userID) {
			level.Debug(logger).Log(
				"msg", "push request failed",
				"code", http.StatusBadRequest,
				"err", err,
			)
		}
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	if d.tenantConfigs.LogPushRequestStreams(userID) {
		var sb strings.Builder
		for _, s := range req.Streams {
			sb.WriteString(s.Labels)
		}
		level.Debug(logger).Log(
			"msg", "push request streams",
			"streams", sb.String(),
		)
	}

	_, err = d.Push(r.Context(), req)
	if err == nil {
		if d.tenantConfigs.LogPushRequest(userID) {
			level.Debug(logger).Log(
				"msg", "push request successful",
			)
		}
		w.WriteHeader(http.StatusNoContent)
		return
	}

	resp, ok := httpgrpc.HTTPResponseFromError(err)
	if ok {
		body := string(resp.Body)
		if d.tenantConfigs.LogPushRequest(userID) {
			level.Debug(logger).Log(
				"msg", "push request failed",
				"code", resp.Code,
				"err", body,
			)
		}
		http.Error(w, body, int(resp.Code))
	} else {
		if d.tenantConfigs.LogPushRequest(userID) {
			level.Debug(logger).Log(
				"msg", "push request failed",
				"code", http.StatusInternalServerError,
				"err", err.Error(),
			)
		}
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func ParseRequest(logger gokit.Logger, userID string, r *http.Request) (*logproto.PushRequest, error) {

	// Body
	var body io.Reader
	// bodySize should always reflect the compressed size of the request body
	bodySize := lokiutil.NewSizeReader(r.Body)
	contentEncoding := r.Header.Get(contentEnc)
	switch contentEncoding {
	case "":
		body = bodySize
	case "snappy":
		// Snappy-decoding is done by `util.ParseProtoReader(..., util.RawSnappy)` below.
		// Pass on body bytes. Note: HTTP clients do not need to set this header,
		// but they sometimes do. See #3407.
		body = bodySize
	case "gzip":
		gzipReader, err := gzip.NewReader(bodySize)
		if err != nil {
			return nil, err
		}
		defer gzipReader.Close()
		body = gzipReader
	default:
		return nil, fmt.Errorf("Content-Encoding %q not supported", contentEncoding)
	}

	contentType := r.Header.Get(contentType)
	var req logproto.PushRequest

	defer func() {
		var (
			entriesSize      int64
			streamLabelsSize int64
			totalEntries     int64
		)

		mostRecentEntry := time.Unix(0, 0)

		for _, s := range req.Streams {
			streamLabelsSize += int64(len(s.Labels))
			for _, e := range s.Entries {
				totalEntries++
				entriesSize += int64(len(e.Line))
				if e.Timestamp.After(mostRecentEntry) {
					mostRecentEntry = e.Timestamp
				}
			}
		}

		// incrementing tenant metrics if we have a tenant.
		if totalEntries != 0 && userID != "" {
			bytesIngested.WithLabelValues(userID).Add(float64(entriesSize))
			linesIngested.WithLabelValues(userID).Add(float64(totalEntries))
		}

		level.Debug(logger).Log(
			"msg", "push request parsed",
			"path", r.URL.Path,
			"contentType", contentType,
			"contentEncoding", contentEncoding,
			"bodySize", humanize.Bytes(uint64(bodySize.Size())),
			"streams", len(req.Streams),
			"entries", totalEntries,
			"streamLabelsSize", humanize.Bytes(uint64(streamLabelsSize)),
			"entriesSize", humanize.Bytes(uint64(entriesSize)),
			"totalSize", humanize.Bytes(uint64(entriesSize+streamLabelsSize)),
			"mostRecentLagMs", time.Since(mostRecentEntry).Milliseconds(),
		)
	}()

	switch contentType {
	case applicationJSON:
		var err error

		// todo once https://github.com/weaveworks/common/commit/73225442af7da93ec8f6a6e2f7c8aafaee3f8840 is in Loki.
		// We can try to pass the body as bytes.buffer instead to avoid reading into another buffer.
		if loghttp.GetVersion(r.RequestURI) == loghttp.VersionV1 {
			err = unmarshal.DecodePushRequest(body, &req)
		} else {
			err = unmarshal_legacy.DecodePushRequest(body, &req)
		}

		if err != nil {
			return nil, err
		}

	default:
		// When no content-type header is set or when it is set to
		// `application/x-protobuf`: expect snappy compression.
		if err := util.ParseProtoReader(r.Context(), body, int(r.ContentLength), math.MaxInt32, &req, util.RawSnappy); err != nil {
			return nil, err
		}
	}
	return &req, nil
}
<|endoftext|>


<commit_before>package dsdk

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"os"
	_path "path"
	"strings"

	uuid "github.com/google/uuid"
)

var (
	EndLogging = make(chan bool)
)

type LogsUpload struct {
	Path string `json:"path,omitempty"`
}

type LogsUploadRequest struct {
	Ctxt  context.Context
	Files []string
}

func newLogsUpload(path string) *LogsUpload {
	return &LogsUpload{
		Path: _path.Join(path, "logs_upload"),
	}
}

func logsUpload(ctxt context.Context, file string) error {
	conn := GetConn(ctxt)
	tid, ok := ctxt.Value("tid").(string)
	if !ok {
		tid = "nil"
	}
	reqId := uuid.Must(uuid.NewRandom()).String()
	var err error
	if conn.apikey == "" {
		if _, err = conn.Login(ctxt); err != nil {
			return err
		}
	}
	key := conn.apikey
	conn.baseUrl.Path = _path.Join(conn.baseUrl.Path, "logs_upload")
	url := conn.baseUrl.String()

	var b bytes.Buffer
	w := multipart.NewWriter(&b)

	f, err := os.Open(file)
	if err != nil {
		return err
	}
	defer f.Close()
	err = w.WriteField("ecosystem", "kubernetes")
	if err != nil {
		return err
	}

	fw, err := w.CreateFormFile("logs.tar.gz", file)
	if err != nil {
		return err
	}
	if _, err = io.Copy(fw, f); err != nil {
		return err
	}
	// Don't forget to close the multipart writer.
	// If you don't close it, your request will be missing the terminating boundary.
	w.Close()

	// Now that you have a form, you can submit it to your handler.
	req, err := http.NewRequest(http.MethodPut, url, &b)
	if err != nil {
		return err
	}
	// Don't forget to set the content type, this will contain the boundary.
	req.Header.Set("Content-Type", w.FormDataContentType())
	req.Header.Set("Auth-Token", key)

	// Submit the request
	client := &http.Client{}
	sheaders, err := json.Marshal(req.Header)
	if err != nil {
		Log().Errorf("Couldn't stringify headers, %s", req.Header)
	}
	Log().Debugf(strings.Join([]string{"\nDatera Trace ID: %s",
		"Datera Request ID: %s",
		"Datera Request URL: %s",
		"Datera Request Method: %s",
		"Datera Request Payload: %s",
		"Datera Request Headers: %s\n"}, "\n"),
		tid, reqId, url, http.MethodPut, req.Body, sheaders)
	res, err := client.Do(req)
	if err != nil {
		return err
	}
	Log().Debugf("Status Code: %d", res.StatusCode)
	// Check the response
	if res.StatusCode != http.StatusOK {
		err = fmt.Errorf("bad status: %s", res.Status)
		bodyBytes, _ := ioutil.ReadAll(res.Body)
		Log().Error(err)
		Log().Error(string(bodyBytes))
		return err
	}
	return nil
}

func rotateLogs(rule string) error {
	if _, err := RunCmd("logrotate", "-f", rule); err != nil {
		return err
	}
	return nil
}

func (e *LogsUpload) Upload(ro *LogsUploadRequest) (*LogsUpload, *ApiErrorResponse, error) {
	return nil, nil, logsUpload(ro.Ctxt, ro.Files[0])
}

func (e *LogsUpload) RotateUploadRemove(ctxt context.Context, rule, rotated string) error {
	if err := rotateLogs(rule); err != nil {
		return err
	}

	// Determine if filtered logs exist
	lf, err := os.Open(rotated)
	if err != nil {
		return err
	}
	defer os.Remove(rotated)
	defer lf.Close()
	fstat, err := lf.Stat()
	if err != nil {
		return err
	}
	// Even a single line of logs will be greater than 100 bytes
	if fstat.Size() > 100 {
		Log().Debug("Uploading logs")
		_, apierr, err := e.Upload(&LogsUploadRequest{
			Ctxt:  ctxt,
			Files: []string{rotated},
		})
		if apierr != nil {
			Log().Errorf("%s\n", Pretty(apierr))
		}
		if err != nil {
			Log().Error(err)
		}
	} else {
		Log().Debugf("No new filtered logs detected.  Size: %d\n", fstat.Size())
	}
	return nil
}
<commit_msg>Removing payload from log upload logs<commit_after>package dsdk

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"os"
	_path "path"
	"strings"

	uuid "github.com/google/uuid"
)

var (
	EndLogging = make(chan bool)
)

type LogsUpload struct {
	Path string `json:"path,omitempty"`
}

type LogsUploadRequest struct {
	Ctxt  context.Context
	Files []string
}

func newLogsUpload(path string) *LogsUpload {
	return &LogsUpload{
		Path: _path.Join(path, "logs_upload"),
	}
}

func logsUpload(ctxt context.Context, file string) error {
	conn := GetConn(ctxt)
	tid, ok := ctxt.Value("tid").(string)
	if !ok {
		tid = "nil"
	}
	reqId := uuid.Must(uuid.NewRandom()).String()
	var err error
	if conn.apikey == "" {
		if _, err = conn.Login(ctxt); err != nil {
			return err
		}
	}
	key := conn.apikey
	conn.baseUrl.Path = _path.Join(conn.baseUrl.Path, "logs_upload")
	url := conn.baseUrl.String()

	var b bytes.Buffer
	w := multipart.NewWriter(&b)

	f, err := os.Open(file)
	if err != nil {
		return err
	}
	defer f.Close()
	err = w.WriteField("ecosystem", "kubernetes")
	if err != nil {
		return err
	}

	fw, err := w.CreateFormFile("logs.tar.gz", file)
	if err != nil {
		return err
	}
	if _, err = io.Copy(fw, f); err != nil {
		return err
	}
	// Don't forget to close the multipart writer.
	// If you don't close it, your request will be missing the terminating boundary.
	w.Close()

	// Now that you have a form, you can submit it to your handler.
	req, err := http.NewRequest(http.MethodPut, url, &b)
	if err != nil {
		return err
	}
	// Don't forget to set the content type, this will contain the boundary.
	req.Header.Set("Content-Type", w.FormDataContentType())
	req.Header.Set("Auth-Token", key)

	// Submit the request
	client := &http.Client{}
	sheaders, err := json.Marshal(req.Header)
	if err != nil {
		Log().Errorf("Couldn't stringify headers, %s", req.Header)
	}
	Log().Debugf(strings.Join([]string{"\nDatera Trace ID: %s",
		"Datera Request ID: %s",
		"Datera Request URL: %s",
		"Datera Request Method: %s",
		"Datera Request Payload: <binary stream>",
		"Datera Request Headers: %s\n"}, "\n"),
		tid, reqId, url, http.MethodPut, sheaders)
	res, err := client.Do(req)
	if err != nil {
		return err
	}
	Log().Debugf("Status Code: %d", res.StatusCode)
	// Check the response
	if res.StatusCode != http.StatusOK {
		err = fmt.Errorf("bad status: %s", res.Status)
		bodyBytes, _ := ioutil.ReadAll(res.Body)
		Log().Error(err)
		Log().Error(string(bodyBytes))
		return err
	}
	return nil
}

func rotateLogs(rule string) error {
	if _, err := RunCmd("logrotate", "-f", rule); err != nil {
		return err
	}
	return nil
}

func (e *LogsUpload) Upload(ro *LogsUploadRequest) (*LogsUpload, *ApiErrorResponse, error) {
	return nil, nil, logsUpload(ro.Ctxt, ro.Files[0])
}

func (e *LogsUpload) RotateUploadRemove(ctxt context.Context, rule, rotated string) error {
	if err := rotateLogs(rule); err != nil {
		return err
	}

	// Determine if filtered logs exist
	lf, err := os.Open(rotated)
	if err != nil {
		return err
	}
	defer os.Remove(rotated)
	defer lf.Close()
	fstat, err := lf.Stat()
	if err != nil {
		return err
	}
	// Even a single line of logs will be greater than 100 bytes
	if fstat.Size() > 100 {
		Log().Debug("Uploading logs")
		_, apierr, err := e.Upload(&LogsUploadRequest{
			Ctxt:  ctxt,
			Files: []string{rotated},
		})
		if apierr != nil {
			Log().Errorf("%s\n", Pretty(apierr))
		}
		if err != nil {
			Log().Error(err)
		}
	} else {
		Log().Debugf("No new filtered logs detected.  Size: %d\n", fstat.Size())
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes 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 cmd

import (
	"fmt"
	"io"
	"net/url"

	dockerterm "github.com/docker/docker/pkg/term"
	"github.com/spf13/cobra"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	coreclient "k8s.io/client-go/kubernetes/typed/core/v1"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/tools/remotecommand"
	api "k8s.io/kubernetes/pkg/apis/core"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/scheme"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
	"k8s.io/kubernetes/pkg/kubectl/util/term"
	"k8s.io/kubernetes/pkg/util/interrupt"
)

var (
	exec_example = templates.Examples(i18n.T(`
		# Get output from running 'date' from pod 123456-7890, using the first container by default
		kubectl exec 123456-7890 date

		# Get output from running 'date' in ruby-container from pod 123456-7890
		kubectl exec 123456-7890 -c ruby-container date

		# Switch to raw terminal mode, sends stdin to 'bash' in ruby-container from pod 123456-7890
		# and sends stdout/stderr from 'bash' back to the client
		kubectl exec 123456-7890 -c ruby-container -i -t -- bash -il

		# List contents of /usr from the first container of pod 123456-7890 and sort by modification time.
		# If the command you want to execute in the pod has any flags in common (e.g. -i),
		# you must use two dashes (--) to separate your command's flags/arguments.
		# Also note, do not surround your command and its flags/arguments with quotes
		# unless that is how you would execute it normally (i.e., do ls -t /usr, not "ls -t /usr").
		kubectl exec 123456-7890 -i -t -- ls -t /usr
		`))
)

const (
	execUsageStr = "expected 'exec POD_NAME COMMAND [ARG1] [ARG2] ... [ARGN]'.\nPOD_NAME and COMMAND are required arguments for the exec command"
)

func NewCmdExec(f cmdutil.Factory, streams genericclioptions.IOStreams) *cobra.Command {
	options := &ExecOptions{
		StreamOptions: StreamOptions{
			IOStreams: streams,
		},

		Executor: &DefaultRemoteExecutor{},
	}
	cmd := &cobra.Command{
		Use: "exec POD [-c CONTAINER] -- COMMAND [args...]",
		DisableFlagsInUseLine: true,
		Short:   i18n.T("Execute a command in a container"),
		Long:    "Execute a command in a container.",
		Example: exec_example,
		Run: func(cmd *cobra.Command, args []string) {
			argsLenAtDash := cmd.ArgsLenAtDash()
			cmdutil.CheckErr(options.Complete(f, cmd, args, argsLenAtDash))
			cmdutil.CheckErr(options.Validate())
			cmdutil.CheckErr(options.Run())
		},
	}
	cmd.Flags().StringVarP(&options.PodName, "pod", "p", options.PodName, "Pod name")
	cmd.Flags().MarkDeprecated("pod", "This flag is deprecated and will be removed in future. Use exec POD_NAME instead.")
	// TODO support UID
	cmd.Flags().StringVarP(&options.ContainerName, "container", "c", options.ContainerName, "Container name. If omitted, the first container in the pod will be chosen")
	cmd.Flags().BoolVarP(&options.Stdin, "stdin", "i", options.Stdin, "Pass stdin to the container")
	cmd.Flags().BoolVarP(&options.TTY, "tty", "t", options.TTY, "Stdin is a TTY")
	return cmd
}

// RemoteExecutor defines the interface accepted by the Exec command - provided for test stubbing
type RemoteExecutor interface {
	Execute(method string, url *url.URL, config *restclient.Config, stdin io.Reader, stdout, stderr io.Writer, tty bool, terminalSizeQueue remotecommand.TerminalSizeQueue) error
}

// DefaultRemoteExecutor is the standard implementation of remote command execution
type DefaultRemoteExecutor struct{}

func (*DefaultRemoteExecutor) Execute(method string, url *url.URL, config *restclient.Config, stdin io.Reader, stdout, stderr io.Writer, tty bool, terminalSizeQueue remotecommand.TerminalSizeQueue) error {
	exec, err := remotecommand.NewSPDYExecutor(config, method, url)
	if err != nil {
		return err
	}
	return exec.Stream(remotecommand.StreamOptions{
		Stdin:             stdin,
		Stdout:            stdout,
		Stderr:            stderr,
		Tty:               tty,
		TerminalSizeQueue: terminalSizeQueue,
	})
}

type StreamOptions struct {
	Namespace     string
	PodName       string
	ContainerName string
	Stdin         bool
	TTY           bool
	// minimize unnecessary output
	Quiet bool
	// InterruptParent, if set, is used to handle interrupts while attached
	InterruptParent *interrupt.Handler

	genericclioptions.IOStreams

	// for testing
	overrideStreams func() (io.ReadCloser, io.Writer, io.Writer)
	isTerminalIn    func(t term.TTY) bool
}

// ExecOptions declare the arguments accepted by the Exec command
type ExecOptions struct {
	StreamOptions

	Command []string

	FullCmdName       string
	SuggestedCmdUsage string

	Executor  RemoteExecutor
	PodClient coreclient.PodsGetter
	Config    *restclient.Config
}

// Complete verifies command line arguments and loads data from the command environment
func (p *ExecOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, argsIn []string, argsLenAtDash int) error {
	// Let kubectl exec follow rules for `--`, see #13004 issue
	if len(p.PodName) == 0 && (len(argsIn) == 0 || argsLenAtDash == 0) {
		return cmdutil.UsageErrorf(cmd, execUsageStr)
	}
	if len(p.PodName) != 0 {
		if len(argsIn) < 1 {
			return cmdutil.UsageErrorf(cmd, execUsageStr)
		}
		p.Command = argsIn
	} else {
		p.PodName = argsIn[0]
		p.Command = argsIn[1:]
		if len(p.Command) < 1 {
			return cmdutil.UsageErrorf(cmd, execUsageStr)
		}
	}

	namespace, _, err := f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}
	p.Namespace = namespace

	cmdParent := cmd.Parent()
	if cmdParent != nil {
		p.FullCmdName = cmdParent.CommandPath()
	}
	if len(p.FullCmdName) > 0 && cmdutil.IsSiblingCommandExists(cmd, "describe") {
		p.SuggestedCmdUsage = fmt.Sprintf("Use '%s describe pod/%s -n %s' to see all of the containers in this pod.", p.FullCmdName, p.PodName, p.Namespace)
	}

	config, err := f.ToRESTConfig()
	if err != nil {
		return err
	}
	p.Config = config

	clientset, err := f.KubernetesClientSet()
	if err != nil {
		return err
	}
	p.PodClient = clientset.CoreV1()

	return nil
}

// Validate checks that the provided exec options are specified.
func (p *ExecOptions) Validate() error {
	if len(p.PodName) == 0 {
		return fmt.Errorf("pod name must be specified")
	}
	if len(p.Command) == 0 {
		return fmt.Errorf("you must specify at least one command for the container")
	}
	if p.Out == nil || p.ErrOut == nil {
		return fmt.Errorf("both output and error output must be provided")
	}
	if p.Executor == nil || p.PodClient == nil || p.Config == nil {
		return fmt.Errorf("client, client config, and executor must be provided")
	}
	return nil
}

func (o *StreamOptions) setupTTY() term.TTY {
	t := term.TTY{
		Parent: o.InterruptParent,
		Out:    o.Out,
	}

	if !o.Stdin {
		// need to nil out o.In to make sure we don't create a stream for stdin
		o.In = nil
		o.TTY = false
		return t
	}

	t.In = o.In
	if !o.TTY {
		return t
	}

	if o.isTerminalIn == nil {
		o.isTerminalIn = func(tty term.TTY) bool {
			return tty.IsTerminalIn()
		}
	}
	if !o.isTerminalIn(t) {
		o.TTY = false

		if o.ErrOut != nil {
			fmt.Fprintln(o.ErrOut, "Unable to use a TTY - input is not a terminal or the right kind of file")
		}

		return t
	}

	// if we get to here, the user wants to attach stdin, wants a TTY, and o.In is a terminal, so we
	// can safely set t.Raw to true
	t.Raw = true

	if o.overrideStreams == nil {
		// use dockerterm.StdStreams() to get the right I/O handles on Windows
		o.overrideStreams = dockerterm.StdStreams
	}
	stdin, stdout, _ := o.overrideStreams()
	o.In = stdin
	t.In = stdin
	if o.Out != nil {
		o.Out = stdout
		t.Out = stdout
	}

	return t
}

// Run executes a validated remote execution against a pod.
func (p *ExecOptions) Run() error {
	pod, err := p.PodClient.Pods(p.Namespace).Get(p.PodName, metav1.GetOptions{})
	if err != nil {
		return err
	}

	if pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed {
		return fmt.Errorf("cannot exec into a container in a completed pod; current phase is %s", pod.Status.Phase)
	}

	containerName := p.ContainerName
	if len(containerName) == 0 {
		if len(pod.Spec.Containers) > 1 {
			usageString := fmt.Sprintf("Defaulting container name to %s.", pod.Spec.Containers[0].Name)
			if len(p.SuggestedCmdUsage) > 0 {
				usageString = fmt.Sprintf("%s\n%s", usageString, p.SuggestedCmdUsage)
			}
			fmt.Fprintf(p.ErrOut, "%s\n", usageString)
		}
		containerName = pod.Spec.Containers[0].Name
	}

	// ensure we can recover the terminal while attached
	t := p.setupTTY()

	var sizeQueue remotecommand.TerminalSizeQueue
	if t.Raw {
		// this call spawns a goroutine to monitor/update the terminal size
		sizeQueue = t.MonitorSize(t.GetSize())

		// unset p.Err if it was previously set because both stdout and stderr go over p.Out when tty is
		// true
		p.ErrOut = nil
	}

	fn := func() error {
		restClient, err := restclient.RESTClientFor(p.Config)
		if err != nil {
			return err
		}

		// TODO: consider abstracting into a client invocation or client helper
		req := restClient.Post().
			Resource("pods").
			Name(pod.Name).
			Namespace(pod.Namespace).
			SubResource("exec").
			Param("container", containerName)
		req.VersionedParams(&api.PodExecOptions{
			Container: containerName,
			Command:   p.Command,
			Stdin:     p.Stdin,
			Stdout:    p.Out != nil,
			Stderr:    p.ErrOut != nil,
			TTY:       t.Raw,
		}, scheme.ParameterCodec)

		return p.Executor.Execute("POST", req.URL(), p.Config, p.In, p.Out, p.ErrOut, t.Raw, sizeQueue)
	}

	if err := t.Safe(fn); err != nil {
		return err
	}

	return nil
}
<commit_msg>Updated PodExecOptions to external version<commit_after>/*
Copyright 2014 The Kubernetes 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 cmd

import (
	"fmt"
	"io"
	"net/url"

	dockerterm "github.com/docker/docker/pkg/term"
	"github.com/spf13/cobra"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	coreclient "k8s.io/client-go/kubernetes/typed/core/v1"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/tools/remotecommand"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/scheme"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
	"k8s.io/kubernetes/pkg/kubectl/util/term"
	"k8s.io/kubernetes/pkg/util/interrupt"
)

var (
	exec_example = templates.Examples(i18n.T(`
		# Get output from running 'date' from pod 123456-7890, using the first container by default
		kubectl exec 123456-7890 date

		# Get output from running 'date' in ruby-container from pod 123456-7890
		kubectl exec 123456-7890 -c ruby-container date

		# Switch to raw terminal mode, sends stdin to 'bash' in ruby-container from pod 123456-7890
		# and sends stdout/stderr from 'bash' back to the client
		kubectl exec 123456-7890 -c ruby-container -i -t -- bash -il

		# List contents of /usr from the first container of pod 123456-7890 and sort by modification time.
		# If the command you want to execute in the pod has any flags in common (e.g. -i),
		# you must use two dashes (--) to separate your command's flags/arguments.
		# Also note, do not surround your command and its flags/arguments with quotes
		# unless that is how you would execute it normally (i.e., do ls -t /usr, not "ls -t /usr").
		kubectl exec 123456-7890 -i -t -- ls -t /usr
		`))
)

const (
	execUsageStr = "expected 'exec POD_NAME COMMAND [ARG1] [ARG2] ... [ARGN]'.\nPOD_NAME and COMMAND are required arguments for the exec command"
)

func NewCmdExec(f cmdutil.Factory, streams genericclioptions.IOStreams) *cobra.Command {
	options := &ExecOptions{
		StreamOptions: StreamOptions{
			IOStreams: streams,
		},

		Executor: &DefaultRemoteExecutor{},
	}
	cmd := &cobra.Command{
		Use: "exec POD [-c CONTAINER] -- COMMAND [args...]",
		DisableFlagsInUseLine: true,
		Short:   i18n.T("Execute a command in a container"),
		Long:    "Execute a command in a container.",
		Example: exec_example,
		Run: func(cmd *cobra.Command, args []string) {
			argsLenAtDash := cmd.ArgsLenAtDash()
			cmdutil.CheckErr(options.Complete(f, cmd, args, argsLenAtDash))
			cmdutil.CheckErr(options.Validate())
			cmdutil.CheckErr(options.Run())
		},
	}
	cmd.Flags().StringVarP(&options.PodName, "pod", "p", options.PodName, "Pod name")
	cmd.Flags().MarkDeprecated("pod", "This flag is deprecated and will be removed in future. Use exec POD_NAME instead.")
	// TODO support UID
	cmd.Flags().StringVarP(&options.ContainerName, "container", "c", options.ContainerName, "Container name. If omitted, the first container in the pod will be chosen")
	cmd.Flags().BoolVarP(&options.Stdin, "stdin", "i", options.Stdin, "Pass stdin to the container")
	cmd.Flags().BoolVarP(&options.TTY, "tty", "t", options.TTY, "Stdin is a TTY")
	return cmd
}

// RemoteExecutor defines the interface accepted by the Exec command - provided for test stubbing
type RemoteExecutor interface {
	Execute(method string, url *url.URL, config *restclient.Config, stdin io.Reader, stdout, stderr io.Writer, tty bool, terminalSizeQueue remotecommand.TerminalSizeQueue) error
}

// DefaultRemoteExecutor is the standard implementation of remote command execution
type DefaultRemoteExecutor struct{}

func (*DefaultRemoteExecutor) Execute(method string, url *url.URL, config *restclient.Config, stdin io.Reader, stdout, stderr io.Writer, tty bool, terminalSizeQueue remotecommand.TerminalSizeQueue) error {
	exec, err := remotecommand.NewSPDYExecutor(config, method, url)
	if err != nil {
		return err
	}
	return exec.Stream(remotecommand.StreamOptions{
		Stdin:             stdin,
		Stdout:            stdout,
		Stderr:            stderr,
		Tty:               tty,
		TerminalSizeQueue: terminalSizeQueue,
	})
}

type StreamOptions struct {
	Namespace     string
	PodName       string
	ContainerName string
	Stdin         bool
	TTY           bool
	// minimize unnecessary output
	Quiet bool
	// InterruptParent, if set, is used to handle interrupts while attached
	InterruptParent *interrupt.Handler

	genericclioptions.IOStreams

	// for testing
	overrideStreams func() (io.ReadCloser, io.Writer, io.Writer)
	isTerminalIn    func(t term.TTY) bool
}

// ExecOptions declare the arguments accepted by the Exec command
type ExecOptions struct {
	StreamOptions

	Command []string

	FullCmdName       string
	SuggestedCmdUsage string

	Executor  RemoteExecutor
	PodClient coreclient.PodsGetter
	Config    *restclient.Config
}

// Complete verifies command line arguments and loads data from the command environment
func (p *ExecOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, argsIn []string, argsLenAtDash int) error {
	// Let kubectl exec follow rules for `--`, see #13004 issue
	if len(p.PodName) == 0 && (len(argsIn) == 0 || argsLenAtDash == 0) {
		return cmdutil.UsageErrorf(cmd, execUsageStr)
	}
	if len(p.PodName) != 0 {
		if len(argsIn) < 1 {
			return cmdutil.UsageErrorf(cmd, execUsageStr)
		}
		p.Command = argsIn
	} else {
		p.PodName = argsIn[0]
		p.Command = argsIn[1:]
		if len(p.Command) < 1 {
			return cmdutil.UsageErrorf(cmd, execUsageStr)
		}
	}

	namespace, _, err := f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}
	p.Namespace = namespace

	cmdParent := cmd.Parent()
	if cmdParent != nil {
		p.FullCmdName = cmdParent.CommandPath()
	}
	if len(p.FullCmdName) > 0 && cmdutil.IsSiblingCommandExists(cmd, "describe") {
		p.SuggestedCmdUsage = fmt.Sprintf("Use '%s describe pod/%s -n %s' to see all of the containers in this pod.", p.FullCmdName, p.PodName, p.Namespace)
	}

	config, err := f.ToRESTConfig()
	if err != nil {
		return err
	}
	p.Config = config

	clientset, err := f.KubernetesClientSet()
	if err != nil {
		return err
	}
	p.PodClient = clientset.CoreV1()

	return nil
}

// Validate checks that the provided exec options are specified.
func (p *ExecOptions) Validate() error {
	if len(p.PodName) == 0 {
		return fmt.Errorf("pod name must be specified")
	}
	if len(p.Command) == 0 {
		return fmt.Errorf("you must specify at least one command for the container")
	}
	if p.Out == nil || p.ErrOut == nil {
		return fmt.Errorf("both output and error output must be provided")
	}
	if p.Executor == nil || p.PodClient == nil || p.Config == nil {
		return fmt.Errorf("client, client config, and executor must be provided")
	}
	return nil
}

func (o *StreamOptions) setupTTY() term.TTY {
	t := term.TTY{
		Parent: o.InterruptParent,
		Out:    o.Out,
	}

	if !o.Stdin {
		// need to nil out o.In to make sure we don't create a stream for stdin
		o.In = nil
		o.TTY = false
		return t
	}

	t.In = o.In
	if !o.TTY {
		return t
	}

	if o.isTerminalIn == nil {
		o.isTerminalIn = func(tty term.TTY) bool {
			return tty.IsTerminalIn()
		}
	}
	if !o.isTerminalIn(t) {
		o.TTY = false

		if o.ErrOut != nil {
			fmt.Fprintln(o.ErrOut, "Unable to use a TTY - input is not a terminal or the right kind of file")
		}

		return t
	}

	// if we get to here, the user wants to attach stdin, wants a TTY, and o.In is a terminal, so we
	// can safely set t.Raw to true
	t.Raw = true

	if o.overrideStreams == nil {
		// use dockerterm.StdStreams() to get the right I/O handles on Windows
		o.overrideStreams = dockerterm.StdStreams
	}
	stdin, stdout, _ := o.overrideStreams()
	o.In = stdin
	t.In = stdin
	if o.Out != nil {
		o.Out = stdout
		t.Out = stdout
	}

	return t
}

// Run executes a validated remote execution against a pod.
func (p *ExecOptions) Run() error {
	pod, err := p.PodClient.Pods(p.Namespace).Get(p.PodName, metav1.GetOptions{})
	if err != nil {
		return err
	}

	if pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed {
		return fmt.Errorf("cannot exec into a container in a completed pod; current phase is %s", pod.Status.Phase)
	}

	containerName := p.ContainerName
	if len(containerName) == 0 {
		if len(pod.Spec.Containers) > 1 {
			usageString := fmt.Sprintf("Defaulting container name to %s.", pod.Spec.Containers[0].Name)
			if len(p.SuggestedCmdUsage) > 0 {
				usageString = fmt.Sprintf("%s\n%s", usageString, p.SuggestedCmdUsage)
			}
			fmt.Fprintf(p.ErrOut, "%s\n", usageString)
		}
		containerName = pod.Spec.Containers[0].Name
	}

	// ensure we can recover the terminal while attached
	t := p.setupTTY()

	var sizeQueue remotecommand.TerminalSizeQueue
	if t.Raw {
		// this call spawns a goroutine to monitor/update the terminal size
		sizeQueue = t.MonitorSize(t.GetSize())

		// unset p.Err if it was previously set because both stdout and stderr go over p.Out when tty is
		// true
		p.ErrOut = nil
	}

	fn := func() error {
		restClient, err := restclient.RESTClientFor(p.Config)
		if err != nil {
			return err
		}

		// TODO: consider abstracting into a client invocation or client helper
		req := restClient.Post().
			Resource("pods").
			Name(pod.Name).
			Namespace(pod.Namespace).
			SubResource("exec").
			Param("container", containerName)
		req.VersionedParams(&corev1.PodExecOptions{
			Container: containerName,
			Command:   p.Command,
			Stdin:     p.Stdin,
			Stdout:    p.Out != nil,
			Stderr:    p.ErrOut != nil,
			TTY:       t.Raw,
		}, scheme.ParameterCodec)

		return p.Executor.Execute("POST", req.URL(), p.Config, p.In, p.Out, p.ErrOut, t.Raw, sizeQueue)
	}

	if err := t.Safe(fn); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package shutdown

import (
	"errors"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
)

var h = newHandler()

type handler struct {
	Active bool

	mtx   sync.Mutex
	stack []func()
}

func newHandler() *handler {
	h := &handler{}
	go h.wait()
	return h
}

func IsActive() bool {
	return h.Active
}

func BeforeExit(f func()) {
	h.BeforeExit(f)
}

func (h *handler) BeforeExit(f func()) {
	h.mtx.Lock()
	h.stack = append(h.stack, f)
	h.mtx.Unlock()
}

func Fatal(v ...interface{}) {
	h.Fatal(v)
}

func (h *handler) Fatal(v ...interface{}) {
	h.exit(errors.New(fmt.Sprint(v...)))
}

func (h *handler) wait() {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt, os.Signal(syscall.SIGTERM))
	<-ch
	h.exit(nil)
}

func (h *handler) exit(err error) {
	h.mtx.Lock()
	h.Active = true
	for i := len(h.stack) - 1; i > 0; i-- {
		h.stack[i]()
	}
	if err != nil {
		log.New(os.Stderr, "", log.Lshortfile|log.Lmicroseconds).Output(3, err.Error())
		os.Exit(1)
	}
	os.Exit(0)
}
<commit_msg>pkg/shutdown: Make shutdown flag check thread-safe<commit_after>package shutdown

import (
	"errors"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"
	"sync/atomic"
	"syscall"
)

var h = newHandler()

type handler struct {
	active atomic.Value
	mtx    sync.Mutex
	stack  []func()
}

func newHandler() *handler {
	h := &handler{}
	h.active.Store(false)
	go h.wait()
	return h
}

func IsActive() bool {
	return h.active.Load().(bool)
}

func BeforeExit(f func()) {
	h.BeforeExit(f)
}

func (h *handler) BeforeExit(f func()) {
	h.mtx.Lock()
	h.stack = append(h.stack, f)
	h.mtx.Unlock()
}

func Fatal(v ...interface{}) {
	h.Fatal(v)
}

func (h *handler) Fatal(v ...interface{}) {
	h.exit(errors.New(fmt.Sprint(v...)))
}

func (h *handler) wait() {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt, os.Signal(syscall.SIGTERM))
	<-ch
	h.exit(nil)
}

func (h *handler) exit(err error) {
	h.mtx.Lock()
	h.active.Store(true)
	for i := len(h.stack) - 1; i > 0; i-- {
		h.stack[i]()
	}
	if err != nil {
		log.New(os.Stderr, "", log.Lshortfile|log.Lmicroseconds).Output(3, err.Error())
		os.Exit(1)
	}
	os.Exit(0)
}
<|endoftext|>


<commit_before>package tracking

import (
	"fmt"
	"time"

	"github.com/SeerUK/tid/pkg/errhandling"
	"github.com/SeerUK/tid/pkg/state"
	"github.com/SeerUK/tid/proto"
)

const (
	// KeyStatus is the key for the current tracking status in the store.
	KeyStatus = "status"
	// KeyEntryFmt is the formatting string for the entry keys in the store.
	KeyEntryFmt = "entry:%s"
	// KeyTimesheetDateFmt is the date formatting string for timesheet keys in the store.
	KeyTimesheetDateFmt = "sheet:2006-01-02"
	// KeyTimesheetFmt is the formatting string for the timesheet keys in the store.
	KeyTimesheetFmt = "sheet:%s"
)

// Gateway provides access to timesheet data in the database.
type Gateway struct {
	// The underlying storage to access.
	store state.Store
}

// NewGateway creates a new timesheet gateway.
func NewGateway(store state.Store) Gateway {
	return Gateway{
		store: store,
	}
}

// FindEntry attempts to find an entry with the given key.
func (g *Gateway) FindEntry(entryKey string) (*Entry, error) {
	entry := NewEntry("")

	return entry, g.store.Read(fmt.Sprintf(KeyEntryFmt, entryKey), entry.Message)
}

// FindOrCreateStatus attempts to find the current status.
func (g *Gateway) FindOrCreateStatus() (*Status, error) {
	status := NewStatus()

	err := g.store.Read(KeyStatus, status.Message)
	if err != nil && err != state.ErrNilResult {
		return nil, err
	}

	return status, nil
}

// FindTimesheet attempts to find a timesheet with the given date.
func (g *Gateway) FindOrCreateTimesheet(sheetKey string) (*Timesheet, error) {
	sheet := NewTimesheet(&proto.TrackingTimesheet{
		Key: sheetKey,
	})

	err := g.store.Read(fmt.Sprintf(KeyTimesheetFmt, sheetKey), sheet.Message)
	if err != nil && err != state.ErrNilResult {
		return nil, err
	}

	return sheet, nil
}

// FindTodaysTimesheet attempts to find the timesheet for the current date.
func (g *Gateway) FindOrCreateTodaysTimesheet() (*Timesheet, error) {
	return g.FindOrCreateTimesheet(time.Now().Local().Format(KeyTimesheetDateFmt))
}

// PersistEntry persists a given entry to the store.
func (g *Gateway) PersistEntry(entry *Entry) error {
	entryRef := &proto.TrackingEntryRef{
		Key:   entry.ShortHash(),
		Entry: entry.Hash(),
	}

	// Persisting an entry is a 2-step process, as we need to also store the short-key so we can
	// look up the long key.
	errs := errhandling.NewErrorStack()
	errs.Add(g.store.Write(fmt.Sprintf(KeyEntryFmt, entry.ShortHash()), entryRef))
	errs.Add(g.store.Write(fmt.Sprintf(KeyEntryFmt, entry.Hash()), entry.Message))

	return errs.Errors()
}

// PersistStatus persists a given status to the store.
func (g *Gateway) PersistStatus(status *Status) error {
	return g.store.Write(KeyStatus, status.Message)
}

// PersistTimesheet persists a given timesheet to the store.
func (g *Gateway) PersistTimesheet(sheet *Timesheet) error {
	return g.store.Write(fmt.Sprintf(KeyTimesheetFmt, sheet.Key()), sheet.Message)
}

// RemoveEntry attempts to remove an entry.
func (g *Gateway) RemoveEntry(entry *Entry) error {
	// Remove from status, if applicable
	status, err := g.FindOrCreateStatus()
	if err != nil {
		return err
	}

	if status.Ref().Entry == entry.Hash() {
		status.StopAndClear()
	}

	// Remove from timesheet
	sheet, err := g.FindOrCreateTimesheet(entry.Message.Timesheet)
	if err != nil {
		return err
	}

	sheet.RemoveEntry(entry)

	errs := errhandling.NewErrorStack()
	errs.Add(g.PersistStatus(status))
	errs.Add(g.PersistTimesheet(sheet))

	if !errs.Empty() {
		return errs.Errors()
	}

	// Remove entry
	errs = errhandling.NewErrorStack()
	errs.Add(g.store.Delete(fmt.Sprintf(KeyEntryFmt, entry.ShortHash())))
	errs.Add(g.store.Delete(fmt.Sprintf(KeyEntryFmt, entry.Hash())))

	return errs.Errors()
}
<commit_msg>You can now remove by short hash.<commit_after>package tracking

import (
	"fmt"
	"time"

	"github.com/SeerUK/tid/pkg/errhandling"
	"github.com/SeerUK/tid/pkg/state"
	"github.com/SeerUK/tid/proto"
)

const (
	// KeyStatus is the key for the current tracking status in the store.
	KeyStatus = "status"
	// KeyEntryFmt is the formatting string for the entry keys in the store.
	KeyEntryFmt = "entry:%s"
	// KeyTimesheetDateFmt is the date formatting string for timesheet keys in the store.
	KeyTimesheetDateFmt = "sheet:2006-01-02"
	// KeyTimesheetFmt is the formatting string for the timesheet keys in the store.
	KeyTimesheetFmt = "sheet:%s"
)

// Gateway provides access to timesheet data in the database.
type Gateway struct {
	// The underlying storage to access.
	store state.Store
}

// NewGateway creates a new timesheet gateway.
func NewGateway(store state.Store) Gateway {
	return Gateway{
		store: store,
	}
}

// FindEntry attempts to find an entry with the given key.
func (g *Gateway) FindEntry(hash string) (*Entry, error) {
	entry := NewEntry("")
	key := hash

	if len(hash) == 7 {
		longHash, err := g.FindEntryHashByShortHash(hash)
		if err != nil {
			return nil, err
		}

		key = longHash
	}

	return entry, g.store.Read(fmt.Sprintf(KeyEntryFmt, key), entry.Message)
}

// FindEntryHashByShortHash attempts to find an entry's hash by it's short hash.
func (g *Gateway) FindEntryHashByShortHash(hash string) (string, error) {
	var ref proto.TrackingEntryRef

	err := g.store.Read(fmt.Sprintf(KeyEntryFmt, hash), &ref)
	if err != nil {
		return "", err
	}

	return ref.Entry, nil
}

// FindOrCreateStatus attempts to find the current status.
func (g *Gateway) FindOrCreateStatus() (*Status, error) {
	status := NewStatus()

	err := g.store.Read(KeyStatus, status.Message)
	if err != nil && err != state.ErrNilResult {
		return nil, err
	}

	return status, nil
}

// FindTimesheet attempts to find a timesheet with the given date.
func (g *Gateway) FindOrCreateTimesheet(sheetKey string) (*Timesheet, error) {
	sheet := NewTimesheet(&proto.TrackingTimesheet{
		Key: sheetKey,
	})

	err := g.store.Read(fmt.Sprintf(KeyTimesheetFmt, sheetKey), sheet.Message)
	if err != nil && err != state.ErrNilResult {
		return nil, err
	}

	return sheet, nil
}

// FindTodaysTimesheet attempts to find the timesheet for the current date.
func (g *Gateway) FindOrCreateTodaysTimesheet() (*Timesheet, error) {
	return g.FindOrCreateTimesheet(time.Now().Local().Format(KeyTimesheetDateFmt))
}

// PersistEntry persists a given entry to the store.
func (g *Gateway) PersistEntry(entry *Entry) error {
	entryRef := &proto.TrackingEntryRef{
		Key:   entry.ShortHash(),
		Entry: entry.Hash(),
	}

	// Persisting an entry is a 2-step process, as we need to also store the short-key so we can
	// look up the long key.
	errs := errhandling.NewErrorStack()
	errs.Add(g.store.Write(fmt.Sprintf(KeyEntryFmt, entry.ShortHash()), entryRef))
	errs.Add(g.store.Write(fmt.Sprintf(KeyEntryFmt, entry.Hash()), entry.Message))

	return errs.Errors()
}

// PersistStatus persists a given status to the store.
func (g *Gateway) PersistStatus(status *Status) error {
	return g.store.Write(KeyStatus, status.Message)
}

// PersistTimesheet persists a given timesheet to the store.
func (g *Gateway) PersistTimesheet(sheet *Timesheet) error {
	return g.store.Write(fmt.Sprintf(KeyTimesheetFmt, sheet.Key()), sheet.Message)
}

// RemoveEntry attempts to remove an entry.
func (g *Gateway) RemoveEntry(entry *Entry) error {
	// Remove from status, if applicable
	status, err := g.FindOrCreateStatus()
	if err != nil {
		return err
	}

	if status.Ref().Entry == entry.Hash() {
		status.StopAndClear()
	}

	// Remove from timesheet
	sheet, err := g.FindOrCreateTimesheet(entry.Message.Timesheet)
	if err != nil {
		return err
	}

	sheet.RemoveEntry(entry)

	errs := errhandling.NewErrorStack()
	errs.Add(g.PersistStatus(status))
	errs.Add(g.PersistTimesheet(sheet))

	if !errs.Empty() {
		return errs.Errors()
	}

	// Remove entry
	errs = errhandling.NewErrorStack()
	errs.Add(g.store.Delete(fmt.Sprintf(KeyEntryFmt, entry.ShortHash())))
	errs.Add(g.store.Delete(fmt.Sprintf(KeyEntryFmt, entry.Hash())))

	return errs.Errors()
}
<|endoftext|>


<commit_before>// Copyright 2018 the LinuxBoot Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package visitors

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/linuxboot/fiano/pkg/uefi"
)

var (
	force  = flag.Bool("force", false, "force extract to non empty directory")
	remove = flag.Bool("remove", false, "remove existing directory before extracting")
)

// Extract extracts any Firmware node to DirPath
type Extract struct {
	BasePath string
	DirPath  string
	Index    *uint64
}

// extractBinary simply dumps the binary to a specified directory and filename.
// It creates the directory if it doesn't already exist, and dumps the buffer to it.
// It returns the filepath of the binary, and an error if it exists.
// This is meant as a helper function for other Extract functions.
func (v *Extract) extractBinary(buf []byte, filename string) (string, error) {
	// Create the directory if it doesn't exist
	dirPath := filepath.Join(v.BasePath, v.DirPath)
	if err := os.MkdirAll(dirPath, 0755); err != nil {
		return "", err
	}

	// Dump the binary.
	fp := filepath.Join(dirPath, filename)
	if err := ioutil.WriteFile(fp, buf, 0666); err != nil {
		// Make sure we return "" since we don't want an invalid path to be serialized out.
		return "", err
	}
	// Return only the relative path from the root of the tree
	return filepath.Join(v.DirPath, filename), nil
}

// Run wraps Visit and performs some setup and teardown tasks.
func (v *Extract) Run(f uefi.Firmware) error {
	// Optionally remove directory if it already exists.
	if *remove {
		if err := os.RemoveAll(v.BasePath); err != nil {
			return err
		}
	}

	if !*force {
		// Check that directory does not exist or is empty.
		files, err := ioutil.ReadDir(v.BasePath)
		if err == nil {
			if len(files) != 0 {
				return errors.New("existing directory not empty, use --force to override")
			}
		} else if !os.IsNotExist(err) {
			// Error was not EEXIST, we do not know what went wrong.
			return err
		}
	}

	// Create the directory if it does not exist.
	if err := os.MkdirAll(v.BasePath, 0755); err != nil {
		return err
	}

	// Reset the index
	*v.Index = 0
	if err := f.Apply(v); err != nil {
		return err
	}

	// Output summary json.
	json, err := uefi.MarshalFirmware(f)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(filepath.Join(v.BasePath, "summary.json"), json, 0666)
}

// Visit applies the Extract visitor to any Firmware type.
func (v *Extract) Visit(f uefi.Firmware) error {
	// The visitor must be cloned before modification; otherwise, the
	// sibling's values are modified.
	v2 := *v

	var err error
	switch f := f.(type) {

	case *uefi.FirmwareVolume:
		v2.DirPath = filepath.Join(v.DirPath, fmt.Sprintf("%#x", f.FVOffset))
		if len(f.Files) == 0 {
			f.ExtractPath, err = v2.extractBinary(f.Buf(), "fv.bin")
		} else {
			f.ExtractPath, err = v2.extractBinary(f.Buf()[:f.DataOffset], "fvh.bin")
		}

	case *uefi.File:
		// For files we use the GUID as the folder name.
		v2.DirPath = filepath.Join(v.DirPath, f.Header.GUID.String())
		// Crappy hack to make unique ids unique
		v2.DirPath = filepath.Join(v2.DirPath, fmt.Sprint(*v.Index))
		*v.Index++
		if len(f.Sections) == 0 && f.NVarStore == nil {
			f.ExtractPath, err = v2.extractBinary(f.Buf(), fmt.Sprintf("%v.ffs", f.Header.GUID))
		}

	case *uefi.Section:
		// For sections we use the file order as the folder name.
		v2.DirPath = filepath.Join(v.DirPath, fmt.Sprint(f.FileOrder))
		if len(f.Encapsulated) == 0 {
			f.ExtractPath, err = v2.extractBinary(f.Buf(), fmt.Sprintf("%v.sec", f.FileOrder))
		}

	case *uefi.NVar:
		// For NVar we use the GUID as the folder name the Name as file name and add the offset to links to make them unique
		v2.DirPath = filepath.Join(v.DirPath, f.GUID.String())
		if f.IsValid() {
			if f.NVarStore == nil {
				if f.Type == uefi.LinkNVarEntry {
					f.ExtractPath, err = v2.extractBinary(f.Buf()[f.DataOffset:], fmt.Sprintf("%v-%#x.bin", f.Name, f.Offset))
				} else {
					f.ExtractPath, err = v2.extractBinary(f.Buf()[f.DataOffset:], fmt.Sprintf("%v.bin", f.Name))
				}
			}
		} else {
			f.ExtractPath, err = v2.extractBinary(f.Buf(), fmt.Sprintf("%#x.nvar", f.Offset))
		}

	case *uefi.FlashDescriptor:
		v2.DirPath = filepath.Join(v.DirPath, "ifd")
		f.ExtractPath, err = v2.extractBinary(f.Buf(), "flashdescriptor.bin")

	case *uefi.BIOSRegion:
		v2.DirPath = filepath.Join(v.DirPath, "bios")
		if len(f.Elements) == 0 {
			f.ExtractPath, err = v2.extractBinary(f.Buf(), "biosregion.bin")
		}

	case *uefi.MERegion:
		v2.DirPath = filepath.Join(v.DirPath, "me")
		f.ExtractPath, err = v2.extractBinary(f.Buf(), "meregion.bin")

	case *uefi.RawRegion:
		v2.DirPath = filepath.Join(v.DirPath, f.Type().String())
		f.ExtractPath, err = v2.extractBinary(f.Buf(), fmt.Sprintf("%#x.bin", f.FlashRegion().BaseOffset()))

	case *uefi.BIOSPadding:
		v2.DirPath = filepath.Join(v.DirPath, fmt.Sprintf("biospad_%#x", f.Offset))
		f.ExtractPath, err = v2.extractBinary(f.Buf(), "pad.bin")
	}
	if err != nil {
		return err
	}

	return f.ApplyChildren(&v2)
}

func init() {
	var fileIndex uint64
	RegisterCLI("extract", "extract the files to a directory", 1, func(args []string) (uefi.Visitor, error) {
		return &Extract{
			BasePath: args[0],
			DirPath:  ".",
			Index:    &fileIndex,
		}, nil
	})
}
<commit_msg>visitors/extract: elaborate on command usage<commit_after>// Copyright 2018 the LinuxBoot Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package visitors

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/linuxboot/fiano/pkg/uefi"
)

var (
	force  = flag.Bool("force", false, "force extract to non empty directory")
	remove = flag.Bool("remove", false, "remove existing directory before extracting")
)

// Extract extracts any Firmware node to DirPath
type Extract struct {
	BasePath string
	DirPath  string
	Index    *uint64
}

// extractBinary simply dumps the binary to a specified directory and filename.
// It creates the directory if it doesn't already exist, and dumps the buffer to it.
// It returns the filepath of the binary, and an error if it exists.
// This is meant as a helper function for other Extract functions.
func (v *Extract) extractBinary(buf []byte, filename string) (string, error) {
	// Create the directory if it doesn't exist
	dirPath := filepath.Join(v.BasePath, v.DirPath)
	if err := os.MkdirAll(dirPath, 0755); err != nil {
		return "", err
	}

	// Dump the binary.
	fp := filepath.Join(dirPath, filename)
	if err := ioutil.WriteFile(fp, buf, 0666); err != nil {
		// Make sure we return "" since we don't want an invalid path to be serialized out.
		return "", err
	}
	// Return only the relative path from the root of the tree
	return filepath.Join(v.DirPath, filename), nil
}

// Run wraps Visit and performs some setup and teardown tasks.
func (v *Extract) Run(f uefi.Firmware) error {
	// Optionally remove directory if it already exists.
	if *remove {
		if err := os.RemoveAll(v.BasePath); err != nil {
			return err
		}
	}

	if !*force {
		// Check that directory does not exist or is empty.
		files, err := ioutil.ReadDir(v.BasePath)
		if err == nil {
			if len(files) != 0 {
				return errors.New("existing directory not empty, use --force to override")
			}
		} else if !os.IsNotExist(err) {
			// Error was not EEXIST, we do not know what went wrong.
			return err
		}
	}

	// Create the directory if it does not exist.
	if err := os.MkdirAll(v.BasePath, 0755); err != nil {
		return err
	}

	// Reset the index
	*v.Index = 0
	if err := f.Apply(v); err != nil {
		return err
	}

	// Output summary json.
	json, err := uefi.MarshalFirmware(f)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(filepath.Join(v.BasePath, "summary.json"), json, 0666)
}

// Visit applies the Extract visitor to any Firmware type.
func (v *Extract) Visit(f uefi.Firmware) error {
	// The visitor must be cloned before modification; otherwise, the
	// sibling's values are modified.
	v2 := *v

	var err error
	switch f := f.(type) {

	case *uefi.FirmwareVolume:
		v2.DirPath = filepath.Join(v.DirPath, fmt.Sprintf("%#x", f.FVOffset))
		if len(f.Files) == 0 {
			f.ExtractPath, err = v2.extractBinary(f.Buf(), "fv.bin")
		} else {
			f.ExtractPath, err = v2.extractBinary(f.Buf()[:f.DataOffset], "fvh.bin")
		}

	case *uefi.File:
		// For files we use the GUID as the folder name.
		v2.DirPath = filepath.Join(v.DirPath, f.Header.GUID.String())
		// Crappy hack to make unique ids unique
		v2.DirPath = filepath.Join(v2.DirPath, fmt.Sprint(*v.Index))
		*v.Index++
		if len(f.Sections) == 0 && f.NVarStore == nil {
			f.ExtractPath, err = v2.extractBinary(f.Buf(), fmt.Sprintf("%v.ffs", f.Header.GUID))
		}

	case *uefi.Section:
		// For sections we use the file order as the folder name.
		v2.DirPath = filepath.Join(v.DirPath, fmt.Sprint(f.FileOrder))
		if len(f.Encapsulated) == 0 {
			f.ExtractPath, err = v2.extractBinary(f.Buf(), fmt.Sprintf("%v.sec", f.FileOrder))
		}

	case *uefi.NVar:
		// For NVar we use the GUID as the folder name the Name as file name and add the offset to links to make them unique
		v2.DirPath = filepath.Join(v.DirPath, f.GUID.String())
		if f.IsValid() {
			if f.NVarStore == nil {
				if f.Type == uefi.LinkNVarEntry {
					f.ExtractPath, err = v2.extractBinary(f.Buf()[f.DataOffset:], fmt.Sprintf("%v-%#x.bin", f.Name, f.Offset))
				} else {
					f.ExtractPath, err = v2.extractBinary(f.Buf()[f.DataOffset:], fmt.Sprintf("%v.bin", f.Name))
				}
			}
		} else {
			f.ExtractPath, err = v2.extractBinary(f.Buf(), fmt.Sprintf("%#x.nvar", f.Offset))
		}

	case *uefi.FlashDescriptor:
		v2.DirPath = filepath.Join(v.DirPath, "ifd")
		f.ExtractPath, err = v2.extractBinary(f.Buf(), "flashdescriptor.bin")

	case *uefi.BIOSRegion:
		v2.DirPath = filepath.Join(v.DirPath, "bios")
		if len(f.Elements) == 0 {
			f.ExtractPath, err = v2.extractBinary(f.Buf(), "biosregion.bin")
		}

	case *uefi.MERegion:
		v2.DirPath = filepath.Join(v.DirPath, "me")
		f.ExtractPath, err = v2.extractBinary(f.Buf(), "meregion.bin")

	case *uefi.RawRegion:
		v2.DirPath = filepath.Join(v.DirPath, f.Type().String())
		f.ExtractPath, err = v2.extractBinary(f.Buf(), fmt.Sprintf("%#x.bin", f.FlashRegion().BaseOffset()))

	case *uefi.BIOSPadding:
		v2.DirPath = filepath.Join(v.DirPath, fmt.Sprintf("biospad_%#x", f.Offset))
		f.ExtractPath, err = v2.extractBinary(f.Buf(), "pad.bin")
	}
	if err != nil {
		return err
	}

	return f.ApplyChildren(&v2)
}

func init() {
	var fileIndex uint64
	RegisterCLI("extract", "extract dir\n extract the files to directory `dir`", 1, func(args []string) (uefi.Visitor, error) {
		return &Extract{
			BasePath: args[0],
			DirPath:  ".",
			Index:    &fileIndex,
		}, nil
	})
}
<|endoftext|>