<commit_before>package cli

import (
	"errors"
	"fmt"
	"gopkg.in/alecthomas/kingpin.v2"
	"net/url"
	"os"
	"strings"
)

var (
	Verbose bool
)

func GetModuleName() (string, error) {
	if len(os.Args) < 2 {
		return "", errors.New(fmt.Sprintf(
			"Must have at least one argument for the CLI module name: %s <modname>", os.Args[0]))
	}
	return os.Args[1], nil
}

func GetArguments() []string {
	// Exercise validation of argument count:
	if len(os.Args) < 2 {
		return make([]string, 0)
	}
	return os.Args[2:]
}

func NewApp(version string, author string, longDescription string) (*kingpin.Application, error) {
	modName, err := GetModuleName()
	if err != nil {
		return nil, err
	}

	app := kingpin.New(modName, longDescription)
	app.Version(version)
	app.Author(author)
	return app, nil
}

// Add all of the below arguments and commands

func HandleCommonArgs(
	app *kingpin.Application,
	defaultServiceName string,
	shortDescription string,
	connectionTypes []string) {
	HandleCommonFlags(app, defaultServiceName, shortDescription)
	HandleCommonSections(app, connectionTypes)
}

// Standard Arguments

func HandleCommonFlags(app *kingpin.Application, defaultServiceName string, shortDescription string) {
	app.HelpFlag.Short('h') // in addition to default '--help'
	app.Flag("verbose", "Enable extra logging of requests/responses").Short('v').BoolVar(&Verbose)

	// This fulfills an interface that's expected by the main DC/OS CLI:
	// Prints a description of the module.
	app.Flag("info", "Show short description.").PreAction(func(*kingpin.ParseContext) error {
		fmt.Fprintf(os.Stdout, "%s\n", shortDescription)
		os.Exit(0)
		return nil
	}).Bool()

	app.Flag("force-insecure", "Allow unverified TLS certificates when querying service").BoolVar(&tlsForceInsecure)

	// Overrides of data that we fetch from DC/OS CLI:

	// Support using "DCOS_AUTH_TOKEN" or "AUTH_TOKEN" when available
	app.Flag("custom-auth-token", "Custom auth token to use when querying service").Envar("DCOS_AUTH_TOKEN").PlaceHolder("DCOS_AUTH_TOKEN").StringVar(&dcosAuthToken)
	// Support using "DCOS_URI" or "DCOS_URL" when available
	app.Flag("custom-dcos-url", "Custom cluster URL to use when querying service").Envar("DCOS_URI").Envar("DCOS_URL").PlaceHolder("DCOS_URI/DCOS_URL").StringVar(&dcosUrl)
	// Support using "DCOS_CA_PATH" or "DCOS_CERT_PATH" when available
	app.Flag("custom-cert-path", "Custom TLS CA certificate file to use when querying service").Envar("DCOS_CA_PATH").Envar("DCOS_CERT_PATH").PlaceHolder("DCOS_CA_PATH/DCOS_CERT_PATH").StringVar(&tlsCACertPath)

	// Default to --name <name> : use provided framework name (default to <modulename>.service_name, if available)
	overrideServiceName := OptionalCLIConfigValue(fmt.Sprintf("%s.service_name", os.Args[1]))
	if len(overrideServiceName) != 0 {
		defaultServiceName = overrideServiceName
	}
	app.Flag("name", "Name of the service instance to query").Default(defaultServiceName).StringVar(&ServiceName)
}

// All sections

func HandleCommonSections(app *kingpin.Application, connectionTypes []string) {
	HandleConfigSection(app)
	HandleConnectionSection(app, connectionTypes)
	HandlePlanSection(app)
	HandleStateSection(app)
	HandleHealthSection(app)
}

// Config section

type ConfigHandler struct {
	ShowId string
}

func (cmd *ConfigHandler) RunList(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("v1/configurations"))
	return nil
}
func (cmd *ConfigHandler) RunShow(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet(fmt.Sprintf("v1/configurations/%s", cmd.ShowId)))
	return nil
}
func (cmd *ConfigHandler) RunTarget(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("v1/configurations/target"))
	return nil
}
func (cmd *ConfigHandler) RunTargetId(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("v1/configurations/targetId"))
	return nil
}

func HandleConfigSection(app *kingpin.Application) {
	// config <list, show, target, target_id>
	cmd := &ConfigHandler{}
	config := app.Command("config", "View persisted configurations")

	config.Command("list", "List IDs of all available configurations").Action(cmd.RunList)

	show := config.Command("show", "Display a specified configuration").Action(cmd.RunShow)
	show.Arg("config_id", "ID of the configuration to display").Required().StringVar(&cmd.ShowId)

	config.Command("target", "Display the target configuration").Action(cmd.RunTarget)

	config.Command("target_id", "List ID of the target configuration").Action(cmd.RunTargetId)
}

// Connection section

type ConnectionHandler struct {
	TypeName string
}

func (cmd *ConnectionHandler) RunConnection(c *kingpin.ParseContext) error {
	if len(cmd.TypeName) == 0 {
		// Root endpoint: Always produce JSON
		PrintJSON(HTTPGet("v1/connection"))
	} else {
		// Any custom type endpoints: May be any format, so just print the raw text
		PrintText(HTTPGet(fmt.Sprintf("v1/connection/%s", cmd.TypeName)))
	}
	return nil
}

func HandleConnectionSection(app *kingpin.Application, connectionTypes []string) {
	// connection [type]
	cmd := &ConnectionHandler{}
	connection := app.Command("connection", fmt.Sprintf("View connection information (custom types: %s)", strings.Join(connectionTypes, ", "))).Action(cmd.RunConnection)
	if len(connectionTypes) != 0 {
		connection.Arg("type", fmt.Sprintf("Custom type of the connection data to display (%s)", strings.Join(connectionTypes, ", "))).StringVar(&cmd.TypeName)
	}
}

// Plan section

type PlanHandler struct {
	PhaseId string
	BlockId string
}

func (cmd *PlanHandler) RunActive(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("v1/plan/status"))
	return nil
}
func (cmd *PlanHandler) RunShow(c *kingpin.ParseContext) error {
	// custom behavior: ignore 503 error
	response := HTTPQuery(CreateHTTPRequest("GET", "v1/plan"))
	if response.StatusCode != 503 {
		CheckHTTPResponse(response)
	}
	PrintJSON(response)
	return nil
}

func (cmd *PlanHandler) RunContinue(c *kingpin.ParseContext) error {
	PrintJSON(HTTPPost("v1/plan/continue"))
	return nil
}
func (cmd *PlanHandler) RunInterrupt(c *kingpin.ParseContext) error {
	PrintJSON(HTTPPost("v1/plan/interrupt"))
	return nil
}

func (cmd *PlanHandler) RunForce(c *kingpin.ParseContext) error {
	query := url.Values{}
	query.Set("phase", cmd.PhaseId)
	query.Set("block", cmd.BlockId)
	PrintJSON(HTTPPostQuery("v1/plan/forceComplete", query.Encode()))
	return nil
}
func (cmd *PlanHandler) RunRestart(c *kingpin.ParseContext) error {
	query := url.Values{}
	query.Set("phase", cmd.PhaseId)
	query.Set("block", cmd.BlockId)
	PrintJSON(HTTPPostQuery("v1/plan/restart", query.Encode()))
	return nil
}

func HandlePlanSection(app *kingpin.Application) {
	// plan <active, continue, force, interrupt, restart, show>
	cmd := &PlanHandler{}
	plan := app.Command("plan", "Query service plans")

	plan.Command("active", "Display the active operation chain, if any").Action(cmd.RunActive)
	plan.Command("show", "Display the full plan").Action(cmd.RunShow)

	plan.Command("continue", "Continue a currently Waiting operation").Action(cmd.RunContinue)
	plan.Command("interrupt", "Interrupt the current InProgress operation").Action(cmd.RunInterrupt)

	force := plan.Command("force", "Force the current operation to complete").Action(cmd.RunForce)
	force.Arg("phase", "UUID of the Phase containing the provided Block").Required().StringVar(&cmd.PhaseId)
	force.Arg("block", "UUID of Block to be restarted").Required().StringVar(&cmd.BlockId)

	restart := plan.Command("restart", "Restart the current operation").Action(cmd.RunRestart)
	restart.Arg("phase", "UUID of the Phase containing the provided Block").Required().StringVar(&cmd.PhaseId)
	restart.Arg("block", "UUID of Block to be restarted").Required().StringVar(&cmd.BlockId)
}

// State section

type StateHandler struct {
	TaskName string
}

func (cmd *StateHandler) RunFrameworkId(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("v1/state/frameworkId"))
	return nil
}
func (cmd *StateHandler) RunStatus(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet(fmt.Sprintf("v1/tasks/status/%s", cmd.TaskName)))
	return nil
}
func (cmd *StateHandler) RunTask(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet(fmt.Sprintf("v1/tasks/info/%s", cmd.TaskName)))
	return nil
}
func (cmd *StateHandler) RunTasks(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("v1/tasks"))
	return nil
}

func HandleStateSection(app *kingpin.Application) {
	// state <framework_id, status, task, tasks>
	cmd := &StateHandler{}
	state := app.Command("state", "View persisted state")

	state.Command("framework_id", "Display the mesos framework ID").Action(cmd.RunFrameworkId)

	status := state.Command("status", "Display the TaskStatus for a task name").Action(cmd.RunStatus)
	status.Arg("name", "Name of the task to display").Required().StringVar(&cmd.TaskName)

	task := state.Command("task", "Display the TaskInfo for a task name").Action(cmd.RunTask)
	task.Arg("name", "Name of the task to display").Required().StringVar(&cmd.TaskName)

	state.Command("tasks", "List names of all persisted tasks").Action(cmd.RunTasks)
}

// Health section

type HealthHandler struct {
}

func (cmd *HealthHandler) RunHealthCheck(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("admin/healthcheck"))
	return nil
}

func HandleHealthSection(app *kingpin.Application) {
	// health
	cmd := &HealthHandler{}
	app.Command("health", fmt.Sprintf("View healthcheck information)")).Action(cmd.RunHealthCheck)
}
<commit_msg>Update default CLI plan commands to use 'Step': catch up with PlanResource<commit_after>package cli

import (
	"errors"
	"fmt"
	"gopkg.in/alecthomas/kingpin.v2"
	"net/url"
	"os"
	"strings"
)

var (
	Verbose bool
)

func GetModuleName() (string, error) {
	if len(os.Args) < 2 {
		return "", errors.New(fmt.Sprintf(
			"Must have at least one argument for the CLI module name: %s <modname>", os.Args[0]))
	}
	return os.Args[1], nil
}

func GetArguments() []string {
	// Exercise validation of argument count:
	if len(os.Args) < 2 {
		return make([]string, 0)
	}
	return os.Args[2:]
}

func NewApp(version string, author string, longDescription string) (*kingpin.Application, error) {
	modName, err := GetModuleName()
	if err != nil {
		return nil, err
	}

	app := kingpin.New(modName, longDescription)
	app.Version(version)
	app.Author(author)
	return app, nil
}

// Add all of the below arguments and commands

func HandleCommonArgs(
	app *kingpin.Application,
	defaultServiceName string,
	shortDescription string,
	connectionTypes []string) {
	HandleCommonFlags(app, defaultServiceName, shortDescription)
	HandleCommonSections(app, connectionTypes)
}

// Standard Arguments

func HandleCommonFlags(app *kingpin.Application, defaultServiceName string, shortDescription string) {
	app.HelpFlag.Short('h') // in addition to default '--help'
	app.Flag("verbose", "Enable extra logging of requests/responses").Short('v').BoolVar(&Verbose)

	// This fulfills an interface that's expected by the main DC/OS CLI:
	// Prints a description of the module.
	app.Flag("info", "Show short description.").PreAction(func(*kingpin.ParseContext) error {
		fmt.Fprintf(os.Stdout, "%s\n", shortDescription)
		os.Exit(0)
		return nil
	}).Bool()

	app.Flag("force-insecure", "Allow unverified TLS certificates when querying service").BoolVar(&tlsForceInsecure)

	// Overrides of data that we fetch from DC/OS CLI:

	// Support using "DCOS_AUTH_TOKEN" or "AUTH_TOKEN" when available
	app.Flag("custom-auth-token", "Custom auth token to use when querying service").Envar("DCOS_AUTH_TOKEN").PlaceHolder("DCOS_AUTH_TOKEN").StringVar(&dcosAuthToken)
	// Support using "DCOS_URI" or "DCOS_URL" when available
	app.Flag("custom-dcos-url", "Custom cluster URL to use when querying service").Envar("DCOS_URI").Envar("DCOS_URL").PlaceHolder("DCOS_URI/DCOS_URL").StringVar(&dcosUrl)
	// Support using "DCOS_CA_PATH" or "DCOS_CERT_PATH" when available
	app.Flag("custom-cert-path", "Custom TLS CA certificate file to use when querying service").Envar("DCOS_CA_PATH").Envar("DCOS_CERT_PATH").PlaceHolder("DCOS_CA_PATH/DCOS_CERT_PATH").StringVar(&tlsCACertPath)

	// Default to --name <name> : use provided framework name (default to <modulename>.service_name, if available)
	overrideServiceName := OptionalCLIConfigValue(fmt.Sprintf("%s.service_name", os.Args[1]))
	if len(overrideServiceName) != 0 {
		defaultServiceName = overrideServiceName
	}
	app.Flag("name", "Name of the service instance to query").Default(defaultServiceName).StringVar(&ServiceName)
}

// All sections

func HandleCommonSections(app *kingpin.Application, connectionTypes []string) {
	HandleConfigSection(app)
	HandleConnectionSection(app, connectionTypes)
	HandlePlanSection(app)
	HandleStateSection(app)
	HandleHealthSection(app)
}

// Config section

type ConfigHandler struct {
	ShowId string
}

func (cmd *ConfigHandler) RunList(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("v1/configurations"))
	return nil
}
func (cmd *ConfigHandler) RunShow(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet(fmt.Sprintf("v1/configurations/%s", cmd.ShowId)))
	return nil
}
func (cmd *ConfigHandler) RunTarget(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("v1/configurations/target"))
	return nil
}
func (cmd *ConfigHandler) RunTargetId(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("v1/configurations/targetId"))
	return nil
}

func HandleConfigSection(app *kingpin.Application) {
	// config <list, show, target, target_id>
	cmd := &ConfigHandler{}
	config := app.Command("config", "View persisted configurations")

	config.Command("list", "List IDs of all available configurations").Action(cmd.RunList)

	show := config.Command("show", "Display a specified configuration").Action(cmd.RunShow)
	show.Arg("config_id", "ID of the configuration to display").Required().StringVar(&cmd.ShowId)

	config.Command("target", "Display the target configuration").Action(cmd.RunTarget)

	config.Command("target_id", "List ID of the target configuration").Action(cmd.RunTargetId)
}

// Connection section

type ConnectionHandler struct {
	TypeName string
}

func (cmd *ConnectionHandler) RunConnection(c *kingpin.ParseContext) error {
	if len(cmd.TypeName) == 0 {
		// Root endpoint: Always produce JSON
		PrintJSON(HTTPGet("v1/connection"))
	} else {
		// Any custom type endpoints: May be any format, so just print the raw text
		PrintText(HTTPGet(fmt.Sprintf("v1/connection/%s", cmd.TypeName)))
	}
	return nil
}

func HandleConnectionSection(app *kingpin.Application, connectionTypes []string) {
	// connection [type]
	cmd := &ConnectionHandler{}
	connection := app.Command("connection", fmt.Sprintf("View connection information (custom types: %s)", strings.Join(connectionTypes, ", "))).Action(cmd.RunConnection)
	if len(connectionTypes) != 0 {
		connection.Arg("type", fmt.Sprintf("Custom type of the connection data to display (%s)", strings.Join(connectionTypes, ", "))).StringVar(&cmd.TypeName)
	}
}

// Plan section

type PlanHandler struct {
	PhaseId string
	StepId  string
}

func (cmd *PlanHandler) RunActive(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("v1/plan/status"))
	return nil
}
func (cmd *PlanHandler) RunShow(c *kingpin.ParseContext) error {
	// custom behavior: ignore 503 error
	response := HTTPQuery(CreateHTTPRequest("GET", "v1/plan"))
	if response.StatusCode != 503 {
		CheckHTTPResponse(response)
	}
	PrintJSON(response)
	return nil
}

func (cmd *PlanHandler) RunContinue(c *kingpin.ParseContext) error {
	PrintJSON(HTTPPost("v1/plan/continue"))
	return nil
}
func (cmd *PlanHandler) RunInterrupt(c *kingpin.ParseContext) error {
	PrintJSON(HTTPPost("v1/plan/interrupt"))
	return nil
}

func (cmd *PlanHandler) RunForce(c *kingpin.ParseContext) error {
	query := url.Values{}
	query.Set("phase", cmd.PhaseId)
	query.Set("step", cmd.StepId)
	PrintJSON(HTTPPostQuery("v1/plan/forceComplete", query.Encode()))
	return nil
}
func (cmd *PlanHandler) RunRestart(c *kingpin.ParseContext) error {
	query := url.Values{}
	query.Set("phase", cmd.PhaseId)
	query.Set("step", cmd.StepId)
	PrintJSON(HTTPPostQuery("v1/plan/restart", query.Encode()))
	return nil
}

func HandlePlanSection(app *kingpin.Application) {
	// plan <active, continue, force, interrupt, restart, show>
	cmd := &PlanHandler{}
	plan := app.Command("plan", "Query service plans")

	plan.Command("active", "Display the active operation chain, if any").Action(cmd.RunActive)
	plan.Command("show", "Display the full plan").Action(cmd.RunShow)

	plan.Command("continue", "Continue a currently Waiting operation").Action(cmd.RunContinue)
	plan.Command("interrupt", "Interrupt the current InProgress operation").Action(cmd.RunInterrupt)

	force := plan.Command("force", "Force the current operation to complete").Action(cmd.RunForce)
	force.Arg("phase", "UUID of the Phase containing the provided Step").Required().StringVar(&cmd.PhaseId)
	force.Arg("step", "UUID of Step to be restarted").Required().StringVar(&cmd.StepId)

	restart := plan.Command("restart", "Restart the current operation").Action(cmd.RunRestart)
	restart.Arg("phase", "UUID of the Phase containing the provided Step").Required().StringVar(&cmd.PhaseId)
	restart.Arg("step", "UUID of Step to be restarted").Required().StringVar(&cmd.StepId)
}

// State section

type StateHandler struct {
	TaskName string
}

func (cmd *StateHandler) RunFrameworkId(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("v1/state/frameworkId"))
	return nil
}
func (cmd *StateHandler) RunStatus(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet(fmt.Sprintf("v1/tasks/status/%s", cmd.TaskName)))
	return nil
}
func (cmd *StateHandler) RunTask(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet(fmt.Sprintf("v1/tasks/info/%s", cmd.TaskName)))
	return nil
}
func (cmd *StateHandler) RunTasks(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("v1/tasks"))
	return nil
}

func HandleStateSection(app *kingpin.Application) {
	// state <framework_id, status, task, tasks>
	cmd := &StateHandler{}
	state := app.Command("state", "View persisted state")

	state.Command("framework_id", "Display the mesos framework ID").Action(cmd.RunFrameworkId)

	status := state.Command("status", "Display the TaskStatus for a task name").Action(cmd.RunStatus)
	status.Arg("name", "Name of the task to display").Required().StringVar(&cmd.TaskName)

	task := state.Command("task", "Display the TaskInfo for a task name").Action(cmd.RunTask)
	task.Arg("name", "Name of the task to display").Required().StringVar(&cmd.TaskName)

	state.Command("tasks", "List names of all persisted tasks").Action(cmd.RunTasks)
}

// Health section

type HealthHandler struct {
}

func (cmd *HealthHandler) RunHealthCheck(c *kingpin.ParseContext) error {
	PrintJSON(HTTPGet("admin/healthcheck"))
	return nil
}

func HandleHealthSection(app *kingpin.Application) {
	// health
	cmd := &HealthHandler{}
	app.Command("health", fmt.Sprintf("View healthcheck information)")).Action(cmd.RunHealthCheck)
}
<|endoftext|>


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

package lsp

import (
	"context"
	"go/token"
	"io/ioutil"
	"path/filepath"
	"reflect"
	"sort"
	"strings"
	"testing"

	"golang.org/x/tools/go/packages"
	"golang.org/x/tools/go/packages/packagestest"
	"golang.org/x/tools/internal/lsp/protocol"
	"golang.org/x/tools/internal/lsp/source"
)

func TestLSP(t *testing.T) {
	packagestest.TestAll(t, testLSP)
}

func testLSP(t *testing.T, exporter packagestest.Exporter) {
	const dir = "testdata"

	files := packagestest.MustCopyFileTree(dir)
	subdirs, err := ioutil.ReadDir(dir)
	if err != nil {
		t.Fatal(err)
	}
	for _, subdir := range subdirs {
		if !subdir.IsDir() {
			continue
		}
		dirpath := filepath.Join(dir, subdir.Name())
		if testFiles, err := ioutil.ReadDir(dirpath); err == nil {
			for _, file := range testFiles {
				if trimmed := strings.TrimSuffix(file.Name(), ".in"); trimmed != file.Name() {
					files[filepath.Join(subdir.Name(), trimmed)] = packagestest.Copy(filepath.Join(dirpath, file.Name()))
				}
			}
		}
	}
	modules := []packagestest.Module{
		{
			Name:  "golang.org/x/tools/internal/lsp",
			Files: files,
		},
	}
	exported := packagestest.Export(t, exporter, modules)
	defer exported.Cleanup()

	dirs := make(map[string]bool)

	// collect results for certain tests
	expectedDiagnostics := make(map[string][]protocol.Diagnostic)
	expectedCompletions := make(map[token.Position]*protocol.CompletionItem)

	s := &server{
		view: source.NewView(),
	}
	// merge the config objects
	cfg := *exported.Config
	cfg.Fset = s.view.Config.Fset
	cfg.Mode = packages.LoadSyntax
	s.view.Config = &cfg

	for _, module := range modules {
		for fragment := range module.Files {
			if !strings.HasSuffix(fragment, ".go") {
				continue
			}
			filename := exporter.Filename(exported, module.Name, fragment)
			expectedDiagnostics[filename] = []protocol.Diagnostic{}
			dirs[filepath.Dir(filename)] = true
		}
	}
	// Do a first pass to collect special markers
	if err := exported.Expect(map[string]interface{}{
		"item": func(name string, r packagestest.Range, _, _ string) {
			exported.Mark(name, r)
		},
	}); err != nil {
		t.Fatal(err)
	}
	// Collect any data that needs to be used by subsequent tests.
	if err := exported.Expect(map[string]interface{}{
		"diag": func(pos token.Position, msg string) {
			line := float64(pos.Line - 1)
			col := float64(pos.Column - 1)
			want := protocol.Diagnostic{
				Range: protocol.Range{
					Start: protocol.Position{
						Line:      line,
						Character: col,
					},
					End: protocol.Position{
						Line:      line,
						Character: col,
					},
				},
				Severity: protocol.SeverityError,
				Source:   "LSP",
				Message:  msg,
			}
			if _, ok := expectedDiagnostics[pos.Filename]; ok {
				expectedDiagnostics[pos.Filename] = append(expectedDiagnostics[pos.Filename], want)
			} else {
				t.Errorf("unexpected filename: %v", pos.Filename)
			}
		},
		"item": func(pos token.Position, label, detail, kind string) {
			var k protocol.CompletionItemKind
			switch kind {
			case "struct":
				k = protocol.StructCompletion
			case "func":
				k = protocol.FunctionCompletion
			case "var":
				k = protocol.VariableCompletion
			case "type":
				k = protocol.TypeParameterCompletion
			case "field":
				k = protocol.FieldCompletion
			case "interface":
				k = protocol.InterfaceCompletion
			case "const":
				k = protocol.ConstantCompletion
			case "method":
				k = protocol.MethodCompletion
			}
			expectedCompletions[pos] = &protocol.CompletionItem{
				Label:  label,
				Detail: detail,
				Kind:   float64(k),
			}
		},
	}); err != nil {
		t.Fatal(err)
	}

	// test completion
	testCompletion(t, exported, s, expectedCompletions)

	// test diagnostics
	var dirList []string
	for dir := range dirs {
		dirList = append(dirList, dir)
	}
	exported.Config.Mode = packages.LoadFiles
	pkgs, err := packages.Load(exported.Config, dirList...)
	if err != nil {
		t.Fatal(err)
	}
	testDiagnostics(t, s.view, pkgs, expectedDiagnostics)
}

func testDiagnostics(t *testing.T, v *source.View, pkgs []*packages.Package, wants map[string][]protocol.Diagnostic) {
	for _, pkg := range pkgs {
		for _, filename := range pkg.GoFiles {
			f := v.GetFile(source.ToURI(filename))
			diagnostics, err := source.Diagnostics(context.Background(), v, f)
			if err != nil {
				t.Fatal(err)
			}
			got := toProtocolDiagnostics(v, diagnostics[filename])
			sort.Slice(got, func(i int, j int) bool {
				return got[i].Range.Start.Line < got[j].Range.Start.Line
			})
			want := wants[filename]
			if equal := reflect.DeepEqual(want, got); !equal {
				t.Errorf("diagnostics failed for %s: (expected: %v), (got: %v)", filepath.Base(filename), want, got)
			}
		}
	}
}

func testCompletion(t *testing.T, exported *packagestest.Exported, s *server, wants map[token.Position]*protocol.CompletionItem) {
	if err := exported.Expect(map[string]interface{}{
		"complete": func(src token.Position, expected []token.Position) {
			var want []protocol.CompletionItem
			for _, pos := range expected {
				want = append(want, *wants[pos])
			}
			list, err := s.Completion(context.Background(), &protocol.CompletionParams{
				TextDocumentPositionParams: protocol.TextDocumentPositionParams{
					TextDocument: protocol.TextDocumentIdentifier{
						URI: protocol.DocumentURI(source.ToURI(src.Filename)),
					},
					Position: protocol.Position{
						Line:      float64(src.Line - 1),
						Character: float64(src.Column - 1),
					},
				},
			})
			if err != nil {
				t.Fatal(err)
			}
			got := list.Items
			if equal := reflect.DeepEqual(want, got); !equal {
				t.Errorf("completion failed for %s:%v:%v: (expected: %v), (got: %v)", filepath.Base(src.Filename), src.Line, src.Column, want, got)
			}
		},
	}); err != nil {
		t.Fatal(err)
	}
}
<commit_msg>internal/lsp: better handling of .in files<commit_after>// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package lsp

import (
	"context"
	"go/token"
	"path/filepath"
	"reflect"
	"sort"
	"strings"
	"testing"

	"golang.org/x/tools/go/packages"
	"golang.org/x/tools/go/packages/packagestest"
	"golang.org/x/tools/internal/lsp/protocol"
	"golang.org/x/tools/internal/lsp/source"
)

func TestLSP(t *testing.T) {
	packagestest.TestAll(t, testLSP)
}

func testLSP(t *testing.T, exporter packagestest.Exporter) {
	const dir = "testdata"

	files := packagestest.MustCopyFileTree(dir)
	for fragment, operation := range files {
		if trimmed := strings.TrimSuffix(fragment, ".in"); trimmed != fragment {
			delete(files, fragment)
			files[trimmed] = operation
		}
	}
	modules := []packagestest.Module{
		{
			Name:  "golang.org/x/tools/internal/lsp",
			Files: files,
		},
	}
	exported := packagestest.Export(t, exporter, modules)
	defer exported.Cleanup()

	dirs := make(map[string]bool)

	// collect results for certain tests
	expectedDiagnostics := make(map[string][]protocol.Diagnostic)
	expectedCompletions := make(map[token.Position]*protocol.CompletionItem)

	s := &server{
		view: source.NewView(),
	}
	// merge the config objects
	cfg := *exported.Config
	cfg.Fset = s.view.Config.Fset
	cfg.Mode = packages.LoadSyntax
	s.view.Config = &cfg

	for _, module := range modules {
		for fragment := range module.Files {
			if !strings.HasSuffix(fragment, ".go") {
				continue
			}
			filename := exporter.Filename(exported, module.Name, fragment)
			expectedDiagnostics[filename] = []protocol.Diagnostic{}
			dirs[filepath.Dir(filename)] = true
		}
	}
	// Do a first pass to collect special markers
	if err := exported.Expect(map[string]interface{}{
		"item": func(name string, r packagestest.Range, _, _ string) {
			exported.Mark(name, r)
		},
	}); err != nil {
		t.Fatal(err)
	}
	// Collect any data that needs to be used by subsequent tests.
	if err := exported.Expect(map[string]interface{}{
		"diag": func(pos token.Position, msg string) {
			line := float64(pos.Line - 1)
			col := float64(pos.Column - 1)
			want := protocol.Diagnostic{
				Range: protocol.Range{
					Start: protocol.Position{
						Line:      line,
						Character: col,
					},
					End: protocol.Position{
						Line:      line,
						Character: col,
					},
				},
				Severity: protocol.SeverityError,
				Source:   "LSP",
				Message:  msg,
			}
			if _, ok := expectedDiagnostics[pos.Filename]; ok {
				expectedDiagnostics[pos.Filename] = append(expectedDiagnostics[pos.Filename], want)
			} else {
				t.Errorf("unexpected filename: %v", pos.Filename)
			}
		},
		"item": func(pos token.Position, label, detail, kind string) {
			var k protocol.CompletionItemKind
			switch kind {
			case "struct":
				k = protocol.StructCompletion
			case "func":
				k = protocol.FunctionCompletion
			case "var":
				k = protocol.VariableCompletion
			case "type":
				k = protocol.TypeParameterCompletion
			case "field":
				k = protocol.FieldCompletion
			case "interface":
				k = protocol.InterfaceCompletion
			case "const":
				k = protocol.ConstantCompletion
			case "method":
				k = protocol.MethodCompletion
			}
			expectedCompletions[pos] = &protocol.CompletionItem{
				Label:  label,
				Detail: detail,
				Kind:   float64(k),
			}
		},
	}); err != nil {
		t.Fatal(err)
	}

	// test completion
	testCompletion(t, exported, s, expectedCompletions)

	// test diagnostics
	var dirList []string
	for dir := range dirs {
		dirList = append(dirList, dir)
	}
	exported.Config.Mode = packages.LoadFiles
	pkgs, err := packages.Load(exported.Config, dirList...)
	if err != nil {
		t.Fatal(err)
	}
	testDiagnostics(t, s.view, pkgs, expectedDiagnostics)
}

func testDiagnostics(t *testing.T, v *source.View, pkgs []*packages.Package, wants map[string][]protocol.Diagnostic) {
	for _, pkg := range pkgs {
		for _, filename := range pkg.GoFiles {
			f := v.GetFile(source.ToURI(filename))
			diagnostics, err := source.Diagnostics(context.Background(), v, f)
			if err != nil {
				t.Fatal(err)
			}
			got := toProtocolDiagnostics(v, diagnostics[filename])
			sort.Slice(got, func(i int, j int) bool {
				return got[i].Range.Start.Line < got[j].Range.Start.Line
			})
			want := wants[filename]
			if equal := reflect.DeepEqual(want, got); !equal {
				t.Errorf("diagnostics failed for %s: (expected: %v), (got: %v)", filepath.Base(filename), want, got)
			}
		}
	}
}

func testCompletion(t *testing.T, exported *packagestest.Exported, s *server, wants map[token.Position]*protocol.CompletionItem) {
	if err := exported.Expect(map[string]interface{}{
		"complete": func(src token.Position, expected []token.Position) {
			var want []protocol.CompletionItem
			for _, pos := range expected {
				want = append(want, *wants[pos])
			}
			list, err := s.Completion(context.Background(), &protocol.CompletionParams{
				TextDocumentPositionParams: protocol.TextDocumentPositionParams{
					TextDocument: protocol.TextDocumentIdentifier{
						URI: protocol.DocumentURI(source.ToURI(src.Filename)),
					},
					Position: protocol.Position{
						Line:      float64(src.Line - 1),
						Character: float64(src.Column - 1),
					},
				},
			})
			if err != nil {
				t.Fatal(err)
			}
			got := list.Items
			if equal := reflect.DeepEqual(want, got); !equal {
				t.Errorf("completion failed for %s:%v:%v: (expected: %v), (got: %v)", filepath.Base(src.Filename), src.Line, src.Column, want, got)
			}
		},
	}); err != nil {
		t.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 CoreOS, Inc.
//
// 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 resource

import (
	"bytes"
	"compress/gzip"
	"context"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"

	configErrors "github.com/coreos/ignition/v2/config/shared/errors"
	"github.com/coreos/ignition/v2/internal/log"
	"github.com/coreos/ignition/v2/internal/util"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/pin/tftp"
	"github.com/vincent-petithory/dataurl"
)

var (
	ErrSchemeUnsupported      = errors.New("unsupported source scheme")
	ErrPathNotAbsolute        = errors.New("path is not absolute")
	ErrNotFound               = errors.New("resource not found")
	ErrFailed                 = errors.New("failed to fetch resource")
	ErrCompressionUnsupported = errors.New("compression is not supported with that scheme")

	// ConfigHeaders are the HTTP headers that should be used when the Ignition
	// config is being fetched
	ConfigHeaders = http.Header{
		"Accept-Encoding": []string{"identity"},
		"Accept":          []string{"application/vnd.coreos.ignition+json; version=2.2.0, application/vnd.coreos.ignition+json; version=1; q=0.5, */*; q=0.1"},
	}
)

// Fetcher holds settings for fetching resources from URLs
type Fetcher struct {
	// The logger object to use when logging information.
	Logger *log.Logger

	// client is the http client that will be used when fetching http(s)
	// resources. If left nil, one will be created and used, but this means any
	// timeouts Ignition was configured to used will be ignored.
	client *HttpClient

	// The AWS Session to use when fetching resources from S3. If left nil, the
	// first S3 object that is fetched will initialize the field. This can be
	// used to set credentials.
	AWSSession *session.Session

	// The region where the AWS machine trying to fetch is.
	// This is used as a hint to fetch the S3 bucket from the right partition and region.
	S3RegionHint string
}

type FetchOptions struct {
	// Headers are the http headers that will be used when fetching http(s)
	// resources. They have no effect on other fetching schemes.
	Headers http.Header

	// Hash is the hash to use when calculating a fetched resource's hash. If
	// left as nil, no hash will be calculated.
	Hash hash.Hash

	// The expected sum to be produced by the given hasher. If the Hash field is
	// nil, this field is ignored.
	ExpectedSum []byte

	// Compression specifies the type of compression to use when decompressing
	// the fetched object. If left empty, no decompression will be used.
	Compression string
}

// FetchToBuffer will fetch the given url into a temporrary file, and then read
// in the contents of the file and delete it. It will return the downloaded
// contents, or an error if one was encountered.
func (f *Fetcher) FetchToBuffer(u url.URL, opts FetchOptions) ([]byte, error) {
	file, err := ioutil.TempFile("", "ignition")
	if err != nil {
		return nil, err
	}
	defer os.Remove(file.Name())
	defer file.Close()
	err = f.Fetch(u, file, opts)
	if err != nil {
		return nil, err
	}
	_, err = file.Seek(0, os.SEEK_SET)
	if err != nil {
		return nil, err
	}
	res, err := ioutil.ReadAll(file)
	if err != nil {
		return nil, err
	}
	return res, nil
}

// Fetch calls the appropriate FetchFrom* function based on the scheme of the
// given URL. The results will be decompressed if compression is set in opts,
// and written into dest. If opts.Hash is set the data stream will also be
// hashed and compared against opts.ExpectedSum, and any match failures will
// result in an error being returned.
//
// Fetch expects dest to be an empty file and for the cursor in the file to be
// at the beginning. Since some url schemes (ex: s3) use chunked downloads and
// fetch chunks out of order, Fetch's behavior when dest is not an empty file is
// undefined.
func (f *Fetcher) Fetch(u url.URL, dest *os.File, opts FetchOptions) error {
	switch u.Scheme {
	case "http", "https":
		return f.FetchFromHTTP(u, dest, opts)
	case "tftp":
		return f.FetchFromTFTP(u, dest, opts)
	case "data":
		return f.FetchFromDataURL(u, dest, opts)
	case "s3":
		return f.FetchFromS3(u, dest, opts)
	case "":
		return nil
	default:
		return ErrSchemeUnsupported
	}
}

// FetchFromTFTP fetches a resource from u via TFTP into dest, returning an
// error if one is encountered.
func (f *Fetcher) FetchFromTFTP(u url.URL, dest *os.File, opts FetchOptions) error {
	if !strings.ContainsRune(u.Host, ':') {
		u.Host = u.Host + ":69"
	}
	c, err := tftp.NewClient(u.Host)
	if err != nil {
		return err
	}
	wt, err := c.Receive(u.Path, "octet")
	if err != nil {
		return err
	}
	// The TFTP library takes an io.Writer to send data in to, but to decompress
	// the stream the gzip library wraps an io.Reader, so let's create a pipe to
	// connect these two things
	pReader, pWriter := io.Pipe()
	doneChan := make(chan error, 2)

	checkForDoneChanErr := func(err error) error {
		// If an error is encountered while decompressing or copying data out of
		// the pipe, there's probably an error from writing into the pipe that
		// will better describe what went wrong. This function does a
		// non-blocking read of doneChan, overriding the returned error val if
		// there's anything in doneChan.
		select {
		case writeErr := <-doneChan:
			if writeErr != nil {
				return writeErr
			}
			return err
		default:
			return err
		}
	}

	// A goroutine is used to handle writing the fetched data into the pipe
	// while also copying it out of the pipe concurrently
	go func() {
		_, err := wt.WriteTo(pWriter)
		doneChan <- err
		err = pWriter.Close()
		doneChan <- err
	}()
	err = f.decompressCopyHashAndVerify(dest, pReader, opts)
	if err != nil {
		return checkForDoneChanErr(err)
	}
	// receive the error from wt.WriteTo()
	err = <-doneChan
	if err != nil {
		return err
	}
	// receive the error from pWriter.Close()
	err = <-doneChan
	if err != nil {
		return err
	}
	return nil
}

// FetchFromHTTP fetches a resource from u via HTTP(S) into dest, returning an
// error if one is encountered.
func (f *Fetcher) FetchFromHTTP(u url.URL, dest *os.File, opts FetchOptions) error {
	// for the case when "config is not valid"
	// this if necessary if not spawned through kola (e.g. Packet Dashboard)
	if f.client == nil {
		logger := log.New(true)
		f.Logger = &logger
		if err := f.newHttpClient(); err != nil {
			return err
		}
	}

	dataReader, status, ctxCancel, err := f.client.getReaderWithHeader(u.String(), opts.Headers)
	if ctxCancel != nil {
		// whatever context getReaderWithHeader created for the request should
		// be cancelled once we're done reading the response
		defer ctxCancel()
	}
	if err != nil {
		return err
	}
	defer dataReader.Close()

	switch status {
	case http.StatusOK, http.StatusNoContent:
		break
	case http.StatusNotFound:
		return ErrNotFound
	default:
		return ErrFailed
	}

	return f.decompressCopyHashAndVerify(dest, dataReader, opts)
}

// FetchFromDataURL writes the data stored in the dataurl u into dest, returning
// an error if one is encountered.
func (f *Fetcher) FetchFromDataURL(u url.URL, dest *os.File, opts FetchOptions) error {
	if opts.Compression != "" {
		return ErrCompressionUnsupported
	}
	url, err := dataurl.DecodeString(u.String())
	if err != nil {
		return err
	}

	return f.decompressCopyHashAndVerify(dest, bytes.NewBuffer(url.Data), opts)
}

// FetchFromS3 gets data from an S3 bucket as described by u and writes it into
// dest, returning an error if one is encountered. It will attempt to acquire
// IAM credentials from the EC2 metadata service, and if this fails will attempt
// to fetch the object with anonymous credentials.
func (f *Fetcher) FetchFromS3(u url.URL, dest *os.File, opts FetchOptions) error {
	if opts.Compression != "" {
		return ErrCompressionUnsupported
	}
	ctx := context.Background()
	if f.client != nil && f.client.timeout != 0 {
		var cancelFn context.CancelFunc
		ctx, cancelFn = context.WithTimeout(ctx, f.client.timeout)
		defer cancelFn()
	}

	if f.AWSSession == nil {
		var err error
		f.AWSSession, err = session.NewSession(&aws.Config{
			Credentials: credentials.AnonymousCredentials,
		})
		if err != nil {
			return err
		}
	}
	sess := f.AWSSession.Copy()

	// Determine the partition and region this bucket is in
	regionHint := "us-east-1"
	if f.S3RegionHint != "" {
		regionHint = f.S3RegionHint
	}
	region, err := s3manager.GetBucketRegion(ctx, sess, u.Host, regionHint)
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok && aerr.Code() == "NotFound" {
			return fmt.Errorf("couldn't determine the region for bucket %q: %v", u.Host, err)
		}
		return err
	}

	sess.Config.Region = aws.String(region)

	var versionId *string
	if v, ok := u.Query()["versionId"]; ok && len(v) > 0 {
		versionId = aws.String(v[0])
	}

	input := &s3.GetObjectInput{
		Bucket:    &u.Host,
		Key:       &u.Path,
		VersionId: versionId,
	}
	err = f.fetchFromS3WithCreds(ctx, dest, input, sess)
	if err != nil {
		return err
	}
	if opts.Hash != nil {
		opts.Hash.Reset()
		_, err = dest.Seek(0, os.SEEK_SET)
		if err != nil {
			return err
		}
		_, err = io.Copy(opts.Hash, dest)
		if err != nil {
			return err
		}
		calculatedSum := opts.Hash.Sum(nil)
		if !bytes.Equal(calculatedSum, opts.ExpectedSum) {
			return util.ErrHashMismatch{
				Calculated: hex.EncodeToString(calculatedSum),
				Expected:   hex.EncodeToString(opts.ExpectedSum),
			}
		}
		f.Logger.Debug("file matches expected sum of: %s", hex.EncodeToString(opts.ExpectedSum))
	}
	return nil
}

func (f *Fetcher) fetchFromS3WithCreds(ctx context.Context, dest *os.File, input *s3.GetObjectInput, sess *session.Session) error {
	httpClient, err := defaultHTTPClient()
	if err != nil {
		return err
	}

	awsConfig := aws.NewConfig().WithHTTPClient(httpClient)
	s3Client := s3.New(sess, awsConfig)
	downloader := s3manager.NewDownloaderWithClient(s3Client)
	if _, err := downloader.DownloadWithContext(ctx, dest, input); err != nil {
		if awserrval, ok := err.(awserr.Error); ok && awserrval.Code() == "EC2RoleRequestError" {
			// If this error was due to an EC2 role request error, try again
			// with the anonymous credentials.
			sess.Config.Credentials = credentials.AnonymousCredentials
			return f.fetchFromS3WithCreds(ctx, dest, input, sess)
		}
		return err
	}
	return nil
}

// uncompress will wrap the given io.Reader in a decompresser specified in the
// FetchOptions, and return an io.ReadCloser with the decompressed data stream.
func (f *Fetcher) uncompress(r io.Reader, opts FetchOptions) (io.ReadCloser, error) {
	switch opts.Compression {
	case "":
		return ioutil.NopCloser(r), nil
	case "gzip":
		return gzip.NewReader(r)
	default:
		return nil, configErrors.ErrCompressionInvalid
	}
}

// decompressCopyHashAndVerify will decompress src if necessary, copy src into
// dest until src returns an io.EOF while also calculating a hash if one is set,
// and will return an error if there's any problems with any of this or if the
// hash doesn't match the expected hash in the opts.
func (f *Fetcher) decompressCopyHashAndVerify(dest io.Writer, src io.Reader, opts FetchOptions) error {
	decompressor, err := f.uncompress(src, opts)
	if err != nil {
		return err
	}
	defer decompressor.Close()
	if opts.Hash != nil {
		opts.Hash.Reset()
		dest = io.MultiWriter(dest, opts.Hash)
	}
	_, err = io.Copy(dest, decompressor)
	if err != nil {
		return err
	}
	if opts.Hash != nil {
		calculatedSum := opts.Hash.Sum(nil)
		if !bytes.Equal(calculatedSum, opts.ExpectedSum) {
			return util.ErrHashMismatch{
				Calculated: hex.EncodeToString(calculatedSum),
				Expected:   hex.EncodeToString(opts.ExpectedSum),
			}
		}
		f.Logger.Debug("file matches expected sum of: %s", hex.EncodeToString(opts.ExpectedSum))
	}
	return nil
}
<commit_msg>resource/url: update schema version in Accept header<commit_after>// Copyright 2016 CoreOS, Inc.
//
// 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 resource

import (
	"bytes"
	"compress/gzip"
	"context"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"

	configErrors "github.com/coreos/ignition/v2/config/shared/errors"
	"github.com/coreos/ignition/v2/internal/log"
	"github.com/coreos/ignition/v2/internal/util"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/pin/tftp"
	"github.com/vincent-petithory/dataurl"
)

var (
	ErrSchemeUnsupported      = errors.New("unsupported source scheme")
	ErrPathNotAbsolute        = errors.New("path is not absolute")
	ErrNotFound               = errors.New("resource not found")
	ErrFailed                 = errors.New("failed to fetch resource")
	ErrCompressionUnsupported = errors.New("compression is not supported with that scheme")

	// ConfigHeaders are the HTTP headers that should be used when the Ignition
	// config is being fetched
	ConfigHeaders = http.Header{
		"Accept-Encoding": []string{"identity"},
		"Accept":          []string{"application/vnd.coreos.ignition+json;version=3.0.0, */*;q=0.1"},
	}
)

// Fetcher holds settings for fetching resources from URLs
type Fetcher struct {
	// The logger object to use when logging information.
	Logger *log.Logger

	// client is the http client that will be used when fetching http(s)
	// resources. If left nil, one will be created and used, but this means any
	// timeouts Ignition was configured to used will be ignored.
	client *HttpClient

	// The AWS Session to use when fetching resources from S3. If left nil, the
	// first S3 object that is fetched will initialize the field. This can be
	// used to set credentials.
	AWSSession *session.Session

	// The region where the AWS machine trying to fetch is.
	// This is used as a hint to fetch the S3 bucket from the right partition and region.
	S3RegionHint string
}

type FetchOptions struct {
	// Headers are the http headers that will be used when fetching http(s)
	// resources. They have no effect on other fetching schemes.
	Headers http.Header

	// Hash is the hash to use when calculating a fetched resource's hash. If
	// left as nil, no hash will be calculated.
	Hash hash.Hash

	// The expected sum to be produced by the given hasher. If the Hash field is
	// nil, this field is ignored.
	ExpectedSum []byte

	// Compression specifies the type of compression to use when decompressing
	// the fetched object. If left empty, no decompression will be used.
	Compression string
}

// FetchToBuffer will fetch the given url into a temporrary file, and then read
// in the contents of the file and delete it. It will return the downloaded
// contents, or an error if one was encountered.
func (f *Fetcher) FetchToBuffer(u url.URL, opts FetchOptions) ([]byte, error) {
	file, err := ioutil.TempFile("", "ignition")
	if err != nil {
		return nil, err
	}
	defer os.Remove(file.Name())
	defer file.Close()
	err = f.Fetch(u, file, opts)
	if err != nil {
		return nil, err
	}
	_, err = file.Seek(0, os.SEEK_SET)
	if err != nil {
		return nil, err
	}
	res, err := ioutil.ReadAll(file)
	if err != nil {
		return nil, err
	}
	return res, nil
}

// Fetch calls the appropriate FetchFrom* function based on the scheme of the
// given URL. The results will be decompressed if compression is set in opts,
// and written into dest. If opts.Hash is set the data stream will also be
// hashed and compared against opts.ExpectedSum, and any match failures will
// result in an error being returned.
//
// Fetch expects dest to be an empty file and for the cursor in the file to be
// at the beginning. Since some url schemes (ex: s3) use chunked downloads and
// fetch chunks out of order, Fetch's behavior when dest is not an empty file is
// undefined.
func (f *Fetcher) Fetch(u url.URL, dest *os.File, opts FetchOptions) error {
	switch u.Scheme {
	case "http", "https":
		return f.FetchFromHTTP(u, dest, opts)
	case "tftp":
		return f.FetchFromTFTP(u, dest, opts)
	case "data":
		return f.FetchFromDataURL(u, dest, opts)
	case "s3":
		return f.FetchFromS3(u, dest, opts)
	case "":
		return nil
	default:
		return ErrSchemeUnsupported
	}
}

// FetchFromTFTP fetches a resource from u via TFTP into dest, returning an
// error if one is encountered.
func (f *Fetcher) FetchFromTFTP(u url.URL, dest *os.File, opts FetchOptions) error {
	if !strings.ContainsRune(u.Host, ':') {
		u.Host = u.Host + ":69"
	}
	c, err := tftp.NewClient(u.Host)
	if err != nil {
		return err
	}
	wt, err := c.Receive(u.Path, "octet")
	if err != nil {
		return err
	}
	// The TFTP library takes an io.Writer to send data in to, but to decompress
	// the stream the gzip library wraps an io.Reader, so let's create a pipe to
	// connect these two things
	pReader, pWriter := io.Pipe()
	doneChan := make(chan error, 2)

	checkForDoneChanErr := func(err error) error {
		// If an error is encountered while decompressing or copying data out of
		// the pipe, there's probably an error from writing into the pipe that
		// will better describe what went wrong. This function does a
		// non-blocking read of doneChan, overriding the returned error val if
		// there's anything in doneChan.
		select {
		case writeErr := <-doneChan:
			if writeErr != nil {
				return writeErr
			}
			return err
		default:
			return err
		}
	}

	// A goroutine is used to handle writing the fetched data into the pipe
	// while also copying it out of the pipe concurrently
	go func() {
		_, err := wt.WriteTo(pWriter)
		doneChan <- err
		err = pWriter.Close()
		doneChan <- err
	}()
	err = f.decompressCopyHashAndVerify(dest, pReader, opts)
	if err != nil {
		return checkForDoneChanErr(err)
	}
	// receive the error from wt.WriteTo()
	err = <-doneChan
	if err != nil {
		return err
	}
	// receive the error from pWriter.Close()
	err = <-doneChan
	if err != nil {
		return err
	}
	return nil
}

// FetchFromHTTP fetches a resource from u via HTTP(S) into dest, returning an
// error if one is encountered.
func (f *Fetcher) FetchFromHTTP(u url.URL, dest *os.File, opts FetchOptions) error {
	// for the case when "config is not valid"
	// this if necessary if not spawned through kola (e.g. Packet Dashboard)
	if f.client == nil {
		logger := log.New(true)
		f.Logger = &logger
		if err := f.newHttpClient(); err != nil {
			return err
		}
	}

	dataReader, status, ctxCancel, err := f.client.getReaderWithHeader(u.String(), opts.Headers)
	if ctxCancel != nil {
		// whatever context getReaderWithHeader created for the request should
		// be cancelled once we're done reading the response
		defer ctxCancel()
	}
	if err != nil {
		return err
	}
	defer dataReader.Close()

	switch status {
	case http.StatusOK, http.StatusNoContent:
		break
	case http.StatusNotFound:
		return ErrNotFound
	default:
		return ErrFailed
	}

	return f.decompressCopyHashAndVerify(dest, dataReader, opts)
}

// FetchFromDataURL writes the data stored in the dataurl u into dest, returning
// an error if one is encountered.
func (f *Fetcher) FetchFromDataURL(u url.URL, dest *os.File, opts FetchOptions) error {
	if opts.Compression != "" {
		return ErrCompressionUnsupported
	}
	url, err := dataurl.DecodeString(u.String())
	if err != nil {
		return err
	}

	return f.decompressCopyHashAndVerify(dest, bytes.NewBuffer(url.Data), opts)
}

// FetchFromS3 gets data from an S3 bucket as described by u and writes it into
// dest, returning an error if one is encountered. It will attempt to acquire
// IAM credentials from the EC2 metadata service, and if this fails will attempt
// to fetch the object with anonymous credentials.
func (f *Fetcher) FetchFromS3(u url.URL, dest *os.File, opts FetchOptions) error {
	if opts.Compression != "" {
		return ErrCompressionUnsupported
	}
	ctx := context.Background()
	if f.client != nil && f.client.timeout != 0 {
		var cancelFn context.CancelFunc
		ctx, cancelFn = context.WithTimeout(ctx, f.client.timeout)
		defer cancelFn()
	}

	if f.AWSSession == nil {
		var err error
		f.AWSSession, err = session.NewSession(&aws.Config{
			Credentials: credentials.AnonymousCredentials,
		})
		if err != nil {
			return err
		}
	}
	sess := f.AWSSession.Copy()

	// Determine the partition and region this bucket is in
	regionHint := "us-east-1"
	if f.S3RegionHint != "" {
		regionHint = f.S3RegionHint
	}
	region, err := s3manager.GetBucketRegion(ctx, sess, u.Host, regionHint)
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok && aerr.Code() == "NotFound" {
			return fmt.Errorf("couldn't determine the region for bucket %q: %v", u.Host, err)
		}
		return err
	}

	sess.Config.Region = aws.String(region)

	var versionId *string
	if v, ok := u.Query()["versionId"]; ok && len(v) > 0 {
		versionId = aws.String(v[0])
	}

	input := &s3.GetObjectInput{
		Bucket:    &u.Host,
		Key:       &u.Path,
		VersionId: versionId,
	}
	err = f.fetchFromS3WithCreds(ctx, dest, input, sess)
	if err != nil {
		return err
	}
	if opts.Hash != nil {
		opts.Hash.Reset()
		_, err = dest.Seek(0, os.SEEK_SET)
		if err != nil {
			return err
		}
		_, err = io.Copy(opts.Hash, dest)
		if err != nil {
			return err
		}
		calculatedSum := opts.Hash.Sum(nil)
		if !bytes.Equal(calculatedSum, opts.ExpectedSum) {
			return util.ErrHashMismatch{
				Calculated: hex.EncodeToString(calculatedSum),
				Expected:   hex.EncodeToString(opts.ExpectedSum),
			}
		}
		f.Logger.Debug("file matches expected sum of: %s", hex.EncodeToString(opts.ExpectedSum))
	}
	return nil
}

func (f *Fetcher) fetchFromS3WithCreds(ctx context.Context, dest *os.File, input *s3.GetObjectInput, sess *session.Session) error {
	httpClient, err := defaultHTTPClient()
	if err != nil {
		return err
	}

	awsConfig := aws.NewConfig().WithHTTPClient(httpClient)
	s3Client := s3.New(sess, awsConfig)
	downloader := s3manager.NewDownloaderWithClient(s3Client)
	if _, err := downloader.DownloadWithContext(ctx, dest, input); err != nil {
		if awserrval, ok := err.(awserr.Error); ok && awserrval.Code() == "EC2RoleRequestError" {
			// If this error was due to an EC2 role request error, try again
			// with the anonymous credentials.
			sess.Config.Credentials = credentials.AnonymousCredentials
			return f.fetchFromS3WithCreds(ctx, dest, input, sess)
		}
		return err
	}
	return nil
}

// uncompress will wrap the given io.Reader in a decompresser specified in the
// FetchOptions, and return an io.ReadCloser with the decompressed data stream.
func (f *Fetcher) uncompress(r io.Reader, opts FetchOptions) (io.ReadCloser, error) {
	switch opts.Compression {
	case "":
		return ioutil.NopCloser(r), nil
	case "gzip":
		return gzip.NewReader(r)
	default:
		return nil, configErrors.ErrCompressionInvalid
	}
}

// decompressCopyHashAndVerify will decompress src if necessary, copy src into
// dest until src returns an io.EOF while also calculating a hash if one is set,
// and will return an error if there's any problems with any of this or if the
// hash doesn't match the expected hash in the opts.
func (f *Fetcher) decompressCopyHashAndVerify(dest io.Writer, src io.Reader, opts FetchOptions) error {
	decompressor, err := f.uncompress(src, opts)
	if err != nil {
		return err
	}
	defer decompressor.Close()
	if opts.Hash != nil {
		opts.Hash.Reset()
		dest = io.MultiWriter(dest, opts.Hash)
	}
	_, err = io.Copy(dest, decompressor)
	if err != nil {
		return err
	}
	if opts.Hash != nil {
		calculatedSum := opts.Hash.Sum(nil)
		if !bytes.Equal(calculatedSum, opts.ExpectedSum) {
			return util.ErrHashMismatch{
				Calculated: hex.EncodeToString(calculatedSum),
				Expected:   hex.EncodeToString(opts.ExpectedSum),
			}
		}
		f.Logger.Debug("file matches expected sum of: %s", hex.EncodeToString(opts.ExpectedSum))
	}
	return nil
}
<|endoftext|>


<commit_before>package cli

import (
	"encoding/json"
	"errors"
	"fmt"

	log "github.com/Sirupsen/logrus"
	"github.com/bitrise-io/go-utils/colorstring"
	"github.com/codegangsta/cli"
)

// ValidationItemModel ...
type ValidationItemModel struct {
	IsValid bool   `json:"is_valid" yaml:"is_valid"`
	Error   string `json:"error,omitempty" yaml:"error,omitempty"`
}

// ValidationModel ...
type ValidationModel struct {
	Config  *ValidationItemModel `json:"config,omitempty" yaml:"config,omitempty"`
	Secrets *ValidationItemModel `json:"secrets,omitempty" yaml:"secrets,omitempty"`
}

func printRawValidation(validation ValidationModel) error {
	validConfig := true
	if validation.Config != nil {
		fmt.Println(colorstring.Blue("Config validation result:"))
		configValidation := *validation.Config
		if configValidation.IsValid {
			fmt.Printf("is valid: %s\n", colorstring.Greenf("%v", configValidation.IsValid))
		} else {
			fmt.Printf("is valid: %s\n", colorstring.Redf("%v", configValidation.IsValid))
			fmt.Printf("error: %s\n", colorstring.Red(configValidation.Error))

			validConfig = false
		}
		fmt.Println()
	}

	validSecrets := true
	if validation.Secrets != nil {
		fmt.Println(colorstring.Blue("Secret validation result:"))
		secretValidation := *validation.Secrets
		if secretValidation.IsValid {
			fmt.Printf("is valid: %s\n", colorstring.Greenf("%v", secretValidation.IsValid))
		} else {
			fmt.Printf("is valid: %s\n", colorstring.Redf("%v", secretValidation.IsValid))
			fmt.Printf("error: %s\n", colorstring.Red(secretValidation.Error))

			validSecrets = false
		}
	}

	if !validConfig && !validSecrets {
		return errors.New("Config and secrets are invalid")
	} else if !validConfig {
		return errors.New("Config is invalid")
	} else if !validSecrets {
		return errors.New("Secret is invalid")
	}
	return nil
}

func printJSONValidation(validation ValidationModel) error {
	bytes, err := json.Marshal(validation)
	if err != nil {
		return err
	}

	fmt.Println(string(bytes))
	if (validation.Config != nil && !validation.Config.IsValid) &&
		(validation.Secrets != nil && !validation.Secrets.IsValid) {
		return errors.New("Config and secrets are invalid")
	} else if validation.Config != nil && !validation.Config.IsValid {
		return errors.New("Config is invalid")
	} else if validation.Secrets != nil && !validation.Secrets.IsValid {
		return errors.New("Secret is invalid")
	}
	return nil
}

func validate(c *cli.Context) {
	format := c.String(OuputFormatKey)
	if format == "" {
		format = OutputFormatRaw
	} else if !(format == OutputFormatRaw || format == OutputFormatJSON) {
		log.Fatalf("Invalid format: %s", format)
	}

	validation := ValidationModel{}

	pth, err := GetBitriseConfigFilePath(c)
	if err != nil && err.Error() != "No workflow yml found" {
		log.Fatalf("Faild to get config path, err: %s", err)
	}
	if pth != "" || (pth == "" && c.String(ConfigBase64Key) != "") {
		// Config validation
		isValid := true
		errMsg := ""

		_, err := CreateBitriseConfigFromCLIParams(c)
		if err != nil {
			isValid = false
			errMsg = err.Error()
		}

		validation.Config = &ValidationItemModel{
			IsValid: isValid,
			Error:   errMsg,
		}
	} else {
		log.Warn("No config found for validation")
	}

	pth, err = GetInventoryFilePath(c)
	if err != nil {
		log.Fatalf("Faild to get secrets path, err: %s", err)
	}
	if pth != "" || c.String(InventoryBase64Key) != "" {
		// Inventory validation
		isValid := true
		errMsg := ""

		_, err := CreateInventoryFromCLIParams(c)
		if err != nil {
			isValid = false
			errMsg = err.Error()
		}

		validation.Secrets = &ValidationItemModel{
			IsValid: isValid,
			Error:   errMsg,
		}
	}

	if validation.Config == nil && validation.Secrets == nil {
		log.Fatal("No config or secrets found for validation")
	}

	switch format {
	case OutputFormatRaw:
		if err := printRawValidation(validation); err != nil {
			log.Fatalf("Validation failed, err: %s", err)
		}
		break
	case OutputFormatJSON:
		if err := printJSONValidation(validation); err != nil {
			log.Fatalf("Validation failed, err: %s", err)
		}
		break
	default:
		log.Fatalf("Invalid format: %s", format)
	}
}
<commit_msg>validation fix<commit_after>package cli

import (
	"encoding/json"
	"errors"
	"fmt"

	log "github.com/Sirupsen/logrus"
	"github.com/bitrise-io/go-utils/colorstring"
	"github.com/codegangsta/cli"
)

// ValidationItemModel ...
type ValidationItemModel struct {
	IsValid bool   `json:"is_valid" yaml:"is_valid"`
	Error   string `json:"error,omitempty" yaml:"error,omitempty"`
}

// ValidationModel ...
type ValidationModel struct {
	Config  *ValidationItemModel `json:"config,omitempty" yaml:"config,omitempty"`
	Secrets *ValidationItemModel `json:"secrets,omitempty" yaml:"secrets,omitempty"`
}

func printRawValidation(validation ValidationModel) error {
	validConfig := true
	if validation.Config != nil {
		fmt.Println(colorstring.Blue("Config validation result:"))
		configValidation := *validation.Config
		if configValidation.IsValid {
			fmt.Printf("is valid: %s\n", colorstring.Greenf("%v", configValidation.IsValid))
		} else {
			fmt.Printf("is valid: %s\n", colorstring.Redf("%v", configValidation.IsValid))
			fmt.Printf("error: %s\n", colorstring.Red(configValidation.Error))

			validConfig = false
		}
		fmt.Println()
	}

	validSecrets := true
	if validation.Secrets != nil {
		fmt.Println(colorstring.Blue("Secret validation result:"))
		secretValidation := *validation.Secrets
		if secretValidation.IsValid {
			fmt.Printf("is valid: %s\n", colorstring.Greenf("%v", secretValidation.IsValid))
		} else {
			fmt.Printf("is valid: %s\n", colorstring.Redf("%v", secretValidation.IsValid))
			fmt.Printf("error: %s\n", colorstring.Red(secretValidation.Error))

			validSecrets = false
		}
	}

	if !validConfig && !validSecrets {
		return errors.New("Config and secrets are invalid")
	} else if !validConfig {
		return errors.New("Config is invalid")
	} else if !validSecrets {
		return errors.New("Secret is invalid")
	}
	return nil
}

func printJSONValidation(validation ValidationModel) error {
	bytes, err := json.Marshal(validation)
	if err != nil {
		return err
	}

	fmt.Println(string(bytes))
	if (validation.Config != nil && !validation.Config.IsValid) &&
		(validation.Secrets != nil && !validation.Secrets.IsValid) {
		return errors.New("Config and secrets are invalid")
	} else if validation.Config != nil && !validation.Config.IsValid {
		return errors.New("Config is invalid")
	} else if validation.Secrets != nil && !validation.Secrets.IsValid {
		return errors.New("Secret is invalid")
	}
	return nil
}

func validate(c *cli.Context) {
	format := c.String(OuputFormatKey)
	if format == "" {
		format = OutputFormatRaw
	} else if !(format == OutputFormatRaw || format == OutputFormatJSON) {
		log.Fatalf("Invalid format: %s", format)
	}

	validation := ValidationModel{}

	pth, err := GetBitriseConfigFilePath(c)
	if err != nil && err.Error() != "No workflow yml found" {
		log.Fatalf("Faild to get config path, err: %s", err)
	}
	if pth != "" || (pth == "" && c.String(ConfigBase64Key) != "") {
		// Config validation
		isValid := true
		errMsg := ""

		_, err := CreateBitriseConfigFromCLIParams(c)
		if err != nil {
			isValid = false
			errMsg = err.Error()
		}

		validation.Config = &ValidationItemModel{
			IsValid: isValid,
			Error:   errMsg,
		}
	} else {
		log.Debug("No config found for validation")
	}

	pth, err = GetInventoryFilePath(c)
	if err != nil {
		log.Fatalf("Faild to get secrets path, err: %s", err)
	}
	if pth != "" || c.String(InventoryBase64Key) != "" {
		// Inventory validation
		isValid := true
		errMsg := ""

		_, err := CreateInventoryFromCLIParams(c)
		if err != nil {
			isValid = false
			errMsg = err.Error()
		}

		validation.Secrets = &ValidationItemModel{
			IsValid: isValid,
			Error:   errMsg,
		}
	}

	if validation.Config == nil && validation.Secrets == nil {
		log.Fatal("No config or secrets found for validation")
	}

	switch format {
	case OutputFormatRaw:
		if err := printRawValidation(validation); err != nil {
			log.Fatalf("Validation failed, err: %s", err)
		}
		break
	case OutputFormatJSON:
		if err := printJSONValidation(validation); err != nil {
			log.Fatalf("Validation failed, err: %s", err)
		}
		break
	default:
		log.Fatalf("Invalid format: %s", format)
	}
}
<|endoftext|>


<commit_before>package service

import (
	"bytes"
	"context"
	"database/sql"
	"errors"
	"fmt"
	"html/template"
	"log"
	"net/url"
	"strings"
	"time"

	"github.com/hako/branca"
	webtemplate "github.com/nicolasparada/nakama/web/template"
)

// KeyAuthUserID to use in context.
const KeyAuthUserID = ctxkey("auth_user_id")

const (
	verificationCodeLifespan = time.Minute * 15
	tokenLifespan            = time.Hour * 24 * 14
)

var (
	// ErrUnimplemented denotes a not implemented functionality.
	ErrUnimplemented = errors.New("unimplemented")
	// ErrUnauthenticated denotes no authenticated user in context.
	ErrUnauthenticated = errors.New("unauthenticated")
	// ErrInvalidRedirectURI denotes an invalid redirect URI.
	ErrInvalidRedirectURI = errors.New("invalid redirect URI")
	// ErrInvalidToken denotes an invalid token.
	ErrInvalidToken = errors.New("invalid token")
	// ErrExpiredToken denotes that the token already expired.
	ErrExpiredToken = errors.New("expired token")
	// ErrInvalidVerificationCode denotes an invalid verification code.
	ErrInvalidVerificationCode = errors.New("invalid verification code")
	// ErrVerificationCodeNotFound denotes a not found verification code.
	ErrVerificationCodeNotFound = errors.New("verification code not found")
)

var magicLinkMailTmpl *template.Template

type ctxkey string

// TokenOutput response.
type TokenOutput struct {
	Token     string    `json:"token"`
	ExpiresAt time.Time `json:"expiresAt"`
}

// DevLoginOutput response.
type DevLoginOutput struct {
	User      User      `json:"user"`
	Token     string    `json:"token"`
	ExpiresAt time.Time `json:"expiresAt"`
}

// SendMagicLink to login without passwords.
func (s *Service) SendMagicLink(ctx context.Context, email, redirectURI string) error {
	email = strings.TrimSpace(email)
	if !reEmail.MatchString(email) {
		return ErrInvalidEmail
	}

	uri, err := url.Parse(redirectURI)
	if err != nil || !uri.IsAbs() {
		return ErrInvalidRedirectURI
	}

	var code string
	err = s.DB.QueryRowContext(ctx, `
		INSERT INTO verification_codes (user_id) VALUES (
			(SELECT id FROM users WHERE email = $1)
		) RETURNING id`, email).Scan(&code)
	if isNotNullViolation(err) {
		return ErrUserNotFound
	}

	if err != nil {
		return fmt.Errorf("could not insert verification code: %w", err)
	}

	defer func() {
		if err != nil {
			go func() {
				_, err := s.DB.Exec("DELETE FROM verification_codes WHERE id = $1", code)
				if err != nil {
					log.Printf("could not delete verification code: %v\n", err)
				}
			}()
		}
	}()

	link := cloneURL(s.Origin)
	link.Path = "/api/auth_redirect"
	q := link.Query()
	q.Set("verification_code", code)
	q.Set("redirect_uri", uri.String())
	link.RawQuery = q.Encode()

	if magicLinkMailTmpl == nil {
		magicLinkMailTmpl, err = template.ParseFS(webtemplate.Files, "mail/magic-link.html")
		if err != nil {
			return fmt.Errorf("could not parse magic link mail template: %w", err)
		}
	}

	var b bytes.Buffer
	if err = magicLinkMailTmpl.Execute(&b, map[string]interface{}{
		"MagicLink": link.String(),
		"Minutes":   int(verificationCodeLifespan.Minutes()),
	}); err != nil {
		return fmt.Errorf("could not execute magic link mail template: %w", err)
	}

	if err = s.Sender.Send(email, "Magic Link", b.String()); err != nil {
		return fmt.Errorf("could not send magic link: %w", err)
	}

	return nil
}

// AuthURI to be redirected to and complete the login flow.
// It contains the token in the hash fragment.
func (s *Service) AuthURI(ctx context.Context, verificationCode, redirectURI string) (string, error) {
	verificationCode = strings.TrimSpace(verificationCode)
	if !reUUID.MatchString(verificationCode) {
		return "", ErrInvalidVerificationCode
	}

	uri, err := url.ParseRequestURI(redirectURI)
	if err != nil {
		return "", ErrInvalidRedirectURI
	}

	var uid string
	var createdAt time.Time
	err = s.DB.QueryRowContext(ctx, `
		DELETE FROM verification_codes WHERE id = $1
		RETURNING user_id, created_at`, verificationCode).Scan(&uid, &createdAt)
	if err == sql.ErrNoRows {
		return "", ErrVerificationCodeNotFound
	}

	if err != nil {
		return "", fmt.Errorf("could not delete verification code: %w", err)
	}

	now := time.Now()
	exp := createdAt.Add(verificationCodeLifespan)
	if exp.Equal(now) || exp.Before(now) {
		return "", ErrExpiredToken
	}

	token, err := s.codec().EncodeToString(uid)
	if err != nil {
		return "", fmt.Errorf("could not create token: %w", err)
	}

	f := url.Values{}
	f.Set("token", token)
	f.Set("expires_at", now.Add(tokenLifespan).Format(time.RFC3339Nano))
	uri.Fragment = f.Encode()

	return uri.String(), nil
}

// DevLogin is a login for development purposes only.
// TODO: disable dev login on production.
func (s *Service) DevLogin(ctx context.Context, email string) (DevLoginOutput, error) {
	var out DevLoginOutput

	email = strings.TrimSpace(email)
	if !reEmail.MatchString(email) {
		return out, ErrInvalidEmail
	}

	var avatar sql.NullString
	query := "SELECT id, username, avatar FROM users WHERE email = $1"
	err := s.DB.QueryRowContext(ctx, query, email).Scan(&out.User.ID, &out.User.Username, &avatar)

	if err == sql.ErrNoRows {
		return out, ErrUserNotFound
	}

	if err != nil {
		return out, fmt.Errorf("could not query select user: %w", err)
	}

	out.User.AvatarURL = s.avatarURL(avatar)

	out.Token, err = s.codec().EncodeToString(out.User.ID)
	if err != nil {
		return out, fmt.Errorf("could not create token: %w", err)
	}

	out.ExpiresAt = time.Now().Add(tokenLifespan)

	return out, nil
}

// AuthUserIDFromToken decodes the token into a user ID.
func (s *Service) AuthUserIDFromToken(token string) (string, error) {
	uid, err := s.codec().DecodeToString(token)
	if err != nil {
		// We check error string because branca doesn't export errors.
		msg := err.Error()
		if msg == "invalid base62 token" || msg == "invalid token version" {
			return "", ErrInvalidToken
		}
		if msg == "token is expired" {
			return "", ErrExpiredToken
		}
		return "", fmt.Errorf("could not decode token: %w", err)
	}

	if !reUUID.MatchString(uid) {
		return "", ErrInvalidUserID
	}

	return uid, nil
}

// AuthUser is the current authenticated user.
func (s *Service) AuthUser(ctx context.Context) (User, error) {
	var u User
	uid, ok := ctx.Value(KeyAuthUserID).(string)
	if !ok {
		return u, ErrUnauthenticated
	}

	return s.userByID(ctx, uid)
}

// Token to authenticate requests.
func (s *Service) Token(ctx context.Context) (TokenOutput, error) {
	var out TokenOutput
	uid, ok := ctx.Value(KeyAuthUserID).(string)
	if !ok {
		return out, ErrUnauthenticated
	}

	var err error
	out.Token, err = s.codec().EncodeToString(uid)
	if err != nil {
		return out, fmt.Errorf("could not create token: %w", err)
	}

	out.ExpiresAt = time.Now().Add(tokenLifespan)

	return out, nil
}

func (s *Service) deleteExpiredVerificationCodesJob(ctx context.Context) {
	ticker := time.NewTicker(time.Hour * 24)
	done := ctx.Done()

loop:
	for {
		select {
		case <-ticker.C:
			if err := s.deleteExpiredVerificationCodes(ctx); err != nil {
				log.Println(err)
			}
		case <-done:
			ticker.Stop()
			break loop
		}
	}
}

func (s *Service) deleteExpiredVerificationCodes(ctx context.Context) error {
	query := fmt.Sprintf("DELETE FROM verification_codes WHERE (created_at - INTERVAL '%dm') <= now()", int64(verificationCodeLifespan.Minutes()))
	if _, err := s.DB.ExecContext(ctx, query); err != nil {
		return fmt.Errorf("could not delete expired verification code: %w", err)
	}
	return nil
}

func (s *Service) codec() *branca.Branca {
	cdc := branca.NewBranca(s.TokenKey)
	cdc.SetTTL(uint32(tokenLifespan.Seconds()))
	return cdc
}
<commit_msg>service: update branca error checking<commit_after>package service

import (
	"bytes"
	"context"
	"database/sql"
	"errors"
	"fmt"
	"html/template"
	"log"
	"net/url"
	"strings"
	"time"

	"github.com/hako/branca"
	webtemplate "github.com/nicolasparada/nakama/web/template"
)

// KeyAuthUserID to use in context.
const KeyAuthUserID = ctxkey("auth_user_id")

const (
	verificationCodeLifespan = time.Minute * 15
	tokenLifespan            = time.Hour * 24 * 14
)

var (
	// ErrUnimplemented denotes a not implemented functionality.
	ErrUnimplemented = errors.New("unimplemented")
	// ErrUnauthenticated denotes no authenticated user in context.
	ErrUnauthenticated = errors.New("unauthenticated")
	// ErrInvalidRedirectURI denotes an invalid redirect URI.
	ErrInvalidRedirectURI = errors.New("invalid redirect URI")
	// ErrInvalidToken denotes an invalid token.
	ErrInvalidToken = errors.New("invalid token")
	// ErrExpiredToken denotes that the token already expired.
	ErrExpiredToken = errors.New("expired token")
	// ErrInvalidVerificationCode denotes an invalid verification code.
	ErrInvalidVerificationCode = errors.New("invalid verification code")
	// ErrVerificationCodeNotFound denotes a not found verification code.
	ErrVerificationCodeNotFound = errors.New("verification code not found")
)

var magicLinkMailTmpl *template.Template

type ctxkey string

// TokenOutput response.
type TokenOutput struct {
	Token     string    `json:"token"`
	ExpiresAt time.Time `json:"expiresAt"`
}

// DevLoginOutput response.
type DevLoginOutput struct {
	User      User      `json:"user"`
	Token     string    `json:"token"`
	ExpiresAt time.Time `json:"expiresAt"`
}

// SendMagicLink to login without passwords.
func (s *Service) SendMagicLink(ctx context.Context, email, redirectURI string) error {
	email = strings.TrimSpace(email)
	if !reEmail.MatchString(email) {
		return ErrInvalidEmail
	}

	uri, err := url.Parse(redirectURI)
	if err != nil || !uri.IsAbs() {
		return ErrInvalidRedirectURI
	}

	var code string
	err = s.DB.QueryRowContext(ctx, `
		INSERT INTO verification_codes (user_id) VALUES (
			(SELECT id FROM users WHERE email = $1)
		) RETURNING id`, email).Scan(&code)
	if isNotNullViolation(err) {
		return ErrUserNotFound
	}

	if err != nil {
		return fmt.Errorf("could not insert verification code: %w", err)
	}

	defer func() {
		if err != nil {
			go func() {
				_, err := s.DB.Exec("DELETE FROM verification_codes WHERE id = $1", code)
				if err != nil {
					log.Printf("could not delete verification code: %v\n", err)
				}
			}()
		}
	}()

	link := cloneURL(s.Origin)
	link.Path = "/api/auth_redirect"
	q := link.Query()
	q.Set("verification_code", code)
	q.Set("redirect_uri", uri.String())
	link.RawQuery = q.Encode()

	if magicLinkMailTmpl == nil {
		magicLinkMailTmpl, err = template.ParseFS(webtemplate.Files, "mail/magic-link.html")
		if err != nil {
			return fmt.Errorf("could not parse magic link mail template: %w", err)
		}
	}

	var b bytes.Buffer
	if err = magicLinkMailTmpl.Execute(&b, map[string]interface{}{
		"MagicLink": link.String(),
		"Minutes":   int(verificationCodeLifespan.Minutes()),
	}); err != nil {
		return fmt.Errorf("could not execute magic link mail template: %w", err)
	}

	if err = s.Sender.Send(email, "Magic Link", b.String()); err != nil {
		return fmt.Errorf("could not send magic link: %w", err)
	}

	return nil
}

// AuthURI to be redirected to and complete the login flow.
// It contains the token in the hash fragment.
func (s *Service) AuthURI(ctx context.Context, verificationCode, redirectURI string) (string, error) {
	verificationCode = strings.TrimSpace(verificationCode)
	if !reUUID.MatchString(verificationCode) {
		return "", ErrInvalidVerificationCode
	}

	uri, err := url.ParseRequestURI(redirectURI)
	if err != nil {
		return "", ErrInvalidRedirectURI
	}

	var uid string
	var createdAt time.Time
	err = s.DB.QueryRowContext(ctx, `
		DELETE FROM verification_codes WHERE id = $1
		RETURNING user_id, created_at`, verificationCode).Scan(&uid, &createdAt)
	if err == sql.ErrNoRows {
		return "", ErrVerificationCodeNotFound
	}

	if err != nil {
		return "", fmt.Errorf("could not delete verification code: %w", err)
	}

	now := time.Now()
	exp := createdAt.Add(verificationCodeLifespan)
	if exp.Equal(now) || exp.Before(now) {
		return "", ErrExpiredToken
	}

	token, err := s.codec().EncodeToString(uid)
	if err != nil {
		return "", fmt.Errorf("could not create token: %w", err)
	}

	f := url.Values{}
	f.Set("token", token)
	f.Set("expires_at", now.Add(tokenLifespan).Format(time.RFC3339Nano))
	uri.Fragment = f.Encode()

	return uri.String(), nil
}

// DevLogin is a login for development purposes only.
// TODO: disable dev login on production.
func (s *Service) DevLogin(ctx context.Context, email string) (DevLoginOutput, error) {
	var out DevLoginOutput

	email = strings.TrimSpace(email)
	if !reEmail.MatchString(email) {
		return out, ErrInvalidEmail
	}

	var avatar sql.NullString
	query := "SELECT id, username, avatar FROM users WHERE email = $1"
	err := s.DB.QueryRowContext(ctx, query, email).Scan(&out.User.ID, &out.User.Username, &avatar)

	if err == sql.ErrNoRows {
		return out, ErrUserNotFound
	}

	if err != nil {
		return out, fmt.Errorf("could not query select user: %w", err)
	}

	out.User.AvatarURL = s.avatarURL(avatar)

	out.Token, err = s.codec().EncodeToString(out.User.ID)
	if err != nil {
		return out, fmt.Errorf("could not create token: %w", err)
	}

	out.ExpiresAt = time.Now().Add(tokenLifespan)

	return out, nil
}

// AuthUserIDFromToken decodes the token into a user ID.
func (s *Service) AuthUserIDFromToken(token string) (string, error) {
	uid, err := s.codec().DecodeToString(token)
	if err != nil {
		// We check error string because branca doesn't export errors.
		if errors.Is(err, branca.ErrInvalidToken) || errors.Is(err, branca.ErrInvalidTokenVersion) {
			return "", ErrInvalidToken
		}
		if _, ok := err.(*branca.ErrExpiredToken); ok {
			return "", ErrExpiredToken
		}
		return "", fmt.Errorf("could not decode token: %w", err)
	}

	if !reUUID.MatchString(uid) {
		return "", ErrInvalidUserID
	}

	return uid, nil
}

// AuthUser is the current authenticated user.
func (s *Service) AuthUser(ctx context.Context) (User, error) {
	var u User
	uid, ok := ctx.Value(KeyAuthUserID).(string)
	if !ok {
		return u, ErrUnauthenticated
	}

	return s.userByID(ctx, uid)
}

// Token to authenticate requests.
func (s *Service) Token(ctx context.Context) (TokenOutput, error) {
	var out TokenOutput
	uid, ok := ctx.Value(KeyAuthUserID).(string)
	if !ok {
		return out, ErrUnauthenticated
	}

	var err error
	out.Token, err = s.codec().EncodeToString(uid)
	if err != nil {
		return out, fmt.Errorf("could not create token: %w", err)
	}

	out.ExpiresAt = time.Now().Add(tokenLifespan)

	return out, nil
}

func (s *Service) deleteExpiredVerificationCodesJob(ctx context.Context) {
	ticker := time.NewTicker(time.Hour * 24)
	done := ctx.Done()

loop:
	for {
		select {
		case <-ticker.C:
			if err := s.deleteExpiredVerificationCodes(ctx); err != nil {
				log.Println(err)
			}
		case <-done:
			ticker.Stop()
			break loop
		}
	}
}

func (s *Service) deleteExpiredVerificationCodes(ctx context.Context) error {
	query := fmt.Sprintf("DELETE FROM verification_codes WHERE (created_at - INTERVAL '%dm') <= now()", int64(verificationCodeLifespan.Minutes()))
	if _, err := s.DB.ExecContext(ctx, query); err != nil {
		return fmt.Errorf("could not delete expired verification code: %w", err)
	}
	return nil
}

func (s *Service) codec() *branca.Branca {
	cdc := branca.NewBranca(s.TokenKey)
	cdc.SetTTL(uint32(tokenLifespan.Seconds()))
	return cdc
}
<|endoftext|>


<commit_before>package cli

import (
	"io/ioutil"
	"os"
	"sort"
	"strings"

	"github.com/spf13/cobra"

	"github.com/giantswarm/inago/controller"
)

var (
	validateCmd = &cobra.Command{
		Use:   "validate [directory]",
		Short: "Validate groups",
		Long:  "Validate group directories on the local filesystem",
		Run:   validateRun,
	}
)

func validateRun(cmd *cobra.Command, args []string) {
	groups := args

	// If no groups are specified, assume all directories in current
	// directory are groups to be checked.
	if len(args) == 0 {
		files, err := ioutil.ReadDir(".")
		if err != nil {
			newLogger.Error(newCtx, "%#v", maskAny(err))
			os.Exit(1)
		}

		for _, file := range files {
			if file.IsDir() && !strings.HasPrefix(file.Name(), ".") {
				groups = append(groups, file.Name())
			}
		}
	}

	sort.Strings(groups)

	requests := []controller.Request{}
	for _, group := range groups {
		newRequestConfig := controller.DefaultRequestConfig()
		newRequestConfig.Group = group
		request := controller.NewRequest(newRequestConfig)

		request, err := extendRequestWithContent(fs, request)
		if err != nil {
			newLogger.Error(newCtx, "%#v", maskAny(err))
			os.Exit(1)
		}
		requests = append(requests, request)
	}

	for _, request := range requests {
		ok, err := controller.ValidateRequest(request)
		if ok {
			newLogger.Info(newCtx, "Group '%v' is valid.", request.Group)
		} else {
			newLogger.Info(newCtx, "Group '%v' not valid: %v.", request.Group, err)
		}
	}

	ok, err := controller.ValidateMultipleRequest(requests)
	if ok {
		newLogger.Info(newCtx, "Groups are valid globally.")
	} else {
		newLogger.Info(newCtx, "Groups are not valid globally:", err)
	}
}
<commit_msg>add dots to show multiple folder are possible<commit_after>package cli

import (
	"io/ioutil"
	"os"
	"sort"
	"strings"

	"github.com/spf13/cobra"

	"github.com/giantswarm/inago/controller"
)

var (
	validateCmd = &cobra.Command{
		Use:   "validate [directory...]",
		Short: "Validate groups",
		Long:  "Validate group directories on the local filesystem",
		Run:   validateRun,
	}
)

func validateRun(cmd *cobra.Command, args []string) {
	groups := args

	// If no groups are specified, assume all directories in current
	// directory are groups to be checked.
	if len(args) == 0 {
		files, err := ioutil.ReadDir(".")
		if err != nil {
			newLogger.Error(newCtx, "%#v", maskAny(err))
			os.Exit(1)
		}

		for _, file := range files {
			if file.IsDir() && !strings.HasPrefix(file.Name(), ".") {
				groups = append(groups, file.Name())
			}
		}
	}

	sort.Strings(groups)

	requests := []controller.Request{}
	for _, group := range groups {
		newRequestConfig := controller.DefaultRequestConfig()
		newRequestConfig.Group = group
		request := controller.NewRequest(newRequestConfig)

		request, err := extendRequestWithContent(fs, request)
		if err != nil {
			newLogger.Error(newCtx, "%#v", maskAny(err))
			os.Exit(1)
		}
		requests = append(requests, request)
	}

	for _, request := range requests {
		ok, err := controller.ValidateRequest(request)
		if ok {
			newLogger.Info(newCtx, "Group '%v' is valid.", request.Group)
		} else {
			newLogger.Info(newCtx, "Group '%v' not valid: %v.", request.Group, err)
		}
	}

	ok, err := controller.ValidateMultipleRequest(requests)
	if ok {
		newLogger.Info(newCtx, "Groups are valid globally.")
	} else {
		newLogger.Info(newCtx, "Groups are not valid globally:", err)
	}
}
<|endoftext|>


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

package worker

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/google/safehtml/template"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/fetch"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/postgres"
	"golang.org/x/sync/errgroup"
)

type annotation struct {
	error
	msg string
}

// doIndexPage writes the status page. On error it returns the error and a short
// string to be written back to the client.
func (s *Server) doIndexPage(w http.ResponseWriter, r *http.Request) (err error) {
	defer derrors.Wrap(&err, "doIndexPage")
	var (
		experiments []*internal.Experiment
		excluded    []string
	)
	if s.getExperiments != nil {
		experiments = s.getExperiments()
	}
	g, ctx := errgroup.WithContext(r.Context())
	g.Go(func() error {
		var err error
		excluded, err = s.db.GetExcludedPrefixes(ctx)
		if err != nil {
			return annotation{err, "error fetching excluded"}
		}
		return nil
	})
	if err := g.Wait(); err != nil {
		var e annotation
		if errors.As(err, &e) {
			log.Errorf(ctx, e.msg, err)
		}
		return err
	}

	var gms runtime.MemStats
	runtime.ReadMemStats(&gms)
	sms, err := getSystemMemStats()
	if err != nil {
		return err
	}
	pms, err := getProcessMemStats()
	if err != nil {
		return err
	}

	page := struct {
		Config                *config.Config
		Env                   string
		ResourcePrefix        string
		LatestTimestamp       *time.Time
		LocationID            string
		Experiments           []*internal.Experiment
		ExperimentsFromConfig bool
		Excluded              []string
		LoadShedStats         fetch.LoadShedStats
		GoMemStats            runtime.MemStats
		ProcessStats          processMemStats
		SystemStats           systemMemStats
	}{
		Config:                s.cfg,
		Env:                   env(s.cfg),
		ResourcePrefix:        strings.ToLower(env(s.cfg)) + "-",
		LocationID:            s.cfg.LocationID,
		Experiments:           experiments,
		ExperimentsFromConfig: os.Getenv("GO_DISCOVERY_EXPERIMENTS_FROM_CONFIG") == "true",
		Excluded:              excluded,
		LoadShedStats:         fetch.ZipLoadShedStats(),
		GoMemStats:            gms,
		ProcessStats:          pms,
		SystemStats:           sms,
	}
	return renderPage(ctx, w, page, s.templates[indexTemplate])
}

func (s *Server) doVersionsPage(w http.ResponseWriter, r *http.Request) (err error) {
	defer derrors.Wrap(&err, "doVersionsPage")
	const pageSize = 20
	g, ctx := errgroup.WithContext(r.Context())
	var (
		next, failures, recents []*internal.ModuleVersionState
		stats                   *postgres.VersionStats
	)
	g.Go(func() error {
		var err error
		next, err = s.db.GetNextModulesToFetch(ctx, pageSize)
		if err != nil {
			return annotation{err, "error fetching next versions"}
		}
		return nil
	})
	g.Go(func() error {
		var err error
		failures, err = s.db.GetRecentFailedVersions(ctx, pageSize)
		if err != nil {
			return annotation{err, "error fetching recent failures"}
		}
		return nil
	})
	g.Go(func() error {
		var err error
		recents, err = s.db.GetRecentVersions(ctx, pageSize)
		if err != nil {
			return annotation{err, "error fetching recent versions"}
		}
		return nil
	})
	g.Go(func() error {
		var err error
		stats, err = s.db.GetVersionStats(ctx)
		if err != nil {
			return annotation{err, "error fetching stats"}
		}
		return nil
	})
	if err := g.Wait(); err != nil {
		var e annotation
		if errors.As(err, &e) {
			log.Errorf(ctx, e.msg, err)
		}
		return err
	}

	type count struct {
		Code  int
		Desc  string
		Count int
	}
	var counts []*count
	for code, n := range stats.VersionCounts {
		c := &count{Code: code, Count: n}
		if e := derrors.FromStatus(code, ""); e != nil && e != derrors.Unknown {
			c.Desc = e.Error()
		}
		counts = append(counts, c)
	}
	sort.Slice(counts, func(i, j int) bool { return counts[i].Code < counts[j].Code })
	page := struct {
		Next, Recent, RecentFailures []*internal.ModuleVersionState
		Config                       *config.Config
		Env                          string
		ResourcePrefix               string
		LatestTimestamp              *time.Time
		Counts                       []*count
	}{
		Next:            next,
		Recent:          recents,
		RecentFailures:  failures,
		Config:          s.cfg,
		Env:             env(s.cfg),
		ResourcePrefix:  strings.ToLower(env(s.cfg)) + "-",
		LatestTimestamp: &stats.LatestTimestamp,
		Counts:          counts,
	}
	return renderPage(ctx, w, page, s.templates[versionsTemplate])
}

func env(cfg *config.Config) string {
	e := cfg.DeploymentEnvironment()
	return strings.ToUpper(e[:1]) + e[1:]
}

func renderPage(ctx context.Context, w http.ResponseWriter, page interface{}, tmpl *template.Template) (err error) {
	defer derrors.Wrap(&err, "renderPage")
	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, page); err != nil {
		return err
	}
	if _, err := io.Copy(w, &buf); err != nil {
		log.Errorf(ctx, "Error copying buffer to ResponseWriter: %v", err)
		return err
	}
	return nil
}

// systemMemStats holds values from the /proc/meminfo
// file, which describes the total system memory.
// All values are in bytes.
type systemMemStats struct {
	Total     uint64
	Free      uint64
	Available uint64
	Used      uint64
	Buffers   uint64
	Cached    uint64
}

// getSystemMemStats reads the /proc/meminfo file.
func getSystemMemStats() (systemMemStats, error) {
	f, err := os.Open("/proc/meminfo")
	if err != nil {
		return systemMemStats{}, err
	}
	defer f.Close()

	// Read the number, convert kibibytes to bytes, and set p.
	set := func(p *uint64, words []string) {
		if len(words) != 3 || words[2] != "kB" {
			err = fmt.Errorf("got %+v, want 3 words, third is 'kB'", words)
			return
		}
		var ki uint64
		ki, err = strconv.ParseUint(words[1], 10, 64)
		if err == nil {
			*p = ki * 1024
		}
	}

	scan := bufio.NewScanner(f)
	var sms systemMemStats
	for scan.Scan() {
		words := strings.Fields(scan.Text())
		switch words[0] {
		case "MemTotal:":
			set(&sms.Total, words)
		case "MemFree:":
			set(&sms.Free, words)
		case "MemAvailable:":
			set(&sms.Available, words)
		case "Buffers:":
			set(&sms.Buffers, words)
		case "Cached:":
			set(&sms.Cached, words)
		}
	}
	if err != nil {
		return systemMemStats{}, err
	}
	sms.Used = sms.Total - sms.Free - sms.Buffers - sms.Cached // see `man free`
	return sms, nil
}

// processMemStats holds values that describe the current process's memory.
// All values are in bytes.
type processMemStats struct {
	VSize uint64 // virtual memory size
	RSS   uint64 // resident set size (physical memory in use)
}

func getProcessMemStats() (processMemStats, error) {
	f, err := os.Open("/proc/self/stat")
	if err != nil {
		return processMemStats{}, err
	}
	defer f.Close()
	// Values from `man proc`.
	var (
		d          int
		s          string
		c          byte
		vsize, rss uint64
	)
	_, err = fmt.Fscanf(f, "%d %s %c %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d",
		&d, &s, &c, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &vsize, &rss)
	if err != nil {
		return processMemStats{}, err
	}
	const pageSize = 4 * 1024 // Linux page size, from `getconf PAGESIZE`
	return processMemStats{
		VSize: vsize,
		RSS:   rss * pageSize,
	}, nil
}
<commit_msg>internal/worker: don't fail if we can't get memory stats<commit_after>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/google/safehtml/template"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/fetch"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/postgres"
	"golang.org/x/sync/errgroup"
)

type annotation struct {
	error
	msg string
}

// doIndexPage writes the status page. On error it returns the error and a short
// string to be written back to the client.
func (s *Server) doIndexPage(w http.ResponseWriter, r *http.Request) (err error) {
	defer derrors.Wrap(&err, "doIndexPage")
	var (
		experiments []*internal.Experiment
		excluded    []string
	)
	if s.getExperiments != nil {
		experiments = s.getExperiments()
	}
	g, ctx := errgroup.WithContext(r.Context())
	g.Go(func() error {
		var err error
		excluded, err = s.db.GetExcludedPrefixes(ctx)
		if err != nil {
			return annotation{err, "error fetching excluded"}
		}
		return nil
	})
	if err := g.Wait(); err != nil {
		var e annotation
		if errors.As(err, &e) {
			log.Errorf(ctx, e.msg, err)
		}
		return err
	}

	var gms runtime.MemStats
	runtime.ReadMemStats(&gms)
	sms, err := getSystemMemStats()
	if err != nil {
		log.Errorf(ctx, "could not get system stats: %v", err)
	}
	pms, err := getProcessMemStats()
	if err != nil {
		log.Errorf(ctx, "could not get process stats: %v", err)
	}

	page := struct {
		Config                *config.Config
		Env                   string
		ResourcePrefix        string
		LatestTimestamp       *time.Time
		LocationID            string
		Experiments           []*internal.Experiment
		ExperimentsFromConfig bool
		Excluded              []string
		LoadShedStats         fetch.LoadShedStats
		GoMemStats            runtime.MemStats
		ProcessStats          processMemStats
		SystemStats           systemMemStats
	}{
		Config:                s.cfg,
		Env:                   env(s.cfg),
		ResourcePrefix:        strings.ToLower(env(s.cfg)) + "-",
		LocationID:            s.cfg.LocationID,
		Experiments:           experiments,
		ExperimentsFromConfig: os.Getenv("GO_DISCOVERY_EXPERIMENTS_FROM_CONFIG") == "true",
		Excluded:              excluded,
		LoadShedStats:         fetch.ZipLoadShedStats(),
		GoMemStats:            gms,
		ProcessStats:          pms,
		SystemStats:           sms,
	}
	return renderPage(ctx, w, page, s.templates[indexTemplate])
}

func (s *Server) doVersionsPage(w http.ResponseWriter, r *http.Request) (err error) {
	defer derrors.Wrap(&err, "doVersionsPage")
	const pageSize = 20
	g, ctx := errgroup.WithContext(r.Context())
	var (
		next, failures, recents []*internal.ModuleVersionState
		stats                   *postgres.VersionStats
	)
	g.Go(func() error {
		var err error
		next, err = s.db.GetNextModulesToFetch(ctx, pageSize)
		if err != nil {
			return annotation{err, "error fetching next versions"}
		}
		return nil
	})
	g.Go(func() error {
		var err error
		failures, err = s.db.GetRecentFailedVersions(ctx, pageSize)
		if err != nil {
			return annotation{err, "error fetching recent failures"}
		}
		return nil
	})
	g.Go(func() error {
		var err error
		recents, err = s.db.GetRecentVersions(ctx, pageSize)
		if err != nil {
			return annotation{err, "error fetching recent versions"}
		}
		return nil
	})
	g.Go(func() error {
		var err error
		stats, err = s.db.GetVersionStats(ctx)
		if err != nil {
			return annotation{err, "error fetching stats"}
		}
		return nil
	})
	if err := g.Wait(); err != nil {
		var e annotation
		if errors.As(err, &e) {
			log.Errorf(ctx, e.msg, err)
		}
		return err
	}

	type count struct {
		Code  int
		Desc  string
		Count int
	}
	var counts []*count
	for code, n := range stats.VersionCounts {
		c := &count{Code: code, Count: n}
		if e := derrors.FromStatus(code, ""); e != nil && e != derrors.Unknown {
			c.Desc = e.Error()
		}
		counts = append(counts, c)
	}
	sort.Slice(counts, func(i, j int) bool { return counts[i].Code < counts[j].Code })
	page := struct {
		Next, Recent, RecentFailures []*internal.ModuleVersionState
		Config                       *config.Config
		Env                          string
		ResourcePrefix               string
		LatestTimestamp              *time.Time
		Counts                       []*count
	}{
		Next:            next,
		Recent:          recents,
		RecentFailures:  failures,
		Config:          s.cfg,
		Env:             env(s.cfg),
		ResourcePrefix:  strings.ToLower(env(s.cfg)) + "-",
		LatestTimestamp: &stats.LatestTimestamp,
		Counts:          counts,
	}
	return renderPage(ctx, w, page, s.templates[versionsTemplate])
}

func env(cfg *config.Config) string {
	e := cfg.DeploymentEnvironment()
	return strings.ToUpper(e[:1]) + e[1:]
}

func renderPage(ctx context.Context, w http.ResponseWriter, page interface{}, tmpl *template.Template) (err error) {
	defer derrors.Wrap(&err, "renderPage")
	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, page); err != nil {
		return err
	}
	if _, err := io.Copy(w, &buf); err != nil {
		log.Errorf(ctx, "Error copying buffer to ResponseWriter: %v", err)
		return err
	}
	return nil
}

// systemMemStats holds values from the /proc/meminfo
// file, which describes the total system memory.
// All values are in bytes.
type systemMemStats struct {
	Total     uint64
	Free      uint64
	Available uint64
	Used      uint64
	Buffers   uint64
	Cached    uint64
}

// getSystemMemStats reads the /proc/meminfo file.
func getSystemMemStats() (systemMemStats, error) {
	f, err := os.Open("/proc/meminfo")
	if err != nil {
		return systemMemStats{}, err
	}
	defer f.Close()

	// Read the number, convert kibibytes to bytes, and set p.
	set := func(p *uint64, words []string) {
		if len(words) != 3 || words[2] != "kB" {
			err = fmt.Errorf("got %+v, want 3 words, third is 'kB'", words)
			return
		}
		var ki uint64
		ki, err = strconv.ParseUint(words[1], 10, 64)
		if err == nil {
			*p = ki * 1024
		}
	}

	scan := bufio.NewScanner(f)
	var sms systemMemStats
	for scan.Scan() {
		words := strings.Fields(scan.Text())
		switch words[0] {
		case "MemTotal:":
			set(&sms.Total, words)
		case "MemFree:":
			set(&sms.Free, words)
		case "MemAvailable:":
			set(&sms.Available, words)
		case "Buffers:":
			set(&sms.Buffers, words)
		case "Cached:":
			set(&sms.Cached, words)
		}
	}
	if err != nil {
		return systemMemStats{}, err
	}
	sms.Used = sms.Total - sms.Free - sms.Buffers - sms.Cached // see `man free`
	return sms, nil
}

// processMemStats holds values that describe the current process's memory.
// All values are in bytes.
type processMemStats struct {
	VSize uint64 // virtual memory size
	RSS   uint64 // resident set size (physical memory in use)
}

func getProcessMemStats() (processMemStats, error) {
	f, err := os.Open("/proc/self/stat")
	if err != nil {
		return processMemStats{}, err
	}
	defer f.Close()
	// Values from `man proc`.
	var (
		d          int
		s          string
		c          byte
		vsize, rss uint64
	)
	_, err = fmt.Fscanf(f, "%d %s %c %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d",
		&d, &s, &c, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &d, &vsize, &rss)
	if err != nil {
		return processMemStats{}, err
	}
	const pageSize = 4 * 1024 // Linux page size, from `getconf PAGESIZE`
	return processMemStats{
		VSize: vsize,
		RSS:   rss * pageSize,
	}, nil
}
<|endoftext|>