<commit_before>// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cli

import (
	"context"
	"encoding/hex"
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"regexp"
	"strings"
	"time"

	"github.com/maruel/subcommands"
	"google.golang.org/grpc/metadata"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/common/cli"
	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/data/strpair"
	"go.chromium.org/luci/common/data/text"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/flag"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/system/exitcode"
	"go.chromium.org/luci/common/system/signals"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/lucictx"
	"go.chromium.org/luci/server/auth/realms"

	"go.chromium.org/luci/resultdb/internal/services/recorder"
	"go.chromium.org/luci/resultdb/pbutil"
	pb "go.chromium.org/luci/resultdb/proto/v1"
	"go.chromium.org/luci/resultdb/sink"
)

var matchInvalidInvocationIDChars = regexp.MustCompile(`[^a-z0-9_\-:.]`)

func cmdStream(p Params) *subcommands.Command {
	return &subcommands.Command{
		UsageLine: `stream [flags] TEST_CMD [TEST_ARG]...`,
		ShortDesc: "Run a given test command and upload the results to ResultDB",
		// TODO(crbug.com/1017288): add a link to ResultSink protocol doc
		LongDesc: text.Doc(`
			Run a given test command, continuously collect the results over IPC, and
			upload them to ResultDB. Either use the current invocation from
			LUCI_CONTEXT or create/finalize a new one. Example:
				rdb stream -new -realm chromium:public ./out/chrome/test/browser_tests
		`),
		CommandRun: func() subcommands.CommandRun {
			r := &streamRun{
				vars: make(map[string]string),
				tags: make(strpair.Map),
			}
			r.baseCommandRun.RegisterGlobalFlags(p)
			r.Flags.BoolVar(&r.isNew, "new", false, text.Doc(`
				If true, create and use a new invocation for the test command.
				If false, use the current invocation, set in LUCI_CONTEXT.
			`))
			r.Flags.StringVar(&r.realm, "realm", "", text.Doc(`
				Realm to create the new invocation in. Required if -new is set,
				ignored otherwise.
				e.g. "chromium:public"
			`))
			r.Flags.StringVar(&r.testIDPrefix, "test-id-prefix", "", text.Doc(`
				Prefix to prepend to the test ID of every test result.
			`))
			r.Flags.Var(flag.StringMap(r.vars), "var", text.Doc(`
				Variant to add to every test result in "key:value" format.
				If the test command adds a variant with the same key, the value given by
				this flag will get overridden.
			`))
			r.Flags.UintVar(&r.artChannelMaxLeases, "max-concurrent-artifact-uploads",
				sink.DefaultArtChannelMaxLeases, text.Doc(`
				The maximum number of goroutines uploading artifacts.
			`))
			r.Flags.UintVar(&r.trChannelMaxLeases, "max-concurrent-test-result-uploads",
				sink.DefaultTestResultChannelMaxLeases, text.Doc(`
				The maximum number of goroutines uploading test results.
			`))
			r.Flags.StringVar(&r.testTestLocationBase, "test-location-base", "", text.Doc(`
				File base to prepend to the test location file name, if the file name is a relative path.
				It must start with "//".
			`))
			r.Flags.Var(flag.StringPairs(r.tags), "tag", text.Doc(`
				Tag to add to every test result in "key:value" format.
				A key can be repeated.
			`))
			r.Flags.BoolVar(&r.coerceNegativeDuration, "coerce-negative-duration",
				false, text.Doc(`
				If true, all negative durations will be coerced to 0.
				If false, test results with negative durations will be rejected.
			`))
			return r
		},
	}
}

type streamRun struct {
	baseCommandRun

	// flags
	isNew                  bool
	realm                  string
	testIDPrefix           string
	testTestLocationBase   string
	vars                   map[string]string
	artChannelMaxLeases    uint
	trChannelMaxLeases     uint
	tags                   strpair.Map
	pbTags                 []*pb.StringPair
	coerceNegativeDuration bool
	// TODO(ddoman): add flags
	// - invocation-tag
	// - log-file

	invocation lucictx.ResultDBInvocation
}

func (r *streamRun) validate(ctx context.Context, args []string) (err error) {
	if len(args) == 0 {
		return errors.Reason("missing a test command to run").Err()
	}
	if err := pbutil.ValidateVariant(&pb.Variant{Def: r.vars}); err != nil {
		return errors.Annotate(err, "invalid variant").Err()
	}
	if r.realm != "" {
		if err := realms.ValidateRealmName(r.realm, realms.GlobalScope); err != nil {
			return errors.Annotate(err, "invalid realm").Err()
		}
	}
	return nil
}

func (r *streamRun) Run(a subcommands.Application, args []string, env subcommands.Env) (ret int) {
	ctx := cli.GetContext(a, r, env)

	if err := r.validate(ctx, args); err != nil {
		return r.done(err)
	}

	loginMode := auth.OptionalLogin
	// login is required only if it creates a new invocation.
	if r.isNew {
		if r.realm == "" {
			return r.done(errors.Reason("-realm is required for new invocations").Err())
		}
		loginMode = auth.SilentLogin
	}
	if err := r.initClients(ctx, loginMode); err != nil {
		return r.done(err)
	}

	// if -new is passed, create a new invocation. If not, use the existing one set in
	// lucictx.
	if r.isNew {
		ninv, err := r.createInvocation(ctx, r.realm)
		if err != nil {
			return r.done(err)
		}
		r.invocation = ninv

		// Update lucictx with the new invocation.
		ctx = lucictx.SetResultDB(ctx, &lucictx.ResultDB{
			Hostname:          r.host,
			CurrentInvocation: &r.invocation,
		})
	} else {
		if r.resultdbCtx == nil {
			return r.done(errors.Reason("the environment does not have an existing invocation; use -new to create a new one").Err())
		}
		if err := r.validateCurrentInvocation(); err != nil {
			return r.done(err)
		}
		r.invocation = *r.resultdbCtx.CurrentInvocation
	}

	defer func() {
		// Finalize the invocation if it was created by -new.
		if r.isNew {
			if err := r.finalizeInvocation(ctx); err != nil {
				logging.Errorf(ctx, "failed to finalize the invocation: %s", err)
				ret = r.done(err)
			}
		}
	}()

	err := r.runTestCmd(ctx, args)
	ec, ok := exitcode.Get(err)
	if !ok {
		logging.Errorf(ctx, "rdb-stream: failed to run the test command: %s", err)
		return r.done(err)
	}
	logging.Infof(ctx, "rdb-stream: exiting with %d", ec)
	return ec
}

func (r *streamRun) runTestCmd(ctx context.Context, args []string) error {
	// Kill the subprocess if rdb-stream is asked to stop.
	// Subprocess exiting will unblock rdb-stream and it will stop soon.
	cmdCtx, cancelCmd := context.WithCancel(ctx)
	defer cancelCmd()
	defer signals.HandleInterrupt(func() {
		logging.Warningf(ctx, "Interrupt signal received; killing the subprocess")
		cancelCmd()
	})()

	cmd := exec.CommandContext(cmdCtx, args[0], args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	// TODO(ddoman): send the logs of SinkServer to --log-file

	cfg := sink.ServerConfig{
		Recorder:                   r.recorder,
		Invocation:                 r.invocation.Name,
		UpdateToken:                r.invocation.UpdateToken,
		TestIDPrefix:               r.testIDPrefix,
		BaseVariant:                &pb.Variant{Def: r.vars},
		ArtifactUploader:           &sink.ArtifactUploader{Client: r.http, Host: r.host},
		ArtChannelMaxLeases:        r.artChannelMaxLeases,
		TestResultChannelMaxLeases: r.trChannelMaxLeases,
		TestLocationBase:           r.testTestLocationBase,
		BaseTags:                   pbutil.FromStrpairMap(r.tags),
		CoerceNegativeDuration:     r.coerceNegativeDuration,
	}
	return sink.Run(ctx, cfg, func(ctx context.Context, cfg sink.ServerConfig) error {
		exported, err := lucictx.Export(ctx)
		if err != nil {
			return err
		}
		defer func() {
			logging.Infof(ctx, "rdb-stream: the test process terminated")
			exported.Close()
		}()
		exported.SetInCmd(cmd)
		logging.Infof(ctx, "rdb-stream: starting the test command - %q", cmd.Args)
		if err := cmd.Start(); err != nil {
			return errors.Annotate(err, "cmd.start").Err()
		}
		return cmd.Wait()
	})
}

func (r *streamRun) createInvocation(ctx context.Context, realm string) (ret lucictx.ResultDBInvocation, err error) {
	invID, err := genInvID(ctx)
	if err != nil {
		return
	}

	md := metadata.MD{}
	resp, err := r.recorder.CreateInvocation(ctx, &pb.CreateInvocationRequest{
		InvocationId: invID,
		Invocation: &pb.Invocation{
			Realm: realm,
		},
	}, prpc.Header(&md))
	if err != nil {
		err = errors.Annotate(err, "failed to create an invocation").Err()
		return
	}
	tks := md.Get(recorder.UpdateTokenMetadataKey)
	if len(tks) == 0 {
		err = errors.Reason("Missing header: update-token").Err()
		return
	}

	ret = lucictx.ResultDBInvocation{Name: resp.Name, UpdateToken: tks[0]}
	fmt.Fprintf(os.Stderr, "rdb-stream: created invocation - https://ci.chromium.org/ui/inv/%s\n", invID)
	return
}

// finalizeInvocation finalizes the invocation.
func (r *streamRun) finalizeInvocation(ctx context.Context) error {
	ctx = metadata.AppendToOutgoingContext(
		ctx, recorder.UpdateTokenMetadataKey, r.invocation.UpdateToken)
	_, err := r.recorder.FinalizeInvocation(ctx, &pb.FinalizeInvocationRequest{
		Name: r.invocation.Name,
	})
	return err
}

// genInvID generates an invocation ID, made of the username, the current timestamp
// in a human-friendly format, and a random suffix.
//
// This can be used to generate a random invocation ID, but the creator and creation time
// can be easily found.
func genInvID(ctx context.Context) (string, error) {
	whoami, err := user.Current()
	if err != nil {
		return "", err
	}
	bytes := make([]byte, 8)
	if _, err := mathrand.Read(ctx, bytes); err != nil {
		return "", err
	}

	username := strings.ToLower(whoami.Username)
	username = matchInvalidInvocationIDChars.ReplaceAllString(username, "")

	suffix := strings.ToLower(fmt.Sprintf(
		"%s-%s", time.Now().UTC().Format("2006-01-02-15-04-00"),
		// Note: cannot use base64 because not all of its characters are allowed
		// in invocation IDs.
		hex.EncodeToString(bytes)))

	// An invocation ID can contain up to 100 ascii characters that conform to the regex,
	return fmt.Sprintf("u-%.*s-%s", 100-len(suffix), username, suffix), nil
}
<commit_msg>[resultdb] print the invocation URL when finalizing it<commit_after>// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cli

import (
	"context"
	"encoding/hex"
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"regexp"
	"strings"
	"time"

	"github.com/maruel/subcommands"
	"google.golang.org/grpc/metadata"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/common/cli"
	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/data/strpair"
	"go.chromium.org/luci/common/data/text"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/flag"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/system/exitcode"
	"go.chromium.org/luci/common/system/signals"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/lucictx"
	"go.chromium.org/luci/server/auth/realms"

	"go.chromium.org/luci/resultdb/internal/services/recorder"
	"go.chromium.org/luci/resultdb/pbutil"
	pb "go.chromium.org/luci/resultdb/proto/v1"
	"go.chromium.org/luci/resultdb/sink"
)

var matchInvalidInvocationIDChars = regexp.MustCompile(`[^a-z0-9_\-:.]`)

func cmdStream(p Params) *subcommands.Command {
	return &subcommands.Command{
		UsageLine: `stream [flags] TEST_CMD [TEST_ARG]...`,
		ShortDesc: "Run a given test command and upload the results to ResultDB",
		// TODO(crbug.com/1017288): add a link to ResultSink protocol doc
		LongDesc: text.Doc(`
			Run a given test command, continuously collect the results over IPC, and
			upload them to ResultDB. Either use the current invocation from
			LUCI_CONTEXT or create/finalize a new one. Example:
				rdb stream -new -realm chromium:public ./out/chrome/test/browser_tests
		`),
		CommandRun: func() subcommands.CommandRun {
			r := &streamRun{
				vars: make(map[string]string),
				tags: make(strpair.Map),
			}
			r.baseCommandRun.RegisterGlobalFlags(p)
			r.Flags.BoolVar(&r.isNew, "new", false, text.Doc(`
				If true, create and use a new invocation for the test command.
				If false, use the current invocation, set in LUCI_CONTEXT.
			`))
			r.Flags.StringVar(&r.realm, "realm", "", text.Doc(`
				Realm to create the new invocation in. Required if -new is set,
				ignored otherwise.
				e.g. "chromium:public"
			`))
			r.Flags.StringVar(&r.testIDPrefix, "test-id-prefix", "", text.Doc(`
				Prefix to prepend to the test ID of every test result.
			`))
			r.Flags.Var(flag.StringMap(r.vars), "var", text.Doc(`
				Variant to add to every test result in "key:value" format.
				If the test command adds a variant with the same key, the value given by
				this flag will get overridden.
			`))
			r.Flags.UintVar(&r.artChannelMaxLeases, "max-concurrent-artifact-uploads",
				sink.DefaultArtChannelMaxLeases, text.Doc(`
				The maximum number of goroutines uploading artifacts.
			`))
			r.Flags.UintVar(&r.trChannelMaxLeases, "max-concurrent-test-result-uploads",
				sink.DefaultTestResultChannelMaxLeases, text.Doc(`
				The maximum number of goroutines uploading test results.
			`))
			r.Flags.StringVar(&r.testTestLocationBase, "test-location-base", "", text.Doc(`
				File base to prepend to the test location file name, if the file name is a relative path.
				It must start with "//".
			`))
			r.Flags.Var(flag.StringPairs(r.tags), "tag", text.Doc(`
				Tag to add to every test result in "key:value" format.
				A key can be repeated.
			`))
			r.Flags.BoolVar(&r.coerceNegativeDuration, "coerce-negative-duration",
				false, text.Doc(`
				If true, all negative durations will be coerced to 0.
				If false, test results with negative durations will be rejected.
			`))
			return r
		},
	}
}

type streamRun struct {
	baseCommandRun

	// flags
	isNew                  bool
	realm                  string
	testIDPrefix           string
	testTestLocationBase   string
	vars                   map[string]string
	artChannelMaxLeases    uint
	trChannelMaxLeases     uint
	tags                   strpair.Map
	pbTags                 []*pb.StringPair
	coerceNegativeDuration bool
	// TODO(ddoman): add flags
	// - invocation-tag
	// - log-file

	invocation lucictx.ResultDBInvocation
}

func (r *streamRun) validate(ctx context.Context, args []string) (err error) {
	if len(args) == 0 {
		return errors.Reason("missing a test command to run").Err()
	}
	if err := pbutil.ValidateVariant(&pb.Variant{Def: r.vars}); err != nil {
		return errors.Annotate(err, "invalid variant").Err()
	}
	if r.realm != "" {
		if err := realms.ValidateRealmName(r.realm, realms.GlobalScope); err != nil {
			return errors.Annotate(err, "invalid realm").Err()
		}
	}
	return nil
}

func (r *streamRun) Run(a subcommands.Application, args []string, env subcommands.Env) (ret int) {
	ctx := cli.GetContext(a, r, env)

	if err := r.validate(ctx, args); err != nil {
		return r.done(err)
	}

	loginMode := auth.OptionalLogin
	// login is required only if it creates a new invocation.
	if r.isNew {
		if r.realm == "" {
			return r.done(errors.Reason("-realm is required for new invocations").Err())
		}
		loginMode = auth.SilentLogin
	}
	if err := r.initClients(ctx, loginMode); err != nil {
		return r.done(err)
	}

	// if -new is passed, create a new invocation. If not, use the existing one set in
	// lucictx.
	if r.isNew {
		ninv, err := r.createInvocation(ctx, r.realm)
		if err != nil {
			return r.done(err)
		}
		r.invocation = ninv

		// Update lucictx with the new invocation.
		ctx = lucictx.SetResultDB(ctx, &lucictx.ResultDB{
			Hostname:          r.host,
			CurrentInvocation: &r.invocation,
		})
	} else {
		if r.resultdbCtx == nil {
			return r.done(errors.Reason("the environment does not have an existing invocation; use -new to create a new one").Err())
		}
		if err := r.validateCurrentInvocation(); err != nil {
			return r.done(err)
		}
		r.invocation = *r.resultdbCtx.CurrentInvocation
	}

	defer func() {
		// Finalize the invocation if it was created by -new.
		if r.isNew {
			if err := r.finalizeInvocation(ctx); err != nil {
				logging.Errorf(ctx, "failed to finalize the invocation: %s", err)
				ret = r.done(err)
			}
		}
	}()

	err := r.runTestCmd(ctx, args)
	ec, ok := exitcode.Get(err)
	if !ok {
		logging.Errorf(ctx, "rdb-stream: failed to run the test command: %s", err)
		return r.done(err)
	}
	logging.Infof(ctx, "rdb-stream: exiting with %d", ec)
	return ec
}

func (r *streamRun) runTestCmd(ctx context.Context, args []string) error {
	// Kill the subprocess if rdb-stream is asked to stop.
	// Subprocess exiting will unblock rdb-stream and it will stop soon.
	cmdCtx, cancelCmd := context.WithCancel(ctx)
	defer cancelCmd()
	defer signals.HandleInterrupt(func() {
		logging.Warningf(ctx, "Interrupt signal received; killing the subprocess")
		cancelCmd()
	})()

	cmd := exec.CommandContext(cmdCtx, args[0], args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	// TODO(ddoman): send the logs of SinkServer to --log-file

	cfg := sink.ServerConfig{
		Recorder:                   r.recorder,
		Invocation:                 r.invocation.Name,
		UpdateToken:                r.invocation.UpdateToken,
		TestIDPrefix:               r.testIDPrefix,
		BaseVariant:                &pb.Variant{Def: r.vars},
		ArtifactUploader:           &sink.ArtifactUploader{Client: r.http, Host: r.host},
		ArtChannelMaxLeases:        r.artChannelMaxLeases,
		TestResultChannelMaxLeases: r.trChannelMaxLeases,
		TestLocationBase:           r.testTestLocationBase,
		BaseTags:                   pbutil.FromStrpairMap(r.tags),
		CoerceNegativeDuration:     r.coerceNegativeDuration,
	}
	return sink.Run(ctx, cfg, func(ctx context.Context, cfg sink.ServerConfig) error {
		exported, err := lucictx.Export(ctx)
		if err != nil {
			return err
		}
		defer func() {
			logging.Infof(ctx, "rdb-stream: the test process terminated")
			exported.Close()
		}()
		exported.SetInCmd(cmd)
		logging.Infof(ctx, "rdb-stream: starting the test command - %q", cmd.Args)
		if err := cmd.Start(); err != nil {
			return errors.Annotate(err, "cmd.start").Err()
		}
		return cmd.Wait()
	})
}

func (r *streamRun) createInvocation(ctx context.Context, realm string) (ret lucictx.ResultDBInvocation, err error) {
	invID, err := genInvID(ctx)
	if err != nil {
		return
	}

	md := metadata.MD{}
	resp, err := r.recorder.CreateInvocation(ctx, &pb.CreateInvocationRequest{
		InvocationId: invID,
		Invocation: &pb.Invocation{
			Realm: realm,
		},
	}, prpc.Header(&md))
	if err != nil {
		err = errors.Annotate(err, "failed to create an invocation").Err()
		return
	}
	tks := md.Get(recorder.UpdateTokenMetadataKey)
	if len(tks) == 0 {
		err = errors.Reason("Missing header: update-token").Err()
		return
	}

	ret = lucictx.ResultDBInvocation{Name: resp.Name, UpdateToken: tks[0]}
	fmt.Fprintf(os.Stderr, "rdb-stream: created invocation - https://ci.chromium.org/ui/inv/%s\n", invID)
	return
}

// finalizeInvocation finalizes the invocation.
func (r *streamRun) finalizeInvocation(ctx context.Context) error {
	id, err := pbutil.ParseInvocationName(r.invocation.Name)
	if err != nil {
		return errors.Reason("failed to parse invocation name(%q): %s", r.invocation.Name, err).Err()
	}

	fmt.Fprintf(os.Stderr, "rdb-stream: finalizing the invocation - https://ci.chromium.org/ui/inv/%s\n", id)
	ctx = metadata.AppendToOutgoingContext(
		ctx, recorder.UpdateTokenMetadataKey, r.invocation.UpdateToken)
	_, err = r.recorder.FinalizeInvocation(ctx, &pb.FinalizeInvocationRequest{
		Name: r.invocation.Name,
	})
	return err
}

// genInvID generates an invocation ID, made of the username, the current timestamp
// in a human-friendly format, and a random suffix.
//
// This can be used to generate a random invocation ID, but the creator and creation time
// can be easily found.
func genInvID(ctx context.Context) (string, error) {
	whoami, err := user.Current()
	if err != nil {
		return "", err
	}
	bytes := make([]byte, 8)
	if _, err := mathrand.Read(ctx, bytes); err != nil {
		return "", err
	}

	username := strings.ToLower(whoami.Username)
	username = matchInvalidInvocationIDChars.ReplaceAllString(username, "")

	suffix := strings.ToLower(fmt.Sprintf(
		"%s-%s", time.Now().UTC().Format("2006-01-02-15-04-00"),
		// Note: cannot use base64 because not all of its characters are allowed
		// in invocation IDs.
		hex.EncodeToString(bytes)))

	// An invocation ID can contain up to 100 ascii characters that conform to the regex,
	return fmt.Sprintf("u-%.*s-%s", 100-len(suffix), username, suffix), nil
}
<|endoftext|>


<commit_before>// Copyright 2015 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package discovery

import (
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/model"
	"golang.org/x/net/context"

	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/util/strutil"
)

const (
	ec2Label              = model.MetaLabelPrefix + "ec2_"
	ec2LabelAZ            = ec2Label + "availability_zone"
	ec2LabelInstanceID    = ec2Label + "instance_id"
	ec2LabelInstanceState = ec2Label + "instance_state"
	ec2LabelPublicDNS     = ec2Label + "public_dns_name"
	ec2LabelPublicIP      = ec2Label + "public_ip"
	ec2LabelPrivateIP     = ec2Label + "private_ip"
	ec2LabelSubnetID      = ec2Label + "subnet_id"
	ec2LabelTag           = ec2Label + "tag_"
	ec2LabelVPCID         = ec2Label + "vpc_id"
	subnetSeparator       = ","
)

// EC2Discovery periodically performs EC2-SD requests. It implements
// the TargetProvider interface.
type EC2Discovery struct {
	aws      *aws.Config
	interval time.Duration
	port     int
}

// NewEC2Discovery returns a new EC2Discovery which periodically refreshes its targets.
func NewEC2Discovery(conf *config.EC2SDConfig) *EC2Discovery {
	creds := credentials.NewStaticCredentials(conf.AccessKey, conf.SecretKey, "")
	if conf.AccessKey == "" && conf.SecretKey == "" {
		creds = defaults.DefaultChainCredentials
	}
	return &EC2Discovery{
		aws: &aws.Config{
			Region:      &conf.Region,
			Credentials: creds,
		},
		interval: time.Duration(conf.RefreshInterval),
		port:     conf.Port,
	}
}

// Run implements the TargetProvider interface.
func (ed *EC2Discovery) Run(ctx context.Context, ch chan<- []*config.TargetGroup) {
	defer close(ch)

	ticker := time.NewTicker(ed.interval)
	defer ticker.Stop()

	// Get an initial set right away.
	tg, err := ed.refresh()
	if err != nil {
		log.Error(err)
	} else {
		ch <- []*config.TargetGroup{tg}
	}

	for {
		select {
		case <-ticker.C:
			tg, err := ed.refresh()
			if err != nil {
				log.Error(err)
			} else {
				ch <- []*config.TargetGroup{tg}
			}
		case <-ctx.Done():
			return
		}
	}
}

func (ed *EC2Discovery) refresh() (*config.TargetGroup, error) {
	ec2s := ec2.New(ed.aws)
	tg := &config.TargetGroup{
		Source: *ed.aws.Region,
	}
	if err := ec2s.DescribeInstancesPages(nil, func(p *ec2.DescribeInstancesOutput, lastPage bool) bool {
		for _, r := range p.Reservations {
			for _, inst := range r.Instances {
				if inst.PrivateIpAddress == nil {
					continue
				}
				labels := model.LabelSet{
					ec2LabelInstanceID: model.LabelValue(*inst.InstanceId),
				}
				labels[ec2LabelPrivateIP] = model.LabelValue(*inst.PrivateIpAddress)
				addr := net.JoinHostPort(*inst.PrivateIpAddress, fmt.Sprintf("%d", ed.port))
				labels[model.AddressLabel] = model.LabelValue(addr)

				if inst.PublicIpAddress != nil {
					labels[ec2LabelPublicIP] = model.LabelValue(*inst.PublicIpAddress)
					labels[ec2LabelPublicDNS] = model.LabelValue(*inst.PublicDnsName)
				}

				labels[ec2LabelAZ] = model.LabelValue(*inst.Placement.AvailabilityZone)
				labels[ec2LabelInstanceState] = model.LabelValue(*inst.State.Name)

				if inst.VpcId != nil {
					labels[ec2LabelVPCID] = model.LabelValue(*inst.VpcId)

					subnetsMap := make(map[string]struct{})
					for _, eni := range inst.NetworkInterfaces {
						subnetsMap[*eni.SubnetId] = struct{}{}
					}
					subnets := []string{}
					for k := range subnetsMap {
						subnets = append(subnets, k)
					}
					labels[ec2LabelSubnetID] = model.LabelValue(
						subnetSeparator +
							strings.Join(subnets, subnetSeparator) +
							subnetSeparator)
				}

				for _, t := range inst.Tags {
					name := strutil.SanitizeLabelName(*t.Key)
					labels[ec2LabelTag+model.LabelName(name)] = model.LabelValue(*t.Value)
				}
				tg.Targets = append(tg.Targets, labels)
			}
		}
		return true
	}); err != nil {
		return nil, fmt.Errorf("could not describe instances: %s", err)
	}
	return tg, nil
}
<commit_msg>Add EC2 SD metrics (#2095)<commit_after>// Copyright 2015 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package discovery

import (
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/model"
	"golang.org/x/net/context"

	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/util/strutil"
)

const (
	ec2Label              = model.MetaLabelPrefix + "ec2_"
	ec2LabelAZ            = ec2Label + "availability_zone"
	ec2LabelInstanceID    = ec2Label + "instance_id"
	ec2LabelInstanceState = ec2Label + "instance_state"
	ec2LabelPublicDNS     = ec2Label + "public_dns_name"
	ec2LabelPublicIP      = ec2Label + "public_ip"
	ec2LabelPrivateIP     = ec2Label + "private_ip"
	ec2LabelSubnetID      = ec2Label + "subnet_id"
	ec2LabelTag           = ec2Label + "tag_"
	ec2LabelVPCID         = ec2Label + "vpc_id"
	subnetSeparator       = ","
)

var (
	ec2SDScrapeFailuresCount = prometheus.NewCounter(
		prometheus.CounterOpts{
			Namespace: namespace,
			Name:      "ec2_sd_scape_failures_total",
			Help:      "The number of EC2-SD scrape failures.",
		})
	ec2SDScrapeDuration = prometheus.NewSummary(
		prometheus.SummaryOpts{
			Namespace: namespace,
			Name:      "ec2_sd_scrape_duration_seconds",
			Help:      "The duration of a EC2-SD scrape in seconds.",
		})
)

func init() {
	prometheus.MustRegister(ec2SDScrapeFailuresCount)
	prometheus.MustRegister(ec2SDScrapeDuration)
}

// EC2Discovery periodically performs EC2-SD requests. It implements
// the TargetProvider interface.
type EC2Discovery struct {
	aws      *aws.Config
	interval time.Duration
	port     int
}

// NewEC2Discovery returns a new EC2Discovery which periodically refreshes its targets.
func NewEC2Discovery(conf *config.EC2SDConfig) *EC2Discovery {
	creds := credentials.NewStaticCredentials(conf.AccessKey, conf.SecretKey, "")
	if conf.AccessKey == "" && conf.SecretKey == "" {
		creds = defaults.DefaultChainCredentials
	}
	return &EC2Discovery{
		aws: &aws.Config{
			Region:      &conf.Region,
			Credentials: creds,
		},
		interval: time.Duration(conf.RefreshInterval),
		port:     conf.Port,
	}
}

// Run implements the TargetProvider interface.
func (ed *EC2Discovery) Run(ctx context.Context, ch chan<- []*config.TargetGroup) {
	defer close(ch)

	ticker := time.NewTicker(ed.interval)
	defer ticker.Stop()

	// Get an initial set right away.
	tg, err := ed.refresh()
	if err != nil {
		log.Error(err)
	} else {
		ch <- []*config.TargetGroup{tg}
	}

	for {
		select {
		case <-ticker.C:
			tg, err := ed.refresh()
			if err != nil {
				log.Error(err)
			} else {
				ch <- []*config.TargetGroup{tg}
			}
		case <-ctx.Done():
			return
		}
	}
}

func (ed *EC2Discovery) refresh() (tg *config.TargetGroup, err error) {
	t0 := time.Now()
	defer func() {
		ec2SDScrapeDuration.Observe(time.Since(t0).Seconds())
		if err != nil {
			ec2SDScrapeFailuresCount.Inc()
		}
	}()

	ec2s := ec2.New(ed.aws)
	tg = &config.TargetGroup{
		Source: *ed.aws.Region,
	}
	if err = ec2s.DescribeInstancesPages(nil, func(p *ec2.DescribeInstancesOutput, lastPage bool) bool {
		for _, r := range p.Reservations {
			for _, inst := range r.Instances {
				if inst.PrivateIpAddress == nil {
					continue
				}
				labels := model.LabelSet{
					ec2LabelInstanceID: model.LabelValue(*inst.InstanceId),
				}
				labels[ec2LabelPrivateIP] = model.LabelValue(*inst.PrivateIpAddress)
				addr := net.JoinHostPort(*inst.PrivateIpAddress, fmt.Sprintf("%d", ed.port))
				labels[model.AddressLabel] = model.LabelValue(addr)

				if inst.PublicIpAddress != nil {
					labels[ec2LabelPublicIP] = model.LabelValue(*inst.PublicIpAddress)
					labels[ec2LabelPublicDNS] = model.LabelValue(*inst.PublicDnsName)
				}

				labels[ec2LabelAZ] = model.LabelValue(*inst.Placement.AvailabilityZone)
				labels[ec2LabelInstanceState] = model.LabelValue(*inst.State.Name)

				if inst.VpcId != nil {
					labels[ec2LabelVPCID] = model.LabelValue(*inst.VpcId)

					subnetsMap := make(map[string]struct{})
					for _, eni := range inst.NetworkInterfaces {
						subnetsMap[*eni.SubnetId] = struct{}{}
					}
					subnets := []string{}
					for k := range subnetsMap {
						subnets = append(subnets, k)
					}
					labels[ec2LabelSubnetID] = model.LabelValue(
						subnetSeparator +
							strings.Join(subnets, subnetSeparator) +
							subnetSeparator)
				}

				for _, t := range inst.Tags {
					name := strutil.SanitizeLabelName(*t.Key)
					labels[ec2LabelTag+model.LabelName(name)] = model.LabelValue(*t.Value)
				}
				tg.Targets = append(tg.Targets, labels)
			}
		}
		return true
	}); err != nil {
		return nil, fmt.Errorf("could not describe instances: %s", err)
	}
	return tg, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math/rand"
	"os"
	"os/exec"
	"os/signal"
	"time"
)

const (
	lose int = iota
	tie
	win
)

var player string
var givenAnswer int
var signalChan = make(chan os.Signal, 1) // channel to catch ctrl-c

// game holds the data collected during game play
type game struct {
	attempts int    // track how many rounds played
	player   string // player name
	pAnswer  *int   // pointer to given answer
	cAnswer  int    // computer's answer
	results  []int  // array to hold wins, loses, and ties
}

// checkValidAnswer makes sure the given answer is valid
func checkValidAnswer(pa *int) bool {
	if *pa == lose || *pa == tie || *pa == win {
		return true
	}
	return false
}

// clearScreen runs a shell clear command
func clearScreen() {
	c := exec.Command("clear")
	c.Stdout = os.Stdout
	c.Run()
}

// genStats outputs the game play statistics
func (g *game) genStats() {
	var wins, loses, ties int
	for _, i := range g.results {
		switch {
		case i == win:
			wins++
		case i == lose:
			loses++
		case i == tie:
			ties++
		}
	}
	fmt.Printf("\n\n%s, here are your stats...\n", g.player)
	fmt.Printf("Rounds: %d, Wins: %d, Loses: %d, Ties: %d\n\n", len(g.results), wins, loses, ties)
	os.Exit(1) // Since it was a ctrl-c, exit non-zero
}

// genComputerAnswer will randomly generate a number used as an answer
func genComputerAnswer() int {
	rand.Seed(time.Now().UTC().UnixNano())
	return rand.Intn(3)
}

func main() {
	clearScreen()
	fmt.Print("+ Rock-Paper-Scissors +\n\n")
	fmt.Println("Enter 0 for rock, 1 for paper, and 2 for scissors")
	fmt.Print("Enter your name: ")
	fmt.Scanf("%s", &player)
	g := game{
		player:   player,
		attempts: 0,
		results:  make([]int, 0),
	}
	signal.Notify(signalChan, os.Interrupt)
	// setup go routine to catch a ctrl-c
	go func() {
		for range signalChan {
			g.genStats()
		}
	}()
	for {
		fmt.Print("Enter answer: ")
		fmt.Scanf("%d", &givenAnswer)
		if !checkValidAnswer(&givenAnswer) {
			fmt.Println("invalid answer, try again")
			continue
		}
		g.attempts = g.attempts + 1
		g.pAnswer = &givenAnswer
		g.cAnswer = genComputerAnswer()
		switch {
		case g.cAnswer%3+1 == *g.pAnswer:
			g.results = append(g.results, win)
			fmt.Println("Win")
		case *g.pAnswer%3+1 == g.cAnswer:
			g.results = append(g.results, lose)
			fmt.Println("lose")
		default:
			g.results = append(g.results, tie)
			fmt.Println("tie")
		}
	}
}
<commit_msg>license addition to code as well as description and small changes<commit_after>// Copyright 2015 Brian J. Downs
//
// 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.

// This is a basic implementation of the game rock, paper, scissors.
package main

import (
	"fmt"
	"math/rand"
	"os"
	"os/exec"
	"os/signal"
	"time"
)

const (
	lose int = iota
	tie
	win
)

var player string
var givenAnswer int
var signalChan = make(chan os.Signal, 1) // channel to catch ctrl-c

// game holds the data collected during game play
type game struct {
	attempts int   // track how many rounds played
	pAnswer  *int  // pointer to given answer
	cAnswer  int   // computer's answer
	results  []int // array to hold wins, loses, and ties
}

// checkValidAnswer makes sure the given answer is valid
func checkValidAnswer(pa *int) bool {
	if *pa == lose || *pa == tie || *pa == win {
		return true
	}
	return false
}

// clearScreen runs a shell clear command
func clearScreen() {
	c := exec.Command("clear")
	c.Stdout = os.Stdout
	c.Run()
}

// genStats outputs the game play statistics
func (g *game) genStats() {
	var wins, loses, ties int
	for _, i := range g.results {
		switch {
		case i == win:
			wins++
		case i == lose:
			loses++
		case i == tie:
			ties++
		}
	}
	fmt.Printf("\n\nRounds: %d, Wins: %d, Loses: %d, Ties: %d\n\n", len(g.results), wins, loses, ties)
	os.Exit(1) // Since it was a ctrl-c, exit non-zero
}

// genComputerAnswer will randomly generate a number used as an answer
func genComputerAnswer() int {
	rand.Seed(time.Now().UTC().UnixNano())
	return rand.Intn(3)
}

func main() {
	clearScreen()
	fmt.Print("+ Rock-Paper-Scissors (Enter 0 for ROCK, 1 for PAPER, and 2 for SCISSORS)\n\n")
	g := game{
		attempts: 0,
		results:  make([]int, 0),
	}
	signal.Notify(signalChan, os.Interrupt)
	// setup go routine to catch a ctrl-c
	go func() {
		for range signalChan {
			g.genStats()
		}
	}()
	for {
		fmt.Print("Enter answer: ")
		fmt.Scanf("%d", &givenAnswer)
		if !checkValidAnswer(&givenAnswer) {
			fmt.Println("invalid answer, try again")
			continue
		}
		g.attempts = g.attempts + 1
		g.pAnswer = &givenAnswer
		g.cAnswer = genComputerAnswer()
		switch {
		case g.cAnswer%3+1 == *g.pAnswer:
			g.results = append(g.results, win)
			fmt.Println("Win")
		case *g.pAnswer%3+1 == g.cAnswer:
			g.results = append(g.results, lose)
			fmt.Println("lose")
		default:
			g.results = append(g.results, tie)
			fmt.Println("tie")
		}
	}
}
<|endoftext|>


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

package user

import (
	"net/http"

	asymkey_model "code.gitea.io/gitea/models/asymkey"
	"code.gitea.io/gitea/models/perm"
	user_model "code.gitea.io/gitea/models/user"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/convert"
	"code.gitea.io/gitea/modules/setting"
	api "code.gitea.io/gitea/modules/structs"
	"code.gitea.io/gitea/modules/web"
	"code.gitea.io/gitea/routers/api/v1/repo"
	"code.gitea.io/gitea/routers/api/v1/utils"
	asymkey_service "code.gitea.io/gitea/services/asymkey"
)

// appendPrivateInformation appends the owner and key type information to api.PublicKey
func appendPrivateInformation(apiKey *api.PublicKey, key *asymkey_model.PublicKey, defaultUser *user_model.User) (*api.PublicKey, error) {
	if key.Type == asymkey_model.KeyTypeDeploy {
		apiKey.KeyType = "deploy"
	} else if key.Type == asymkey_model.KeyTypeUser {
		apiKey.KeyType = "user"

		if defaultUser.ID == key.OwnerID {
			apiKey.Owner = convert.ToUser(defaultUser, defaultUser)
		} else {
			user, err := user_model.GetUserByID(key.OwnerID)
			if err != nil {
				return apiKey, err
			}
			apiKey.Owner = convert.ToUser(user, user)
		}
	} else {
		apiKey.KeyType = "unknown"
	}
	apiKey.ReadOnly = key.Mode == perm.AccessModeRead
	return apiKey, nil
}

func composePublicKeysAPILink() string {
	return setting.AppURL + "api/v1/user/keys/"
}

func listPublicKeys(ctx *context.APIContext, user *user_model.User) {
	var keys []*asymkey_model.PublicKey
	var err error
	var count int

	fingerprint := ctx.FormString("fingerprint")
	username := ctx.Params("username")

	if fingerprint != "" {
		// Querying not just listing
		if username != "" {
			// Restrict to provided uid
			keys, err = asymkey_model.SearchPublicKey(user.ID, fingerprint)
		} else {
			// Unrestricted
			keys, err = asymkey_model.SearchPublicKey(0, fingerprint)
		}
		count = len(keys)
	} else {
		total, err2 := asymkey_model.CountPublicKeys(user.ID)
		if err2 != nil {
			ctx.InternalServerError(err)
			return
		}
		count = int(total)

		// Use ListPublicKeys
		keys, err = asymkey_model.ListPublicKeys(user.ID, utils.GetListOptions(ctx))
	}

	if err != nil {
		ctx.Error(http.StatusInternalServerError, "ListPublicKeys", err)
		return
	}

	apiLink := composePublicKeysAPILink()
	apiKeys := make([]*api.PublicKey, len(keys))
	for i := range keys {
		apiKeys[i] = convert.ToPublicKey(apiLink, keys[i])
		if ctx.Doer.IsAdmin || ctx.Doer.ID == keys[i].OwnerID {
			apiKeys[i], _ = appendPrivateInformation(apiKeys[i], keys[i], user)
		}
	}

	ctx.SetTotalCountHeader(int64(count))
	ctx.JSON(http.StatusOK, &apiKeys)
}

// ListMyPublicKeys list all of the authenticated user's public keys
func ListMyPublicKeys(ctx *context.APIContext) {
	// swagger:operation GET /user/keys user userCurrentListKeys
	// ---
	// summary: List the authenticated user's public keys
	// parameters:
	// - name: fingerprint
	//   in: query
	//   description: fingerprint of the key
	//   type: string
	// - name: page
	//   in: query
	//   description: page number of results to return (1-based)
	//   type: integer
	// - name: limit
	//   in: query
	//   description: page size of results
	//   type: integer
	// produces:
	// - application/json
	// responses:
	//   "200":
	//     "$ref": "#/responses/PublicKeyList"

	listPublicKeys(ctx, ctx.Doer)
}

// ListPublicKeys list the given user's public keys
func ListPublicKeys(ctx *context.APIContext) {
	// swagger:operation GET /users/{username}/keys user userListKeys
	// ---
	// summary: List the given user's public keys
	// produces:
	// - application/json
	// parameters:
	// - name: username
	//   in: path
	//   description: username of user
	//   type: string
	//   required: true
	// - name: fingerprint
	//   in: query
	//   description: fingerprint of the key
	//   type: string
	// - name: page
	//   in: query
	//   description: page number of results to return (1-based)
	//   type: integer
	// - name: limit
	//   in: query
	//   description: page size of results
	//   type: integer
	// responses:
	//   "200":
	//     "$ref": "#/responses/PublicKeyList"

	listPublicKeys(ctx, ctx.ContextUser)
}

// GetPublicKey get a public key
func GetPublicKey(ctx *context.APIContext) {
	// swagger:operation GET /user/keys/{id} user userCurrentGetKey
	// ---
	// summary: Get a public key
	// produces:
	// - application/json
	// parameters:
	// - name: id
	//   in: path
	//   description: id of key to get
	//   type: integer
	//   format: int64
	//   required: true
	// responses:
	//   "200":
	//     "$ref": "#/responses/PublicKey"
	//   "404":
	//     "$ref": "#/responses/notFound"

	key, err := asymkey_model.GetPublicKeyByID(ctx.ParamsInt64(":id"))
	if err != nil {
		if asymkey_model.IsErrKeyNotExist(err) {
			ctx.NotFound()
		} else {
			ctx.Error(http.StatusInternalServerError, "GetPublicKeyByID", err)
		}
		return
	}

	apiLink := composePublicKeysAPILink()
	apiKey := convert.ToPublicKey(apiLink, key)
	if ctx.Doer.IsAdmin || ctx.Doer.ID == key.OwnerID {
		apiKey, _ = appendPrivateInformation(apiKey, key, ctx.Doer)
	}
	ctx.JSON(http.StatusOK, apiKey)
}

// CreateUserPublicKey creates new public key to given user by ID.
func CreateUserPublicKey(ctx *context.APIContext, form api.CreateKeyOption, uid int64) {
	content, err := asymkey_model.CheckPublicKeyString(form.Key)
	if err != nil {
		repo.HandleCheckKeyStringError(ctx, err)
		return
	}

	key, err := asymkey_model.AddPublicKey(uid, form.Title, content, 0)
	if err != nil {
		repo.HandleAddKeyError(ctx, err)
		return
	}
	apiLink := composePublicKeysAPILink()
	apiKey := convert.ToPublicKey(apiLink, key)
	if ctx.Doer.IsAdmin || ctx.Doer.ID == key.OwnerID {
		apiKey, _ = appendPrivateInformation(apiKey, key, ctx.Doer)
	}
	ctx.JSON(http.StatusCreated, apiKey)
}

// CreatePublicKey create one public key for me
func CreatePublicKey(ctx *context.APIContext) {
	// swagger:operation POST /user/keys user userCurrentPostKey
	// ---
	// summary: Create a public key
	// consumes:
	// - application/json
	// produces:
	// - application/json
	// parameters:
	// - name: body
	//   in: body
	//   schema:
	//     "$ref": "#/definitions/CreateKeyOption"
	// responses:
	//   "201":
	//     "$ref": "#/responses/PublicKey"
	//   "422":
	//     "$ref": "#/responses/validationError"

	form := web.GetForm(ctx).(*api.CreateKeyOption)
	CreateUserPublicKey(ctx, *form, ctx.Doer.ID)
}

// DeletePublicKey delete one public key
func DeletePublicKey(ctx *context.APIContext) {
	// swagger:operation DELETE /user/keys/{id} user userCurrentDeleteKey
	// ---
	// summary: Delete a public key
	// produces:
	// - application/json
	// parameters:
	// - name: id
	//   in: path
	//   description: id of key to delete
	//   type: integer
	//   format: int64
	//   required: true
	// responses:
	//   "204":
	//     "$ref": "#/responses/empty"
	//   "403":
	//     "$ref": "#/responses/forbidden"
	//   "404":
	//     "$ref": "#/responses/notFound"

	id := ctx.ParamsInt64(":id")
	externallyManaged, err := asymkey_model.PublicKeyIsExternallyManaged(id)
	if err != nil {
		ctx.Error(http.StatusInternalServerError, "PublicKeyIsExternallyManaged", err)
	}
	if externallyManaged {
		ctx.Error(http.StatusForbidden, "", "SSH Key is externally managed for this user")
	}

	if err := asymkey_service.DeletePublicKey(ctx.Doer, id); err != nil {
		if asymkey_model.IsErrKeyNotExist(err) {
			ctx.NotFound()
		} else if asymkey_model.IsErrKeyAccessDenied(err) {
			ctx.Error(http.StatusForbidden, "", "You do not have access to this key")
		} else {
			ctx.Error(http.StatusInternalServerError, "DeletePublicKey", err)
		}
		return
	}

	ctx.Status(http.StatusNoContent)
}
<commit_msg>Fix DELETE request for non-existent public key (#19443)<commit_after>// Copyright 2015 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package user

import (
	"net/http"

	asymkey_model "code.gitea.io/gitea/models/asymkey"
	"code.gitea.io/gitea/models/perm"
	user_model "code.gitea.io/gitea/models/user"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/convert"
	"code.gitea.io/gitea/modules/setting"
	api "code.gitea.io/gitea/modules/structs"
	"code.gitea.io/gitea/modules/web"
	"code.gitea.io/gitea/routers/api/v1/repo"
	"code.gitea.io/gitea/routers/api/v1/utils"
	asymkey_service "code.gitea.io/gitea/services/asymkey"
)

// appendPrivateInformation appends the owner and key type information to api.PublicKey
func appendPrivateInformation(apiKey *api.PublicKey, key *asymkey_model.PublicKey, defaultUser *user_model.User) (*api.PublicKey, error) {
	if key.Type == asymkey_model.KeyTypeDeploy {
		apiKey.KeyType = "deploy"
	} else if key.Type == asymkey_model.KeyTypeUser {
		apiKey.KeyType = "user"

		if defaultUser.ID == key.OwnerID {
			apiKey.Owner = convert.ToUser(defaultUser, defaultUser)
		} else {
			user, err := user_model.GetUserByID(key.OwnerID)
			if err != nil {
				return apiKey, err
			}
			apiKey.Owner = convert.ToUser(user, user)
		}
	} else {
		apiKey.KeyType = "unknown"
	}
	apiKey.ReadOnly = key.Mode == perm.AccessModeRead
	return apiKey, nil
}

func composePublicKeysAPILink() string {
	return setting.AppURL + "api/v1/user/keys/"
}

func listPublicKeys(ctx *context.APIContext, user *user_model.User) {
	var keys []*asymkey_model.PublicKey
	var err error
	var count int

	fingerprint := ctx.FormString("fingerprint")
	username := ctx.Params("username")

	if fingerprint != "" {
		// Querying not just listing
		if username != "" {
			// Restrict to provided uid
			keys, err = asymkey_model.SearchPublicKey(user.ID, fingerprint)
		} else {
			// Unrestricted
			keys, err = asymkey_model.SearchPublicKey(0, fingerprint)
		}
		count = len(keys)
	} else {
		total, err2 := asymkey_model.CountPublicKeys(user.ID)
		if err2 != nil {
			ctx.InternalServerError(err)
			return
		}
		count = int(total)

		// Use ListPublicKeys
		keys, err = asymkey_model.ListPublicKeys(user.ID, utils.GetListOptions(ctx))
	}

	if err != nil {
		ctx.Error(http.StatusInternalServerError, "ListPublicKeys", err)
		return
	}

	apiLink := composePublicKeysAPILink()
	apiKeys := make([]*api.PublicKey, len(keys))
	for i := range keys {
		apiKeys[i] = convert.ToPublicKey(apiLink, keys[i])
		if ctx.Doer.IsAdmin || ctx.Doer.ID == keys[i].OwnerID {
			apiKeys[i], _ = appendPrivateInformation(apiKeys[i], keys[i], user)
		}
	}

	ctx.SetTotalCountHeader(int64(count))
	ctx.JSON(http.StatusOK, &apiKeys)
}

// ListMyPublicKeys list all of the authenticated user's public keys
func ListMyPublicKeys(ctx *context.APIContext) {
	// swagger:operation GET /user/keys user userCurrentListKeys
	// ---
	// summary: List the authenticated user's public keys
	// parameters:
	// - name: fingerprint
	//   in: query
	//   description: fingerprint of the key
	//   type: string
	// - name: page
	//   in: query
	//   description: page number of results to return (1-based)
	//   type: integer
	// - name: limit
	//   in: query
	//   description: page size of results
	//   type: integer
	// produces:
	// - application/json
	// responses:
	//   "200":
	//     "$ref": "#/responses/PublicKeyList"

	listPublicKeys(ctx, ctx.Doer)
}

// ListPublicKeys list the given user's public keys
func ListPublicKeys(ctx *context.APIContext) {
	// swagger:operation GET /users/{username}/keys user userListKeys
	// ---
	// summary: List the given user's public keys
	// produces:
	// - application/json
	// parameters:
	// - name: username
	//   in: path
	//   description: username of user
	//   type: string
	//   required: true
	// - name: fingerprint
	//   in: query
	//   description: fingerprint of the key
	//   type: string
	// - name: page
	//   in: query
	//   description: page number of results to return (1-based)
	//   type: integer
	// - name: limit
	//   in: query
	//   description: page size of results
	//   type: integer
	// responses:
	//   "200":
	//     "$ref": "#/responses/PublicKeyList"

	listPublicKeys(ctx, ctx.ContextUser)
}

// GetPublicKey get a public key
func GetPublicKey(ctx *context.APIContext) {
	// swagger:operation GET /user/keys/{id} user userCurrentGetKey
	// ---
	// summary: Get a public key
	// produces:
	// - application/json
	// parameters:
	// - name: id
	//   in: path
	//   description: id of key to get
	//   type: integer
	//   format: int64
	//   required: true
	// responses:
	//   "200":
	//     "$ref": "#/responses/PublicKey"
	//   "404":
	//     "$ref": "#/responses/notFound"

	key, err := asymkey_model.GetPublicKeyByID(ctx.ParamsInt64(":id"))
	if err != nil {
		if asymkey_model.IsErrKeyNotExist(err) {
			ctx.NotFound()
		} else {
			ctx.Error(http.StatusInternalServerError, "GetPublicKeyByID", err)
		}
		return
	}

	apiLink := composePublicKeysAPILink()
	apiKey := convert.ToPublicKey(apiLink, key)
	if ctx.Doer.IsAdmin || ctx.Doer.ID == key.OwnerID {
		apiKey, _ = appendPrivateInformation(apiKey, key, ctx.Doer)
	}
	ctx.JSON(http.StatusOK, apiKey)
}

// CreateUserPublicKey creates new public key to given user by ID.
func CreateUserPublicKey(ctx *context.APIContext, form api.CreateKeyOption, uid int64) {
	content, err := asymkey_model.CheckPublicKeyString(form.Key)
	if err != nil {
		repo.HandleCheckKeyStringError(ctx, err)
		return
	}

	key, err := asymkey_model.AddPublicKey(uid, form.Title, content, 0)
	if err != nil {
		repo.HandleAddKeyError(ctx, err)
		return
	}
	apiLink := composePublicKeysAPILink()
	apiKey := convert.ToPublicKey(apiLink, key)
	if ctx.Doer.IsAdmin || ctx.Doer.ID == key.OwnerID {
		apiKey, _ = appendPrivateInformation(apiKey, key, ctx.Doer)
	}
	ctx.JSON(http.StatusCreated, apiKey)
}

// CreatePublicKey create one public key for me
func CreatePublicKey(ctx *context.APIContext) {
	// swagger:operation POST /user/keys user userCurrentPostKey
	// ---
	// summary: Create a public key
	// consumes:
	// - application/json
	// produces:
	// - application/json
	// parameters:
	// - name: body
	//   in: body
	//   schema:
	//     "$ref": "#/definitions/CreateKeyOption"
	// responses:
	//   "201":
	//     "$ref": "#/responses/PublicKey"
	//   "422":
	//     "$ref": "#/responses/validationError"

	form := web.GetForm(ctx).(*api.CreateKeyOption)
	CreateUserPublicKey(ctx, *form, ctx.Doer.ID)
}

// DeletePublicKey delete one public key
func DeletePublicKey(ctx *context.APIContext) {
	// swagger:operation DELETE /user/keys/{id} user userCurrentDeleteKey
	// ---
	// summary: Delete a public key
	// produces:
	// - application/json
	// parameters:
	// - name: id
	//   in: path
	//   description: id of key to delete
	//   type: integer
	//   format: int64
	//   required: true
	// responses:
	//   "204":
	//     "$ref": "#/responses/empty"
	//   "403":
	//     "$ref": "#/responses/forbidden"
	//   "404":
	//     "$ref": "#/responses/notFound"

	id := ctx.ParamsInt64(":id")
	externallyManaged, err := asymkey_model.PublicKeyIsExternallyManaged(id)
	if err != nil {
		if asymkey_model.IsErrKeyNotExist(err) {
			ctx.NotFound()
		} else {
			ctx.Error(http.StatusInternalServerError, "PublicKeyIsExternallyManaged", err)
		}
		return
	}

	if externallyManaged {
		ctx.Error(http.StatusForbidden, "", "SSH Key is externally managed for this user")
		return
	}

	if err := asymkey_service.DeletePublicKey(ctx.Doer, id); err != nil {
		if asymkey_model.IsErrKeyAccessDenied(err) {
			ctx.Error(http.StatusForbidden, "", "You do not have access to this key")
		} else {
			ctx.Error(http.StatusInternalServerError, "DeletePublicKey", err)
		}
		return
	}

	ctx.Status(http.StatusNoContent)
}
<|endoftext|>


<commit_before>package genetics

import (
	"testing"
	"fmt"
	"github.com/yaricom/goNEAT/neat/network"
	"bytes"
	"github.com/yaricom/goNEAT/neat"
)

// Tests Gene WriteGene
func TestGene_WriteGene(t *testing.T) {
	// gene  1 1 4 1.1983046913458986 0 1.0 1.1983046913458986 0
	traitId, inNodeId, outNodeId, innov_num := 1, 1, 4, int64(1)
	weight, mut_num := 1.1983046913458986, 1.1983046913458986
	recurrent, enabled := false, false
	gene_str := fmt.Sprintf("%d %d %d %g %t %d %g %t",
		traitId, inNodeId, outNodeId, weight, recurrent, innov_num, mut_num, enabled)

	trait := neat.NewTrait()
	trait.Id = traitId
	gene := NewGeneWithTrait(trait, weight, network.NewNNode(1, network.InputNeuron),
		network.NewNNode(4, network.HiddenNeuron), recurrent, innov_num, mut_num)
	gene.IsEnabled = enabled

	out_buf := bytes.NewBufferString("")
	gene.Write(out_buf)

	out_str := out_buf.String()
	if gene_str != out_str {
		t.Errorf("Wrong Gene serialization\n[%s]\n[%s]", gene_str, out_str)
	}
}
<commit_msg>Implemented copy constructor test and removed write test<commit_after>package genetics

import (
	"testing"
	"github.com/yaricom/goNEAT/neat/network"
	"github.com/yaricom/goNEAT/neat"
	"reflect"
)

// Tests Gene WriteGene
func TestNewGeneCopy(t *testing.T) {
	nodes := []*network.NNode{
		{Id:1, NeuronType: network.InputNeuron, ActivationType: network.NullActivation, Incoming:make([]*network.Link, 0), Outgoing:make([]*network.Link, 0)},
		{Id:2, NeuronType: network.OutputNeuron, ActivationType: network.SigmoidSteepenedActivation, Incoming:make([]*network.Link, 0), Outgoing:make([]*network.Link, 0)},
	}
	trait := &neat.Trait{Id:1, Params:[]float64{0.1, 0, 0, 0, 0, 0, 0, 0}}
	g1 := NewGeneWithTrait(trait, 3.2, nodes[0], nodes[1], true, 42, 5.2)

	// test
	g := NewGeneCopy(g1, trait, nodes[0], nodes[1])
	if g.Link.InNode.Id != nodes[0].Id {
		t.Error("g.Link.InNode.Id != nodes[0].Id", g.Link.InNode.Id)
	}
	if g.Link.OutNode.Id != nodes[1].Id {
		t.Error("g.Link.OutNode.Id != nodes[1].Id", g.Link.OutNode.Id)
	}
	if g.Link.Trait.Id != trait.Id {
		t.Error("g.Link.Trait.Id != trait.Id", g.Link.Trait.Id)
	}
	if reflect.DeepEqual(g.Link.Trait.Params, trait.Params) == false {
		t.Error("reflect.DeepEqual(g.Link.Trait.Params, trait.Params) == false")
	}
	if g.InnovationNum != g1.InnovationNum {
		t.Error("g.InnovationNum != g1.InnovationNum", g.InnovationNum, g1.InnovationNum)
	}
	if g.MutationNum != g1.MutationNum {
		t.Error("g.MutationNum != g1.MutationNum", g.MutationNum, g1.MutationNum)
	}
	if g.IsEnabled != g1.IsEnabled {
		t.Error("g.IsEnabled != g1.IsEnabled", g.IsEnabled, g1.IsEnabled)
	}
}
<|endoftext|>


<commit_before>package zookeeper

import (
	"bufio"
	"context"
	"crypto/tls"
	"fmt"
	"net"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	tlsint "github.com/influxdata/telegraf/plugins/common/tls"
	"github.com/influxdata/telegraf/plugins/inputs"
)

var zookeeperFormatRE = regexp.MustCompile(`^zk_(\w+)\s+([\w\.\-]+)`)

// Zookeeper is a zookeeper plugin
type Zookeeper struct {
	Servers []string
	Timeout internal.Duration

	EnableTLS bool `toml:"enable_tls"`
	EnableSSL bool `toml:"enable_ssl"` // deprecated in 1.7; use enable_tls
	tlsint.ClientConfig

	initialized bool
	tlsConfig   *tls.Config
}

var sampleConfig = `
  ## An array of address to gather stats about. Specify an ip or hostname
  ## with port. ie localhost:2181, 10.0.0.1:2181, etc.

  ## If no servers are specified, then localhost is used as the host.
  ## If no port is specified, 2181 is used
  servers = [":2181"]

  ## Timeout for metric collections from all servers.  Minimum timeout is "1s".
  # timeout = "5s"

  ## Optional TLS Config
  # enable_tls = true
  # tls_ca = "/etc/telegraf/ca.pem"
  # tls_cert = "/etc/telegraf/cert.pem"
  # tls_key = "/etc/telegraf/key.pem"
  ## If false, skip chain & host verification
  # insecure_skip_verify = true
`

var defaultTimeout = 5 * time.Second

// SampleConfig returns sample configuration message
func (z *Zookeeper) SampleConfig() string {
	return sampleConfig
}

// Description returns description of Zookeeper plugin
func (z *Zookeeper) Description() string {
	return `Reads 'mntr' stats from one or many zookeeper servers`
}

func (z *Zookeeper) dial(ctx context.Context, addr string) (net.Conn, error) {
	var dialer net.Dialer
	if z.EnableTLS || z.EnableSSL {
		deadline, ok := ctx.Deadline()
		if ok {
			dialer.Deadline = deadline
		}
		return tls.DialWithDialer(&dialer, "tcp", addr, z.tlsConfig)
	} else {
		return dialer.DialContext(ctx, "tcp", addr)
	}
}

// Gather reads stats from all configured servers accumulates stats
func (z *Zookeeper) Gather(acc telegraf.Accumulator) error {
	ctx := context.Background()

	if !z.initialized {
		tlsConfig, err := z.ClientConfig.TLSConfig()
		if err != nil {
			return err
		}
		z.tlsConfig = tlsConfig
		z.initialized = true
	}

	if z.Timeout.Duration < 1*time.Second {
		z.Timeout.Duration = defaultTimeout
	}

	ctx, cancel := context.WithTimeout(ctx, z.Timeout.Duration)
	defer cancel()

	if len(z.Servers) == 0 {
		z.Servers = []string{":2181"}
	}

	for _, serverAddress := range z.Servers {
		acc.AddError(z.gatherServer(ctx, serverAddress, acc))
	}
	return nil
}

func (z *Zookeeper) gatherServer(ctx context.Context, address string, acc telegraf.Accumulator) error {
	var zookeeper_state string
	_, _, err := net.SplitHostPort(address)
	if err != nil {
		address = address + ":2181"
	}

	c, err := z.dial(ctx, address)
	if err != nil {
		return err
	}
	defer c.Close()

	// Apply deadline to connection
	deadline, ok := ctx.Deadline()
	if ok {
		c.SetDeadline(deadline)
	}

	fmt.Fprintf(c, "%s\n", "mntr")
	rdr := bufio.NewReader(c)
	scanner := bufio.NewScanner(rdr)

	service := strings.Split(address, ":")
	if len(service) != 2 {
		return fmt.Errorf("Invalid service address: %s", address)
	}

	fields := make(map[string]interface{})
	for scanner.Scan() {
		line := scanner.Text()
		parts := zookeeperFormatRE.FindStringSubmatch(string(line))

		if len(parts) != 3 {
			return fmt.Errorf("unexpected line in mntr response: %q", line)
		}

		measurement := strings.TrimPrefix(parts[1], "zk_")
		if measurement == "server_state" {
			zookeeper_state = parts[2]
		} else {
			sValue := string(parts[2])

			iVal, err := strconv.ParseInt(sValue, 10, 64)
			if err == nil {
				fields[measurement] = iVal
			} else {
				fields[measurement] = sValue
			}
		}
	}

	srv := "localhost"
	if service[0] != "" {
		srv = service[0]
	}

	tags := map[string]string{
		"server": srv,
		"port":   service[1],
		"state":  zookeeper_state,
	}
	acc.AddFields("zookeeper", fields, tags)

	return nil
}

func init() {
	inputs.Add("zookeeper", func() telegraf.Input {
		return &Zookeeper{}
	})
}
<commit_msg>improve mntr regex to match user specific keys. (#7533)<commit_after>package zookeeper

import (
	"bufio"
	"context"
	"crypto/tls"
	"fmt"
	"net"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	tlsint "github.com/influxdata/telegraf/plugins/common/tls"
	"github.com/influxdata/telegraf/plugins/inputs"
)

var zookeeperFormatRE = regexp.MustCompile(`^zk_(\w[\w\.\-]*)\s+([\w\.\-]+)`)

// Zookeeper is a zookeeper plugin
type Zookeeper struct {
	Servers []string
	Timeout internal.Duration

	EnableTLS bool `toml:"enable_tls"`
	EnableSSL bool `toml:"enable_ssl"` // deprecated in 1.7; use enable_tls
	tlsint.ClientConfig

	initialized bool
	tlsConfig   *tls.Config
}

var sampleConfig = `
  ## An array of address to gather stats about. Specify an ip or hostname
  ## with port. ie localhost:2181, 10.0.0.1:2181, etc.

  ## If no servers are specified, then localhost is used as the host.
  ## If no port is specified, 2181 is used
  servers = [":2181"]

  ## Timeout for metric collections from all servers.  Minimum timeout is "1s".
  # timeout = "5s"

  ## Optional TLS Config
  # enable_tls = true
  # tls_ca = "/etc/telegraf/ca.pem"
  # tls_cert = "/etc/telegraf/cert.pem"
  # tls_key = "/etc/telegraf/key.pem"
  ## If false, skip chain & host verification
  # insecure_skip_verify = true
`

var defaultTimeout = 5 * time.Second

// SampleConfig returns sample configuration message
func (z *Zookeeper) SampleConfig() string {
	return sampleConfig
}

// Description returns description of Zookeeper plugin
func (z *Zookeeper) Description() string {
	return `Reads 'mntr' stats from one or many zookeeper servers`
}

func (z *Zookeeper) dial(ctx context.Context, addr string) (net.Conn, error) {
	var dialer net.Dialer
	if z.EnableTLS || z.EnableSSL {
		deadline, ok := ctx.Deadline()
		if ok {
			dialer.Deadline = deadline
		}
		return tls.DialWithDialer(&dialer, "tcp", addr, z.tlsConfig)
	} else {
		return dialer.DialContext(ctx, "tcp", addr)
	}
}

// Gather reads stats from all configured servers accumulates stats
func (z *Zookeeper) Gather(acc telegraf.Accumulator) error {
	ctx := context.Background()

	if !z.initialized {
		tlsConfig, err := z.ClientConfig.TLSConfig()
		if err != nil {
			return err
		}
		z.tlsConfig = tlsConfig
		z.initialized = true
	}

	if z.Timeout.Duration < 1*time.Second {
		z.Timeout.Duration = defaultTimeout
	}

	ctx, cancel := context.WithTimeout(ctx, z.Timeout.Duration)
	defer cancel()

	if len(z.Servers) == 0 {
		z.Servers = []string{":2181"}
	}

	for _, serverAddress := range z.Servers {
		acc.AddError(z.gatherServer(ctx, serverAddress, acc))
	}
	return nil
}

func (z *Zookeeper) gatherServer(ctx context.Context, address string, acc telegraf.Accumulator) error {
	var zookeeper_state string
	_, _, err := net.SplitHostPort(address)
	if err != nil {
		address = address + ":2181"
	}

	c, err := z.dial(ctx, address)
	if err != nil {
		return err
	}
	defer c.Close()

	// Apply deadline to connection
	deadline, ok := ctx.Deadline()
	if ok {
		c.SetDeadline(deadline)
	}

	fmt.Fprintf(c, "%s\n", "mntr")
	rdr := bufio.NewReader(c)
	scanner := bufio.NewScanner(rdr)

	service := strings.Split(address, ":")
	if len(service) != 2 {
		return fmt.Errorf("Invalid service address: %s", address)
	}

	fields := make(map[string]interface{})
	for scanner.Scan() {
		line := scanner.Text()
		parts := zookeeperFormatRE.FindStringSubmatch(string(line))

		if len(parts) != 3 {
			return fmt.Errorf("unexpected line in mntr response: %q", line)
		}

		measurement := strings.TrimPrefix(parts[1], "zk_")
		if measurement == "server_state" {
			zookeeper_state = parts[2]
		} else {
			sValue := string(parts[2])

			iVal, err := strconv.ParseInt(sValue, 10, 64)
			if err == nil {
				fields[measurement] = iVal
			} else {
				fields[measurement] = sValue
			}
		}
	}

	srv := "localhost"
	if service[0] != "" {
		srv = service[0]
	}

	tags := map[string]string{
		"server": srv,
		"port":   service[1],
		"state":  zookeeper_state,
	}
	acc.AddFields("zookeeper", fields, tags)

	return nil
}

func init() {
	inputs.Add("zookeeper", func() telegraf.Input {
		return &Zookeeper{}
	})
}
<|endoftext|>


<commit_before>// Copyright 2015, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// 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 ctxjwt_test

import (
	"testing"

	"bytes"

	"github.com/corestoreio/csfw/config"
	"github.com/corestoreio/csfw/net/ctxjwt"
	"github.com/stretchr/testify/assert"
)

func TestPasswordFromConfig(t *testing.T) {

	cfg := config.NewMockGetter(
		config.WithMockValues(config.MockPV{
			config.MockPathScopeDefault(ctxjwt.PathJWTPassword): `Rump3lst!lzch3n`,
		}),
	)

	jm, err := ctxjwt.NewService(
		ctxjwt.WithPasswordFromConfig(cfg),
	)
	assert.NoError(t, err)

	theToken, _, err := jm.GenerateToken(nil)
	assert.NoError(t, err)
	assert.NotEmpty(t, theToken)

}

func TestWithRSAReaderFail(t *testing.T) {

	jm, err := ctxjwt.NewService(
		ctxjwt.WithRSA(bytes.NewReader([]byte(`invalid pem data`))),
	)
	assert.Nil(t, jm)
	assert.Equal(t, "Private Key from io.Reader no found", err.Error())

}
<commit_msg>net/ctxjwt: Remove MockPathScopeDefault<commit_after>// Copyright 2015, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// 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 ctxjwt_test

import (
	"testing"

	"bytes"

	"github.com/corestoreio/csfw/config"
	"github.com/corestoreio/csfw/config/scope"
	"github.com/corestoreio/csfw/net/ctxjwt"
	"github.com/stretchr/testify/assert"
)

func TestPasswordFromConfig(t *testing.T) {

	cfg := config.NewMockGetter(
		config.WithMockValues(config.MockPV{
			scope.StrDefault.FQPathInt64(ctxjwt.PathJWTPassword): `Rump3lst!lzch3n`,
		}),
	)

	jm, err := ctxjwt.NewService(
		ctxjwt.WithPasswordFromConfig(cfg),
	)
	assert.NoError(t, err)

	theToken, _, err := jm.GenerateToken(nil)
	assert.NoError(t, err)
	assert.NotEmpty(t, theToken)

}

func TestWithRSAReaderFail(t *testing.T) {

	jm, err := ctxjwt.NewService(
		ctxjwt.WithRSA(bytes.NewReader([]byte(`invalid pem data`))),
	)
	assert.Nil(t, jm)
	assert.Equal(t, "Private Key from io.Reader no found", err.Error())

}
<|endoftext|>


<commit_before>package network

import (
	"testing"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/utils/constants"
	"github.com/stretchr/testify/assert"
)

func TestPeersData(t *testing.T) {
	data := peersData{}
	data.initialize()

	peer1 := peer{
		nodeID: ids.ShortID{0x01},
	}

	// add of first peer is handled
	data.add(&peer1)
	retrievedPeer1, peer1Found := data.getByID(peer1.nodeID)
	assert.True(t, peer1Found)
	assert.True(t, &peer1 == retrievedPeer1)
	assert.True(t, data.size() == 1)

	// re-addition of peer works as update
	updatedPeer1 := peer{
		nodeID: ids.ShortID{0x01},
	}
	data.add(&updatedPeer1)
	retrievedPeer1, peer1Found = data.getByID(peer1.nodeID)
	assert.True(t, peer1Found)
	assert.True(t, &updatedPeer1 == retrievedPeer1)
	assert.True(t, data.size() == 1)

	peer2 := peer{
		nodeID: ids.ShortID{0x02},
	}

	// add of another peer is handled
	data.add(&peer2)
	retrievedPeer2, peer2Found := data.getByID(peer2.nodeID)
	assert.True(t, peer2Found)
	assert.True(t, &peer2 == retrievedPeer2)
	assert.True(t, data.size() == 2)

	// removal of added peer is handled
	data.remove(&peer1)
	retrievedPeer1, peer1Found = data.getByID(peer1.nodeID)
	assert.False(t, peer1Found)
	assert.True(t, retrievedPeer1 == nil)
	retrievedPeer2, peer2Found = data.getByID(peer2.nodeID)
	assert.True(t, peer2Found)
	assert.True(t, &peer2 == retrievedPeer2)
	assert.True(t, data.size() == 1)

	unknownPeer := peer{
		nodeID: ids.ShortID{0xff},
	}

	// query for unknown peer is handled
	retrievedUnknownPeer, unknownPeerfound := data.getByID(unknownPeer.nodeID)
	assert.False(t, unknownPeerfound)
	assert.True(t, retrievedUnknownPeer == nil)

	// removal of unknown peer is handled
	data.remove(&unknownPeer)
	retrievedPeer2, peer2Found = data.getByID(peer2.nodeID)
	assert.True(t, peer2Found)
	assert.True(t, &peer2 == retrievedPeer2)
	assert.True(t, data.size() == 1)

	// retrival by inbound index is handled
	peer3 := peer{
		nodeID: ids.ShortID{0x03},
	}
	peer4 := peer{
		nodeID: ids.ShortID{0x04},
	}
	data.add(&peer3)
	data.add(&peer4)
	assert.True(t, data.size() == 3)

	thirdPeer, ok := data.getByIdx(1)
	assert.True(t, ok)
	assert.True(t, &peer3 == thirdPeer)

	// retrival by outbound index is handled
	outOfIndexPeer, ok := data.getByIdx(data.size())
	assert.False(t, ok)
	assert.True(t, outOfIndexPeer == nil)

	// reset is idempotent
	data.reset()
	assert.True(t, data.size() == 0)

	data.reset()
	assert.True(t, data.size() == 0)
}

func TestPeersDataSample(t *testing.T) {
	data := peersData{}
	data.initialize()
	trackedSubnetIDs := ids.Set{}
	trackedSubnetIDs.Add(constants.PrimaryNetworkID)
	// Case: Empty
	peers, err := data.sample(constants.PrimaryNetworkID, 0)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	peers, err = data.sample(constants.PrimaryNetworkID, 1)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	// Case: 1 peer who hasn't finished handshake
	peer1 := peer{
		nodeID:         ids.ShortID{0x01},
		trackedSubnets: trackedSubnetIDs,
	}
	data.add(&peer1)
	peers, err = data.sample(constants.PrimaryNetworkID, 0)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	peers, err = data.sample(constants.PrimaryNetworkID, 1)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	// Case: 1 peer who hasn't finished handshake, 1 who has
	peer2 := peer{
		nodeID:         ids.ShortID{0x02},
		trackedSubnets: trackedSubnetIDs,
	}
	peer2.finishedHandshake.SetValue(true)
	data.add(&peer2)

	peers, err = data.sample(constants.PrimaryNetworkID, 0)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	peers, err = data.sample(constants.PrimaryNetworkID, 1)
	assert.NoError(t, err)
	assert.Len(t, peers, 1)
	assert.EqualValues(t, peers[0].nodeID, peer2.nodeID)

	peers, err = data.sample(constants.PrimaryNetworkID, 2)
	assert.NoError(t, err)
	assert.Len(t, peers, 1)
	assert.EqualValues(t, peers[0].nodeID, peer2.nodeID)

	// Case: 2 peers who have finished handshake
	peer1.finishedHandshake.SetValue(true)
	peers, err = data.sample(constants.PrimaryNetworkID, 0)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	peers, err = data.sample(constants.PrimaryNetworkID, 1)
	assert.NoError(t, err)
	assert.Len(t, peers, 1)

	peers, err = data.sample(constants.PrimaryNetworkID, 2)
	assert.NoError(t, err)
	assert.Len(t, peers, 2)
	// Ensure both peers are sampled once
	assert.True(t,
		(peers[0].nodeID == peer1.nodeID && peers[1].nodeID == peer2.nodeID) ||
			(peers[0].nodeID == peer2.nodeID && peers[1].nodeID == peer1.nodeID),
	)

	peers, err = data.sample(constants.PrimaryNetworkID, 3)
	assert.NoError(t, err)
	assert.Len(t, peers, 2)
	// Ensure both peers are sampled once
	assert.True(t,
		(peers[0].nodeID == peer1.nodeID && peers[1].nodeID == peer2.nodeID) ||
			(peers[0].nodeID == peer2.nodeID && peers[1].nodeID == peer1.nodeID),
	)

	// peer with additional subnet sampled
	testID := ids.GenerateTestID()

	// no peers has this subnet
	peers, err = data.sample(testID, 3)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	trackedSubnetIDs.Add(testID)
	peer3 := peer{
		nodeID:         ids.ShortID{0x03},
		trackedSubnets: trackedSubnetIDs,
	}
	peer3.finishedHandshake.SetValue(true)
	data.add(&peer3)
	peers, err = data.sample(testID, 3)
	assert.NoError(t, err)
	assert.Len(t, peers, 3)
	// Ensure peer is sampled
	assert.Equal(t, peers[0].nodeID, peer3.nodeID)
}
<commit_msg>fix failing test<commit_after>package network

import (
	"testing"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/utils/constants"
	"github.com/stretchr/testify/assert"
)

func TestPeersData(t *testing.T) {
	data := peersData{}
	data.initialize()

	peer1 := peer{
		nodeID: ids.ShortID{0x01},
	}

	// add of first peer is handled
	data.add(&peer1)
	retrievedPeer1, peer1Found := data.getByID(peer1.nodeID)
	assert.True(t, peer1Found)
	assert.True(t, &peer1 == retrievedPeer1)
	assert.True(t, data.size() == 1)

	// re-addition of peer works as update
	updatedPeer1 := peer{
		nodeID: ids.ShortID{0x01},
	}
	data.add(&updatedPeer1)
	retrievedPeer1, peer1Found = data.getByID(peer1.nodeID)
	assert.True(t, peer1Found)
	assert.True(t, &updatedPeer1 == retrievedPeer1)
	assert.True(t, data.size() == 1)

	peer2 := peer{
		nodeID: ids.ShortID{0x02},
	}

	// add of another peer is handled
	data.add(&peer2)
	retrievedPeer2, peer2Found := data.getByID(peer2.nodeID)
	assert.True(t, peer2Found)
	assert.True(t, &peer2 == retrievedPeer2)
	assert.True(t, data.size() == 2)

	// removal of added peer is handled
	data.remove(&peer1)
	retrievedPeer1, peer1Found = data.getByID(peer1.nodeID)
	assert.False(t, peer1Found)
	assert.True(t, retrievedPeer1 == nil)
	retrievedPeer2, peer2Found = data.getByID(peer2.nodeID)
	assert.True(t, peer2Found)
	assert.True(t, &peer2 == retrievedPeer2)
	assert.True(t, data.size() == 1)

	unknownPeer := peer{
		nodeID: ids.ShortID{0xff},
	}

	// query for unknown peer is handled
	retrievedUnknownPeer, unknownPeerfound := data.getByID(unknownPeer.nodeID)
	assert.False(t, unknownPeerfound)
	assert.True(t, retrievedUnknownPeer == nil)

	// removal of unknown peer is handled
	data.remove(&unknownPeer)
	retrievedPeer2, peer2Found = data.getByID(peer2.nodeID)
	assert.True(t, peer2Found)
	assert.True(t, &peer2 == retrievedPeer2)
	assert.True(t, data.size() == 1)

	// retrival by inbound index is handled
	peer3 := peer{
		nodeID: ids.ShortID{0x03},
	}
	peer4 := peer{
		nodeID: ids.ShortID{0x04},
	}
	data.add(&peer3)
	data.add(&peer4)
	assert.True(t, data.size() == 3)

	thirdPeer, ok := data.getByIdx(1)
	assert.True(t, ok)
	assert.True(t, &peer3 == thirdPeer)

	// retrival by outbound index is handled
	outOfIndexPeer, ok := data.getByIdx(data.size())
	assert.False(t, ok)
	assert.True(t, outOfIndexPeer == nil)

	// reset is idempotent
	data.reset()
	assert.True(t, data.size() == 0)

	data.reset()
	assert.True(t, data.size() == 0)
}

func TestPeersDataSample(t *testing.T) {
	data := peersData{}
	data.initialize()
	trackedSubnetIDs := ids.Set{}
	trackedSubnetIDs.Add(constants.PrimaryNetworkID)
	// Case: Empty
	peers, err := data.sample(constants.PrimaryNetworkID, 0)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	peers, err = data.sample(constants.PrimaryNetworkID, 1)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	// Case: 1 peer who hasn't finished handshake
	peer1 := peer{
		nodeID:         ids.ShortID{0x01},
		trackedSubnets: trackedSubnetIDs,
	}
	data.add(&peer1)
	peers, err = data.sample(constants.PrimaryNetworkID, 0)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	peers, err = data.sample(constants.PrimaryNetworkID, 1)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	// Case: 1 peer who hasn't finished handshake, 1 who has
	peer2 := peer{
		nodeID:         ids.ShortID{0x02},
		trackedSubnets: trackedSubnetIDs,
	}
	peer2.finishedHandshake.SetValue(true)
	data.add(&peer2)

	peers, err = data.sample(constants.PrimaryNetworkID, 0)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	peers, err = data.sample(constants.PrimaryNetworkID, 1)
	assert.NoError(t, err)
	assert.Len(t, peers, 1)
	assert.EqualValues(t, peers[0].nodeID, peer2.nodeID)

	peers, err = data.sample(constants.PrimaryNetworkID, 2)
	assert.NoError(t, err)
	assert.Len(t, peers, 1)
	assert.EqualValues(t, peers[0].nodeID, peer2.nodeID)

	// Case: 2 peers who have finished handshake
	peer1.finishedHandshake.SetValue(true)
	peers, err = data.sample(constants.PrimaryNetworkID, 0)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	peers, err = data.sample(constants.PrimaryNetworkID, 1)
	assert.NoError(t, err)
	assert.Len(t, peers, 1)

	peers, err = data.sample(constants.PrimaryNetworkID, 2)
	assert.NoError(t, err)
	assert.Len(t, peers, 2)
	// Ensure both peers are sampled once
	assert.True(t,
		(peers[0].nodeID == peer1.nodeID && peers[1].nodeID == peer2.nodeID) ||
			(peers[0].nodeID == peer2.nodeID && peers[1].nodeID == peer1.nodeID),
	)

	peers, err = data.sample(constants.PrimaryNetworkID, 3)
	assert.NoError(t, err)
	assert.Len(t, peers, 2)
	// Ensure both peers are sampled once
	assert.True(t,
		(peers[0].nodeID == peer1.nodeID && peers[1].nodeID == peer2.nodeID) ||
			(peers[0].nodeID == peer2.nodeID && peers[1].nodeID == peer1.nodeID),
	)

	testSubnetID := ids.GenerateTestID()

	// no peers has this subnet
	peers, err = data.sample(testSubnetID, 3)
	assert.NoError(t, err)
	assert.Len(t, peers, 0)

	// peer with additional subnet sampled
	newSubnetSet := ids.Set{}
	newSubnetSet.Add(constants.PrimaryNetworkID, testSubnetID)

	peer3 := peer{
		nodeID:         ids.ShortID{0x03},
		trackedSubnets: newSubnetSet,
	}
	peer3.finishedHandshake.SetValue(true)
	data.add(&peer3)

	peers, err = data.sample(testSubnetID, 3)
	assert.NoError(t, err)
	assert.Len(t, peers, 1)

	// Ensure peer is sampled
	assert.Equal(t, peer3.nodeID, peers[0].nodeID)
}
<|endoftext|>


<commit_before>// Thread safe engine for MyMySQL
//
// In contrast to native engine:
// - one connection can be used by multiple gorutines,
// - if connection is idle pings are sent to the server (once per minute) to
//   avoid timeout.
//
// See documentation of mymysql/native for details
package thrsafe

import (
	"github.com/ziutek/mymysql/mysql"
	_ "github.com/ziutek/mymysql/native"
	"io"
	"sync"
	"time"
)

type Conn struct {
	mysql.Conn
	mutex *sync.Mutex

	stopPinger chan struct{}
	lastUsed   time.Time
}

func (c *Conn) lock() {
	//log.Println(c, ":: lock @", c.mutex)
	c.mutex.Lock()
}

func (c *Conn) unlock() {
	//log.Println(c, ":: unlock @", c.mutex)
	c.lastUsed = time.Now()
	c.mutex.Unlock()
}

type Result struct {
	mysql.Result
	conn *Conn
}

type Stmt struct {
	mysql.Stmt
	conn *Conn
}

type Transaction struct {
	*Conn
	conn *Conn
}

func New(proto, laddr, raddr, user, passwd string, db ...string) mysql.Conn {
	return &Conn{
		Conn:  orgNew(proto, laddr, raddr, user, passwd, db...),
		mutex: new(sync.Mutex),
	}
}

func (c *Conn) Clone() mysql.Conn {
	return &Conn{
		Conn:  c.Conn.Clone(),
		mutex: new(sync.Mutex),
	}
}

func (c *Conn) pinger() {
	c.stopPinger = make(chan struct{})
	defer func() { c.stopPinger = nil }()

	const to = 60 * time.Second
	sleep := to
	for {
		timer := time.After(sleep)
		select {
		case <-c.stopPinger:
			return
		case t := <-timer:
			sleep := to - t.Sub(c.lastUsed)
			if sleep <= 0 {
				if c.Ping() != nil {
					return
				}
				sleep = to
			}
		}
	}
}

func (c *Conn) Connect() error {
	//log.Println("Connect")
	c.lock()
	defer c.unlock()
	go c.pinger()
	return c.Conn.Connect()
}

func (c *Conn) Close() error {
	//log.Println("Close")
	close(c.stopPinger) // Stop pinger before lock connection
	c.lock()
	defer c.unlock()
	return c.Conn.Close()
}

func (c *Conn) Reconnect() error {
	//log.Println("Reconnect")
	c.lock()
	defer c.unlock()
	if c.stopPinger == nil {
		go c.pinger()
	}
	return c.Conn.Reconnect()
}

func (c *Conn) Use(dbname string) error {
	//log.Println("Use")
	c.lock()
	defer c.unlock()
	return c.Conn.Use(dbname)
}

func (c *Conn) Start(sql string, params ...interface{}) (mysql.Result, error) {
	//log.Println("Start")
	c.lock()
	res, err := c.Conn.Start(sql, params...)
	// Unlock if error or OK result (which doesn't provide any fields)
	if err != nil {
		c.unlock()
		return nil, err
	}
	if res.StatusOnly() && !res.MoreResults() {
		c.unlock()
	}
	return &Result{Result: res, conn: c}, err
}

func (res *Result) ScanRow(row mysql.Row) error {
	//log.Println("ScanRow")
	err := res.Result.ScanRow(row)
	if err == nil {
		// There are more rows to read
		return nil
	}
	if err != io.EOF || !res.StatusOnly() && !res.MoreResults() {
		// Error or no more rows in not empty result set and no more resutls.
		// In case if empty result set and no more resutls Start have unlocked
		// it before.
		res.conn.unlock()
	}
	return err
}

func (res *Result) GetRow() (mysql.Row, error) {
	return mysql.GetRow(res)
}

func (res *Result) NextResult() (mysql.Result, error) {
	//log.Println("NextResult")
	next, err := res.Result.NextResult()
	if err != nil {
		return nil, err
	}
	if next == nil {
		return nil, nil
	}
	if next.StatusOnly() && !next.MoreResults() {
		res.conn.unlock()
	}
	return &Result{next, res.conn}, nil
}

func (c *Conn) Ping() error {
	c.lock()
	defer c.unlock()
	return c.Conn.Ping()
}

func (c *Conn) Prepare(sql string) (mysql.Stmt, error) {
	//log.Println("Prepare")
	c.lock()
	defer c.unlock()
	stmt, err := c.Conn.Prepare(sql)
	if err != nil {
		return nil, err
	}
	return &Stmt{Stmt: stmt, conn: c}, nil
}

func (stmt *Stmt) Run(params ...interface{}) (mysql.Result, error) {
	//log.Println("Run")
	stmt.conn.lock()
	res, err := stmt.Stmt.Run(params...)
	// Unlock if error or OK result (which doesn't provide any fields)
	if err != nil {
		stmt.conn.unlock()
		return nil, err
	}
	if res.StatusOnly() && !res.MoreResults() {
		stmt.conn.unlock()
	}
	return &Result{Result: res, conn: stmt.conn}, nil
}

func (stmt *Stmt) Delete() error {
	//log.Println("Delete")
	stmt.conn.lock()
	defer stmt.conn.unlock()
	return stmt.Stmt.Delete()
}

func (stmt *Stmt) Reset() error {
	//log.Println("Reset")
	stmt.conn.lock()
	defer stmt.conn.unlock()
	return stmt.Stmt.Reset()
}

func (stmt *Stmt) SendLongData(pnum int, data interface{}, pkt_size int) error {
	//log.Println("SendLongData")
	stmt.conn.lock()
	defer stmt.conn.unlock()
	return stmt.Stmt.SendLongData(pnum, data, pkt_size)
}

// See mysql.Query
func (c *Conn) Query(sql string, params ...interface{}) ([]mysql.Row, mysql.Result, error) {
	return mysql.Query(c, sql, params...)
}

// See mysql.QueryFirst
func (my *Conn) QueryFirst(sql string, params ...interface{}) (mysql.Row, mysql.Result, error) {
	return mysql.QueryFirst(my, sql, params...)
}

// See mysql.QueryLast
func (my *Conn) QueryLast(sql string, params ...interface{}) (mysql.Row, mysql.Result, error) {
	return mysql.QueryLast(my, sql, params...)
}

// See mysql.Exec
func (stmt *Stmt) Exec(params ...interface{}) ([]mysql.Row, mysql.Result, error) {
	return mysql.Exec(stmt, params...)
}

// See mysql.ExecFirst
func (stmt *Stmt) ExecFirst(params ...interface{}) (mysql.Row, mysql.Result, error) {
	return mysql.ExecFirst(stmt, params...)
}

// See mysql.ExecLast
func (stmt *Stmt) ExecLast(params ...interface{}) (mysql.Row, mysql.Result, error) {
	return mysql.ExecLast(stmt, params...)
}

// See mysql.End
func (res *Result) End() error {
	return mysql.End(res)
}

// See mysql.GetFirstRow
func (res *Result) GetFirstRow() (mysql.Row, error) {
	return mysql.GetFirstRow(res)
}

// See mysql.GetLastRow
func (res *Result) GetLastRow() (mysql.Row, error) {
	return mysql.GetLastRow(res)
}

// See mysql.GetRows
func (res *Result) GetRows() ([]mysql.Row, error) {
	return mysql.GetRows(res)
}

// Begins a new transaction. No any other thread can send command on this
// connection until Commit or Rollback will be called.
// Periodical pinging the server is disabled during transaction.

func (c *Conn) Begin() (mysql.Transaction, error) {
	//log.Println("Begin")
	c.lock()
	tr := Transaction{
		&Conn{Conn: c.Conn, mutex: new(sync.Mutex)},
		c,
	}
	_, err := c.Conn.Start("START TRANSACTION")
	if err != nil {
		c.unlock()
		return nil, err
	}
	return &tr, nil
}

func (tr *Transaction) end(cr string) error {
	tr.lock()
	_, err := tr.conn.Conn.Start(cr)
	tr.conn.unlock()
	// Invalidate this transaction
	m := tr.Conn.mutex
	tr.Conn = nil
	tr.conn = nil
	m.Unlock() // One goorutine which still uses this transaction will panic
	return err
}

func (tr *Transaction) Commit() error {
	//log.Println("Commit")
	return tr.end("COMMIT")
}

func (tr *Transaction) Rollback() error {
	//log.Println("Rollback")
	return tr.end("ROLLBACK")
}

func (tr *Transaction) IsValid() bool {
	return tr.Conn != nil
}

func (tr *Transaction) Do(st mysql.Stmt) mysql.Stmt {
	if s, ok := st.(*Stmt); ok && s.conn == tr.conn {
		// Returns new statement which uses statement mutexes
		return &Stmt{s.Stmt, tr.Conn}
	}
	panic("Transaction and statement doesn't belong to the same connection")
}

var orgNew func(proto, laddr, raddr, user, passwd string, db ...string) mysql.Conn

func init() {
	orgNew = mysql.New
	mysql.New = New
}
<commit_msg>Fixes issue #53<commit_after>// Thread safe engine for MyMySQL
//
// In contrast to native engine:
// - one connection can be used by multiple gorutines,
// - if connection is idle pings are sent to the server (once per minute) to
//   avoid timeout.
//
// See documentation of mymysql/native for details
package thrsafe

import (
	"github.com/ziutek/mymysql/mysql"
	_ "github.com/ziutek/mymysql/native"
	"io"
	"log"
	"sync"
	"time"
)

type Conn struct {
	mysql.Conn
	mutex *sync.Mutex

	stopPinger chan struct{}
	lastUsed   time.Time
}

func (c *Conn) lock() {
	//log.Println(c, ":: lock @", c.mutex)
	c.mutex.Lock()
}

func (c *Conn) unlock() {
	//log.Println(c, ":: unlock @", c.mutex)
	c.lastUsed = time.Now()
	c.mutex.Unlock()
}

type Result struct {
	mysql.Result
	conn *Conn
}

type Stmt struct {
	mysql.Stmt
	conn *Conn
}

type Transaction struct {
	*Conn
	conn *Conn
}

func New(proto, laddr, raddr, user, passwd string, db ...string) mysql.Conn {
	return &Conn{
		Conn:  orgNew(proto, laddr, raddr, user, passwd, db...),
		mutex: new(sync.Mutex),
	}
}

func (c *Conn) Clone() mysql.Conn {
	return &Conn{
		Conn:  c.Conn.Clone(),
		mutex: new(sync.Mutex),
	}
}

func (c *Conn) pinger() {
	c.stopPinger = make(chan struct{})
	defer func() { c.stopPinger = nil }()

	const to = 60 * time.Second
	sleep := to
	for {
		timer := time.After(sleep)
		select {
		case <-c.stopPinger:
			return
		case t := <-timer:
			sleep := to - t.Sub(c.lastUsed)
			if sleep <= 0 {
				if c.Ping() != nil {
					return
				}
				sleep = to
			}
		}
	}
}

func (c *Conn) Connect() error {
	//log.Println("Connect")
	c.lock()
	defer c.unlock()
	go c.pinger()
	return c.Conn.Connect()
}

func (c *Conn) Close() error {
	//log.Println("Close")
	close(c.stopPinger) // Stop pinger before lock connection
	c.lock()
	defer c.unlock()
	return c.Conn.Close()
}

func (c *Conn) Reconnect() error {
	//log.Println("Reconnect")
	c.lock()
	defer c.unlock()
	if c.stopPinger == nil {
		go c.pinger()
	}
	return c.Conn.Reconnect()
}

func (c *Conn) Use(dbname string) error {
	//log.Println("Use")
	c.lock()
	defer c.unlock()
	return c.Conn.Use(dbname)
}

func (c *Conn) Start(sql string, params ...interface{}) (mysql.Result, error) {
	//log.Println("Start")
	c.lock()
	res, err := c.Conn.Start(sql, params...)
	// Unlock if error or OK result (which doesn't provide any fields)
	if err != nil {
		c.unlock()
		return nil, err
	}
	if res.StatusOnly() && !res.MoreResults() {
		c.unlock()
	}
	return &Result{Result: res, conn: c}, err
}

func (res *Result) ScanRow(row mysql.Row) error {
	//log.Println("ScanRow")
	err := res.Result.ScanRow(row)
	if err == nil {
		// There are more rows to read
		return nil
	}
	if err == mysql.ErrReadAfterEOR {
		// Trying read after EOR - connection unlocked before
		return err
	}
	if err != io.EOF || !res.StatusOnly() && !res.MoreResults() {
		log.Println("Debug ***", res, err != io.EOF && err != mysql.ErrReadAfterEOR, res.StatusOnly(), res.MoreResults())
		// Error or no more rows in not empty result set and no more resutls.
		// In case if empty result set and no more resutls Start has unlocked
		// it before.
		res.conn.unlock()
	}
	return err
}

func (res *Result) GetRow() (mysql.Row, error) {
	return mysql.GetRow(res)
}

func (res *Result) NextResult() (mysql.Result, error) {
	//log.Println("NextResult")
	next, err := res.Result.NextResult()
	if err != nil {
		return nil, err
	}
	if next == nil {
		return nil, nil
	}
	if next.StatusOnly() && !next.MoreResults() {
		res.conn.unlock()
	}
	return &Result{next, res.conn}, nil
}

func (c *Conn) Ping() error {
	c.lock()
	defer c.unlock()
	return c.Conn.Ping()
}

func (c *Conn) Prepare(sql string) (mysql.Stmt, error) {
	//log.Println("Prepare")
	c.lock()
	defer c.unlock()
	stmt, err := c.Conn.Prepare(sql)
	if err != nil {
		return nil, err
	}
	return &Stmt{Stmt: stmt, conn: c}, nil
}

func (stmt *Stmt) Run(params ...interface{}) (mysql.Result, error) {
	//log.Println("Run")
	stmt.conn.lock()
	res, err := stmt.Stmt.Run(params...)
	// Unlock if error or OK result (which doesn't provide any fields)
	if err != nil {
		stmt.conn.unlock()
		return nil, err
	}
	if res.StatusOnly() && !res.MoreResults() {
		stmt.conn.unlock()
	}
	return &Result{Result: res, conn: stmt.conn}, nil
}

func (stmt *Stmt) Delete() error {
	//log.Println("Delete")
	stmt.conn.lock()
	defer stmt.conn.unlock()
	return stmt.Stmt.Delete()
}

func (stmt *Stmt) Reset() error {
	//log.Println("Reset")
	stmt.conn.lock()
	defer stmt.conn.unlock()
	return stmt.Stmt.Reset()
}

func (stmt *Stmt) SendLongData(pnum int, data interface{}, pkt_size int) error {
	//log.Println("SendLongData")
	stmt.conn.lock()
	defer stmt.conn.unlock()
	return stmt.Stmt.SendLongData(pnum, data, pkt_size)
}

// See mysql.Query
func (c *Conn) Query(sql string, params ...interface{}) ([]mysql.Row, mysql.Result, error) {
	return mysql.Query(c, sql, params...)
}

// See mysql.QueryFirst
func (my *Conn) QueryFirst(sql string, params ...interface{}) (mysql.Row, mysql.Result, error) {
	return mysql.QueryFirst(my, sql, params...)
}

// See mysql.QueryLast
func (my *Conn) QueryLast(sql string, params ...interface{}) (mysql.Row, mysql.Result, error) {
	return mysql.QueryLast(my, sql, params...)
}

// See mysql.Exec
func (stmt *Stmt) Exec(params ...interface{}) ([]mysql.Row, mysql.Result, error) {
	return mysql.Exec(stmt, params...)
}

// See mysql.ExecFirst
func (stmt *Stmt) ExecFirst(params ...interface{}) (mysql.Row, mysql.Result, error) {
	return mysql.ExecFirst(stmt, params...)
}

// See mysql.ExecLast
func (stmt *Stmt) ExecLast(params ...interface{}) (mysql.Row, mysql.Result, error) {
	return mysql.ExecLast(stmt, params...)
}

// See mysql.End
func (res *Result) End() error {
	return mysql.End(res)
}

// See mysql.GetFirstRow
func (res *Result) GetFirstRow() (mysql.Row, error) {
	return mysql.GetFirstRow(res)
}

// See mysql.GetLastRow
func (res *Result) GetLastRow() (mysql.Row, error) {
	return mysql.GetLastRow(res)
}

// See mysql.GetRows
func (res *Result) GetRows() ([]mysql.Row, error) {
	return mysql.GetRows(res)
}

// Begins a new transaction. No any other thread can send command on this
// connection until Commit or Rollback will be called.
// Periodical pinging the server is disabled during transaction.

func (c *Conn) Begin() (mysql.Transaction, error) {
	//log.Println("Begin")
	c.lock()
	tr := Transaction{
		&Conn{Conn: c.Conn, mutex: new(sync.Mutex)},
		c,
	}
	_, err := c.Conn.Start("START TRANSACTION")
	if err != nil {
		c.unlock()
		return nil, err
	}
	return &tr, nil
}

func (tr *Transaction) end(cr string) error {
	tr.lock()
	_, err := tr.conn.Conn.Start(cr)
	tr.conn.unlock()
	// Invalidate this transaction
	m := tr.Conn.mutex
	tr.Conn = nil
	tr.conn = nil
	m.Unlock() // One goorutine which still uses this transaction will panic
	return err
}

func (tr *Transaction) Commit() error {
	//log.Println("Commit")
	return tr.end("COMMIT")
}

func (tr *Transaction) Rollback() error {
	//log.Println("Rollback")
	return tr.end("ROLLBACK")
}

func (tr *Transaction) IsValid() bool {
	return tr.Conn != nil
}

func (tr *Transaction) Do(st mysql.Stmt) mysql.Stmt {
	if s, ok := st.(*Stmt); ok && s.conn == tr.conn {
		// Returns new statement which uses statement mutexes
		return &Stmt{s.Stmt, tr.Conn}
	}
	panic("Transaction and statement doesn't belong to the same connection")
}

var orgNew func(proto, laddr, raddr, user, passwd string, db ...string) mysql.Conn

func init() {
	orgNew = mysql.New
	mysql.New = New
}
<|endoftext|>