<commit_before>package clicommand

import (
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/buildkite/agent/v3/api"
	"github.com/buildkite/agent/v3/cliconfig"
	"github.com/buildkite/roko"
	"github.com/urfave/cli"
)

type OIDCTokenConfig struct {
	Audience string `cli:"audience"`
	Job      string `cli:"job"      validate:"required"`

	// Global flags
	Debug       bool     `cli:"debug"`
	LogLevel    string   `cli:"log-level"`
	NoColor     bool     `cli:"no-color"`
	Experiments []string `cli:"experiment" normalize:"list"`
	Profile     string   `cli:"profile"`

	// API config
	DebugHTTP        bool   `cli:"debug-http"`
	AgentAccessToken string `cli:"agent-access-token" validate:"required"`
	Endpoint         string `cli:"endpoint"           validate:"required"`
	NoHTTP2          bool   `cli:"no-http2"`
}

const (
	oidcTokenDescription = `Usage:

   buildkite-agent oidc request-token [options...]

Description:
   Requests and prints an OIDC token from Buildkite that claims the Job ID
   (amongst other things) and the specified audience. If no audience is
   specified, the endpoint's default audience will be claimed.

Example:
   $ buildkite-agent oidc request-token --audience sts.amazonaws.com

   Requests and prints an OIDC token from Buildkite that claims the Job ID
   (amongst other things) and the audience "sts.amazonaws.com".
`
	backoffSeconds = 2
	maxAttempts    = 5
)

var OIDCRequestTokenCommand = cli.Command{
	Name:        "request-token",
	Usage:       "Requests and prints an OIDC token from Buildkite with the specified audience,",
	Description: oidcTokenDescription,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "audience",
			Value: "",
			Usage: "The audience that will consume the OIDC token",
		},
		cli.StringFlag{
			Name:   "job",
			Value:  "",
			Usage:  "Buildkite Job Id to claim in the OIDC token",
			EnvVar: "BUILDKITE_JOB_ID",
		},

		// API Flags
		AgentAccessTokenFlag,
		EndpointFlag,
		NoHTTP2Flag,
		DebugHTTPFlag,

		// Global flags
		NoColorFlag,
		DebugFlag,
		LogLevelFlag,
		ExperimentsFlag,
		ProfileFlag,
	},
	Action: func(c *cli.Context) error {
		// The configuration will be loaded into this struct
		cfg := OIDCTokenConfig{}

		loader := cliconfig.Loader{CLI: c, Config: &cfg}
		warnings, err := loader.Load()
		if err != nil {
			fmt.Fprintf(c.App.ErrWriter, "%s\n", err)
			os.Exit(1)
		}

		l := CreateLogger(&cfg)

		// Now that we have a logger, log out the warnings that loading config generated
		for _, warning := range warnings {
			l.Warn("%s", warning)
		}

		// Setup any global configuration options
		done := HandleGlobalFlags(l, cfg)
		defer done()

		// Create the API client
		client := api.NewClient(l, loadAPIClientConfig(cfg, "AgentAccessToken"))

		// Find the meta data value
		var token *api.OIDCToken
		var resp *api.Response

		if err := roko.NewRetrier(
			roko.WithMaxAttempts(maxAttempts),
			roko.WithStrategy(roko.Exponential(backoffSeconds*time.Second, 0)),
		).Do(func(r *roko.Retrier) error {
			req := &api.OIDCTokenRequest{
				Job:      cfg.Job,
				Audience: cfg.Audience,
			}

			token, resp, err = client.OIDCToken(req)
			if resp != nil {
				switch resp.StatusCode {
				// Don't bother retrying if the response was one of these statuses
				case http.StatusBadRequest, http.StatusUnauthorized, http.StatusForbidden, http.StatusUnprocessableEntity:
					r.Break()
					return err
				}
			}

			if err != nil {
				l.Warn("%s (%s)", err, r)
			}

			return err
		}); err != nil {
			if len(cfg.Audience) > 0 {
				l.Error("Could not obtain OIDC token for audience %s", cfg.Audience)
			} else {
				l.Error("Could not obtain OIDC token for default audience")
			}
			return err
		}

		fmt.Println(token.Token)
		return nil
	},
}
<commit_msg>Fix a comment<commit_after>package clicommand

import (
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/buildkite/agent/v3/api"
	"github.com/buildkite/agent/v3/cliconfig"
	"github.com/buildkite/roko"
	"github.com/urfave/cli"
)

type OIDCTokenConfig struct {
	Audience string `cli:"audience"`
	Job      string `cli:"job"      validate:"required"`

	// Global flags
	Debug       bool     `cli:"debug"`
	LogLevel    string   `cli:"log-level"`
	NoColor     bool     `cli:"no-color"`
	Experiments []string `cli:"experiment" normalize:"list"`
	Profile     string   `cli:"profile"`

	// API config
	DebugHTTP        bool   `cli:"debug-http"`
	AgentAccessToken string `cli:"agent-access-token" validate:"required"`
	Endpoint         string `cli:"endpoint"           validate:"required"`
	NoHTTP2          bool   `cli:"no-http2"`
}

const (
	oidcTokenDescription = `Usage:

   buildkite-agent oidc request-token [options...]

Description:
   Requests and prints an OIDC token from Buildkite that claims the Job ID
   (amongst other things) and the specified audience. If no audience is
   specified, the endpoint's default audience will be claimed.

Example:
   $ buildkite-agent oidc request-token --audience sts.amazonaws.com

   Requests and prints an OIDC token from Buildkite that claims the Job ID
   (amongst other things) and the audience "sts.amazonaws.com".
`
	backoffSeconds = 2
	maxAttempts    = 5
)

var OIDCRequestTokenCommand = cli.Command{
	Name:        "request-token",
	Usage:       "Requests and prints an OIDC token from Buildkite with the specified audience,",
	Description: oidcTokenDescription,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "audience",
			Value: "",
			Usage: "The audience that will consume the OIDC token",
		},
		cli.StringFlag{
			Name:   "job",
			Value:  "",
			Usage:  "Buildkite Job Id to claim in the OIDC token",
			EnvVar: "BUILDKITE_JOB_ID",
		},

		// API Flags
		AgentAccessTokenFlag,
		EndpointFlag,
		NoHTTP2Flag,
		DebugHTTPFlag,

		// Global flags
		NoColorFlag,
		DebugFlag,
		LogLevelFlag,
		ExperimentsFlag,
		ProfileFlag,
	},
	Action: func(c *cli.Context) error {
		// The configuration will be loaded into this struct
		cfg := OIDCTokenConfig{}

		loader := cliconfig.Loader{CLI: c, Config: &cfg}
		warnings, err := loader.Load()
		if err != nil {
			fmt.Fprintf(c.App.ErrWriter, "%s\n", err)
			os.Exit(1)
		}

		l := CreateLogger(&cfg)

		// Now that we have a logger, log out the warnings that loading config generated
		for _, warning := range warnings {
			l.Warn("%s", warning)
		}

		// Setup any global configuration options
		done := HandleGlobalFlags(l, cfg)
		defer done()

		// Create the API client
		client := api.NewClient(l, loadAPIClientConfig(cfg, "AgentAccessToken"))

		// Request the token
		var token *api.OIDCToken
		var resp *api.Response

		if err := roko.NewRetrier(
			roko.WithMaxAttempts(maxAttempts),
			roko.WithStrategy(roko.Exponential(backoffSeconds*time.Second, 0)),
		).Do(func(r *roko.Retrier) error {
			req := &api.OIDCTokenRequest{
				Job:      cfg.Job,
				Audience: cfg.Audience,
			}

			token, resp, err = client.OIDCToken(req)
			if resp != nil {
				switch resp.StatusCode {
				// Don't bother retrying if the response was one of these statuses
				case http.StatusBadRequest, http.StatusUnauthorized, http.StatusForbidden, http.StatusUnprocessableEntity:
					r.Break()
					return err
				}
			}

			if err != nil {
				l.Warn("%s (%s)", err, r)
			}

			return err
		}); err != nil {
			if len(cfg.Audience) > 0 {
				l.Error("Could not obtain OIDC token for audience %s", cfg.Audience)
			} else {
				l.Error("Could not obtain OIDC token for default audience")
			}
			return err
		}

		fmt.Println(token.Token)
		return nil
	},
}
<|endoftext|>


<commit_before>// Package jsonfilelog provides the default Logger implementation for
// Docker logging. This logger logs to files on the host server in the
// JSON format.
package jsonfilelog

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/daemon/logger"
	"github.com/docker/docker/daemon/logger/loggerutils"
	"github.com/docker/docker/pkg/jsonlog"
	"github.com/docker/go-units"
)

// Name is the name of the file that the jsonlogger logs to.
const Name = "json-file"

// JSONFileLogger is Logger implementation for default Docker logging.
type JSONFileLogger struct {
	buf     *bytes.Buffer
	writer  *loggerutils.RotateFileWriter
	mu      sync.Mutex
	ctx     logger.Context
	readers map[*logger.LogWatcher]struct{} // stores the active log followers
	extra   []byte                          // json-encoded extra attributes
}

func init() {
	if err := logger.RegisterLogDriver(Name, New); err != nil {
		logrus.Fatal(err)
	}
	if err := logger.RegisterLogOptValidator(Name, ValidateLogOpt); err != nil {
		logrus.Fatal(err)
	}
}

// New creates new JSONFileLogger which writes to filename passed in
// on given context.
func New(ctx logger.Context) (logger.Logger, error) {
	var capval int64 = -1
	if capacity, ok := ctx.Config["max-size"]; ok {
		var err error
		capval, err = units.FromHumanSize(capacity)
		if err != nil {
			return nil, err
		}
	}
	var maxFiles = 1
	if maxFileString, ok := ctx.Config["max-file"]; ok {
		var err error
		maxFiles, err = strconv.Atoi(maxFileString)
		if err != nil {
			return nil, err
		}
		if maxFiles < 1 {
			return nil, fmt.Errorf("max-file cannot be less than 1")
		}
	}

	writer, err := loggerutils.NewRotateFileWriter(ctx.LogPath, capval, maxFiles)
	if err != nil {
		return nil, err
	}

	var extra []byte
	if attrs := ctx.ExtraAttributes(nil); len(attrs) > 0 {
		var err error
		extra, err = json.Marshal(attrs)
		if err != nil {
			return nil, err
		}
	}

	return &JSONFileLogger{
		buf:     bytes.NewBuffer(nil),
		writer:  writer,
		readers: make(map[*logger.LogWatcher]struct{}),
		extra:   extra,
	}, nil
}

// Log converts logger.Message to jsonlog.JSONLog and serializes it to file.
func (l *JSONFileLogger) Log(msg *logger.Message) error {
	timestamp, err := jsonlog.FastTimeMarshalJSON(msg.Timestamp)
	if err != nil {
		return err
	}
	err = (&jsonlog.JSONLogs{
		Log:      append(msg.Line, '\n'),
		Stream:   msg.Source,
		Created:  timestamp,
		RawAttrs: l.extra,
	}).MarshalJSONBuf(l.buf)
	if err != nil {
		return err
	}

	l.buf.WriteByte('\n')
	_, err = l.writer.Write(l.buf.Bytes())
	l.buf.Reset()

	return err
}

// ValidateLogOpt looks for json specific log options max-file & max-size.
func ValidateLogOpt(cfg map[string]string) error {
	for key := range cfg {
		switch key {
		case "max-file":
		case "max-size":
		case "labels":
		case "env":
		default:
			return fmt.Errorf("unknown log opt '%s' for json-file log driver", key)
		}
	}
	return nil
}

// LogPath returns the location the given json logger logs to.
func (l *JSONFileLogger) LogPath() string {
	return l.writer.LogPath()
}

// Close closes underlying file and signals all readers to stop.
func (l *JSONFileLogger) Close() error {
	l.mu.Lock()
	err := l.writer.Close()
	for r := range l.readers {
		r.Close()
		delete(l.readers, r)
	}
	l.mu.Unlock()
	return err
}

// Name returns name of this logger.
func (l *JSONFileLogger) Name() string {
	return Name
}
<commit_msg>Fix race condition in JSONFileLogger.Log<commit_after>// Package jsonfilelog provides the default Logger implementation for
// Docker logging. This logger logs to files on the host server in the
// JSON format.
package jsonfilelog

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/daemon/logger"
	"github.com/docker/docker/daemon/logger/loggerutils"
	"github.com/docker/docker/pkg/jsonlog"
	"github.com/docker/go-units"
)

// Name is the name of the file that the jsonlogger logs to.
const Name = "json-file"

// JSONFileLogger is Logger implementation for default Docker logging.
type JSONFileLogger struct {
	buf     *bytes.Buffer
	writer  *loggerutils.RotateFileWriter
	mu      sync.Mutex
	ctx     logger.Context
	readers map[*logger.LogWatcher]struct{} // stores the active log followers
	extra   []byte                          // json-encoded extra attributes
}

func init() {
	if err := logger.RegisterLogDriver(Name, New); err != nil {
		logrus.Fatal(err)
	}
	if err := logger.RegisterLogOptValidator(Name, ValidateLogOpt); err != nil {
		logrus.Fatal(err)
	}
}

// New creates new JSONFileLogger which writes to filename passed in
// on given context.
func New(ctx logger.Context) (logger.Logger, error) {
	var capval int64 = -1
	if capacity, ok := ctx.Config["max-size"]; ok {
		var err error
		capval, err = units.FromHumanSize(capacity)
		if err != nil {
			return nil, err
		}
	}
	var maxFiles = 1
	if maxFileString, ok := ctx.Config["max-file"]; ok {
		var err error
		maxFiles, err = strconv.Atoi(maxFileString)
		if err != nil {
			return nil, err
		}
		if maxFiles < 1 {
			return nil, fmt.Errorf("max-file cannot be less than 1")
		}
	}

	writer, err := loggerutils.NewRotateFileWriter(ctx.LogPath, capval, maxFiles)
	if err != nil {
		return nil, err
	}

	var extra []byte
	if attrs := ctx.ExtraAttributes(nil); len(attrs) > 0 {
		var err error
		extra, err = json.Marshal(attrs)
		if err != nil {
			return nil, err
		}
	}

	return &JSONFileLogger{
		buf:     bytes.NewBuffer(nil),
		writer:  writer,
		readers: make(map[*logger.LogWatcher]struct{}),
		extra:   extra,
	}, nil
}

// Log converts logger.Message to jsonlog.JSONLog and serializes it to file.
func (l *JSONFileLogger) Log(msg *logger.Message) error {
	timestamp, err := jsonlog.FastTimeMarshalJSON(msg.Timestamp)
	if err != nil {
		return err
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	err = (&jsonlog.JSONLogs{
		Log:      append(msg.Line, '\n'),
		Stream:   msg.Source,
		Created:  timestamp,
		RawAttrs: l.extra,
	}).MarshalJSONBuf(l.buf)
	if err != nil {
		return err
	}

	l.buf.WriteByte('\n')
	_, err = l.writer.Write(l.buf.Bytes())
	l.buf.Reset()

	return err
}

// ValidateLogOpt looks for json specific log options max-file & max-size.
func ValidateLogOpt(cfg map[string]string) error {
	for key := range cfg {
		switch key {
		case "max-file":
		case "max-size":
		case "labels":
		case "env":
		default:
			return fmt.Errorf("unknown log opt '%s' for json-file log driver", key)
		}
	}
	return nil
}

// LogPath returns the location the given json logger logs to.
func (l *JSONFileLogger) LogPath() string {
	return l.writer.LogPath()
}

// Close closes underlying file and signals all readers to stop.
func (l *JSONFileLogger) Close() error {
	l.mu.Lock()
	err := l.writer.Close()
	for r := range l.readers {
		r.Close()
		delete(l.readers, r)
	}
	l.mu.Unlock()
	return err
}

// Name returns name of this logger.
func (l *JSONFileLogger) Name() string {
	return Name
}
<|endoftext|>


<commit_before>// Copyright 2015 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 lib

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/command"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/digest"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/filemetadata"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/uploadinfo"
	"github.com/maruel/subcommands"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/auth/client/authcli"
	"go.chromium.org/luci/cipd/version"
	"go.chromium.org/luci/client/archiver"
	"go.chromium.org/luci/client/cas"
	"go.chromium.org/luci/client/internal/common"
	"go.chromium.org/luci/client/isolate"
	"go.chromium.org/luci/common/data/text/units"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/isolatedclient"
	"go.chromium.org/luci/common/system/filesystem"
)

type commonFlags struct {
	subcommands.CommandRunBase
	defaultFlags common.Flags
}

func (c *commonFlags) Init() {
	c.defaultFlags.Init(&c.Flags)
}

func (c *commonFlags) Parse() error {
	return c.defaultFlags.Parse()
}

type commonServerFlags struct {
	commonFlags
	isolatedFlags isolatedclient.Flags
	authFlags     authcli.Flags

	parsedAuthOpts auth.Options
}

func (c *commonServerFlags) Init(authOpts auth.Options) {
	c.commonFlags.Init()
	c.isolatedFlags.Init(&c.Flags)
	c.authFlags.Register(&c.Flags, authOpts)
}

func (c *commonServerFlags) Parse() error {
	var err error
	if err = c.commonFlags.Parse(); err != nil {
		return err
	}
	if err = c.isolatedFlags.Parse(); err != nil {
		return err
	}
	c.parsedAuthOpts, err = c.authFlags.Options()
	return err
}

func (c *commonServerFlags) createAuthClient(ctx context.Context) (*http.Client, error) {
	// Don't enforce authentication by using OptionalLogin mode. This is needed
	// for IP-allowed bots: they have NO credentials to send.
	return auth.NewAuthenticator(ctx, auth.OptionalLogin, c.parsedAuthOpts).Client()
}

func (c *commonServerFlags) createIsolatedClient(authCl *http.Client) (*isolatedclient.Client, error) {
	userAgent := "isolate-go/" + IsolateVersion
	if ver, err := version.GetStartupVersion(); err == nil && ver.InstanceID != "" {
		userAgent += fmt.Sprintf(" (%s@%s)", ver.PackageName, ver.InstanceID)
	}
	return c.isolatedFlags.NewClient(isolatedclient.WithAuthClient(authCl), isolatedclient.WithUserAgent(userAgent))
}

type isolateFlags struct {
	// TODO(tandrii): move ArchiveOptions from isolate pkg to here.
	isolate.ArchiveOptions
}

func (c *isolateFlags) Init(f *flag.FlagSet) {
	c.ArchiveOptions.Init()
	f.StringVar(&c.Isolate, "isolate", "", ".isolate file to load the dependency data from")
	f.StringVar(&c.Isolate, "i", "", "Alias for -isolate")
	f.StringVar(&c.IgnoredPathFilterRe, "ignored-path-filter-re", "", "A regular expression for filtering away the paths to be ignored. Note that this regexp matches ANY part of the path. So if you want to match the beginning of a path, you need to explicitly prepend ^ (same for $). Please use the Go regexp syntax. I.e. use double backslack \\\\ if you need a backslash literal.")
	f.Var(&c.ConfigVariables, "config-variable", "Config variables are used to determine which conditions should be matched when loading a .isolate file, default: [].")
	f.Var(&c.PathVariables, "path-variable", "Path variables are used to replace file paths when loading a .isolate file, default: {}")
	f.BoolVar(&c.AllowMissingFileDir, "allow-missing-file-dir", false, "If this flag is true, invalid entries in the isolated file are only logged, but won't stop it from being processed.")
}

// RequiredIsolateFlags specifies which flags are required on the command line
// being parsed.
type RequiredIsolateFlags uint

const (
	// RequireIsolateFile means the -isolate flag is required.
	RequireIsolateFile RequiredIsolateFlags = 1 << iota
	// RequireIsolatedFile means the -isolated flag is required.
	RequireIsolatedFile
)

func (c *isolateFlags) Parse(cwd string, flags RequiredIsolateFlags) error {
	if !filepath.IsAbs(cwd) {
		return errors.Reason("cwd must be absolute path").Err()
	}
	for _, vars := range [](map[string]string){c.ConfigVariables, c.PathVariables} {
		for k := range vars {
			if !isolate.IsValidVariable(k) {
				return fmt.Errorf("invalid key %s", k)
			}
		}
	}
	// Account for EXECUTABLE_SUFFIX.
	if len(c.ConfigVariables) != 0 || len(c.PathVariables) > 1 {
		os.Stderr.WriteString(
			"WARNING: -config-variables and -path-variables\n" +
				"         will be unsupported soon. Please contact the LUCI team.\n" +
				"         https://crbug.com/907880\n")
	}

	if c.Isolate == "" {
		if flags&RequireIsolateFile != 0 {
			return errors.Reason("-isolate must be specified").Err()
		}
	} else {
		if !filepath.IsAbs(c.Isolate) {
			c.Isolate = filepath.Clean(filepath.Join(cwd, c.Isolate))
		}
	}

	if c.Isolated == "" {
		if flags&RequireIsolatedFile != 0 {
			return errors.Reason("-isolated must be specified").Err()
		}
	} else {
		if !filepath.IsAbs(c.Isolated) {
			c.Isolated = filepath.Clean(filepath.Join(cwd, c.Isolated))
		}
	}
	return nil
}

// loggingFlags configures eventlog logging.
type loggingFlags struct {
	EventlogEndpoint string
}

func (lf *loggingFlags) Init(f *flag.FlagSet) {
	f.StringVar(&lf.EventlogEndpoint, "eventlog-endpoint", "", `The URL destination for eventlogs. The special values "prod" or "test" may be used to target the standard prod or test urls respectively. An empty string disables eventlogging.`)
}

func elideNestedPaths(deps []string, pathSep string) []string {
	// For |deps| having a pattern like below:
	// "ab/"
	// "ab/cd/"
	// "ab/foo.txt"
	//
	// We need to elide the nested paths under "ab/" to make HardlinkRecursively
	// work. Without this step, all files have already been hard linked when
	// processing "ab/", so "ab/cd/" would lead to an error.
	sort.Strings(deps)
	prefixDir := ""
	var result []string
	for _, dep := range deps {
		if len(prefixDir) > 0 && strings.HasPrefix(dep, prefixDir) {
			continue
		}
		// |dep| can be either an unseen directory, or an individual file
		result = append(result, dep)
		prefixDir = ""
		if strings.HasSuffix(dep, pathSep) {
			// |dep| is a directory
			prefixDir = dep
		}
	}
	return result
}

func recreateTree(outDir string, rootDir string, deps []string) error {
	if err := filesystem.MakeDirs(outDir); err != nil {
		return errors.Annotate(err, "failed to create directory: %s", outDir).Err()
	}
	deps = elideNestedPaths(deps, string(os.PathSeparator))
	createdDirs := make(map[string]struct{})
	for _, dep := range deps {
		dst := filepath.Join(outDir, dep[len(rootDir):])
		dstDir := filepath.Dir(dst)
		if _, ok := createdDirs[dstDir]; !ok {
			if err := filesystem.MakeDirs(dstDir); err != nil {
				return errors.Annotate(err, "failed to call MakeDirs(%s)", dstDir).Err()
			}
			createdDirs[dstDir] = struct{}{}
		}

		err := filesystem.HardlinkRecursively(dep, dst)
		if err != nil {
			return errors.Annotate(err, "failed to call HardlinkRecursively(%s, %s)", dep, dst).Err()
		}
	}
	return nil
}

// archiveLogger reports stats to stderr.
type archiveLogger struct {
	start time.Time
	quiet bool
}

// LogSummary logs (to eventlog and stderr) a high-level summary of archive operations(s).
func (al *archiveLogger) LogSummary(ctx context.Context, hits, misses int64, bytesHit, bytesPushed units.Size, digests []string) {
	end := time.Now()

	if !al.quiet {
		duration := end.Sub(al.start)
		fmt.Fprintf(os.Stderr, "Hits    : %5d (%s)\n", hits, bytesHit)
		fmt.Fprintf(os.Stderr, "Misses  : %5d (%s)\n", misses, bytesPushed)
		fmt.Fprintf(os.Stderr, "Duration: %s\n", duration.Round(time.Millisecond))
	}
}

// Print acts like fmt.Printf, but may prepend a prefix to format, depending on the value of al.quiet.
func (al *archiveLogger) Printf(format string, a ...interface{}) (n int, err error) {
	return al.Fprintf(os.Stdout, format, a...)
}

// Print acts like fmt.fprintf, but may prepend a prefix to format, depending on the value of al.quiet.
func (al *archiveLogger) Fprintf(w io.Writer, format string, a ...interface{}) (n int, err error) {
	prefix := "\n"
	if al.quiet {
		prefix = ""
	}
	args := []interface{}{prefix}
	args = append(args, a...)
	return fmt.Printf("%s"+format, args...)
}

func (al *archiveLogger) printSummary(summary archiver.IsolatedSummary) {
	al.Printf("%s\t%s\n", summary.Digest, summary.Name)
}

func buildCASInputSpec(opts *isolate.ArchiveOptions) (string, *command.InputSpec, error) {
	inputPaths, execRoot, err := isolate.ProcessIsolateForCAS(opts)
	if err != nil {
		return "", nil, err
	}

	inputSpec := &command.InputSpec{
		Inputs: inputPaths,
	}
	if opts.IgnoredPathFilterRe != "" {
		inputSpec.InputExclusions = []*command.InputExclusion{
			{
				Regex: opts.IgnoredPathFilterRe,
				Type:  command.UnspecifiedInputType,
			},
		}
	}

	return execRoot, inputSpec, nil
}

func uploadToCAS(ctx context.Context, dumpJSON string, authOpts auth.Options, fl *cas.Flags, al *archiveLogger, opts ...*isolate.ArchiveOptions) ([]digest.Digest, error) {
	cl, err := newCasClient(ctx, fl.Instance, authOpts, false)
	if err != nil {
		return nil, err
	}
	defer cl.Close()

	fmCache := filemetadata.NewSingleFlightCache()
	var rootDgs []digest.Digest
	var entries []*uploadinfo.Entry
	for _, o := range opts {
		execRoot, is, err := buildCASInputSpec(o)
		if err != nil {
			return nil, errors.Annotate(err, "failed to call buildCASInputSpec").Err()
		}
		rootDg, entrs, _, err := cl.ComputeMerkleTree(execRoot, is, fmCache)
		if err != nil {
			return nil, errors.Annotate(err, "failed to call ComputeMerkleTree").Err()
		}
		rootDgs = append(rootDgs, rootDg)
		entries = append(entries, entrs...)
	}
	uploadedDgs, err := cl.UploadIfMissing(ctx, entries...)
	if err != nil {
		return nil, err
	}

	if al != nil {
		missing := int64(len(uploadedDgs))
		hits := int64(len(entries)) - missing
		bytesPushed := int64(0)
		bytesTotal := int64(0)
		for _, e := range entries {
			bytesTotal += e.Digest.Size
		}
		for _, dg := range uploadedDgs {
			bytesPushed += dg.Size
		}

		var dgsStr []string
		for _, dg := range rootDgs {
			dgsStr = append(dgsStr, dg.String())
		}
		al.LogSummary(ctx, missing, hits, units.Size(bytesTotal-bytesPushed), units.Size(bytesPushed), dgsStr)
	}

	if dumpJSON == "" {
		return rootDgs, nil
	}

	f, err := os.Create(dumpJSON)
	if err != nil {
		return rootDgs, err
	}
	defer f.Close()

	m := make(map[string]string)
	for i, o := range opts {
		m[filesystem.GetFilenameNoExt(o.Isolate)] = rootDgs[i].String()
	}
	return rootDgs, json.NewEncoder(f).Encode(m)
}

// This is overwritten in test.
var newCasClient = cas.NewClient
<commit_msg>[isolate] fix upload stats arguments<commit_after>// Copyright 2015 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 lib

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/command"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/digest"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/filemetadata"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/uploadinfo"
	"github.com/maruel/subcommands"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/auth/client/authcli"
	"go.chromium.org/luci/cipd/version"
	"go.chromium.org/luci/client/archiver"
	"go.chromium.org/luci/client/cas"
	"go.chromium.org/luci/client/internal/common"
	"go.chromium.org/luci/client/isolate"
	"go.chromium.org/luci/common/data/text/units"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/isolatedclient"
	"go.chromium.org/luci/common/system/filesystem"
)

type commonFlags struct {
	subcommands.CommandRunBase
	defaultFlags common.Flags
}

func (c *commonFlags) Init() {
	c.defaultFlags.Init(&c.Flags)
}

func (c *commonFlags) Parse() error {
	return c.defaultFlags.Parse()
}

type commonServerFlags struct {
	commonFlags
	isolatedFlags isolatedclient.Flags
	authFlags     authcli.Flags

	parsedAuthOpts auth.Options
}

func (c *commonServerFlags) Init(authOpts auth.Options) {
	c.commonFlags.Init()
	c.isolatedFlags.Init(&c.Flags)
	c.authFlags.Register(&c.Flags, authOpts)
}

func (c *commonServerFlags) Parse() error {
	var err error
	if err = c.commonFlags.Parse(); err != nil {
		return err
	}
	if err = c.isolatedFlags.Parse(); err != nil {
		return err
	}
	c.parsedAuthOpts, err = c.authFlags.Options()
	return err
}

func (c *commonServerFlags) createAuthClient(ctx context.Context) (*http.Client, error) {
	// Don't enforce authentication by using OptionalLogin mode. This is needed
	// for IP-allowed bots: they have NO credentials to send.
	return auth.NewAuthenticator(ctx, auth.OptionalLogin, c.parsedAuthOpts).Client()
}

func (c *commonServerFlags) createIsolatedClient(authCl *http.Client) (*isolatedclient.Client, error) {
	userAgent := "isolate-go/" + IsolateVersion
	if ver, err := version.GetStartupVersion(); err == nil && ver.InstanceID != "" {
		userAgent += fmt.Sprintf(" (%s@%s)", ver.PackageName, ver.InstanceID)
	}
	return c.isolatedFlags.NewClient(isolatedclient.WithAuthClient(authCl), isolatedclient.WithUserAgent(userAgent))
}

type isolateFlags struct {
	// TODO(tandrii): move ArchiveOptions from isolate pkg to here.
	isolate.ArchiveOptions
}

func (c *isolateFlags) Init(f *flag.FlagSet) {
	c.ArchiveOptions.Init()
	f.StringVar(&c.Isolate, "isolate", "", ".isolate file to load the dependency data from")
	f.StringVar(&c.Isolate, "i", "", "Alias for -isolate")
	f.StringVar(&c.IgnoredPathFilterRe, "ignored-path-filter-re", "", "A regular expression for filtering away the paths to be ignored. Note that this regexp matches ANY part of the path. So if you want to match the beginning of a path, you need to explicitly prepend ^ (same for $). Please use the Go regexp syntax. I.e. use double backslack \\\\ if you need a backslash literal.")
	f.Var(&c.ConfigVariables, "config-variable", "Config variables are used to determine which conditions should be matched when loading a .isolate file, default: [].")
	f.Var(&c.PathVariables, "path-variable", "Path variables are used to replace file paths when loading a .isolate file, default: {}")
	f.BoolVar(&c.AllowMissingFileDir, "allow-missing-file-dir", false, "If this flag is true, invalid entries in the isolated file are only logged, but won't stop it from being processed.")
}

// RequiredIsolateFlags specifies which flags are required on the command line
// being parsed.
type RequiredIsolateFlags uint

const (
	// RequireIsolateFile means the -isolate flag is required.
	RequireIsolateFile RequiredIsolateFlags = 1 << iota
	// RequireIsolatedFile means the -isolated flag is required.
	RequireIsolatedFile
)

func (c *isolateFlags) Parse(cwd string, flags RequiredIsolateFlags) error {
	if !filepath.IsAbs(cwd) {
		return errors.Reason("cwd must be absolute path").Err()
	}
	for _, vars := range [](map[string]string){c.ConfigVariables, c.PathVariables} {
		for k := range vars {
			if !isolate.IsValidVariable(k) {
				return fmt.Errorf("invalid key %s", k)
			}
		}
	}
	// Account for EXECUTABLE_SUFFIX.
	if len(c.ConfigVariables) != 0 || len(c.PathVariables) > 1 {
		os.Stderr.WriteString(
			"WARNING: -config-variables and -path-variables\n" +
				"         will be unsupported soon. Please contact the LUCI team.\n" +
				"         https://crbug.com/907880\n")
	}

	if c.Isolate == "" {
		if flags&RequireIsolateFile != 0 {
			return errors.Reason("-isolate must be specified").Err()
		}
	} else {
		if !filepath.IsAbs(c.Isolate) {
			c.Isolate = filepath.Clean(filepath.Join(cwd, c.Isolate))
		}
	}

	if c.Isolated == "" {
		if flags&RequireIsolatedFile != 0 {
			return errors.Reason("-isolated must be specified").Err()
		}
	} else {
		if !filepath.IsAbs(c.Isolated) {
			c.Isolated = filepath.Clean(filepath.Join(cwd, c.Isolated))
		}
	}
	return nil
}

// loggingFlags configures eventlog logging.
type loggingFlags struct {
	EventlogEndpoint string
}

func (lf *loggingFlags) Init(f *flag.FlagSet) {
	f.StringVar(&lf.EventlogEndpoint, "eventlog-endpoint", "", `The URL destination for eventlogs. The special values "prod" or "test" may be used to target the standard prod or test urls respectively. An empty string disables eventlogging.`)
}

func elideNestedPaths(deps []string, pathSep string) []string {
	// For |deps| having a pattern like below:
	// "ab/"
	// "ab/cd/"
	// "ab/foo.txt"
	//
	// We need to elide the nested paths under "ab/" to make HardlinkRecursively
	// work. Without this step, all files have already been hard linked when
	// processing "ab/", so "ab/cd/" would lead to an error.
	sort.Strings(deps)
	prefixDir := ""
	var result []string
	for _, dep := range deps {
		if len(prefixDir) > 0 && strings.HasPrefix(dep, prefixDir) {
			continue
		}
		// |dep| can be either an unseen directory, or an individual file
		result = append(result, dep)
		prefixDir = ""
		if strings.HasSuffix(dep, pathSep) {
			// |dep| is a directory
			prefixDir = dep
		}
	}
	return result
}

func recreateTree(outDir string, rootDir string, deps []string) error {
	if err := filesystem.MakeDirs(outDir); err != nil {
		return errors.Annotate(err, "failed to create directory: %s", outDir).Err()
	}
	deps = elideNestedPaths(deps, string(os.PathSeparator))
	createdDirs := make(map[string]struct{})
	for _, dep := range deps {
		dst := filepath.Join(outDir, dep[len(rootDir):])
		dstDir := filepath.Dir(dst)
		if _, ok := createdDirs[dstDir]; !ok {
			if err := filesystem.MakeDirs(dstDir); err != nil {
				return errors.Annotate(err, "failed to call MakeDirs(%s)", dstDir).Err()
			}
			createdDirs[dstDir] = struct{}{}
		}

		err := filesystem.HardlinkRecursively(dep, dst)
		if err != nil {
			return errors.Annotate(err, "failed to call HardlinkRecursively(%s, %s)", dep, dst).Err()
		}
	}
	return nil
}

// archiveLogger reports stats to stderr.
type archiveLogger struct {
	start time.Time
	quiet bool
}

// LogSummary logs (to eventlog and stderr) a high-level summary of archive operations(s).
func (al *archiveLogger) LogSummary(ctx context.Context, hits, misses int64, bytesHit, bytesPushed units.Size, digests []string) {
	end := time.Now()

	if !al.quiet {
		duration := end.Sub(al.start)
		fmt.Fprintf(os.Stderr, "Hits    : %5d (%s)\n", hits, bytesHit)
		fmt.Fprintf(os.Stderr, "Misses  : %5d (%s)\n", misses, bytesPushed)
		fmt.Fprintf(os.Stderr, "Duration: %s\n", duration.Round(time.Millisecond))
	}
}

// Print acts like fmt.Printf, but may prepend a prefix to format, depending on the value of al.quiet.
func (al *archiveLogger) Printf(format string, a ...interface{}) (n int, err error) {
	return al.Fprintf(os.Stdout, format, a...)
}

// Print acts like fmt.fprintf, but may prepend a prefix to format, depending on the value of al.quiet.
func (al *archiveLogger) Fprintf(w io.Writer, format string, a ...interface{}) (n int, err error) {
	prefix := "\n"
	if al.quiet {
		prefix = ""
	}
	args := []interface{}{prefix}
	args = append(args, a...)
	return fmt.Printf("%s"+format, args...)
}

func (al *archiveLogger) printSummary(summary archiver.IsolatedSummary) {
	al.Printf("%s\t%s\n", summary.Digest, summary.Name)
}

func buildCASInputSpec(opts *isolate.ArchiveOptions) (string, *command.InputSpec, error) {
	inputPaths, execRoot, err := isolate.ProcessIsolateForCAS(opts)
	if err != nil {
		return "", nil, err
	}

	inputSpec := &command.InputSpec{
		Inputs: inputPaths,
	}
	if opts.IgnoredPathFilterRe != "" {
		inputSpec.InputExclusions = []*command.InputExclusion{
			{
				Regex: opts.IgnoredPathFilterRe,
				Type:  command.UnspecifiedInputType,
			},
		}
	}

	return execRoot, inputSpec, nil
}

func uploadToCAS(ctx context.Context, dumpJSON string, authOpts auth.Options, fl *cas.Flags, al *archiveLogger, opts ...*isolate.ArchiveOptions) ([]digest.Digest, error) {
	cl, err := newCasClient(ctx, fl.Instance, authOpts, false)
	if err != nil {
		return nil, err
	}
	defer cl.Close()

	fmCache := filemetadata.NewSingleFlightCache()
	var rootDgs []digest.Digest
	var entries []*uploadinfo.Entry
	for _, o := range opts {
		execRoot, is, err := buildCASInputSpec(o)
		if err != nil {
			return nil, errors.Annotate(err, "failed to call buildCASInputSpec").Err()
		}
		rootDg, entrs, _, err := cl.ComputeMerkleTree(execRoot, is, fmCache)
		if err != nil {
			return nil, errors.Annotate(err, "failed to call ComputeMerkleTree").Err()
		}
		rootDgs = append(rootDgs, rootDg)
		entries = append(entries, entrs...)
	}
	uploadedDgs, err := cl.UploadIfMissing(ctx, entries...)
	if err != nil {
		return nil, err
	}

	if al != nil {
		missing := int64(len(uploadedDgs))
		hits := int64(len(entries)) - missing
		bytesPushed := int64(0)
		bytesTotal := int64(0)
		for _, e := range entries {
			bytesTotal += e.Digest.Size
		}
		for _, dg := range uploadedDgs {
			bytesPushed += dg.Size
		}

		var dgsStr []string
		for _, dg := range rootDgs {
			dgsStr = append(dgsStr, dg.String())
		}
		al.LogSummary(ctx, hits, missing, units.Size(bytesTotal-bytesPushed), units.Size(bytesPushed), dgsStr)
	}

	if dumpJSON == "" {
		return rootDgs, nil
	}

	f, err := os.Create(dumpJSON)
	if err != nil {
		return rootDgs, err
	}
	defer f.Close()

	m := make(map[string]string)
	for i, o := range opts {
		m[filesystem.GetFilenameNoExt(o.Isolate)] = rootDgs[i].String()
	}
	return rootDgs, json.NewEncoder(f).Encode(m)
}

// This is overwritten in test.
var newCasClient = cas.NewClient
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2017-2018 Alibaba Group Holding Limited
 */
package config

import (
	"fmt"
	"github.com/aliyun/alibaba-cloud-sdk-go/sdk"
	"github.com/aliyun/alibaba-cloud-sdk-go/sdk/auth/credentials"
	"github.com/aliyun/aliyun-cli/cli"
)

type AuthenticateMode string

const (
	AK         = AuthenticateMode("AK")
	StsToken   = AuthenticateMode("StsToken")
	RamRoleArn = AuthenticateMode("RamRoleArn")
	EcsRamRole = AuthenticateMode("EcsRamRole")
	RsaKeyPair = AuthenticateMode("RsaKeyPair")
)

type Profile struct {
	Name            string           `json:"name"`
	Mode            AuthenticateMode `json:"mode"`
	AccessKeyId     string           `json:"access_key_id"`
	AccessKeySecret string           `json:"access_key_secret"`
	StsToken        string           `json:"sts_token"`
	RamRoleName     string           `json:"ram_role_name"`
	RamRoleArn      string           `json:"ram_role_arn"`
	RoleSessionName string           `json:"ram_session_name"`
	PrivateKey      string           `json:"private_key"`
	KeyPairName     string           `json:"key_pair_name"`
	ExpiredSeconds  int              `json:"expired_seconds"`
	Verified        string           `json:"verified"`
	RegionId        string           `json:"region_id"`
	OutputFormat    string           `json:"output_format"`
	Language        string           `json:"language"`
	Site            string           `json:"Site"`
}

func NewProfile(name string) (Profile) {
	return Profile{
		Name:         name,
		Mode:         AK,
		OutputFormat: "json",
		Language:     "en",
	}
}

func (cp *Profile) Validate() error {
	if cp.Mode == "" {
		return fmt.Errorf("profile %s is not configure yet, run `aliyun configure --profile %s` first", cp.Name, cp.Name)
	}
	switch cp.Mode {
	case AK:
		return cp.ValidateAK()
	case StsToken:
		err := cp.ValidateAK()
		if err != nil {
			return err
		}
		if cp.StsToken == "" {
			return fmt.Errorf("invailed sts_token")
		}
	case RamRoleArn:
		err := cp.ValidateAK()
		if err != nil {
			return err
		}
		if cp.RamRoleArn == "" {
			return fmt.Errorf("invailed ram_role_arn")
		}
		if cp.RoleSessionName == "" {
			return fmt.Errorf("invailed role_session_name")
		}
	case EcsRamRole:
		if cp.RamRoleName == "" {
			return fmt.Errorf("invailed ram_role_name")
		}
	case RsaKeyPair:
		if cp.PrivateKey == "" {
			return fmt.Errorf("invailed private_key")
		}
		if cp.KeyPairName == "" {
			return fmt.Errorf("invailed key_pair_name")
		}
	default:
		return fmt.Errorf("invailed mode: %s", cp.Mode)
	}
	if cp.RegionId == "" {
		return fmt.Errorf("region can't be empty")
	}
	return nil
}

func (cp *Profile) OverwriteWithFlags(ctx *cli.Context) {
	cp.Mode = AuthenticateMode(ModeFlag.GetValueOrDefault(ctx, string(cp.Mode)))
	cp.AccessKeyId = AccessKeyIdFlag.GetValueOrDefault(ctx, cp.AccessKeyId)
	cp.AccessKeySecret = AccessKeySecretFlag.GetValueOrDefault(ctx, cp.AccessKeySecret)
	cp.StsToken = StsTokenFlag.GetValueOrDefault(ctx, cp.StsToken)
	cp.RamRoleName = RamRoleNameFlag.GetValueOrDefault(ctx, cp.RamRoleName)
	cp.RamRoleArn = RamRoleArnFlag.GetValueOrDefault(ctx, cp.RamRoleArn)
	cp.RoleSessionName = RoleSessionNameFlag.GetValueOrDefault(ctx, cp.RoleSessionName)
	cp.KeyPairName = KeyPairNameFlag.GetValueOrDefault(ctx, cp.KeyPairName)
	cp.PrivateKey = PrivateKeyFlag.GetValueOrDefault(ctx, cp.PrivateKey)
	cp.RegionId = RegionFlag.GetValueOrDefault(ctx, cp.RegionId)
	cp.Language = LanguageFlag.GetValueOrDefault(ctx, cp.Language)

	if cp.AccessKeyId != "" && cp.AccessKeySecret != "" {
		cp.Mode = AK
		if cp.StsToken != "" {
			cp.Mode = StsToken
		} else if cp.RamRoleArn != "" {
			cp.Mode = RamRoleArn
		}
	}
	if cp.PrivateKey != "" && cp.KeyPairName != "" {
		cp.Mode = RsaKeyPair
	}
}

func (cp *Profile) ValidateAK() error {
	if len(cp.AccessKeyId) == 0 {
		return fmt.Errorf("invalid access_key_id: %s", cp.AccessKeyId)
	}
	if len(cp.AccessKeySecret) == 0 {
		return fmt.Errorf("invaild access_key_secret: %s", cp.AccessKeySecret)
	}
	return nil
}

func (cp *Profile) GetClient() (*sdk.Client, error) {
	switch cp.Mode {
	case AK:
		return cp.GetClientByAK()
	case StsToken:
		return cp.GetClientBySts()
	case RamRoleArn:
		return cp.GetClientByRoleArn()
	case EcsRamRole:
		return cp.GetClientByEcsRamRole()
	case RsaKeyPair:
		return cp.GetClientByPrivateKey()
	default:
		return nil, fmt.Errorf("unexcepted certificate mode: %s", cp.Mode)
	}
}

func (cp *Profile) GetClientByAK() (*sdk.Client, error) {
	if cp.AccessKeyId == "" || cp.AccessKeySecret == "" {
		return nil, fmt.Errorf("AccessKeyId/AccessKeySecret is empty! run `aliyun configure` first")
	}

	if cp.RegionId == "" {
		return nil, fmt.Errorf("default RegionId is empty! run `aliyun configure` first")
	}

	client, err := sdk.NewClientWithAccessKey(cp.RegionId, cp.AccessKeyId, cp.AccessKeySecret)
	return client, err
}

func (cp *Profile) GetClientBySts() (*sdk.Client, error) {
	cred := credentials.NewStsTokenCredential(cp.AccessKeyId, cp.AccessKeySecret, cp.StsToken)
	config := sdk.NewConfig()
	client, err := sdk.NewClientWithOptions(cp.RegionId, config, cred)
	return client, err
}

func (cp *Profile) GetClientByEcsRamRole() (*sdk.Client, error) {
	if cp.RamRoleName == "" {
		return nil, fmt.Errorf("RamRole is empty! run `aliyun configure` first")
	}

	cred := credentials.NewEcsRamRoleCredential(cp.RamRoleName)
	config := sdk.NewConfig()
	// keyId, keySecret, stsToken := singer.GetSessionKey(cred)

	client, err := sdk.NewClientWithOptions(cp.RegionId, config, cred)
	return client, err
}

func (cp *Profile) GetClientByRoleArn() (*sdk.Client, error) {
	cred := credentials.NewRamRoleArnCredential(cp.AccessKeyId, cp.AccessKeySecret, cp.RamRoleArn, cp.RoleSessionName, cp.ExpiredSeconds)
	config := sdk.NewConfig()
	client, err := sdk.NewClientWithOptions(cp.RegionId, config, cred)
	return client, err
}

func (cp *Profile) GetClientByPrivateKey() (*sdk.Client, error) {
	cred := credentials.NewRsaKeyPairCredential(cp.PrivateKey, cp.KeyPairName, cp.ExpiredSeconds)
	config := sdk.NewConfig()
	client, err := sdk.NewClientWithOptions(cp.RegionId, config, cred)
	return client, err
}
<commit_msg>migrate RamRoleArn|EcsRamRole session credentials fetch code with CLI, not integrated with SDK.<commit_after>/*
 * Copyright (C) 2017-2018 Alibaba Group Holding Limited
 */
package config

import (
	"fmt"
	"github.com/aliyun/alibaba-cloud-sdk-go/sdk"
	"github.com/aliyun/alibaba-cloud-sdk-go/sdk/auth/credentials"
	"github.com/aliyun/aliyun-cli/cli"
	"github.com/aliyun/alibaba-cloud-sdk-go/services/sts"
	"github.com/aliyun/alibaba-cloud-sdk-go/sdk/requests"
	"net/http"
	"strings"
	"github.com/aliyun/alibaba-cloud-sdk-go/sdk/responses"
	"encoding/json"
	"github.com/jmespath/go-jmespath"
	"github.com/aliyun/alibaba-cloud-sdk-go/sdk/auth/signers"
)

type AuthenticateMode string

const (
	AK         = AuthenticateMode("AK")
	StsToken   = AuthenticateMode("StsToken")
	RamRoleArn = AuthenticateMode("RamRoleArn")
	EcsRamRole = AuthenticateMode("EcsRamRole")
	RsaKeyPair = AuthenticateMode("RsaKeyPair")
)

type Profile struct {
	Name            string           `json:"name"`
	Mode            AuthenticateMode `json:"mode"`
	AccessKeyId     string           `json:"access_key_id"`
	AccessKeySecret string           `json:"access_key_secret"`
	StsToken        string           `json:"sts_token"`
	RamRoleName     string           `json:"ram_role_name"`
	RamRoleArn      string           `json:"ram_role_arn"`
	RoleSessionName string           `json:"ram_session_name"`
	PrivateKey      string           `json:"private_key"`
	KeyPairName     string           `json:"key_pair_name"`
	ExpiredSeconds  int              `json:"expired_seconds"`
	Verified        string           `json:"verified"`
	RegionId        string           `json:"region_id"`
	OutputFormat    string           `json:"output_format"`
	Language        string           `json:"language"`
	Site            string           `json:"Site"`
}

func NewProfile(name string) (Profile) {
	return Profile{
		Name:         name,
		Mode:         AK,
		OutputFormat: "json",
		Language:     "en",
	}
}

func (cp *Profile) Validate() error {
	if cp.Mode == "" {
		return fmt.Errorf("profile %s is not configure yet, run `aliyun configure --profile %s` first", cp.Name, cp.Name)
	}
	switch cp.Mode {
	case AK:
		return cp.ValidateAK()
	case StsToken:
		err := cp.ValidateAK()
		if err != nil {
			return err
		}
		if cp.StsToken == "" {
			return fmt.Errorf("invailed sts_token")
		}
	case RamRoleArn:
		err := cp.ValidateAK()
		if err != nil {
			return err
		}
		if cp.RamRoleArn == "" {
			return fmt.Errorf("invailed ram_role_arn")
		}
		if cp.RoleSessionName == "" {
			return fmt.Errorf("invailed role_session_name")
		}
	case EcsRamRole:
		if cp.RamRoleName == "" {
			return fmt.Errorf("invailed ram_role_name")
		}
	case RsaKeyPair:
		if cp.PrivateKey == "" {
			return fmt.Errorf("invailed private_key")
		}
		if cp.KeyPairName == "" {
			return fmt.Errorf("invailed key_pair_name")
		}
	default:
		return fmt.Errorf("invailed mode: %s", cp.Mode)
	}
	if cp.RegionId == "" {
		return fmt.Errorf("region can't be empty")
	}
	return nil
}

func (cp *Profile) OverwriteWithFlags(ctx *cli.Context) {
	cp.Mode = AuthenticateMode(ModeFlag.GetValueOrDefault(ctx, string(cp.Mode)))
	cp.AccessKeyId = AccessKeyIdFlag.GetValueOrDefault(ctx, cp.AccessKeyId)
	cp.AccessKeySecret = AccessKeySecretFlag.GetValueOrDefault(ctx, cp.AccessKeySecret)
	cp.StsToken = StsTokenFlag.GetValueOrDefault(ctx, cp.StsToken)
	cp.RamRoleName = RamRoleNameFlag.GetValueOrDefault(ctx, cp.RamRoleName)
	cp.RamRoleArn = RamRoleArnFlag.GetValueOrDefault(ctx, cp.RamRoleArn)
	cp.RoleSessionName = RoleSessionNameFlag.GetValueOrDefault(ctx, cp.RoleSessionName)
	cp.KeyPairName = KeyPairNameFlag.GetValueOrDefault(ctx, cp.KeyPairName)
	cp.PrivateKey = PrivateKeyFlag.GetValueOrDefault(ctx, cp.PrivateKey)
	cp.RegionId = RegionFlag.GetValueOrDefault(ctx, cp.RegionId)
	cp.Language = LanguageFlag.GetValueOrDefault(ctx, cp.Language)

	if cp.AccessKeyId != "" && cp.AccessKeySecret != "" {
		cp.Mode = AK
		if cp.StsToken != "" {
			cp.Mode = StsToken
		} else if cp.RamRoleArn != "" {
			cp.Mode = RamRoleArn
		}
	}
	if cp.PrivateKey != "" && cp.KeyPairName != "" {
		cp.Mode = RsaKeyPair
	}
}

func (cp *Profile) ValidateAK() error {
	if len(cp.AccessKeyId) == 0 {
		return fmt.Errorf("invalid access_key_id: %s", cp.AccessKeyId)
	}
	if len(cp.AccessKeySecret) == 0 {
		return fmt.Errorf("invaild access_key_secret: %s", cp.AccessKeySecret)
	}
	return nil
}

func (cp *Profile) GetClient() (*sdk.Client, error) {
	switch cp.Mode {
	case AK:
		return cp.GetClientByAK()
	case StsToken:
		return cp.GetClientBySts()
	case RamRoleArn:
		return cp.GetClientByRoleArn()
	case EcsRamRole:
		return cp.GetClientByEcsRamRole()
	case RsaKeyPair:
		return cp.GetClientByPrivateKey()
	default:
		return nil, fmt.Errorf("unexcepted certificate mode: %s", cp.Mode)
	}
}

func (cp *Profile) GetSessionCredential() (*signers.SessionCredential, error) {
	switch cp.Mode {
	case AK:
		return &signers.SessionCredential{
			AccessKeyId:cp.AccessKeyId,
			AccessKeySecret:cp.AccessKeySecret,
		}, nil
	case StsToken:
		return &signers.SessionCredential{
			AccessKeyId:cp.AccessKeyId,
			AccessKeySecret:cp.AccessKeySecret,
			StsToken:cp.StsToken,
		}, nil
	case RamRoleArn:
		return cp.GetSessionCredentialByRoleArn()
	case EcsRamRole:
		return cp.GetSessionCredentialByEcsRamRole()
	default:
		return nil, fmt.Errorf("unsupported mode '%s' to GetSessionCredential", cp.Mode)
	}
}

func (cp *Profile) GetClientByAK() (*sdk.Client, error) {
	if cp.AccessKeyId == "" || cp.AccessKeySecret == "" {
		return nil, fmt.Errorf("AccessKeyId/AccessKeySecret is empty! run `aliyun configure` first")
	}

	if cp.RegionId == "" {
		return nil, fmt.Errorf("default RegionId is empty! run `aliyun configure` first")
	}

	client, err := sdk.NewClientWithAccessKey(cp.RegionId, cp.AccessKeyId, cp.AccessKeySecret)
	return client, err
}

func (cp *Profile) GetClientBySts() (*sdk.Client, error) {
	cred := credentials.NewStsTokenCredential(cp.AccessKeyId, cp.AccessKeySecret, cp.StsToken)
	config := sdk.NewConfig()
	client, err := sdk.NewClientWithOptions(cp.RegionId, config, cred)
	return client, err
}

func (cp *Profile) GetSessionCredentialByRoleArn() (*signers.SessionCredential, error) {
	client, err := sts.NewClientWithAccessKey(cp.RegionId, cp.AccessKeyId, cp.AccessKeySecret)
	if err != nil {
		return nil, fmt.Errorf("new sts client failed %s", err)
	}
	if client == nil {
		return nil, fmt.Errorf("new sts client with nil")
	}

	request := sts.CreateAssumeRoleRequest()
	request.RoleArn = cp.RamRoleArn
	request.RoleSessionName = cp.RoleSessionName
	request.DurationSeconds = requests.NewInteger(600)
	request.Scheme = "https"

	response, err := client.AssumeRole(request)
	if err != nil {
		return nil, fmt.Errorf("sts:AssumeRole() failed %s", err)
	}

	return &signers.SessionCredential{
		AccessKeyId: response.Credentials.AccessKeyId,
		AccessKeySecret: response.Credentials.AccessKeySecret,
		StsToken: response.Credentials.AccessKeySecret,
	}, nil
}

func (cp *Profile) GetClientByRoleArn() (*sdk.Client, error) {
	sc, err := cp.GetSessionCredentialByRoleArn()
	if err != nil {
		return nil, fmt.Errorf("get session credential failed %s", err)
	}
	config := sdk.NewConfig()
	cred := credentials.NewStsTokenCredential(sc.AccessKeyId, sc.AccessKeySecret, sc.StsToken)
	client, err := sdk.NewClientWithOptions(cp.RegionId, config, cred)
	return client, err
}

func (cp *Profile) GetSessionCredentialByEcsRamRole() (*signers.SessionCredential, error) {
	if cp.RamRoleName == "" {
		return nil, fmt.Errorf("RamRole is empty! run `aliyun configure` first")
	}

	requestUrl := "http://100.100.100.200/latest/meta-data/ram/security-credentials/" + cp.RamRoleName
	httpRequest, err := http.NewRequest(requests.GET, requestUrl, strings.NewReader(""))
	if err != nil {
		return nil, fmt.Errorf("new http request failed %s", err)
	}

	httpClient := &http.Client{}
	httpResponse, err := httpClient.Do(httpRequest)
	if err != nil {
		return nil, fmt.Errorf("failed get credentials from meta-data %s, please check RAM settings", err)
	}

	response := responses.NewCommonResponse()
	err = responses.Unmarshal(response, httpResponse, "")

	if response.GetHttpStatus() != http.StatusOK {
		return nil, fmt.Errorf("get meta-data status=%d please check RAM settings", response.GetHttpStatus())
	}

	var data interface{}
	err = json.Unmarshal(response.GetHttpContentBytes(), &data)
	if err != nil {
		return nil, fmt.Errorf("unmarshal meta-data fail %s", err)
	}

	code, err := jmespath.Search("Code", data)
	if err != nil {
		return nil, fmt.Errorf("read code from meta-data failed %s", err)
	}
	if code.(string) != "Success" {
		return nil, fmt.Errorf("unexcepted code = %s", code)
	}
	accessKeyId, err := jmespath.Search("AccessKeyId", data)
	if err != nil || accessKeyId == "" {
		return nil, fmt.Errorf("read AccessKeyId from meta-data failed %s", err)
	}
	accessKeySecret, err := jmespath.Search("AccessKeySecret", data)
	if err != nil || accessKeySecret == "" {
		return nil, fmt.Errorf("read AccessKeySecret from meta-data failed %s", err)
	}
	securityToken, err := jmespath.Search("SecurityToken", data)
	if err != nil || securityToken == "" {
		return nil, fmt.Errorf("read SecurityToken from meta-data failed %s", err)
	}
	return &signers.SessionCredential{
		AccessKeyId: accessKeyId.(string),
		AccessKeySecret: accessKeySecret.(string),
		StsToken: securityToken.(string),
	}, nil
}

func (cp *Profile) GetClientByEcsRamRole() (*sdk.Client, error) {
	sc, err := cp.GetSessionCredentialByEcsRamRole()
	if err != nil {
		return nil, fmt.Errorf("get session credential failed %s", err)
	}

	cred := credentials.NewStsTokenCredential(sc.AccessKeyId, sc.AccessKeySecret, sc.StsToken)
	config := sdk.NewConfig()
	client, err := sdk.NewClientWithOptions(cp.RegionId, config, cred)
	return client, err
}

func (cp *Profile) GetClientByPrivateKey() (*sdk.Client, error) {
	cred := credentials.NewRsaKeyPairCredential(cp.PrivateKey, cp.KeyPairName, cp.ExpiredSeconds)
	config := sdk.NewConfig()
	client, err := sdk.NewClientWithOptions(cp.RegionId, config, cred)
	return client, err
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC
//
// 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
//
//     https://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 driver defines a set of interfaces that the blob package uses to interact
// with the underlying blob services.
package driver

import (
	"context"
	"io"
)

// Reader reads an object from the blob.
type Reader interface {
	io.ReadCloser
	Size() int64
}

// Writer writes an object to the blob.
type Writer interface {
	io.WriteCloser
}

// WriterOptions controls behaviors of Writer.
type WriterOptions struct {
	BufferSize int
}

// Bucket provides read, write and delete operations on objects within it on the
// blob service.
type Bucket interface {
	// NewRangeReader returns a Reader that reads part of an object, reading at most
	// length bytes starting at the given offset. If length is 0, it will read only
	// the metadata. If length is negative, it will read till the end of the object.
	NewRangeReader(ctx context.Context, key string, offset, length int64) (Reader, error)

	// NewWriter returns Writer that writes to an object associated with key.
	//
	// A new object will be created unless an object with this key already exists.
	// Otherwise any previous object with the same name will be replaced.
	// The object will not be available (and any previous object will remain)
	// until Close has been called.
	//
	// The caller must call Close on the returned Writer when done writing.
	NewWriter(ctx context.Context, key string, opt *WriterOptions) (Writer, error)

	// Delete deletes the object associated with key. It is a no-op if that object
	// does not exist.
	Delete(ctx context.Context, key string) error
}
<commit_msg>blob/driver: document behavior of Size and Write<commit_after>// Copyright 2018 Google LLC
//
// 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
//
//     https://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 driver defines a set of interfaces that the blob package uses to interact
// with the underlying blob services.
package driver

import (
	"context"
	"io"
)

// Reader reads an object from the blob.
type Reader interface {
	io.ReadCloser

	// Size returns the number of bytes in the whole blob. It must always return
	// the same value.
	Size() int64
}

// Writer writes an object to the blob.
type Writer interface {
	io.WriteCloser
}

// WriterOptions controls behaviors of Writer.
type WriterOptions struct {
	BufferSize int
}

// Bucket provides read, write and delete operations on objects within it on the
// blob service.
type Bucket interface {
	// NewRangeReader returns a Reader that reads part of an object, reading at most
	// length bytes starting at the given offset. If length is 0, it will read only
	// the metadata. If length is negative, it will read till the end of the object.
	NewRangeReader(ctx context.Context, key string, offset, length int64) (Reader, error)

	// NewWriter returns Writer that writes to an object associated with key.
	//
	// A new object will be created unless an object with this key already exists.
	// Otherwise any previous object with the same name will be replaced.
	// The object may not be available (and any previous object will remain)
	// until Close has been called.
	//
	// The caller must call Close on the returned Writer when done writing.
	NewWriter(ctx context.Context, key string, opt *WriterOptions) (Writer, error)

	// Delete deletes the object associated with key. It is a no-op if that object
	// does not exist.
	Delete(ctx context.Context, key string) error
}
<|endoftext|>


<commit_before>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package config

import (
	"bufio"
	"compress/gzip"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	amasshttp "github.com/OWASP/Amass/v3/net/http"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/OWASP/Amass/v3/stringset"
	"github.com/gobuffalo/packr/v2"
)

const (
	outputDirectoryName = "amass"
)

var (
	// BoxOfDefaultFiles is the ./resources project directory embedded into the binary.
	BoxOfDefaultFiles *packr.Box
	boxOnce           sync.Once
)

// LookupASNsByName returns requests.ASNRequest objects for autonomous systems with
// descriptions that contain the string provided by the parameter.
func LookupASNsByName(s string) ([]*requests.ASNRequest, error) {
	var records []*requests.ASNRequest

	if BoxOfDefaultFiles == nil {
		boxOnce.Do(openTheBox)
	}

	content, err := BoxOfDefaultFiles.FindString("asnlist.txt")
	if err != nil {
		return records, fmt.Errorf("Failed to obtain the embedded ASN information: asnlist.txt: %v", err)
	}

	s = strings.ToLower(s)
	scanner := bufio.NewScanner(strings.NewReader(content))
	for scanner.Scan() {
		line := scanner.Text()

		if err := scanner.Err(); err == nil {
			parts := strings.Split(strings.TrimSpace(line), ",")

			if strings.Contains(strings.ToLower(parts[1]), s) {
				a, err := strconv.Atoi(parts[0])
				if err == nil {
					records = append(records, &requests.ASNRequest{
						ASN:         a,
						Description: parts[1],
					})
				}
			}
		}
	}
	return records, nil
}

func openTheBox() {
	BoxOfDefaultFiles = packr.New("Amass Box", "../resources")
}

// AcquireConfig populates the Config struct provided by the config argument.
// The configuration file path and a bool indicating the settings were
// successfully loaded are returned.
func AcquireConfig(dir, file string, config *Config) error {
	var err error

	if file != "" {
		err = config.LoadSettings(file)
		if err == nil {
			return nil
		}
	}
	if dir = OutputDirectory(dir); dir != "" {
		if finfo, err := os.Stat(dir); !os.IsNotExist(err) && finfo.IsDir() {
			file := filepath.Join(dir, "config.ini")

			err = config.LoadSettings(file)
			if err == nil {
				return nil
			}
		}
	}
	return err
}

// OutputDirectory returns the file path of the Amass output directory. A suitable
// path provided will be used as the output directory instead.
func OutputDirectory(dir string) string {
	if dir == "" {
		if path, err := os.UserConfigDir(); err == nil {
			dir = filepath.Join(path, outputDirectoryName)
		}
	}
	return dir
}

// GetListFromFile reads a wordlist text or gzip file and returns the slice of words.
func GetListFromFile(path string) ([]string, error) {
	var reader io.Reader

	file, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("Error opening the file %s: %v", path, err)
	}
	defer file.Close()
	reader = file

	// We need to determine if this is a gzipped file or a plain text file, so we
	// first read the first 512 bytes to pass them down to http.DetectContentType
	// for mime detection. The file is rewinded before passing it along to the
	// next reader
	head := make([]byte, 512)
	if _, err = file.Read(head); err != nil {
		return nil, fmt.Errorf("Error reading the first 512 bytes from %s: %s", path, err)
	}
	if _, err = file.Seek(0, 0); err != nil {
		return nil, fmt.Errorf("Error rewinding the file %s: %s", path, err)
	}

	// Read the file as gzip if it's actually compressed
	if mt := http.DetectContentType(head); mt == "application/gzip" || mt == "application/x-gzip" {
		gzReader, err := gzip.NewReader(file)
		if err != nil {
			return nil, fmt.Errorf("Error gz-reading the file %s: %v", path, err)
		}
		defer gzReader.Close()
		reader = gzReader
	}

	s, err := getWordList(reader)
	return s, err
}

func getWordlistByURL(url string) ([]string, error) {
	page, err := amasshttp.RequestWebPage(url, nil, nil, "", "")
	if err != nil {
		return nil, fmt.Errorf("Failed to obtain the wordlist at %s: %v", url, err)
	}
	return getWordList(strings.NewReader(page))
}

func getWordlistByBox(path string) ([]string, error) {
	if BoxOfDefaultFiles == nil {
		boxOnce.Do(openTheBox)
	}

	content, err := BoxOfDefaultFiles.FindString(path)
	if err != nil {
		return nil, fmt.Errorf("Failed to obtain the embedded wordlist: %s: %v", path, err)
	}
	return getWordList(strings.NewReader(content))
}

func getWordList(reader io.Reader) ([]string, error) {
	var words []string

	scanner := bufio.NewScanner(reader)
	for scanner.Scan() {
		// Get the next word in the list
		w := strings.TrimSpace(scanner.Text())
		if err := scanner.Err(); err == nil && w != "" {
			words = append(words, w)
		}
	}
	return stringset.Deduplicate(words), nil
}

func uniqueIntAppend(s []int, e string) []int {
	if a1, err := strconv.Atoi(e); err == nil {
		var found bool

		for _, a2 := range s {
			if a1 == a2 {
				found = true
				break
			}
		}
		if !found {
			s = append(s, a1)
		}
	}
	return s
}
<commit_msg>small changes to code that opens the output directory<commit_after>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package config

import (
	"bufio"
	"compress/gzip"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	amasshttp "github.com/OWASP/Amass/v3/net/http"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/OWASP/Amass/v3/stringset"
	"github.com/gobuffalo/packr/v2"
)

const (
	outputDirectoryName = "amass"
)

var (
	// BoxOfDefaultFiles is the ./resources project directory embedded into the binary.
	BoxOfDefaultFiles *packr.Box
	boxOnce           sync.Once
)

// LookupASNsByName returns requests.ASNRequest objects for autonomous systems with
// descriptions that contain the string provided by the parameter.
func LookupASNsByName(s string) ([]*requests.ASNRequest, error) {
	var records []*requests.ASNRequest

	if BoxOfDefaultFiles == nil {
		boxOnce.Do(openTheBox)
	}

	content, err := BoxOfDefaultFiles.FindString("asnlist.txt")
	if err != nil {
		return records, fmt.Errorf("Failed to obtain the embedded ASN information: asnlist.txt: %v", err)
	}

	s = strings.ToLower(s)
	scanner := bufio.NewScanner(strings.NewReader(content))
	for scanner.Scan() {
		line := scanner.Text()

		if err := scanner.Err(); err == nil {
			parts := strings.Split(strings.TrimSpace(line), ",")

			if strings.Contains(strings.ToLower(parts[1]), s) {
				a, err := strconv.Atoi(parts[0])
				if err == nil {
					records = append(records, &requests.ASNRequest{
						ASN:         a,
						Description: parts[1],
					})
				}
			}
		}
	}
	return records, nil
}

func openTheBox() {
	BoxOfDefaultFiles = packr.New("Amass Box", "../resources")
}

// AcquireConfig populates the Config struct provided by the config argument.
func AcquireConfig(dir, file string, config *Config) error {
	var err error

	if file != "" {
		err = config.LoadSettings(file)
		if err == nil {
			return nil
		}
	}
	if dir = OutputDirectory(dir); dir != "" {
		if finfo, err := os.Stat(dir); !os.IsNotExist(err) && finfo.IsDir() {
			file := filepath.Join(dir, "config.ini")

			err = config.LoadSettings(file)
			if err == nil {
				return nil
			}
		}
	}
	return err
}

// OutputDirectory returns the file path of the Amass output directory. A suitable
// path provided will be used as the output directory instead.
func OutputDirectory(dir ...string) string {
	if len(dir) > 0 && dir[0] != "" {
		return dir[0]
	}

	if path, err := os.UserConfigDir(); err == nil {
		return filepath.Join(path, outputDirectoryName)
	}

	return ""
}

// GetListFromFile reads a wordlist text or gzip file and returns the slice of words.
func GetListFromFile(path string) ([]string, error) {
	var reader io.Reader

	file, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("Error opening the file %s: %v", path, err)
	}
	defer file.Close()
	reader = file

	// We need to determine if this is a gzipped file or a plain text file, so we
	// first read the first 512 bytes to pass them down to http.DetectContentType
	// for mime detection. The file is rewinded before passing it along to the
	// next reader
	head := make([]byte, 512)
	if _, err = file.Read(head); err != nil {
		return nil, fmt.Errorf("Error reading the first 512 bytes from %s: %s", path, err)
	}
	if _, err = file.Seek(0, 0); err != nil {
		return nil, fmt.Errorf("Error rewinding the file %s: %s", path, err)
	}

	// Read the file as gzip if it's actually compressed
	if mt := http.DetectContentType(head); mt == "application/gzip" || mt == "application/x-gzip" {
		gzReader, err := gzip.NewReader(file)
		if err != nil {
			return nil, fmt.Errorf("Error gz-reading the file %s: %v", path, err)
		}
		defer gzReader.Close()
		reader = gzReader
	}

	s, err := getWordList(reader)
	return s, err
}

func getWordlistByURL(url string) ([]string, error) {
	page, err := amasshttp.RequestWebPage(url, nil, nil, "", "")
	if err != nil {
		return nil, fmt.Errorf("Failed to obtain the wordlist at %s: %v", url, err)
	}
	return getWordList(strings.NewReader(page))
}

func getWordlistByBox(path string) ([]string, error) {
	if BoxOfDefaultFiles == nil {
		boxOnce.Do(openTheBox)
	}

	content, err := BoxOfDefaultFiles.FindString(path)
	if err != nil {
		return nil, fmt.Errorf("Failed to obtain the embedded wordlist: %s: %v", path, err)
	}
	return getWordList(strings.NewReader(content))
}

func getWordList(reader io.Reader) ([]string, error) {
	var words []string

	scanner := bufio.NewScanner(reader)
	for scanner.Scan() {
		// Get the next word in the list
		w := strings.TrimSpace(scanner.Text())
		if err := scanner.Err(); err == nil && w != "" {
			words = append(words, w)
		}
	}
	return stringset.Deduplicate(words), nil
}

func uniqueIntAppend(s []int, e string) []int {
	if a1, err := strconv.Atoi(e); err == nil {
		var found bool

		for _, a2 := range s {
			if a1 == a2 {
				found = true
				break
			}
		}
		if !found {
			s = append(s, a1)
		}
	}
	return s
}
<|endoftext|>


<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 siesta

import (
	"fmt"
	"math/rand"
	"net"
	"os"
	"testing"
	"time"
)

var ci = os.Getenv("TRAVIS_CI") != ""
var brokerUp = true
var brokerAddr = "localhost:9092"

func init() {
	conn, err := net.Dial("tcp", brokerAddr)
	if err == nil {
		brokerUp = true
		err = conn.Close()
		if err != nil {
			panic(err)
		}
	}
}

func TestDefaultConnectorFunctional(t *testing.T) {
	if !brokerUp && !ci {
		t.Skip("Broker is not running. Please spin up the broker at localhost:9092 for this test to work.")
	}

	//numMessages := 1000
	topicName := fmt.Sprintf("siesta-%d", time.Now().Unix())

	connector := testConnector(t)
	testTopicMetadata(t, topicName, connector)
	testOffsetStorage(t, topicName, connector)
	//testProduce(t, topicName, numMessages, connector)
	//testConsume(t, topicName, numMessages, connector)
	//closeWithin(t, time.Second, connector)

	//anotherConnector := testConnector(t)
	//should also work fine - must get topic metadata before consuming
	//testConsume(t, topicName, numMessages, anotherConnector)
	//closeWithin(t, time.Second, anotherConnector)
}

func testTopicMetadata(t *testing.T, topicName string, connector *DefaultConnector) {
	metadata, err := connector.GetTopicMetadata([]string{topicName})
	assertFatal(t, err, nil)

	assertNot(t, len(metadata.Brokers), 0)
	assertNot(t, len(metadata.TopicsMetadata), 0)
	if len(metadata.Brokers) > 1 {
		t.Skip("Cluster should consist only of one broker for this test to run.")
	}

	broker := metadata.Brokers[0]
	assert(t, broker.ID, int32(0))
	if ci {
		// this can be asserted on Travis only as we are guaranteed to advertise the broker as localhost
		assert(t, broker.Host, "localhost")
	}
	assert(t, broker.Port, int32(9092))

	topicMetadata := findTopicMetadata(t, metadata.TopicsMetadata, topicName)
	assert(t, topicMetadata.Error, ErrNoError)
	assert(t, topicMetadata.Topic, topicName)
	assertFatal(t, len(topicMetadata.PartitionsMetadata), 1)

	partitionMetadata := topicMetadata.PartitionsMetadata[0]
	assert(t, partitionMetadata.Error, ErrNoError)
	assert(t, partitionMetadata.ISR, []int32{0})
	assert(t, partitionMetadata.Leader, int32(0))
	assert(t, partitionMetadata.PartitionID, int32(0))
	assert(t, partitionMetadata.Replicas, []int32{0})
}

func testOffsetStorage(t *testing.T, topicName string, connector *DefaultConnector) {
	group := fmt.Sprintf("test-%d", time.Now().Unix())
	targetOffset := rand.Int63()

	offset, err := connector.GetOffset(group, topicName, 0)
	assertFatal(t, err, ErrUnknownTopicOrPartition)
	assert(t, offset, int64(-1))

	err = connector.CommitOffset(group, topicName, 0, targetOffset)
	assertFatal(t, err, nil)

	offset, err = connector.GetOffset(group, topicName, 0)
	assertFatal(t, err, nil)
	assert(t, offset, targetOffset)
}

func testProduce(t *testing.T, topicName string, numMessages int, connector *DefaultConnector) {
	produceRequest := new(ProduceRequest)
	produceRequest.AckTimeoutMs = 1000
	produceRequest.RequiredAcks = 1
	for i := 0; i < numMessages; i++ {
		produceRequest.AddMessage(topicName, 0, &Message{
			Key:   []byte(fmt.Sprintf("%d", numMessages-i)),
			Value: []byte(fmt.Sprintf("%d", i)),
		})
	}

	leader, err := connector.GetLeader(topicName, 0)
	assert(t, err, nil)
	assertNot(t, leader, (*BrokerConnection)(nil))
	bytes, err := connector.syncSendAndReceive(leader, produceRequest)
	assertFatal(t, err, nil)

	produceResponse := new(ProduceResponse)
	decodingErr := connector.decode(bytes, produceResponse)
	assertFatal(t, decodingErr, (*DecodingError)(nil))

	topicBlock, exists := produceResponse.Status[topicName]
	assertFatal(t, exists, true)
	partitionBlock, exists := topicBlock[int32(0)]
	assertFatal(t, exists, true)

	assert(t, partitionBlock.Error, ErrNoError)
	assert(t, partitionBlock.Offset, int64(0))
}

func testConsume(t *testing.T, topicName string, numMessages int, connector *DefaultConnector) {
	response, err := connector.Fetch(topicName, 0, 0)
	assertFatal(t, response.Error(topicName, 0), ErrNoError)
	assertFatal(t, err, nil)
	messages, err := response.GetMessages()
	assertFatal(t, err, nil)
	assertFatal(t, len(messages), numMessages)
	for i := 0; i < numMessages; i++ {
		message := messages[i]
		assert(t, message.Topic, topicName)
		assert(t, message.Partition, int32(0))
		assert(t, message.Offset, int64(i))
		assert(t, message.Key, []byte(fmt.Sprintf("%d", numMessages-i)))
		assert(t, message.Value, []byte(fmt.Sprintf("%d", i)))
	}
}

func findTopicMetadata(t *testing.T, metadata []*TopicMetadata, topic string) *TopicMetadata {
	for _, topicMetadata := range metadata {
		if topicMetadata.Topic == topic {
			return topicMetadata
		}
	}

	t.Fatalf("TopicMetadata for topic %s not found", topic)
	return nil
}
<commit_msg>Removed commented snippets in test<commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 siesta

import (
	"fmt"
	"math/rand"
	"net"
	"os"
	"testing"
	"time"
)

var ci = os.Getenv("TRAVIS_CI") != ""
var brokerUp = true
var brokerAddr = "localhost:9092"

func init() {
	conn, err := net.Dial("tcp", brokerAddr)
	if err == nil {
		brokerUp = true
		err = conn.Close()
		if err != nil {
			panic(err)
		}
	}
}

func TestDefaultConnectorFunctional(t *testing.T) {
	if !brokerUp && !ci {
		t.Skip("Broker is not running. Please spin up the broker at localhost:9092 for this test to work.")
	}

	numMessages := 1000
	topicName := fmt.Sprintf("siesta-%d", time.Now().Unix())

	connector := testConnector(t)
	testTopicMetadata(t, topicName, connector)
	testOffsetStorage(t, topicName, connector)
	testProduce(t, topicName, numMessages, connector)
	testConsume(t, topicName, numMessages, connector)
	closeWithin(t, time.Second, connector)

	anotherConnector := testConnector(t)
	//should also work fine - must get topic metadata before consuming
	testConsume(t, topicName, numMessages, anotherConnector)
	closeWithin(t, time.Second, anotherConnector)
}

func testTopicMetadata(t *testing.T, topicName string, connector *DefaultConnector) {
	metadata, err := connector.GetTopicMetadata([]string{topicName})
	assertFatal(t, err, nil)

	assertNot(t, len(metadata.Brokers), 0)
	assertNot(t, len(metadata.TopicsMetadata), 0)
	if len(metadata.Brokers) > 1 {
		t.Skip("Cluster should consist only of one broker for this test to run.")
	}

	broker := metadata.Brokers[0]
	assert(t, broker.ID, int32(0))
	if ci {
		// this can be asserted on Travis only as we are guaranteed to advertise the broker as localhost
		assert(t, broker.Host, "localhost")
	}
	assert(t, broker.Port, int32(9092))

	topicMetadata := findTopicMetadata(t, metadata.TopicsMetadata, topicName)
	assert(t, topicMetadata.Error, ErrNoError)
	assert(t, topicMetadata.Topic, topicName)
	assertFatal(t, len(topicMetadata.PartitionsMetadata), 1)

	partitionMetadata := topicMetadata.PartitionsMetadata[0]
	assert(t, partitionMetadata.Error, ErrNoError)
	assert(t, partitionMetadata.ISR, []int32{0})
	assert(t, partitionMetadata.Leader, int32(0))
	assert(t, partitionMetadata.PartitionID, int32(0))
	assert(t, partitionMetadata.Replicas, []int32{0})
}

func testOffsetStorage(t *testing.T, topicName string, connector *DefaultConnector) {
	group := fmt.Sprintf("test-%d", time.Now().Unix())
	targetOffset := rand.Int63()

	offset, err := connector.GetOffset(group, topicName, 0)
	assertFatal(t, err, ErrUnknownTopicOrPartition)
	assert(t, offset, int64(-1))

	err = connector.CommitOffset(group, topicName, 0, targetOffset)
	assertFatal(t, err, nil)

	offset, err = connector.GetOffset(group, topicName, 0)
	assertFatal(t, err, nil)
	assert(t, offset, targetOffset)
}

func testProduce(t *testing.T, topicName string, numMessages int, connector *DefaultConnector) {
	produceRequest := new(ProduceRequest)
	produceRequest.AckTimeoutMs = 1000
	produceRequest.RequiredAcks = 1
	for i := 0; i < numMessages; i++ {
		produceRequest.AddMessage(topicName, 0, &Message{
			Key:   []byte(fmt.Sprintf("%d", numMessages-i)),
			Value: []byte(fmt.Sprintf("%d", i)),
		})
	}

	leader, err := connector.GetLeader(topicName, 0)
	assert(t, err, nil)
	assertNot(t, leader, (*BrokerConnection)(nil))
	bytes, err := connector.syncSendAndReceive(leader, produceRequest)
	assertFatal(t, err, nil)

	produceResponse := new(ProduceResponse)
	decodingErr := connector.decode(bytes, produceResponse)
	assertFatal(t, decodingErr, (*DecodingError)(nil))

	topicBlock, exists := produceResponse.Status[topicName]
	assertFatal(t, exists, true)
	partitionBlock, exists := topicBlock[int32(0)]
	assertFatal(t, exists, true)

	assert(t, partitionBlock.Error, ErrNoError)
	assert(t, partitionBlock.Offset, int64(0))
}

func testConsume(t *testing.T, topicName string, numMessages int, connector *DefaultConnector) {
	response, err := connector.Fetch(topicName, 0, 0)
	assertFatal(t, response.Error(topicName, 0), ErrNoError)
	assertFatal(t, err, nil)
	messages, err := response.GetMessages()
	assertFatal(t, err, nil)
	assertFatal(t, len(messages), numMessages)
	for i := 0; i < numMessages; i++ {
		message := messages[i]
		assert(t, message.Topic, topicName)
		assert(t, message.Partition, int32(0))
		assert(t, message.Offset, int64(i))
		assert(t, message.Key, []byte(fmt.Sprintf("%d", numMessages-i)))
		assert(t, message.Value, []byte(fmt.Sprintf("%d", i)))
	}
}

func findTopicMetadata(t *testing.T, metadata []*TopicMetadata, topic string) *TopicMetadata {
	for _, topicMetadata := range metadata {
		if topicMetadata.Topic == topic {
			return topicMetadata
		}
	}

	t.Fatalf("TopicMetadata for topic %s not found", topic)
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2014, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package table

import (
	"bytes"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"github.com/conformal/goleveldb/leveldb/iterator"
	"github.com/conformal/goleveldb/leveldb/opt"
	"github.com/conformal/goleveldb/leveldb/testutil"
	"github.com/conformal/goleveldb/leveldb/util"
)

type tableWrapper struct {
	*Reader
}

func (t tableWrapper) Find(key []byte) (rkey, rvalue []byte, err error) {
	return t.Reader.Find(key, nil)
}

func (t tableWrapper) Get(key []byte) (value []byte, err error) {
	return t.Reader.Get(key, nil)
}

func (t tableWrapper) NewIterator(slice *util.Range) iterator.Iterator {
	return t.Reader.NewIterator(slice, nil)
}

var _ = testutil.Defer(func() {
	Describe("Table", func() {
		Describe("approximate offset test", func() {
			var (
				buf = &bytes.Buffer{}
				o   = &opt.Options{
					BlockSize:   1024,
					Compression: opt.NoCompression,
				}
			)

			// Building the table.
			tw := NewWriter(buf, o)
			tw.Append([]byte("k01"), []byte("hello"))
			tw.Append([]byte("k02"), []byte("hello2"))
			tw.Append([]byte("k03"), bytes.Repeat([]byte{'x'}, 10000))
			tw.Append([]byte("k04"), bytes.Repeat([]byte{'x'}, 200000))
			tw.Append([]byte("k05"), bytes.Repeat([]byte{'x'}, 300000))
			tw.Append([]byte("k06"), []byte("hello3"))
			tw.Append([]byte("k07"), bytes.Repeat([]byte{'x'}, 100000))
			err := tw.Close()

			It("Should be able to approximate offset of a key correctly", func() {
				Expect(err).ShouldNot(HaveOccurred())

				tr := NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()), nil, o)
				CheckOffset := func(key string, expect, threshold int) {
					offset, err := tr.GetApproximateOffset([]byte(key))
					Expect(err).ShouldNot(HaveOccurred())
					Expect(offset).Should(BeNumerically("~", expect, threshold), "Offset of key %q", key)
				}

				CheckOffset("k0", 0, 0)
				CheckOffset("k01a", 0, 0)
				CheckOffset("k02", 0, 0)
				CheckOffset("k03", 0, 0)
				CheckOffset("k04", 10000, 1000)
				CheckOffset("k04a", 210000, 1000)
				CheckOffset("k05", 210000, 1000)
				CheckOffset("k06", 510000, 1000)
				CheckOffset("k07", 510000, 1000)
				CheckOffset("xyz", 610000, 2000)
			})
		})

		Describe("read test", func() {
			testutil.AllKeyValueTesting(nil, func(kv testutil.KeyValue) testutil.DB {
				o := &opt.Options{
					BlockSize:            512,
					BlockRestartInterval: 3,
				}
				buf := &bytes.Buffer{}

				// Building the table.
				tw := NewWriter(buf, o)
				kv.Iterate(func(i int, key, value []byte) {
					tw.Append(key, value)
				})
				tw.Close()

				// Opening the table.
				tr := NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()), nil, o)
				return tableWrapper{tr}
			})
		})
	})
})
<commit_msg>table: Add 'with one key per block' test spec<commit_after>// Copyright (c) 2014, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package table

import (
	"bytes"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"github.com/conformal/goleveldb/leveldb/iterator"
	"github.com/conformal/goleveldb/leveldb/opt"
	"github.com/conformal/goleveldb/leveldb/testutil"
	"github.com/conformal/goleveldb/leveldb/util"
)

type tableWrapper struct {
	*Reader
}

func (t tableWrapper) Find(key []byte) (rkey, rvalue []byte, err error) {
	return t.Reader.Find(key, nil)
}

func (t tableWrapper) Get(key []byte) (value []byte, err error) {
	return t.Reader.Get(key, nil)
}

func (t tableWrapper) NewIterator(slice *util.Range) iterator.Iterator {
	return t.Reader.NewIterator(slice, nil)
}

var _ = testutil.Defer(func() {
	Describe("Table", func() {
		Describe("approximate offset test", func() {
			var (
				buf = &bytes.Buffer{}
				o   = &opt.Options{
					BlockSize:   1024,
					Compression: opt.NoCompression,
				}
			)

			// Building the table.
			tw := NewWriter(buf, o)
			tw.Append([]byte("k01"), []byte("hello"))
			tw.Append([]byte("k02"), []byte("hello2"))
			tw.Append([]byte("k03"), bytes.Repeat([]byte{'x'}, 10000))
			tw.Append([]byte("k04"), bytes.Repeat([]byte{'x'}, 200000))
			tw.Append([]byte("k05"), bytes.Repeat([]byte{'x'}, 300000))
			tw.Append([]byte("k06"), []byte("hello3"))
			tw.Append([]byte("k07"), bytes.Repeat([]byte{'x'}, 100000))
			err := tw.Close()

			It("Should be able to approximate offset of a key correctly", func() {
				Expect(err).ShouldNot(HaveOccurred())

				tr := NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()), nil, o)
				CheckOffset := func(key string, expect, threshold int) {
					offset, err := tr.GetApproximateOffset([]byte(key))
					Expect(err).ShouldNot(HaveOccurred())
					Expect(offset).Should(BeNumerically("~", expect, threshold), "Offset of key %q", key)
				}

				CheckOffset("k0", 0, 0)
				CheckOffset("k01a", 0, 0)
				CheckOffset("k02", 0, 0)
				CheckOffset("k03", 0, 0)
				CheckOffset("k04", 10000, 1000)
				CheckOffset("k04a", 210000, 1000)
				CheckOffset("k05", 210000, 1000)
				CheckOffset("k06", 510000, 1000)
				CheckOffset("k07", 510000, 1000)
				CheckOffset("xyz", 610000, 2000)
			})
		})

		Describe("read test", func() {
			Build := func(kv testutil.KeyValue) testutil.DB {
				o := &opt.Options{
					BlockSize:            512,
					BlockRestartInterval: 3,
				}
				buf := &bytes.Buffer{}

				// Building the table.
				tw := NewWriter(buf, o)
				kv.Iterate(func(i int, key, value []byte) {
					tw.Append(key, value)
				})
				tw.Close()

				// Opening the table.
				tr := NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()), nil, o)
				return tableWrapper{tr}
			}
			Test := func(kv *testutil.KeyValue, body func(r *Reader)) func() {
				return func() {
					db := Build(*kv)
					if body != nil {
						body(db.(tableWrapper).Reader)
					}
					testutil.KeyValueTesting(nil, db, *kv)
				}
			}

			testutil.AllKeyValueTesting(nil, Build)
			Describe("with one key per block", Test(testutil.KeyValue_Generate(nil, 9, 1, 10, 512, 512), func(r *Reader) {
				It("should have correct blocks number", func() {
					Expect(r.indexBlock.restartsLen).Should(Equal(9))
				})
			}))
		})
	})
})
<|endoftext|>


<commit_before>package taskserver

import (
	"database/sql"
	"time"

	"github.com/go-errors/errors"
	_ "github.com/jackc/pgx/stdlib"
	"github.com/privacybydesign/irmago/server/keyshare"
)

type taskHandler struct {
	conf *Configuration
	db   keyshare.DB
}

func newHandler(conf *Configuration) (*taskHandler, error) {
	err := processConfiguration(conf)
	if err != nil {
		return nil, err
	}
	db, err := sql.Open("pgx", conf.DBConnstring)
	if err != nil {
		return nil, err
	}
	if err = db.Ping(); err != nil {
		return nil, errors.Errorf("failed to connect to database: %v", err)
	}

	task := &taskHandler{db: keyshare.DB{DB: db}, conf: conf}
	return task, nil
}

func Do(conf *Configuration) error {
	task, err := newHandler(conf)
	if err != nil {
		return err
	}

	task.cleanupEmails()
	task.cleanupTokens()
	task.cleanupAccounts()
	task.expireAccounts()

	return nil
}

// Remove email addresses marked for deletion long enough ago
func (t *taskHandler) cleanupEmails() {
	_, err := t.db.Exec("DELETE FROM irma.emails WHERE delete_on < $1", time.Now().Unix())
	if err != nil {
		t.conf.Logger.WithField("error", err).Error("Could not remove email addresses marked for deletion")
	}
}

// Remove old login and email verifciation tokens
func (t *taskHandler) cleanupTokens() {
	_, err := t.db.Exec("DELETE FROM irma.email_login_tokens WHERE expiry < $1", time.Now().Unix())
	if err != nil {
		t.conf.Logger.WithField("error", err).Error("Could not remove email login tokens that have expired")
		return
	}
	_, err = t.db.Exec("DELETE FROM irma.email_verification_tokens WHERE expiry < $1", time.Now().Unix())
	if err != nil {
		t.conf.Logger.WithField("error", err).Error("Could not remove email verification tokens that have expired")
	}
}

// Cleanup accounts disabled long enough ago.
func (t *taskHandler) cleanupAccounts() {
	_, err := t.db.Exec("DELETE FROM irma.users WHERE delete_on < $1 AND (coredata IS NULL OR last_seen < delete_on - $2)",
		time.Now().Unix(),
		t.conf.DeleteDelay*24*60*60)
	if err != nil {
		t.conf.Logger.WithField("error", err).Error("Could not remove accounts scheduled for deletion")
	}
}

func (t *taskHandler) sendExpiryEmails(id int64, username, lang string) error {
	// Fetch user's email addresses
	err := t.db.QueryIterate("SELECT email FROM irma.emails WHERE user_id = $1",
		func(emailRes *sql.Rows) error {
			var email string
			err := emailRes.Scan(&email)
			if err != nil {
				return err
			}

			// And send
			err = t.conf.SendEmail(
				t.conf.deleteExpiredAccountTemplate,
				t.conf.DeleteExpiredAccountSubject,
				map[string]string{"Username": username, "Email": email},
				[]string{email},
				lang,
			)
			return err
		},
		id,
	)
	if err != nil {
		t.conf.Logger.WithField("error", err).Error("Could not retrieve user's email addresses")
		return err
	}
	return nil
}

// Mark old unused accounts for deletion, and inform their owners.
func (t *taskHandler) expireAccounts() {
	// Disable this task when email server is not given
	if t.conf.EmailServer == "" {
		t.conf.Logger.Warning("Expiring accounts is disabled, as no email server is configured")
		return
	}

	// Iterate over users we havent seen in ExpiryDelay days, and which have a registered email.
	// We ignore (and thus keep alive) accounts without email addresses, as we can't inform their owners.
	// (Note that for such accounts we store no email addresses, i.e. no personal data whatsoever.)
	err := t.db.QueryIterate(`SELECT id, username, language
							FROM irma.users
							WHERE last_seen < $1
								AND (
										SELECT count(*)
										FROM irma.emails
										WHERE irma.users.id = irma.emails.user_id
									) > 0
							LIMIT 10`,
		func(res *sql.Rows) error {
			var id int64
			var username string
			var lang string
			err := res.Scan(&id, &username, &lang)
			if err != nil {
				return err
			}

			// Send emails
			err = t.sendExpiryEmails(id, username, lang)
			if err != nil {
				return err // already logged, just abort
			}

			// Finally, do marking for deletion
			err = t.db.ExecUser("UPDATE irma.users SET delete_on = $2 WHERE id = $1", id,
				time.Now().Add(time.Duration(24*t.conf.DeleteDelay)*time.Hour).Unix())
			if err != nil {
				return err
			}
			return nil
		},
		time.Now().Add(time.Duration(-24*t.conf.ExpiryDelay)*time.Hour).Unix(),
	)
	if err != nil {
		t.conf.Logger.WithField("error", err).Error("Could not query for accounts that have expired")
		return
	}
}
<commit_msg>refactor: reformat multiline SQL query<commit_after>package taskserver

import (
	"database/sql"
	"time"

	"github.com/go-errors/errors"
	_ "github.com/jackc/pgx/stdlib"
	"github.com/privacybydesign/irmago/server/keyshare"
)

type taskHandler struct {
	conf *Configuration
	db   keyshare.DB
}

func newHandler(conf *Configuration) (*taskHandler, error) {
	err := processConfiguration(conf)
	if err != nil {
		return nil, err
	}
	db, err := sql.Open("pgx", conf.DBConnstring)
	if err != nil {
		return nil, err
	}
	if err = db.Ping(); err != nil {
		return nil, errors.Errorf("failed to connect to database: %v", err)
	}

	task := &taskHandler{db: keyshare.DB{DB: db}, conf: conf}
	return task, nil
}

func Do(conf *Configuration) error {
	task, err := newHandler(conf)
	if err != nil {
		return err
	}

	task.cleanupEmails()
	task.cleanupTokens()
	task.cleanupAccounts()
	task.expireAccounts()

	return nil
}

// Remove email addresses marked for deletion long enough ago
func (t *taskHandler) cleanupEmails() {
	_, err := t.db.Exec("DELETE FROM irma.emails WHERE delete_on < $1", time.Now().Unix())
	if err != nil {
		t.conf.Logger.WithField("error", err).Error("Could not remove email addresses marked for deletion")
	}
}

// Remove old login and email verifciation tokens
func (t *taskHandler) cleanupTokens() {
	_, err := t.db.Exec("DELETE FROM irma.email_login_tokens WHERE expiry < $1", time.Now().Unix())
	if err != nil {
		t.conf.Logger.WithField("error", err).Error("Could not remove email login tokens that have expired")
		return
	}
	_, err = t.db.Exec("DELETE FROM irma.email_verification_tokens WHERE expiry < $1", time.Now().Unix())
	if err != nil {
		t.conf.Logger.WithField("error", err).Error("Could not remove email verification tokens that have expired")
	}
}

// Cleanup accounts disabled long enough ago.
func (t *taskHandler) cleanupAccounts() {
	_, err := t.db.Exec("DELETE FROM irma.users WHERE delete_on < $1 AND (coredata IS NULL OR last_seen < delete_on - $2)",
		time.Now().Unix(),
		t.conf.DeleteDelay*24*60*60)
	if err != nil {
		t.conf.Logger.WithField("error", err).Error("Could not remove accounts scheduled for deletion")
	}
}

func (t *taskHandler) sendExpiryEmails(id int64, username, lang string) error {
	// Fetch user's email addresses
	err := t.db.QueryIterate("SELECT email FROM irma.emails WHERE user_id = $1",
		func(emailRes *sql.Rows) error {
			var email string
			err := emailRes.Scan(&email)
			if err != nil {
				return err
			}

			// And send
			err = t.conf.SendEmail(
				t.conf.deleteExpiredAccountTemplate,
				t.conf.DeleteExpiredAccountSubject,
				map[string]string{"Username": username, "Email": email},
				[]string{email},
				lang,
			)
			return err
		},
		id,
	)
	if err != nil {
		t.conf.Logger.WithField("error", err).Error("Could not retrieve user's email addresses")
		return err
	}
	return nil
}

// Mark old unused accounts for deletion, and inform their owners.
func (t *taskHandler) expireAccounts() {
	// Disable this task when email server is not given
	if t.conf.EmailServer == "" {
		t.conf.Logger.Warning("Expiring accounts is disabled, as no email server is configured")
		return
	}

	// Iterate over users we havent seen in ExpiryDelay days, and which have a registered email.
	// We ignore (and thus keep alive) accounts without email addresses, as we can't inform their owners.
	// (Note that for such accounts we store no email addresses, i.e. no personal data whatsoever.)
	err := t.db.QueryIterate(`
		SELECT id, username, language
		FROM irma.users
		WHERE last_seen < $1 AND (
			SELECT count(*)
			FROM irma.emails
			WHERE irma.users.id = irma.emails.user_id
		) > 0
		LIMIT 10`,
		func(res *sql.Rows) error {
			var id int64
			var username string
			var lang string
			err := res.Scan(&id, &username, &lang)
			if err != nil {
				return err
			}

			// Send emails
			err = t.sendExpiryEmails(id, username, lang)
			if err != nil {
				return err // already logged, just abort
			}

			// Finally, do marking for deletion
			err = t.db.ExecUser("UPDATE irma.users SET delete_on = $2 WHERE id = $1", id,
				time.Now().Add(time.Duration(24*t.conf.DeleteDelay)*time.Hour).Unix())
			if err != nil {
				return err
			}
			return nil
		},
		time.Now().Add(time.Duration(-24*t.conf.ExpiryDelay)*time.Hour).Unix(),
	)
	if err != nil {
		t.conf.Logger.WithField("error", err).Error("Could not query for accounts that have expired")
		return
	}
}
<|endoftext|>