<commit_before>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

// Keybase file system

package main

import (
	"flag"
	"fmt"
	"os"

	"github.com/keybase/client/go/kbfs/dokan"
	"github.com/keybase/client/go/kbfs/env"
	"github.com/keybase/client/go/kbfs/libdokan"
	"github.com/keybase/client/go/kbfs/libfs"
	"github.com/keybase/client/go/kbfs/libkbfs"
	"github.com/keybase/client/go/libkb"
)

var runtimeDir = flag.String("runtime-dir", os.Getenv("KEYBASE_RUNTIME_DIR"), "runtime directory")
var label = flag.String("label", os.Getenv("KEYBASE_LABEL"), "label to help identify if running as a service")
var mountType = flag.String("mount-type", defaultMountType, "mount type: default, force, none")
var version = flag.Bool("version", false, "Print version")
var mountFlags = flag.Int64("mount-flags", int64(libdokan.DefaultMountFlags), "Dokan mount flags")
var dokandll = flag.String("dokan-dll", "", "Absolute path of dokan dll to load")
var servicemount = flag.Bool("mount-from-service", false, "get mount path from service")

const usageFormatStr = `Usage:
  kbfsdokan -version

To run against remote KBFS servers:
  kbfsdokan
    [-runtime-dir=path/to/dir] [-label=label] [-mount-type=force]
    [-mount-flags=n] [-dokan-dll=path/to/dokan.dll]
%s
    -mount-from-service | /path/to/mountpoint

To run in a local testing environment:
  kbfsdokan
    [-runtime-dir=path/to/dir] [-label=label] [-mount-type=force]
    [-mount-flags=n] [-dokan-dll=path/to/dokan.dll]
%s
    -mount-from-service | /path/to/mountpoint

Defaults:
%s
`

func getUsageString(ctx libkbfs.Context) string {
	remoteUsageStr := libkbfs.GetRemoteUsageString()
	localUsageStr := libkbfs.GetLocalUsageString()
	defaultUsageStr := libkbfs.GetDefaultsUsageString(ctx)
	return fmt.Sprintf(usageFormatStr, remoteUsageStr,
		localUsageStr, defaultUsageStr)
}

func start() *libfs.Error {
	err := libkb.SaferDLLLoading()
	if err != nil {
		fmt.Printf("SaferDLLLoading failed: %v\n", err)
	}

	ctx := env.NewContext()
	kbfsParams := libkbfs.AddFlags(flag.CommandLine, ctx)

	flag.Parse()

	if *version {
		fmt.Printf("%s\n", libkbfs.VersionString())
		return nil
	}

	var mountpoint string
	if len(flag.Args()) < 1 {
		if !*servicemount {
			fmt.Print(getUsageString(ctx))
			return libfs.InitError("no mount specified")
		}
	} else {
		mountpoint = flag.Arg(0)
	}

	if len(flag.Args()) > 1 {
		fmt.Print(getUsageString(ctx))
		return libfs.InitError("extra arguments specified (flags go before the first argument)")
	}

	options := libdokan.StartOptions{
		KbfsParams: *kbfsParams,
		RuntimeDir: *runtimeDir,
		Label:      *label,
		DokanConfig: dokan.Config{
			MountFlags: dokan.MountFlag(*mountFlags),
			DllPath:    *dokandll,
		},
		ForceMount: *mountType == "force",
		SkipMount:  *mountType == "none",
		MountPoint: mountpoint,
	}

	return libdokan.Start(options, ctx)
}

func main() {
	err := start()
	if err != nil {
		fmt.Fprintf(os.Stderr, "kbfsdokan error: (%d) %s\n", err.Code, err.Message)

		os.Exit(err.Code)
	}
	os.Exit(0)
}
<commit_msg>kbfsdokan: Use KEYBASE_MOUNT_FLAGS if set for -mount-flags (#16234)<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

// Keybase file system

package main

import (
	"flag"
	"fmt"
	"os"
	"strconv"

	"github.com/keybase/client/go/kbfs/dokan"
	"github.com/keybase/client/go/kbfs/env"
	"github.com/keybase/client/go/kbfs/libdokan"
	"github.com/keybase/client/go/kbfs/libfs"
	"github.com/keybase/client/go/kbfs/libkbfs"
	"github.com/keybase/client/go/libkb"
)

func getDefaultMountFlags() int64 {
	mf := os.Getenv("KEYBASE_MOUNT_FLAGS")
	if mf == "" {
		return int64(libdokan.DefaultMountFlags)
	}
	iv, _ := strconv.Atoi(mf)
	return int64(iv)
}

var runtimeDir = flag.String("runtime-dir", os.Getenv("KEYBASE_RUNTIME_DIR"), "runtime directory")
var label = flag.String("label", os.Getenv("KEYBASE_LABEL"), "label to help identify if running as a service")
var mountType = flag.String("mount-type", defaultMountType, "mount type: default, force, none")
var version = flag.Bool("version", false, "Print version")
var mountFlags = flag.Int64("mount-flags", getDefaultMountFlags(), "Dokan mount flags")
var dokandll = flag.String("dokan-dll", "", "Absolute path of dokan dll to load")
var servicemount = flag.Bool("mount-from-service", false, "get mount path from service")

const usageFormatStr = `Usage:
  kbfsdokan -version

To run against remote KBFS servers:
  kbfsdokan
    [-runtime-dir=path/to/dir] [-label=label] [-mount-type=force]
    [-mount-flags=n] [-dokan-dll=path/to/dokan.dll]
%s
    -mount-from-service | /path/to/mountpoint

To run in a local testing environment:
  kbfsdokan
    [-runtime-dir=path/to/dir] [-label=label] [-mount-type=force]
    [-mount-flags=n] [-dokan-dll=path/to/dokan.dll]
%s
    -mount-from-service | /path/to/mountpoint

Defaults:
%s
`

func getUsageString(ctx libkbfs.Context) string {
	remoteUsageStr := libkbfs.GetRemoteUsageString()
	localUsageStr := libkbfs.GetLocalUsageString()
	defaultUsageStr := libkbfs.GetDefaultsUsageString(ctx)
	return fmt.Sprintf(usageFormatStr, remoteUsageStr,
		localUsageStr, defaultUsageStr)
}

func start() *libfs.Error {
	err := libkb.SaferDLLLoading()
	if err != nil {
		fmt.Printf("SaferDLLLoading failed: %v\n", err)
	}

	ctx := env.NewContext()
	kbfsParams := libkbfs.AddFlags(flag.CommandLine, ctx)

	flag.Parse()

	if *version {
		fmt.Printf("%s\n", libkbfs.VersionString())
		return nil
	}

	var mountpoint string
	if len(flag.Args()) < 1 {
		if !*servicemount {
			fmt.Print(getUsageString(ctx))
			return libfs.InitError("no mount specified")
		}
	} else {
		mountpoint = flag.Arg(0)
	}

	if len(flag.Args()) > 1 {
		fmt.Print(getUsageString(ctx))
		return libfs.InitError("extra arguments specified (flags go before the first argument)")
	}

	options := libdokan.StartOptions{
		KbfsParams: *kbfsParams,
		RuntimeDir: *runtimeDir,
		Label:      *label,
		DokanConfig: dokan.Config{
			MountFlags: dokan.MountFlag(*mountFlags),
			DllPath:    *dokandll,
		},
		ForceMount: *mountType == "force",
		SkipMount:  *mountType == "none",
		MountPoint: mountpoint,
	}

	return libdokan.Start(options, ctx)
}

func main() {
	err := start()
	if err != nil {
		fmt.Fprintf(os.Stderr, "kbfsdokan error: (%d) %s\n", err.Code, err.Message)

		os.Exit(err.Code)
	}
	os.Exit(0)
}
<|endoftext|>


<commit_before>//  Copyright (c) 2014, Google Inc.  All rights reserved.
//
// 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 tao

import (
	"bytes"
	"io/ioutil"
	"os"
	"testing"

	"github.com/jlmucb/cloudproxy/go/tao/auth"
)

func testNewStackedLinuxHost() (*LinuxHost, error) {
	tmpdir, err := ioutil.TempDir("/tmp", "test_new_stacked_linux_host")
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(tmpdir)

	ft, err := NewSoftTao("", nil)
	if err != nil {
		return nil, err
	}

	tg := LiberalGuard
	lh, err := NewStackedLinuxHost(tmpdir, &tg, ft, nil)
	if err != nil {
		return nil, err
	}

	return lh, nil
}

func testNewRootLinuxHost() (*LinuxHost, error) {
	tmpdir, err := ioutil.TempDir("/tmp", "test_new_root_linux_host")
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(tmpdir)

	tg := LiberalGuard
	password := []byte("bad password")
	lh, err := NewRootLinuxHost(tmpdir, &tg, password, nil)
	if err != nil {
		return nil, err
	}

	return lh, nil
}

func TestNewStackedLinuxHost(t *testing.T) {
	if _, err := testNewStackedLinuxHost(); err != nil {
		t.Fatal(err)
	}
}

func TestNewRootLinuxHost(t *testing.T) {
	if _, err := testNewRootLinuxHost(); err != nil {
		t.Fatal(err)
	}
}

// Test the methods directly instead of testing them across a channel.

var testChildLH = &LinuxHostChild{
	channel:      nil,
	Cmd:          nil,
	ChildSubprin: []auth.PrinExt{auth.PrinExt{Name: "TestChild"}},
}

func testLinuxHostHandleGetTaoName(t *testing.T, lh *LinuxHost) {
	if !lh.GetTaoName(testChildLH).Identical(lh.taoHost.HostName().MakeSubprincipal(testChildLH.ChildSubprin)) {
		t.Fatal("Incorrect construction of Tao name")
	}
}

func testLinuxHostHandleGetRandomBytes(t *testing.T, lh *LinuxHost) {
	b, err := lh.GetRandomBytes(testChildLH, 10)
	if err != nil {
		t.Fatal("Failed to get random bytes from the Linux host:", err)
	}

	if len(b) != 10 {
		t.Fatal("Linux host returned the incorrect number of random bytes")
	}
}

func testLinuxHostHandleGetSharedSecret(t *testing.T, lh *LinuxHost) {
	b, err := lh.GetSharedSecret(testChildLH, 10, SharedSecretPolicyDefault)
	if err != nil {
		t.Fatal("Couldn't get a shared secret from the Linux host:", err)
	}

	b2, err := lh.GetSharedSecret(testChildLH, 10, SharedSecretPolicyDefault)
	if err != nil {
		t.Fatal("Couldn't get a second shared secret from the Linux host:", err)
	}

	if len(b) == 0 || !bytes.Equal(b, b2) {
		t.Fatal("Invalid or inconsistent secrets returned from HandleGetSharedSecret in the Linux host")
	}
}

func testLinuxHostHandleSealUnseal(t *testing.T, lh *LinuxHost) {
	data := []byte{1, 2, 3, 4, 5, 6, 7}

	// `in` will be zeroed-out by LinuxHost.Seal(). Make a copy
	// to compare to the result of LinuxHost.Unseal().
	in := make([]byte, len(data))
	copy(in, data)

	b, err := lh.Seal(testChildLH, in, SharedSecretPolicyDefault)
	if err != nil {
		t.Fatal("Couldn't seal the data:", err)
	}

	d, policy, err := lh.Unseal(testChildLH, b)
	if err != nil {
		t.Fatal("Couldn't unseal the sealed data:", err)
	}

	if !bytes.Equal(d, data) {
		t.Fatal("Incorrect unsealed data:", d, data)
	}

	if policy != SharedSecretPolicyDefault {
		t.Fatal("Wrong policy returned by Unseal:", err)
	}
}

func testLinuxHostHandleAttest(t *testing.T, lh *LinuxHost) {
	stmt := auth.Pred{Name: "FakePredicate"}

	a, err := lh.Attest(testChildLH, nil, nil, nil, stmt)
	if err != nil {
		t.Fatal("Couldn't create Attestation")
	}

	if a == nil {
		t.Fatal("Returned invalid Attestation from Attest")
	}

	// TODO(tmroeder): verify this attestation.
}

func testRootLinuxHostHandleGetTaoName(t *testing.T) {
	lh, _ := testNewRootLinuxHost()
	testLinuxHostHandleGetTaoName(t, lh)
}

func testRootLinuxHostHandleGetRandomBytes(t *testing.T) {
	lh, _ := testNewRootLinuxHost()
	testLinuxHostHandleGetRandomBytes(t, lh)
}

func testRootLinuxHostHandleGetSharedSecret(t *testing.T) {
	lh, _ := testNewRootLinuxHost()
	testLinuxHostHandleGetSharedSecret(t, lh)
}

func testRootLinuxHostHandleSealUnseal(t *testing.T) {
	lh, _ := testNewRootLinuxHost()
	testLinuxHostHandleSealUnseal(t, lh)
}

func testRootLinuxHostHandleAttest(t *testing.T) {
	lh, _ := testNewRootLinuxHost()
	testLinuxHostHandleAttest(t, lh)
}

func testStackedLinuxHostHandleGetTaoName(t *testing.T) {
	lh, _ := testNewStackedLinuxHost()
	testLinuxHostHandleGetTaoName(t, lh)
}

func testStackedLinuxHostHandleGetRandomBytes(t *testing.T) {
	lh, _ := testNewStackedLinuxHost()
	testLinuxHostHandleGetRandomBytes(t, lh)
}

func testStackedLinuxHostHandleGetSharedSecret(t *testing.T) {
	lh, _ := testNewStackedLinuxHost()
	testLinuxHostHandleGetSharedSecret(t, lh)
}

func testStackedLinuxHostHandleSealUnseal(t *testing.T) {
	lh, _ := testNewStackedLinuxHost()
	testLinuxHostHandleSealUnseal(t, lh)
}

func testStackedLinuxHostHandleAttest(t *testing.T) {
	lh, _ := testNewStackedLinuxHost()
	testLinuxHostHandleAttest(t, lh)
}
<commit_msg>Export common linux host test.<commit_after>//  Copyright (c) 2014, Google Inc.  All rights reserved.
//
// 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 tao

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"testing"

	"github.com/jlmucb/cloudproxy/go/tao/auth"
)

func testNewStackedLinuxHost() (*LinuxHost, error) {
	tmpdir, err := ioutil.TempDir("/tmp", "test_new_stacked_linux_host")
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(tmpdir)

	ft, err := NewSoftTao("", nil)
	if err != nil {
		return nil, err
	}

	tg := LiberalGuard
	lh, err := NewStackedLinuxHost(tmpdir, &tg, ft, nil)
	if err != nil {
		return nil, err
	}

	return lh, nil
}

func testNewRootLinuxHost() (*LinuxHost, error) {
	tmpdir, err := ioutil.TempDir("/tmp", "test_new_root_linux_host")
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(tmpdir)

	tg := LiberalGuard
	password := []byte("bad password")
	lh, err := NewRootLinuxHost(tmpdir, &tg, password, nil)
	if err != nil {
		return nil, err
	}

	return lh, nil
}

func TestNewStackedLinuxHost(t *testing.T) {
	if _, err := testNewStackedLinuxHost(); err != nil {
		t.Fatal(err)
	}
}

func TestNewRootLinuxHost(t *testing.T) {
	if _, err := testNewRootLinuxHost(); err != nil {
		t.Fatal(err)
	}
}

// Test the methods directly instead of testing them across a channel.

var testChildLH = &LinuxHostChild{
	channel:      nil,
	Cmd:          nil,
	ChildSubprin: []auth.PrinExt{auth.PrinExt{Name: "TestChild"}},
}

func DoTestLinuxHostHandleGetTaoName(lh *LinuxHost) error {
	if !lh.GetTaoName(testChildLH).Identical(lh.taoHost.HostName().MakeSubprincipal(testChildLH.ChildSubprin)) {
		return fmt.Errorf("Incorrect construction of Tao name")
	}

	return nil
}

func DoTestLinuxHostHandleGetRandomBytes(lh *LinuxHost) error {
	b, err := lh.GetRandomBytes(testChildLH, 10)
	if err != nil {
		return fmt.Errorf("Failed to get random bytes from the Linux host: %s", err)
	}

	if len(b) != 10 {
		return fmt.Errorf("Linux host returned the incorrect number of random bytes")
	}

	return nil
}

func DoTestLinuxHostHandleGetSharedSecret(lh *LinuxHost) error {
	b, err := lh.GetSharedSecret(testChildLH, 10, SharedSecretPolicyDefault)
	if err != nil {
		return fmt.Errorf("Couldn't get a shared secret from the Linux host: %s", err)
	}

	b2, err := lh.GetSharedSecret(testChildLH, 10, SharedSecretPolicyDefault)
	if err != nil {
		return fmt.Errorf("Couldn't get a second shared secret from the Linux host: %s", err)
	}

	if len(b) == 0 || !bytes.Equal(b, b2) {
		return fmt.Errorf("Invalid or inconsistent secrets returned from HandleGetSharedSecret in the Linux host")
	}

	return nil
}

func DoTestLinuxHostHandleSealUnseal(lh *LinuxHost) error {
	data := []byte{1, 2, 3, 4, 5, 6, 7}

	// `in` will be zeroed-out by LinuxHost.Seal(). Make a copy
	// to compare to the result of LinuxHost.Unseal().
	in := make([]byte, len(data))
	copy(in, data)

	b, err := lh.Seal(testChildLH, in, SharedSecretPolicyDefault)
	if err != nil {
		return fmt.Errorf("Couldn't seal the data: %s", err)
	}

	d, policy, err := lh.Unseal(testChildLH, b)
	if err != nil {
		return fmt.Errorf("Couldn't unseal the sealed data: %s", err)
	}

	if !bytes.Equal(d, data) {
		return fmt.Errorf("Incorrect unsealed data: %s", d)
	}

	if policy != SharedSecretPolicyDefault {
		return fmt.Errorf("Wrong policy returned by Unseal: %s", policy)
	}

	return nil
}

func DoTestLinuxHostHandleAttest(lh *LinuxHost) error {
	stmt := auth.Pred{Name: "FakePredicate"}

	a, err := lh.Attest(testChildLH, nil, nil, nil, stmt)
	if err != nil {
		return fmt.Errorf("Couldn't create Attestation")
	}

	if a == nil {
		return fmt.Errorf("Returned invalid Attestation from Attest")
	}

	// TODO(tmroeder): verify this attestation.
	return nil
}

func testRootLinuxHostHandleGetTaoName(t *testing.T) {
	lh, _ := testNewRootLinuxHost()
	if err := DoTestLinuxHostHandleGetTaoName(lh); err != nil {
		t.Error(err)
	}
}

func testRootLinuxHostHandleGetRandomBytes(t *testing.T) {
	lh, _ := testNewRootLinuxHost()
	if err := DoTestLinuxHostHandleGetRandomBytes(lh); err != nil {
		t.Error(err)
	}
}

func testRootLinuxHostHandleGetSharedSecret(t *testing.T) {
	lh, _ := testNewRootLinuxHost()
	if err := DoTestLinuxHostHandleGetSharedSecret(lh); err != nil {
		t.Error(err)
	}
}

func testRootLinuxHostHandleSealUnseal(t *testing.T) {
	lh, _ := testNewRootLinuxHost()
	if err := DoTestLinuxHostHandleSealUnseal(lh); err != nil {
		t.Error(err)
	}
}

func testRootLinuxHostHandleAttest(t *testing.T) {
	lh, _ := testNewRootLinuxHost()
	if err := DoTestLinuxHostHandleAttest(lh); err != nil {
		t.Error(err)
	}
}

func testStackedLinuxHostHandleGetTaoName(t *testing.T) {
	lh, _ := testNewStackedLinuxHost()
	if err := DoTestLinuxHostHandleGetTaoName(lh); err != nil {
		t.Error(err)
	}
}

func testStackedLinuxHostHandleGetRandomBytes(t *testing.T) {
	lh, _ := testNewStackedLinuxHost()
	if err := DoTestLinuxHostHandleGetRandomBytes(lh); err != nil {
		t.Error(err)
	}
}

func testStackedLinuxHostHandleGetSharedSecret(t *testing.T) {
	lh, _ := testNewStackedLinuxHost()
	if err := DoTestLinuxHostHandleGetSharedSecret(lh); err != nil {
		t.Error(err)
	}
}

func testStackedLinuxHostHandleSealUnseal(t *testing.T) {
	lh, _ := testNewStackedLinuxHost()
	if err := DoTestLinuxHostHandleSealUnseal(lh); err != nil {
		t.Error(err)
	}
}

func testStackedLinuxHostHandleAttest(t *testing.T) {
	lh, _ := testNewStackedLinuxHost()
	if err := DoTestLinuxHostHandleAttest(lh); err != nil {
		t.Error(err)
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2015-2016 Thomas de Zeeuw.
//
// Licensed under the MIT license that can be found in the LICENSE file.

package grpclogger

import (
	"bytes"
	"reflect"
	"testing"
	"time"

	"github.com/Thomasdezeeuw/logger"
	"google.golang.org/grpc/grpclog"
)

// EventWriter that collects the events and errors.
type eventWriter struct {
	events []logger.Event
	errors []error
	closed bool
}

func (ew *eventWriter) Write(event logger.Event) error {
	ew.events = append(ew.events, event)
	return nil
}

func (ew *eventWriter) HandleError(err error) {
	ew.errors = append(ew.errors, err)
}

func (ew *eventWriter) Close() error {
	ew.closed = true
	return nil
}

func TestGrpcLogger(t *testing.T) {
	oldExit := exit
	closedCalled := 0
	exit = func(closeFn func()) {
		closeFn()
		closedCalled++
	}
	defer func() {
		exit = oldExit
	}()
	closeFn := func() {
		closedCalled++
	}

	var ew eventWriter
	logger.Start(&ew)

	tags := logger.Tags{"TestGrpcLogger"}
	now := time.Now()

	grpclog.SetLogger(CreateLogger(tags, closeFn))
	grpclog.Print("Error message")
	grpclog.Printf("Error %s message", "formatted")
	grpclog.Println("Error message")
	grpclog.Fatal("Fatal message")
	grpclog.Fatalf("Fatal %s message", "formatted")
	grpclog.Fatalln("Fatal message")

	if err := logger.Close(); err != nil {
		t.Fatal("Unexpected error closing: " + err.Error())
	}

	expected := []logger.Event{
		{Type: logger.ErrorEvent, Timestamp: now, Tags: tags, Message: "Error message"},
		{Type: logger.ErrorEvent, Timestamp: now, Tags: tags, Message: "Error formatted message"},
		{Type: logger.ErrorEvent, Timestamp: now, Tags: tags, Message: "Error message"},
		{Type: logger.FatalEvent, Timestamp: now, Tags: tags, Message: "Fatal message"},
		{Type: logger.FatalEvent, Timestamp: now, Tags: tags, Message: "Fatal formatted message"},
		{Type: logger.FatalEvent, Timestamp: now, Tags: tags, Message: "Fatal message"},
	}

	if expectedN, got := len(expected), len(ew.events); expectedN != got {
		t.Fatalf("Expected %d events, but got only got %d", expectedN, got)
	}

	const margin = 100 * time.Millisecond
	for i, event := range ew.events {
		expectedEvent := expected[i]

		// Can't mock time in the logger package, so we have a truncate it.
		if !event.Timestamp.Truncate(margin).Equal(expectedEvent.Timestamp.Truncate(margin)) {
			t.Errorf("Expected event #%d to be %v, but got %v", i, expectedEvent, event)
			continue
		}
		event.Timestamp = expectedEvent.Timestamp

		if expectedEvent.Type == logger.FatalEvent {
			// sortof test the stacktrace, best we can do.
			stacktrace := event.Data.([]byte)
			if !bytes.HasPrefix(stacktrace, []byte("goroutine")) {
				t.Errorf("Expected a stacktrace as data for a Fatal event, but got %s ",
					string(stacktrace))
			}
			event.Data = nil
		}

		if expected, got := expectedEvent, event; !reflect.DeepEqual(expected, got) {
			t.Errorf("Expected event #%d to be %v, but got %v", i, expected, got)
		}
	}

	if closedCalled != 6 {
		t.Fatalf("Expected the exit and close function to be called three times, but got %d", closedCalled/2)
	}
}

func TestExit(t *testing.T) {
	oldExit := osExit
	defer func() {
		osExit = oldExit
	}()

	var exitCode int
	var closedCalled bool
	osExit = func(n int) {
		exitCode = n
	}
	closeFn := func() {
		closedCalled = true
	}

	exit(closeFn)

	if !closedCalled {
		t.Fatal("Close function not called")
	} else if exitCode != 1 {
		t.Fatalf("Expceted exit to be called with 1, but got %d", exitCode)
	}
}
<commit_msg>grpclogger: Clean modified exit function<commit_after>// Copyright (C) 2015-2016 Thomas de Zeeuw.
//
// Licensed under the MIT license that can be found in the LICENSE file.

package grpclogger

import (
	"bytes"
	"reflect"
	"testing"
	"time"

	"github.com/Thomasdezeeuw/logger"
	"google.golang.org/grpc/grpclog"
)

// EventWriter that collects the events and errors.
type eventWriter struct {
	events []logger.Event
	errors []error
	closed bool
}

func (ew *eventWriter) Write(event logger.Event) error {
	ew.events = append(ew.events, event)
	return nil
}

func (ew *eventWriter) HandleError(err error) {
	ew.errors = append(ew.errors, err)
}

func (ew *eventWriter) Close() error {
	ew.closed = true
	return nil
}

func TestGrpcLogger(t *testing.T) {
	closedCalled := setupExitCounter()
	defer resetExitFns()
	closeFn := func() {
		*closedCalled++
	}

	var ew eventWriter
	logger.Start(&ew)

	tags := logger.Tags{"TestGrpcLogger"}
	now := time.Now()

	grpclog.SetLogger(CreateLogger(tags, closeFn))
	grpclog.Print("Error message")
	grpclog.Printf("Error %s message", "formatted")
	grpclog.Println("Error message")
	grpclog.Fatal("Fatal message")
	grpclog.Fatalf("Fatal %s message", "formatted")
	grpclog.Fatalln("Fatal message")

	if err := logger.Close(); err != nil {
		t.Fatal("Unexpected error closing: " + err.Error())
	}

	expected := []logger.Event{
		{Type: logger.ErrorEvent, Timestamp: now, Tags: tags, Message: "Error message"},
		{Type: logger.ErrorEvent, Timestamp: now, Tags: tags, Message: "Error formatted message"},
		{Type: logger.ErrorEvent, Timestamp: now, Tags: tags, Message: "Error message"},
		{Type: logger.FatalEvent, Timestamp: now, Tags: tags, Message: "Fatal message"},
		{Type: logger.FatalEvent, Timestamp: now, Tags: tags, Message: "Fatal formatted message"},
		{Type: logger.FatalEvent, Timestamp: now, Tags: tags, Message: "Fatal message"},
	}

	if expectedN, got := len(expected), len(ew.events); expectedN != got {
		t.Fatalf("Expected %d events, but got only got %d", expectedN, got)
	}

	const margin = 100 * time.Millisecond
	for i, event := range ew.events {
		expectedEvent := expected[i]

		// Can't mock time in the logger package, so we have a truncate it.
		if !event.Timestamp.Truncate(margin).Equal(expectedEvent.Timestamp.Truncate(margin)) {
			t.Errorf("Expected event #%d to be %v, but got %v", i, expectedEvent, event)
			continue
		}
		event.Timestamp = expectedEvent.Timestamp

		if expectedEvent.Type == logger.FatalEvent {
			// sortof test the stacktrace, best we can do.
			stacktrace := event.Data.([]byte)
			if !bytes.HasPrefix(stacktrace, []byte("goroutine")) {
				t.Errorf("Expected a stacktrace as data for a Fatal event, but got %s ",
					string(stacktrace))
			}
			event.Data = nil
		}

		if expected, got := expectedEvent, event; !reflect.DeepEqual(expected, got) {
			t.Errorf("Expected event #%d to be %v, but got %v", i, expected, got)
		}
	}

	if *closedCalled != 6 {
		t.Fatalf("Expected the exit and close function to be called three times, but got %d",
			*closedCalled/2)
	}
}

func TestExit(t *testing.T) {
	defer resetExitFns()

	var exitCode int
	var closedCalled bool
	osExit = func(n int) {
		exitCode = n
	}
	closeFn := func() {
		closedCalled = true
	}

	exit(closeFn)

	if !closedCalled {
		t.Fatal("Close function not called")
	} else if exitCode != 1 {
		t.Fatalf("Expceted exit to be called with 1, but got %d", exitCode)
	}
}

var (
	oldExit   = exit
	oldOSExit = osExit
)

func setupExitCounter() (counter *int) {
	var cnt int
	exit = func(closeFn func()) {
		closeFn()
		cnt++
	}
	return &cnt
}

func resetExitFns() {
	exit = oldExit
	osExit = oldOSExit
}
<|endoftext|>


<commit_before>package confoo

import (
	"fmt"
	"io/ioutil"
	"os"
	"reflect"
	"strings"

	"gopkg.in/yaml.v2"
)

const confVar = "CONFOO_CONFIG_FILE"

func errorPanic(format string, a ...interface{}) {
	m := fmt.Sprintf(format, a...)
	m2 := "CONFOO - " + m
	panic(m2)
}

var confData map[interface{}]interface{}

func init() {
	confFile := os.Getenv("CONFOO_CONFIG_FILE")
	if confFile == "" {
		errorPanic(confVar + " is not set")
	}
	data, err := ioutil.ReadFile(confFile)
	if err != nil {
		errorPanic(err.Error())
	}

	confData = make(map[interface{}]interface{})
	err = yaml.Unmarshal(data, &confData)
	if err != nil {
		errorPanic("cannot decode yaml data")
	}
}

func Configure(path string, target interface{}) {
	subConf := getSubConf(path, confData)
	if subConf != nil {
		configPath(path, reflect.ValueOf(target), subConf)
	}
}

func getSubConf(path string, conf interface{}) interface{} {
	subConf := conf
	for _, p := range strings.Split(path, ".") {
		m, ok := subConf.(map[interface{}]interface{})
		if !ok {
			errorPanic("%s: path not found", path)
		}

		subConf, ok = m[p]
		if !ok {
			return nil
		}
	}
	return subConf
}

func normalizeKey(s string) string {
	parts := strings.Split(s, "_")
	for i, p := range parts {
		parts[i] = strings.Title(p)
	}
	return strings.Join(parts, "")
}

func configStruct(path string, dest reflect.Value, conf interface{}) {
	if conf == nil {
		return
	}
	confMap, ok := conf.(map[interface{}]interface{})
	if !ok {
		errorPanic("%s: expected map not found", path)
	}
	for k, subConf := range confMap {
		kk, ok := k.(string)
		if !ok {
			//FIXME log event if CONFOO_DEBUG is set
			//errorPanic("%s.%v: map key is not a string", path, k)
			continue
		}
		fieldName := normalizeKey(kk)
		fieldVal := dest.FieldByName(fieldName)
		if fieldVal.Kind() == reflect.Invalid {
			//FIXME log event if CONFOO_DEBUG is set
			//errorPanic("%s.%v: field not present in target struct", path, k)
			continue
		}
		configPath(path+"."+kk, dest.FieldByName(fieldName), subConf)
	}
}

func configPath(path string, dest reflect.Value, conf interface{}) {
	destKind := dest.Kind()
	switch destKind {
	case reflect.Ptr:
		configPath(path, dest.Elem(), conf)
	case reflect.Struct:
		configStruct(path, dest, conf)
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
		reflect.String, reflect.Bool, reflect.Map:
		confValue := reflect.ValueOf(conf)
		confKind := confValue.Kind()
		if confKind != destKind {
			dType := dest.Type()
			cType := confValue.Type()
			errorPanic("%s: target type %v != conf type %v", path, dType, cType)
		}
		dest.Set(confValue)
	default:
		errorPanic("%s: conf type %v not handled", path, dest.Type())
	}
}
<commit_msg>handling of yaml lists<commit_after>package confoo

import (
	"fmt"
	"io/ioutil"
	"os"
	"reflect"
	"strconv"
	"strings"

	yaml "gopkg.in/yaml.v2"
)

const confVar = "CONFOO_CONFIG_FILE"

func errorPanic(format string, a ...interface{}) {
	m := fmt.Sprintf(format, a...)
	m2 := "CONFOO - " + m
	panic(m2)
}

var confData map[interface{}]interface{}

func init() {
	confFile := os.Getenv("CONFOO_CONFIG_FILE")
	if confFile == "" {
		errorPanic(confVar + " is not set")
	}
	data, err := ioutil.ReadFile(confFile)
	if err != nil {
		errorPanic(err.Error())
	}

	confData = make(map[interface{}]interface{})
	err = yaml.Unmarshal(data, &confData)
	if err != nil {
		errorPanic("cannot decode yaml data")
	}
}

func Configure(path string, target interface{}) {
	subConf := getSubConf(path, confData)
	if subConf != nil {
		configPath(path, reflect.ValueOf(target), subConf)
	}
}

func getSubConf(path string, conf interface{}) interface{} {
	subConf := conf
	for _, p := range strings.Split(path, ".") {
		m, ok := subConf.(map[interface{}]interface{})
		if !ok {
			errorPanic("%s: path not found", path)
		}

		subConf, ok = m[p]
		if !ok {
			return nil
		}
	}
	return subConf
}

func normalizeKey(s string) string {
	parts := strings.Split(s, "_")
	for i, p := range parts {
		parts[i] = strings.Title(p)
	}
	return strings.Join(parts, "")
}

func configStruct(path string, dest reflect.Value, conf interface{}) {
	if conf == nil {
		return
	}
	confMap, ok := conf.(map[interface{}]interface{})
	if !ok {
		errorPanic("%s: expected map not found", path)
	}
	for k, subConf := range confMap {
		kk, ok := k.(string)
		if !ok {
			//FIXME log event if CONFOO_DEBUG is set
			//errorPanic("%s.%v: map key is not a string", path, k)
			continue
		}
		fieldName := normalizeKey(kk)
		fieldVal := dest.FieldByName(fieldName)
		if fieldVal.Kind() == reflect.Invalid {
			//FIXME log event if CONFOO_DEBUG is set
			//errorPanic("%s.%v: field not present in target struct", path, k)
			continue
		}
		configPath(path+"."+kk, dest.FieldByName(fieldName), subConf)
	}
}

func configPath(path string, dest reflect.Value, conf interface{}) {
	destKind := dest.Kind()
	switch destKind {
	case reflect.Ptr:
		configPath(path, dest.Elem(), conf)
	case reflect.Struct:
		configStruct(path, dest, conf)
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
		reflect.String, reflect.Bool, reflect.Map:
		confValue := reflect.ValueOf(conf)
		confKind := confValue.Kind()
		if confKind != destKind {
			dType := dest.Type()
			cType := confValue.Type()
			errorPanic("%s: target type %v != conf type %v", path, dType, cType)
		}
		dest.Set(confValue)
	case reflect.Slice:
		for i, el := range conf.([]interface{}) {
			idx := strconv.Itoa(i)
			elVal := reflect.New(dest.Type().Elem())
			configPath(path+"."+idx, elVal, el)
			dest.Set(reflect.Append(dest, elVal.Elem()))
		}
	default:
		errorPanic("%s: conf type %v not handled", path, dest.Type())
	}
}
<|endoftext|>


<commit_before>package cfutil

import (
	"errors"
	"fmt"
	consul "github.com/hashicorp/consul/api"
	"net/url"
	"os"
)

// Services() returns the list of services available from the
// Consul cluster
func Services() ([]string, error) {
	client, err := NewConsulClient()
	if err != nil {
		return nil, err
	}
	catalogServices, _, err := client.Catalog().Services(nil)
	if err != nil {
		return []string{}, err
	}
	var services []string
	for k := range catalogServices {
		services = append(services, k)
	}
	return services, nil
}

func DiscoverServiceURL(serviceName, tags string) (string, error) {
	client, err := NewConsulClient()
	if err != nil {
		return "", err
	}
	services, _, err := client.Catalog().Service(serviceName, tags, nil)
	if err != nil {
		return "", fmt.Errorf("Service `%s` not found: %s", serviceName, err)
	}
	if len(services) > 0 {
		return CreateURLFromServiceCatalog(services[0])
	}
	return "", fmt.Errorf("Service `%s` not found", serviceName)

}

func CreateURLFromServiceCatalog(catalog *consul.CatalogService) (string, error) {
	var serviceURL url.URL
	if catalog.ServicePort == 443 {
		serviceURL.Scheme = "https"
		serviceURL.Host = catalog.ServiceAddress
	} else {
		serviceURL.Scheme = "http"
		serviceURL.Host = fmt.Sprintf("%s:%d", catalog.ServiceAddress, catalog.ServicePort)
	}
	return serviceURL.String(), nil
}

// Use ServiceRegister() to register your app in the Consul cluster
// Optionally you can provide a health endpoint on your URL and
// a number of tags to make your service more discoverable
func ServiceRegister(name string, path string, tags ...string) error {
	appEnv, _ := Current()
	client, err := NewConsulClient()
	if err != nil {
		return err
	}
	schema, port := schemaAndPortForServices()
	err = client.Agent().ServiceRegister(&consul.AgentServiceRegistration{
		Name:    name,
		Address: appEnv.ApplicationURIs[0],
		Port:    port,
		Tags:    tags,
		Check: &consul.AgentServiceCheck{
			HTTP:     fmt.Sprintf(schema + "://" + appEnv.ApplicationURIs[0] + path),
			Interval: "60s",
		},
	})
	if err != nil {
		return err
	}
	return nil
}

// NewConsulClient() returns a new consul client which you can use to
// access the Consul cluster HTTP API. It uses `CONSUL_MASTER` and
// `CONSUL_TOKEN` environment variables to set up the HTTP API connection.
func NewConsulClient() (*consul.Client, error) {
	dialScheme, dialHost, err := consulDialstring("consul")
	if err != nil {
		return nil, err
	}
	client, consulErr := consul.NewClient(&consul.Config{
		Address: dialHost,
		Scheme:  dialScheme,
		Token:   os.Getenv("CONSUL_TOKEN"),
	})
	if consulErr != nil {
		return nil, consulErr
	}
	return client, nil
}

func GetConsulKey(mooncoreKey string) (string, error) {
	ns := ConsulNamespace()
	key := "mooncore/" + ns + "/" + mooncoreKey
	client, err := NewConsulClient()
	if err != nil {
		return "", err
	}
	kvPair, _, err := client.KV().Get(key, nil)
	if err != nil {
		return "", err
	}
	if kvPair == nil || kvPair.Value == nil {
		return "", fmt.Errorf("Key not found")
	}
	return string(kvPair.Value), nil
}

func ConsulNamespace() string {
	return os.Getenv("CONSUL_NAMESPACE")
}

func consulDialstring(serviceName string) (string, string, error) {
	consulMaster := ""
	if consulMaster = os.Getenv("CONSUL_MASTER"); consulMaster != "" {
		parsed, err := url.Parse(consulMaster)
		if err == nil {
			return parsed.Scheme, parsed.Host, nil
		}
	}
	return "", "", errors.New("CONSUL_MASTER not found or invalid url")
}

func schemaAndPortForServices() (string, int) {
	if ForceHTTP() {
		return "http", 80
	}
	return "https", 443
}
<commit_msg>Do not ignore supplied port<commit_after>package cfutil

import (
	"errors"
	"fmt"
	consul "github.com/hashicorp/consul/api"
	"net/url"
	"os"
	"strconv"
	"strings"
)

// Services() returns the list of services available from the
// Consul cluster
func Services() ([]string, error) {
	client, err := NewConsulClient()
	if err != nil {
		return nil, err
	}
	catalogServices, _, err := client.Catalog().Services(nil)
	if err != nil {
		return []string{}, err
	}
	var services []string
	for k := range catalogServices {
		services = append(services, k)
	}
	return services, nil
}

func DiscoverServiceURL(serviceName, tags string) (string, error) {
	client, err := NewConsulClient()
	if err != nil {
		return "", err
	}
	services, _, err := client.Catalog().Service(serviceName, tags, nil)
	if err != nil {
		return "", fmt.Errorf("Service `%s` not found: %s", serviceName, err)
	}
	if len(services) > 0 {
		return CreateURLFromServiceCatalog(services[0])
	}
	return "", fmt.Errorf("Service `%s` not found", serviceName)

}

func CreateURLFromServiceCatalog(catalog *consul.CatalogService) (string, error) {
	var serviceURL url.URL
	if catalog.ServicePort == 443 {
		serviceURL.Scheme = "https"
		serviceURL.Host = catalog.ServiceAddress
	} else {
		serviceURL.Scheme = "http"
		serviceURL.Host = fmt.Sprintf("%s:%d", catalog.ServiceAddress, catalog.ServicePort)
	}
	return serviceURL.String(), nil
}

// Use ServiceRegister() to register your app in the Consul cluster
// Optionally you can provide a health endpoint on your URL and
// a number of tags to make your service more discoverable
func ServiceRegister(name string, path string, tags ...string) error {
	appEnv, _ := Current()
	client, err := NewConsulClient()
	if err != nil {
		return err
	}
	schema, port := schemaAndPortForServices()

	appURL, _ := url.Parse(appEnv.ApplicationURIs[0])
	splitted := strings.Split(appURL.Host, ":")
	hostWithoutPort := splitted[0]
	if len(splitted) > 1 {
		if addedPort, err := strconv.Atoi(splitted[1]); err == nil && addedPort != port {
			port = addedPort
		}
	}
	err = client.Agent().ServiceRegister(&consul.AgentServiceRegistration{
		Name:    name,
		Address: hostWithoutPort,
		Port:    port,
		Tags:    tags,
		Check: &consul.AgentServiceCheck{
			HTTP:     fmt.Sprintf(schema + "://" + appURL.Host + path),
			Interval: "60s",
		},
	})
	if err != nil {
		return err
	}
	return nil
}

// NewConsulClient() returns a new consul client which you can use to
// access the Consul cluster HTTP API. It uses `CONSUL_MASTER` and
// `CONSUL_TOKEN` environment variables to set up the HTTP API connection.
func NewConsulClient() (*consul.Client, error) {
	dialScheme, dialHost, err := consulDialstring("consul")
	if err != nil {
		return nil, err
	}
	client, consulErr := consul.NewClient(&consul.Config{
		Address: dialHost,
		Scheme:  dialScheme,
		Token:   os.Getenv("CONSUL_TOKEN"),
	})
	if consulErr != nil {
		return nil, consulErr
	}
	return client, nil
}

func GetConsulKey(mooncoreKey string) (string, error) {
	ns := ConsulNamespace()
	key := "mooncore/" + ns + "/" + mooncoreKey
	client, err := NewConsulClient()
	if err != nil {
		return "", err
	}
	kvPair, _, err := client.KV().Get(key, nil)
	if err != nil {
		return "", err
	}
	if kvPair == nil || kvPair.Value == nil {
		return "", fmt.Errorf("Key not found")
	}
	return string(kvPair.Value), nil
}

func ConsulNamespace() string {
	return os.Getenv("CONSUL_NAMESPACE")
}

func consulDialstring(serviceName string) (string, string, error) {
	consulMaster := ""
	if consulMaster = os.Getenv("CONSUL_MASTER"); consulMaster != "" {
		parsed, err := url.Parse(consulMaster)
		if err == nil {
			return parsed.Scheme, parsed.Host, nil
		}
	}
	return "", "", errors.New("CONSUL_MASTER not found or invalid url")
}

func schemaAndPortForServices() (string, int) {
	if ForceHTTP() {
		return "http", 80
	}
	return "https", 443
}
<|endoftext|>


<commit_before>/*
    Author: Pavel Abramov
    Examples taken from golang docs and Mozilla's
    nodejs clientside session library
*/

package sessioncookie

import (
    "io"
    "time"
    "errors"
    "strings"
    "net/http"
    "crypto/aes"
    "crypto/rand"
    "crypto/hmac"
    "crypto/cipher"
    "crypto/sha256"
    "encoding/base64"
)

func concatValue(hash_mac string, timestamp string, data string, sep string) string {
    if sep == "" {
        sep = "."
    }
    return hash_mac + sep+ timestamp + sep + data
}

func encryptData(secret_key []byte, data []byte) ([]byte, error)  {
    // add padding to block size
    remainder := len(data) % aes.BlockSize
    if remainder != 0 {
        padding_length := aes.BlockSize - remainder
        padding := ""
        for i := 0; i < padding_length; i++ {
            padding += "="
        }
        data = append(data, padding...)
    }

    block, err := aes.NewCipher(secret_key)
    if err != nil {
        return []byte{}, err
    }

    // empty array of size block + data
    ciphertext := make([]byte, aes.BlockSize+len(data))

    // iv is appended to ciphertext
    iv := ciphertext[:aes.BlockSize]

    // reads from rand.Reader global singleton PRNG into iv
    _, err = io.ReadFull(rand.Reader, iv)
    if err != nil {
        return []byte{}, err
    }

    // data is encrypted
    mode := cipher.NewCBCEncrypter(block, iv)
    mode.CryptBlocks(ciphertext[aes.BlockSize:], data)

    return ciphertext, nil
}

func decryptData(secret_key []byte, encrypted_data []byte) (string, error) {
    block, err := aes.NewCipher(secret_key)
    if err != nil {
        return "", err
    }

    if len(encrypted_data) < aes.BlockSize {
        return "", errors.New("encrypted data is less than BlockSize")
    }

    iv := encrypted_data[:aes.BlockSize]
    data := encrypted_data[aes.BlockSize:]

    if len(data) % aes.BlockSize != 0 {
        return "", errors.New("encrypted data is not a multiple of block size")
    }

    mode := cipher.NewCBCDecrypter(block, iv)

    mode.CryptBlocks(data, data)

    // remove padding
    trimmed_data := strings.TrimRight(string(data), "=")

    return trimmed_data, nil
}

func generateHMAC(secret_key []byte, timestamp []byte, encrypted_data []byte) []byte {
    message := append(timestamp, encrypted_data...)

    // create hash mac
    hash_mac := hmac.New(sha256.New, secret_key)
    hash_mac.Write(message)

    // return byte array
    return hash_mac.Sum(nil)
}

func EncryptCookieValue(secret_key []byte, plaintext_data []byte, sep string) (string, error) {
    // generate timestamp
    now := time.Now()
    timestamp := []byte(now.Format(time.UnixDate))
    encoded_timestamp := base64.URLEncoding.EncodeToString(timestamp)

    // generate data blob
    encrypted_data, err := encryptData(secret_key, plaintext_data)
    if err != nil {
        return "", err
    }

    encoded_data := base64.URLEncoding.EncodeToString(encrypted_data)

    // generate hmac
    hash_mac := generateHMAC(secret_key, timestamp, encrypted_data)
    encoded_hash_mac := base64.URLEncoding.EncodeToString(hash_mac)

    // return encoded values
    return concatValue(encoded_hash_mac, encoded_timestamp, encoded_data, sep), nil
}

func DecryptCookieValue(secret_key []byte, cookie_value string, sep string) (string, error) {
    if sep == "" {
        sep = "."
    }

    // extract value
    value := strings.Split(cookie_value, sep)
    encoded_hash_mac := value[0]
    encoded_timestamp := value[1]
    encoded_data := value[2]

    //decode
    encrypted_data, err := base64.URLEncoding.DecodeString(encoded_data)
    if err != nil {
        return "", err
    }

    timestamp, err := base64.URLEncoding.DecodeString(encoded_timestamp)
    if err != nil {
        return "", err
    }

    returned_hmac, err := base64.URLEncoding.DecodeString(encoded_hash_mac)
    if err != nil {
        return "", err
    }

    // compute and compare
    computed_hmac := generateHMAC(secret_key, timestamp, encrypted_data)

    if !hmac.Equal(returned_hmac, computed_hmac) {
        return "", errors.New("Returned HMAC does not match")
    }

    // hmac matches, now decrypt data.
    data, err := decryptData(secret_key, encrypted_data)
    if err != nil {
        return "", err
    }

    return string(data), nil
}

// convencience method to create new encrypted cookie
// TODO: Resolve optionals
func NewSessionCookie(key []byte, name string, plaintext_value []byte, domain, path string, expires time.Time, maxage int, secure, httponly bool) (http.Cookie, error){
    if name == "" || len(plaintext_value) == 0 || len(key) == 0 {
        return http.Cookie{}, errors.New("name, value and key are mandatory")
    }

    data, err := EncryptCookieValue(key, plaintext_value, ".")
    if err != nil {
        return http.Cookie{}, err
    }
    raw := name + "=" + data

    // Provide both Expires and Max-Age
    newcookie := http.Cookie{
        Name: name,
        Value: data,
        Path: path,
        Expires: expires,
        RawExpires: expires.String(),
        MaxAge: maxage,
        Secure: secure,
        HttpOnly: httponly,
        Raw: raw,
        Unparsed: []string{raw},
    }

    return newcookie, nil
}

<commit_msg>added key gen<commit_after>/*
    Author: Pavel Abramov
    Examples taken from golang docs and Mozilla's
    nodejs clientside session library
*/

package sessioncookie

import (
    "io"
    "time"
    "errors"
    "strings"
    "net/http"
    "crypto/aes"
    "crypto/rand"
    "crypto/hmac"
    "crypto/cipher"
    "crypto/sha256"
    "encoding/base64"
)

func GenerateKey(bytesize int) ([]byte, error) {
    if bytesize < 1 {
        return []byte{}, errors.New("Keylength must be > 0")
    }

    key := make([]byte, bytesize)
    _, err := io.ReadFull(rand.Reader, key)

    if err != nil {
        return []byte{}, _
    }

    return key, _
}

func concatValue(hash_mac string, timestamp string, data string, sep string) string {
    if sep == "" {
        sep = "."
    }
    return hash_mac + sep+ timestamp + sep + data
}

func encryptData(secret_key []byte, data []byte) ([]byte, error)  {
    // add padding to block size
    remainder := len(data) % aes.BlockSize
    if remainder != 0 {
        padding_length := aes.BlockSize - remainder
        padding := ""
        for i := 0; i < padding_length; i++ {
            padding += "="
        }
        data = append(data, padding...)
    }

    // aes-128, aes-192, aes-256 determined based on keysize
    block, err := aes.NewCipher(secret_key)
    if err != nil {
        return []byte{}, err
    }

    // empty array of size block + data
    ciphertext := make([]byte, aes.BlockSize+len(data))

    // iv is appended to ciphertext
    iv := ciphertext[:aes.BlockSize]

    // reads from rand.Reader global singleton PRNG into iv
    _, err = io.ReadFull(rand.Reader, iv)
    if err != nil {
        return []byte{}, err
    }

    // data is encrypted
    mode := cipher.NewCBCEncrypter(block, iv)
    mode.CryptBlocks(ciphertext[aes.BlockSize:], data)

    return ciphertext, nil
}

func decryptData(secret_key []byte, encrypted_data []byte) (string, error) {
    block, err := aes.NewCipher(secret_key)
    if err != nil {
        return "", err
    }

    if len(encrypted_data) < aes.BlockSize {
        return "", errors.New("encrypted data is less than BlockSize")
    }

    iv := encrypted_data[:aes.BlockSize]
    data := encrypted_data[aes.BlockSize:]

    if len(data) % aes.BlockSize != 0 {
        return "", errors.New("encrypted data is not a multiple of block size")
    }

    mode := cipher.NewCBCDecrypter(block, iv)

    mode.CryptBlocks(data, data)

    // remove padding
    trimmed_data := strings.TrimRight(string(data), "=")

    return trimmed_data, nil
}

func generateHMAC(secret_key []byte, timestamp []byte, encrypted_data []byte) []byte {
    message := append(timestamp, encrypted_data...)

    // create hash mac
    // TODO: Implement differet hashing algorithms
    // i.e. sha256, 384, 512, etc
    hash_mac := hmac.New(sha256.New, secret_key)
    hash_mac.Write(message)

    // return byte array
    return hash_mac.Sum(nil)
}

func EncryptCookieValue(secret_key []byte, plaintext_data []byte, sep string) (string, error) {
    // generate timestamp
    now := time.Now()
    timestamp := []byte(now.Format(time.UnixDate))
    encoded_timestamp := base64.URLEncoding.EncodeToString(timestamp)

    // check if key is multiple of 16, 24 or 32?

    // generate data blob
    encrypted_data, err := encryptData(secret_key, plaintext_data)
    if err != nil {
        return "", err
    }

    encoded_data := base64.URLEncoding.EncodeToString(encrypted_data)

    // generate hmac
    hash_mac := generateHMAC(secret_key, timestamp, encrypted_data)
    encoded_hash_mac := base64.URLEncoding.EncodeToString(hash_mac)

    // return encoded values
    return concatValue(encoded_hash_mac, encoded_timestamp, encoded_data, sep), nil
}

func DecryptCookieValue(secret_key []byte, cookie_value string, sep string) (string, error) {
    if sep == "" {
        sep = "."
    }

    // extract value
    value := strings.Split(cookie_value, sep)
    encoded_hash_mac := value[0]
    encoded_timestamp := value[1]
    encoded_data := value[2]

    //decode
    encrypted_data, err := base64.URLEncoding.DecodeString(encoded_data)
    if err != nil {
        return "", err
    }

    timestamp, err := base64.URLEncoding.DecodeString(encoded_timestamp)
    if err != nil {
        return "", err
    }

    returned_hmac, err := base64.URLEncoding.DecodeString(encoded_hash_mac)
    if err != nil {
        return "", err
    }

    // compute and compare
    computed_hmac := generateHMAC(secret_key, timestamp, encrypted_data)

    if !hmac.Equal(returned_hmac, computed_hmac) {
        return "", errors.New("Returned HMAC does not match")
    }

    // hmac matches, now decrypt data.
    data, err := decryptData(secret_key, encrypted_data)
    if err != nil {
        return "", err
    }

    return string(data), nil
}

// convencience method to create new encrypted cookie
// TODO: Resolve optionals
func NewSessionCookie(key []byte, name string, plaintext_value []byte, domain, path string, expires time.Time, maxage int, secure, httponly bool) (http.Cookie, error){
    if name == "" || len(plaintext_value) == 0 || len(key) == 0 {
        return http.Cookie{}, errors.New("name, value and key are mandatory")
    }

    data, err := EncryptCookieValue(key, plaintext_value, ".")
    if err != nil {
        return http.Cookie{}, err
    }
    raw := name + "=" + data

    // Provide both Expires and Max-Age
    newcookie := http.Cookie{
        Name: name,
        Value: data,
        Path: path,
        Expires: expires,
        RawExpires: expires.String(),
        MaxAge: maxage,
        Secure: secure,
        HttpOnly: httponly,
        Raw: raw,
        Unparsed: []string{raw},
    }

    return newcookie, nil
}

<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package avalanche

import (
	"fmt"

	"github.com/ava-labs/gecko/cache"
	"github.com/ava-labs/gecko/ids"
	"github.com/ava-labs/gecko/snow/choices"
	"github.com/ava-labs/gecko/snow/consensus/avalanche"
	"github.com/ava-labs/gecko/snow/engine/common"
	"github.com/ava-labs/gecko/snow/engine/common/queue"
	"github.com/ava-labs/gecko/utils/formatting"
	"github.com/prometheus/client_golang/prometheus"
)

const (
	cacheSize = 100000
)

// BootstrapConfig ...
type BootstrapConfig struct {
	common.Config

	// VtxBlocked tracks operations that are blocked on vertices
	// TxBlocked tracks operations that are blocked on transactions
	VtxBlocked, TxBlocked *queue.Jobs

	State State
	VM    DAGVM
}

type bootstrapper struct {
	BootstrapConfig
	metrics
	common.Bootstrapper

	// true if all of the vertices in the original accepted frontier have been processed
	processedStartingAcceptedFrontier bool

	// number of vertices fetched so far
	numFetched uint32

	// tracks which validators were asked for which containers in which requests
	outstandingRequests common.Requests

	// Contains IDs of vertices that have recently been processed
	processedCache *cache.LRU

	// true if bootstrapping is done
	finished bool

	// Called when bootstrapping is done
	onFinished func() error
}

// Initialize this engine.
func (b *bootstrapper) Initialize(config BootstrapConfig) error {
	b.BootstrapConfig = config
	b.processedCache = &cache.LRU{Size: cacheSize}

	b.VtxBlocked.SetParser(&vtxParser{
		log:         config.Context.Log,
		numAccepted: b.numBSVtx,
		numDropped:  b.numBSDroppedVtx,
		state:       b.State,
	})

	b.TxBlocked.SetParser(&txParser{
		log:         config.Context.Log,
		numAccepted: b.numBSTx,
		numDropped:  b.numBSDroppedTx,
		vm:          b.VM,
	})

	config.Bootstrapable = b
	b.Bootstrapper.Initialize(config.Config)
	return nil
}

// CurrentAcceptedFrontier ...
func (b *bootstrapper) CurrentAcceptedFrontier() ids.Set {
	acceptedFrontier := ids.Set{}
	acceptedFrontier.Add(b.State.Edge()...)
	return acceptedFrontier
}

// FilterAccepted ...
func (b *bootstrapper) FilterAccepted(containerIDs ids.Set) ids.Set {
	acceptedVtxIDs := ids.Set{}
	for _, vtxID := range containerIDs.List() {
		if vtx, err := b.State.GetVertex(vtxID); err == nil && vtx.Status() == choices.Accepted {
			acceptedVtxIDs.Add(vtxID)
		}
	}
	return acceptedVtxIDs
}

// Get vertex [vtxID] and its ancestors
func (b *bootstrapper) fetch(vtxID ids.ID) error {
	// Make sure we haven't already requested this block
	if b.outstandingRequests.Contains(vtxID) {
		return nil
	}

	// Make sure we don't already have this vertex
	if _, err := b.State.GetVertex(vtxID); err == nil {
		return nil
	}

	validators := b.BootstrapConfig.Validators.Sample(1) // validator to send request to
	if len(validators) == 0 {
		return fmt.Errorf("Dropping request for %s as there are no validators", vtxID)
	}
	validatorID := validators[0].ID()
	b.RequestID++

	b.outstandingRequests.Add(validatorID, b.RequestID, vtxID)
	b.BootstrapConfig.Sender.GetAncestors(validatorID, b.RequestID, vtxID) // request vertex and ancestors
	return nil
}

// Process vertices
func (b *bootstrapper) process(vtx avalanche.Vertex) error {
	toProcess := newMaxVertexHeap()
	if _, ok := b.processedCache.Get(vtx.ID()); !ok { // only process if we haven't already
		toProcess.Push(vtx)
	}
	for toProcess.Len() > 0 {
		vtx := toProcess.Pop()
		switch vtx.Status() {
		case choices.Unknown:
			if err := b.fetch(vtx.ID()); err != nil {
				return err
			}
		case choices.Rejected:
			return fmt.Errorf("tried to accept %s even though it was previously rejected", vtx.ID())
		case choices.Processing:
			if err := b.VtxBlocked.Push(&vertexJob{
				log:         b.BootstrapConfig.Context.Log,
				numAccepted: b.numBSVtx,
				numDropped:  b.numBSDroppedVtx,
				vtx:         vtx,
			}); err == nil {
				b.numBSBlockedVtx.Inc()
				b.numFetched++ // Progress tracker
				if b.numFetched%common.StatusUpdateFrequency == 0 {
					b.BootstrapConfig.Context.Log.Info("fetched %d vertices", b.numFetched)
				}
			} else {
				b.BootstrapConfig.Context.Log.Verbo("couldn't push to vtxBlocked: %s", err)
			}
			for _, tx := range vtx.Txs() {
				if err := b.TxBlocked.Push(&txJob{
					log:         b.BootstrapConfig.Context.Log,
					numAccepted: b.numBSTx,
					numDropped:  b.numBSDroppedTx,
					tx:          tx,
				}); err == nil {
					b.numBSBlockedTx.Inc()
				} else {
					b.BootstrapConfig.Context.Log.Verbo("couldn't push to txBlocked: %s", err)
				}
			}
			for _, parent := range vtx.Parents() {
				if _, ok := b.processedCache.Get(parent.ID()); ok { // already processed this
					continue
				}
				toProcess.Push(parent)
			}
			b.processedCache.Put(vtx.ID(), nil)
		}
	}

	if err := b.VtxBlocked.Commit(); err != nil {
		return err
	}
	if err := b.TxBlocked.Commit(); err != nil {
		return err
	}

	if numPending := b.outstandingRequests.Len(); numPending == 0 && b.processedStartingAcceptedFrontier {
		return b.finish()
	}
	return nil
}

// MultiPut handles the receipt of multiple containers. Should be received in response to a GetAncestors message to [vdr]
// with request ID [requestID]. Expects vtxs[0] to be the vertex requested in the corresponding GetAncestors.
func (b *bootstrapper) MultiPut(vdr ids.ShortID, requestID uint32, vtxs [][]byte) error {
	if lenVtxs := len(vtxs); lenVtxs > common.MaxContainersPerMultiPut {
		b.BootstrapConfig.Context.Log.Debug("MultiPut(%s, %d) contains more than maximum number of vertices", vdr, requestID)
		return b.GetAncestorsFailed(vdr, requestID)
	} else if lenVtxs == 0 {
		b.BootstrapConfig.Context.Log.Debug("MultiPut(%s, %d) contains no vertices", vdr, requestID)
		return b.GetAncestorsFailed(vdr, requestID)
	}

	// Make sure this is in response to a request we made
	neededVtxID, needed := b.outstandingRequests.Remove(vdr, requestID)
	if !needed { // this message isn't in response to a request we made
		b.BootstrapConfig.Context.Log.Debug("received unexpected MultiPut from %s with ID %d", vdr, requestID)
		return nil
	}

	neededVtx, err := b.State.ParseVertex(vtxs[0]) // the vertex we requested
	if err != nil {
		b.BootstrapConfig.Context.Log.Debug("Failed to parse requested vertex %s: %w", neededVtxID, err)
		b.BootstrapConfig.Context.Log.Verbo("vertex: %s", formatting.DumpBytes{Bytes: vtxs[0]})
		return b.fetch(neededVtxID)
	} else if actualID := neededVtx.ID(); !actualID.Equals(neededVtxID) {
		b.BootstrapConfig.Context.Log.Debug("expected the first block to be the requested block, %s, but is %s", neededVtxID, actualID)
		return b.fetch(neededVtxID)
	}

	for _, vtxBytes := range vtxs { // Parse/persist all the vertices
		if _, err := b.State.ParseVertex(vtxBytes); err != nil { // Persists the vtx
			b.BootstrapConfig.Context.Log.Debug("Failed to parse vertex: %w", err)
			b.BootstrapConfig.Context.Log.Verbo("vertex: %s", formatting.DumpBytes{Bytes: vtxBytes})
		}
	}

	return b.process(neededVtx)
}

// GetAncestorsFailed is called when a GetAncestors message we sent fails
func (b *bootstrapper) GetAncestorsFailed(vdr ids.ShortID, requestID uint32) error {
	vtxID, ok := b.outstandingRequests.Remove(vdr, requestID)
	if !ok {
		b.BootstrapConfig.Context.Log.Debug("GetAncestorsFailed(%s, %d) called but there was no outstanding request to this validator with this ID", vdr, requestID)
		return nil
	}
	// Send another request for the vertex
	return b.fetch(vtxID)
}

// ForceAccepted ...
func (b *bootstrapper) ForceAccepted(acceptedContainerIDs ids.Set) error {
	if err := b.VM.Bootstrapping(); err != nil {
		return fmt.Errorf("failed to notify VM that bootstrapping has started: %w",
			err)
	}

	for _, vtxID := range acceptedContainerIDs.List() {
		if vtx, err := b.State.GetVertex(vtxID); err == nil {
			if err := b.process(vtx); err != nil {
				return err
			}
		} else if err := b.fetch(vtxID); err != nil {
			return err
		}
	}
	b.processedStartingAcceptedFrontier = true

	if numPending := b.outstandingRequests.Len(); numPending == 0 {
		return b.finish()
	}
	return nil
}

// Finish bootstrapping
func (b *bootstrapper) finish() error {
	if b.finished {
		return nil
	}
	b.BootstrapConfig.Context.Log.Info("finished fetching vertices. executing transaction state transitions...")

	if err := b.executeAll(b.TxBlocked, b.numBSBlockedTx); err != nil {
		return err
	}

	b.BootstrapConfig.Context.Log.Info("executing vertex state transitions...")

	if err := b.executeAll(b.VtxBlocked, b.numBSBlockedVtx); err != nil {
		return err
	}

	if err := b.VM.Bootstrapped(); err != nil {
		return fmt.Errorf("failed to notify VM that bootstrapping has finished: %w",
			err)
	}

	// Start consensus
	if err := b.onFinished(); err != nil {
		return err
	}
	b.finished = true
	return nil
}

func (b *bootstrapper) executeAll(jobs *queue.Jobs, numBlocked prometheus.Gauge) error {
	numExecuted := 0
	for job, err := jobs.Pop(); err == nil; job, err = jobs.Pop() {
		numBlocked.Dec()
		b.BootstrapConfig.Context.Log.Debug("Executing: %s", job.ID())
		if err := jobs.Execute(job); err != nil {
			b.BootstrapConfig.Context.Log.Error("Error executing: %s", err)
			return err
		}
		if err := jobs.Commit(); err != nil {
			return err
		}
		numExecuted++
		if numExecuted%common.StatusUpdateFrequency == 0 { // Periodically print progress
			b.BootstrapConfig.Context.Log.Info("executed %d operations", numExecuted)
		}
	}
	return nil
}
<commit_msg>style fix<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package avalanche

import (
	"fmt"

	"github.com/ava-labs/gecko/cache"
	"github.com/ava-labs/gecko/ids"
	"github.com/ava-labs/gecko/snow/choices"
	"github.com/ava-labs/gecko/snow/consensus/avalanche"
	"github.com/ava-labs/gecko/snow/engine/common"
	"github.com/ava-labs/gecko/snow/engine/common/queue"
	"github.com/ava-labs/gecko/utils/formatting"
	"github.com/prometheus/client_golang/prometheus"
)

const (
	cacheSize = 100000
)

// BootstrapConfig ...
type BootstrapConfig struct {
	common.Config

	// VtxBlocked tracks operations that are blocked on vertices
	// TxBlocked tracks operations that are blocked on transactions
	VtxBlocked, TxBlocked *queue.Jobs

	State State
	VM    DAGVM
}

type bootstrapper struct {
	BootstrapConfig
	metrics
	common.Bootstrapper

	// true if all of the vertices in the original accepted frontier have been processed
	processedStartingAcceptedFrontier bool

	// number of vertices fetched so far
	numFetched uint32

	// tracks which validators were asked for which containers in which requests
	outstandingRequests common.Requests

	// Contains IDs of vertices that have recently been processed
	processedCache *cache.LRU

	// true if bootstrapping is done
	finished bool

	// Called when bootstrapping is done
	onFinished func() error
}

// Initialize this engine.
func (b *bootstrapper) Initialize(config BootstrapConfig) error {
	b.BootstrapConfig = config
	b.processedCache = &cache.LRU{Size: cacheSize}

	b.VtxBlocked.SetParser(&vtxParser{
		log:         config.Context.Log,
		numAccepted: b.numBSVtx,
		numDropped:  b.numBSDroppedVtx,
		state:       b.State,
	})

	b.TxBlocked.SetParser(&txParser{
		log:         config.Context.Log,
		numAccepted: b.numBSTx,
		numDropped:  b.numBSDroppedTx,
		vm:          b.VM,
	})

	config.Bootstrapable = b
	b.Bootstrapper.Initialize(config.Config)
	return nil
}

// CurrentAcceptedFrontier ...
func (b *bootstrapper) CurrentAcceptedFrontier() ids.Set {
	acceptedFrontier := ids.Set{}
	acceptedFrontier.Add(b.State.Edge()...)
	return acceptedFrontier
}

// FilterAccepted ...
func (b *bootstrapper) FilterAccepted(containerIDs ids.Set) ids.Set {
	acceptedVtxIDs := ids.Set{}
	for _, vtxID := range containerIDs.List() {
		if vtx, err := b.State.GetVertex(vtxID); err == nil && vtx.Status() == choices.Accepted {
			acceptedVtxIDs.Add(vtxID)
		}
	}
	return acceptedVtxIDs
}

// Get vertex [vtxID] and its ancestors
func (b *bootstrapper) fetch(vtxID ids.ID) error {
	// Make sure we haven't already requested this block
	if b.outstandingRequests.Contains(vtxID) {
		return nil
	}

	// Make sure we don't already have this vertex
	if _, err := b.State.GetVertex(vtxID); err == nil {
		return nil
	}

	validators := b.BootstrapConfig.Validators.Sample(1) // validator to send request to
	if len(validators) == 0 {
		return fmt.Errorf("Dropping request for %s as there are no validators", vtxID)
	}
	validatorID := validators[0].ID()
	b.RequestID++

	b.outstandingRequests.Add(validatorID, b.RequestID, vtxID)
	b.BootstrapConfig.Sender.GetAncestors(validatorID, b.RequestID, vtxID) // request vertex and ancestors
	return nil
}

// Process vertices
func (b *bootstrapper) process(vtx avalanche.Vertex) error {
	toProcess := newMaxVertexHeap()
	if _, ok := b.processedCache.Get(vtx.ID()); !ok { // only process if we haven't already
		toProcess.Push(vtx)
	}
	for toProcess.Len() > 0 {
		vtx := toProcess.Pop()
		switch vtx.Status() {
		case choices.Unknown:
			if err := b.fetch(vtx.ID()); err != nil {
				return err
			}
		case choices.Rejected:
			return fmt.Errorf("tried to accept %s even though it was previously rejected", vtx.ID())
		case choices.Processing:
			if err := b.VtxBlocked.Push(&vertexJob{
				log:         b.BootstrapConfig.Context.Log,
				numAccepted: b.numBSVtx,
				numDropped:  b.numBSDroppedVtx,
				vtx:         vtx,
			}); err == nil {
				b.numBSBlockedVtx.Inc()
				b.numFetched++ // Progress tracker
				if b.numFetched%common.StatusUpdateFrequency == 0 {
					b.BootstrapConfig.Context.Log.Info("fetched %d vertices", b.numFetched)
				}
			} else {
				b.BootstrapConfig.Context.Log.Verbo("couldn't push to vtxBlocked: %s", err)
			}
			for _, tx := range vtx.Txs() {
				if err := b.TxBlocked.Push(&txJob{
					log:         b.BootstrapConfig.Context.Log,
					numAccepted: b.numBSTx,
					numDropped:  b.numBSDroppedTx,
					tx:          tx,
				}); err == nil {
					b.numBSBlockedTx.Inc()
				} else {
					b.BootstrapConfig.Context.Log.Verbo("couldn't push to txBlocked: %s", err)
				}
			}
			for _, parent := range vtx.Parents() {
				if _, ok := b.processedCache.Get(parent.ID()); !ok { // already processed this
					toProcess.Push(parent)
				}
			}
			b.processedCache.Put(vtx.ID(), nil)
		}
	}

	if err := b.VtxBlocked.Commit(); err != nil {
		return err
	}
	if err := b.TxBlocked.Commit(); err != nil {
		return err
	}

	if numPending := b.outstandingRequests.Len(); numPending == 0 && b.processedStartingAcceptedFrontier {
		return b.finish()
	}
	return nil
}

// MultiPut handles the receipt of multiple containers. Should be received in response to a GetAncestors message to [vdr]
// with request ID [requestID]. Expects vtxs[0] to be the vertex requested in the corresponding GetAncestors.
func (b *bootstrapper) MultiPut(vdr ids.ShortID, requestID uint32, vtxs [][]byte) error {
	if lenVtxs := len(vtxs); lenVtxs > common.MaxContainersPerMultiPut {
		b.BootstrapConfig.Context.Log.Debug("MultiPut(%s, %d) contains more than maximum number of vertices", vdr, requestID)
		return b.GetAncestorsFailed(vdr, requestID)
	} else if lenVtxs == 0 {
		b.BootstrapConfig.Context.Log.Debug("MultiPut(%s, %d) contains no vertices", vdr, requestID)
		return b.GetAncestorsFailed(vdr, requestID)
	}

	// Make sure this is in response to a request we made
	neededVtxID, needed := b.outstandingRequests.Remove(vdr, requestID)
	if !needed { // this message isn't in response to a request we made
		b.BootstrapConfig.Context.Log.Debug("received unexpected MultiPut from %s with ID %d", vdr, requestID)
		return nil
	}

	neededVtx, err := b.State.ParseVertex(vtxs[0]) // the vertex we requested
	if err != nil {
		b.BootstrapConfig.Context.Log.Debug("Failed to parse requested vertex %s: %w", neededVtxID, err)
		b.BootstrapConfig.Context.Log.Verbo("vertex: %s", formatting.DumpBytes{Bytes: vtxs[0]})
		return b.fetch(neededVtxID)
	} else if actualID := neededVtx.ID(); !actualID.Equals(neededVtxID) {
		b.BootstrapConfig.Context.Log.Debug("expected the first block to be the requested block, %s, but is %s", neededVtxID, actualID)
		return b.fetch(neededVtxID)
	}

	for _, vtxBytes := range vtxs { // Parse/persist all the vertices
		if _, err := b.State.ParseVertex(vtxBytes); err != nil { // Persists the vtx
			b.BootstrapConfig.Context.Log.Debug("Failed to parse vertex: %w", err)
			b.BootstrapConfig.Context.Log.Verbo("vertex: %s", formatting.DumpBytes{Bytes: vtxBytes})
		}
	}

	return b.process(neededVtx)
}

// GetAncestorsFailed is called when a GetAncestors message we sent fails
func (b *bootstrapper) GetAncestorsFailed(vdr ids.ShortID, requestID uint32) error {
	vtxID, ok := b.outstandingRequests.Remove(vdr, requestID)
	if !ok {
		b.BootstrapConfig.Context.Log.Debug("GetAncestorsFailed(%s, %d) called but there was no outstanding request to this validator with this ID", vdr, requestID)
		return nil
	}
	// Send another request for the vertex
	return b.fetch(vtxID)
}

// ForceAccepted ...
func (b *bootstrapper) ForceAccepted(acceptedContainerIDs ids.Set) error {
	if err := b.VM.Bootstrapping(); err != nil {
		return fmt.Errorf("failed to notify VM that bootstrapping has started: %w",
			err)
	}

	for _, vtxID := range acceptedContainerIDs.List() {
		if vtx, err := b.State.GetVertex(vtxID); err == nil {
			if err := b.process(vtx); err != nil {
				return err
			}
		} else if err := b.fetch(vtxID); err != nil {
			return err
		}
	}
	b.processedStartingAcceptedFrontier = true

	if numPending := b.outstandingRequests.Len(); numPending == 0 {
		return b.finish()
	}
	return nil
}

// Finish bootstrapping
func (b *bootstrapper) finish() error {
	if b.finished {
		return nil
	}
	b.BootstrapConfig.Context.Log.Info("finished fetching vertices. executing transaction state transitions...")

	if err := b.executeAll(b.TxBlocked, b.numBSBlockedTx); err != nil {
		return err
	}

	b.BootstrapConfig.Context.Log.Info("executing vertex state transitions...")

	if err := b.executeAll(b.VtxBlocked, b.numBSBlockedVtx); err != nil {
		return err
	}

	if err := b.VM.Bootstrapped(); err != nil {
		return fmt.Errorf("failed to notify VM that bootstrapping has finished: %w",
			err)
	}

	// Start consensus
	if err := b.onFinished(); err != nil {
		return err
	}
	b.finished = true
	return nil
}

func (b *bootstrapper) executeAll(jobs *queue.Jobs, numBlocked prometheus.Gauge) error {
	numExecuted := 0
	for job, err := jobs.Pop(); err == nil; job, err = jobs.Pop() {
		numBlocked.Dec()
		b.BootstrapConfig.Context.Log.Debug("Executing: %s", job.ID())
		if err := jobs.Execute(job); err != nil {
			b.BootstrapConfig.Context.Log.Error("Error executing: %s", err)
			return err
		}
		if err := jobs.Commit(); err != nil {
			return err
		}
		numExecuted++
		if numExecuted%common.StatusUpdateFrequency == 0 { // Periodically print progress
			b.BootstrapConfig.Context.Log.Info("executed %d operations", numExecuted)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package db

import (
	"errors"
	"fmt"
	"strings"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

//MongodbConfig DB配置
type MongodbConfig struct {
	URL    string
	DBName string
}

//MongoDataInterface 存储的MongoData
type MongoDataInterface interface {
	GenObjectID() bson.ObjectId
	GetObjectID() bson.ObjectId
}

var mgoSession *mgo.Session
var dbConfig *MongodbConfig

//ConnectMongoDB 连接MongoDB
func ConnectMongoDB(url string, dbName string) error {
	dbConfig = &MongodbConfig{}
	dbConfig.URL = url
	dbConfig.DBName = dbName
	var err error
	mgoSession, err = mgo.Dial(url)

	if err != nil {
		return err
	}
	return nil
}

//ConnectMongoDB 连接MongoDB
func ConnectMongoDBWithPasswd(url string, dbName string, user string, passwd string) error {
	if strings.HasPrefix(url, "mongodb://") {
		url = url[10:]
	}
	url = user + ":" + passwd + "@" + url + "/" + dbName
	fmt.Println("url: " + url)
	var err error
	mgoSession, err = mgo.DialWithTimeout(url, 60000000)
	//mgoSession.Ping()
	dbConfig = &MongodbConfig{}
	dbConfig.URL = url
	dbConfig.DBName = dbName

	if err != nil {
		return err
	}
	return nil
}

//GetMongoSession 获取session， 必须手动Close
func getMongoSession() *mgo.Session {
	if mgoSession == nil {
		return nil
	}
	//最大连接池默认为4096
	return mgoSession.Clone()
}

//QueryWithCollection 获取collection对象
func QueryWithCollection(collection string, s func(*mgo.Collection) error) error {
	session := getMongoSession()
	if session != nil {
		defer session.Close()
		c := session.DB(dbConfig.DBName).C(collection)
		return s(c)
	}

	return errors.New("session is nil")
}

func SetNextSeqBaseNum(collection string, key string, base int) int {
	doc := struct{ Seq int }{}
	query := func(c *mgo.Collection) error {
		change := mgo.Change{
			Update: bson.M{
				"$setOnInsert": bson.M{"seq": base},
			},
			Upsert:    true,
			ReturnNew: true,
		}
		_, err := c.Find(bson.M{"_id": key}).Apply(change, &doc)
		if err != nil {
			panic(fmt.Errorf("set counter failed:", err.Error()))
		}
		return err
	}
	err := QueryWithCollection(collection, query)
	if err != nil {
		return -1
	}
	return doc.Seq
}

//GetNextSeq 获取自增id
func GetNextSeq(collection string, key string, num int) int {
	doc := struct{ Seq int }{}
	cid := key
	query := func(c *mgo.Collection) error {
		change := mgo.Change{
			Update:    bson.M{"$inc": bson.M{"seq": num}},
			Upsert:    true,
			ReturnNew: true,
		}
		_, err := c.Find(bson.M{"_id": cid}).Apply(change, &doc)
		if err != nil {
			panic(fmt.Errorf("get counter failed:", err.Error()))
		}
		return err
	}

	err := QueryWithCollection(collection, query)
	if err != nil {
		fmt.Println(err.Error())
		return -1
	}
	return doc.Seq
}
<commit_msg>fix(connectdb):add time param<commit_after>package db

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

//MongodbConfig DB配置
type MongodbConfig struct {
	URL    string
	DBName string
}

//MongoDataInterface 存储的MongoData
type MongoDataInterface interface {
	GenObjectID() bson.ObjectId
	GetObjectID() bson.ObjectId
}

var mgoSession *mgo.Session
var dbConfig *MongodbConfig

//ConnectMongoDB 连接MongoDB
func ConnectMongoDB(url string, dbName string) error {
	dbConfig = &MongodbConfig{}
	dbConfig.URL = url
	dbConfig.DBName = dbName
	var err error
	mgoSession, err = mgo.Dial(url)

	if err != nil {
		return err
	}
	return nil
}

//ConnectMongoDB 连接MongoDB
func ConnectMongoDBWithPasswd(url string, dbName string, user string, passwd string, timeout time.Duration) error {
	if strings.HasPrefix(url, "mongodb://") {
		url = url[10:]
	}
	url = user + ":" + passwd + "@" + url + "/" + dbName
	fmt.Println("url: " + url)
	var err error
	mgoSession, err = mgo.DialWithTimeout(url, timeout)
	//mgoSession.Ping()
	dbConfig = &MongodbConfig{}
	dbConfig.URL = url
	dbConfig.DBName = dbName

	if err != nil {
		return err
	}
	return nil
}

//GetMongoSession 获取session， 必须手动Close
func getMongoSession() *mgo.Session {
	if mgoSession == nil {
		return nil
	}
	//最大连接池默认为4096
	return mgoSession.Clone()
}

//QueryWithCollection 获取collection对象
func QueryWithCollection(collection string, s func(*mgo.Collection) error) error {
	session := getMongoSession()
	if session != nil {
		defer session.Close()
		c := session.DB(dbConfig.DBName).C(collection)
		return s(c)
	}

	return errors.New("session is nil")
}

func SetNextSeqBaseNum(collection string, key string, base int) int {
	doc := struct{ Seq int }{}
	query := func(c *mgo.Collection) error {
		change := mgo.Change{
			Update: bson.M{
				"$setOnInsert": bson.M{"seq": base},
			},
			Upsert:    true,
			ReturnNew: true,
		}
		_, err := c.Find(bson.M{"_id": key}).Apply(change, &doc)
		if err != nil {
			panic(fmt.Errorf("set counter failed:", err.Error()))
		}
		return err
	}
	err := QueryWithCollection(collection, query)
	if err != nil {
		return -1
	}
	return doc.Seq
}

//GetNextSeq 获取自增id
func GetNextSeq(collection string, key string, num int) int {
	doc := struct{ Seq int }{}
	cid := key
	query := func(c *mgo.Collection) error {
		change := mgo.Change{
			Update:    bson.M{"$inc": bson.M{"seq": num}},
			Upsert:    true,
			ReturnNew: true,
		}
		_, err := c.Find(bson.M{"_id": cid}).Apply(change, &doc)
		if err != nil {
			panic(fmt.Errorf("get counter failed:", err.Error()))
		}
		return err
	}

	err := QueryWithCollection(collection, query)
	if err != nil {
		fmt.Println(err.Error())
		return -1
	}
	return doc.Seq
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package avalanche

import (
	"fmt"

	"github.com/ava-labs/gecko/cache"
	"github.com/ava-labs/gecko/ids"
	"github.com/ava-labs/gecko/snow/choices"
	"github.com/ava-labs/gecko/snow/consensus/avalanche"
	"github.com/ava-labs/gecko/snow/engine/common"
	"github.com/ava-labs/gecko/snow/engine/common/queue"
	"github.com/ava-labs/gecko/utils/formatting"
	"github.com/prometheus/client_golang/prometheus"
)

const (
	// We cache processed vertices where height = c * stripeDistance for c = {1,2,3...}
	// This forms a "stripe" of cached DAG vertices at height stripeDistance, 2*stripeDistance, etc.
	// This helps to limit the number of repeated DAG traversals performed
	stripeDistance = 2000
	stripeWidth    = 5
	cacheSize      = 100000
)

// BootstrapConfig ...
type BootstrapConfig struct {
	common.Config

	// VtxBlocked tracks operations that are blocked on vertices
	// TxBlocked tracks operations that are blocked on transactions
	VtxBlocked, TxBlocked *queue.Jobs

	State State
	VM    DAGVM
}

type bootstrapper struct {
	BootstrapConfig
	metrics
	common.Bootstrapper

	// true if all of the vertices in the original accepted frontier have been processed
	processedStartingAcceptedFrontier bool

	// number of vertices fetched so far
	numFetched uint32

	// tracks which validators were asked for which containers in which requests
	outstandingRequests common.Requests

	// Contains IDs of vertices that have recently been processed
	processedCache *cache.LRU

	// true if bootstrapping is done
	finished bool

	// Called when bootstrapping is done
	onFinished func() error
}

// Initialize this engine.
func (b *bootstrapper) Initialize(config BootstrapConfig) error {
	b.BootstrapConfig = config
	b.processedCache = &cache.LRU{Size: cacheSize}

	b.VtxBlocked.SetParser(&vtxParser{
		log:         config.Context.Log,
		numAccepted: b.numBSVtx,
		numDropped:  b.numBSDroppedVtx,
		state:       b.State,
	})

	b.TxBlocked.SetParser(&txParser{
		log:         config.Context.Log,
		numAccepted: b.numBSTx,
		numDropped:  b.numBSDroppedTx,
		vm:          b.VM,
	})

	config.Bootstrapable = b
	b.Bootstrapper.Initialize(config.Config)
	return nil
}

// CurrentAcceptedFrontier ...
func (b *bootstrapper) CurrentAcceptedFrontier() ids.Set {
	acceptedFrontier := ids.Set{}
	acceptedFrontier.Add(b.State.Edge()...)
	return acceptedFrontier
}

// FilterAccepted ...
func (b *bootstrapper) FilterAccepted(containerIDs ids.Set) ids.Set {
	acceptedVtxIDs := ids.Set{}
	for _, vtxID := range containerIDs.List() {
		if vtx, err := b.State.GetVertex(vtxID); err == nil && vtx.Status() == choices.Accepted {
			acceptedVtxIDs.Add(vtxID)
		}
	}
	return acceptedVtxIDs
}

// Get vertex [vtxID] and its ancestors
func (b *bootstrapper) fetch(vtxID ids.ID) error {
	// Make sure we haven't already requested this block
	if b.outstandingRequests.Contains(vtxID) {
		return nil
	}

	// Make sure we don't already have this vertex
	if _, err := b.State.GetVertex(vtxID); err == nil {
		if numPending := b.outstandingRequests.Len(); numPending == 0 && b.processedStartingAcceptedFrontier {
			return b.finish()
		}
		return nil
	}

	validators := b.BootstrapConfig.Validators.Sample(1) // validator to send request to
	if len(validators) == 0 {
		return fmt.Errorf("Dropping request for %s as there are no validators", vtxID)
	}
	validatorID := validators[0].ID()
	b.RequestID++

	b.outstandingRequests.Add(validatorID, b.RequestID, vtxID)
	b.BootstrapConfig.Sender.GetAncestors(validatorID, b.RequestID, vtxID) // request vertex and ancestors
	return nil
}

// Process vertices
func (b *bootstrapper) process(vtxs ...avalanche.Vertex) error {
	toProcess := newMaxVertexHeap()
	for _, vtx := range vtxs {
		if _, ok := b.processedCache.Get(vtx.ID()); !ok { // only process if we haven't already
			toProcess.Push(vtx)
		}
	}

	for toProcess.Len() > 0 {
		vtx := toProcess.Pop()
		switch vtx.Status() {
		case choices.Unknown:
			if err := b.fetch(vtx.ID()); err != nil {
				return err
			}
		case choices.Rejected:
			return fmt.Errorf("tried to accept %s even though it was previously rejected", vtx.ID())
		case choices.Processing:
			if err := b.VtxBlocked.Push(&vertexJob{
				log:         b.BootstrapConfig.Context.Log,
				numAccepted: b.numBSVtx,
				numDropped:  b.numBSDroppedVtx,
				vtx:         vtx,
			}); err == nil {
				b.numBSBlockedVtx.Inc()
				b.numFetched++ // Progress tracker
				if b.numFetched%common.StatusUpdateFrequency == 0 {
					b.BootstrapConfig.Context.Log.Info("fetched %d vertices", b.numFetched)
				}
			} else {
				b.BootstrapConfig.Context.Log.Verbo("couldn't push to vtxBlocked: %s", err)
			}
			for _, tx := range vtx.Txs() {
				if err := b.TxBlocked.Push(&txJob{
					log:         b.BootstrapConfig.Context.Log,
					numAccepted: b.numBSTx,
					numDropped:  b.numBSDroppedTx,
					tx:          tx,
				}); err == nil {
					b.numBSBlockedTx.Inc()
				} else {
					b.BootstrapConfig.Context.Log.Verbo("couldn't push to txBlocked: %s", err)
				}
			}
			for _, parent := range vtx.Parents() {
				if _, ok := b.processedCache.Get(parent.ID()); !ok { // already processed this
					toProcess.Push(parent)
				}
			}
			if vtx.Height()%stripeDistance < stripeWidth {
				b.processedCache.Put(vtx.ID(), nil)
			}
		}
	}

	if err := b.VtxBlocked.Commit(); err != nil {
		return err
	}
	if err := b.TxBlocked.Commit(); err != nil {
		return err
	}

	if numPending := b.outstandingRequests.Len(); numPending == 0 && b.processedStartingAcceptedFrontier {
		return b.finish()
	}
	return nil
}

// MultiPut handles the receipt of multiple containers. Should be received in response to a GetAncestors message to [vdr]
// with request ID [requestID]. Expects vtxs[0] to be the vertex requested in the corresponding GetAncestors.
func (b *bootstrapper) MultiPut(vdr ids.ShortID, requestID uint32, vtxs [][]byte) error {
	if lenVtxs := len(vtxs); lenVtxs > common.MaxContainersPerMultiPut {
		b.BootstrapConfig.Context.Log.Debug("MultiPut(%s, %d) contains more than maximum number of vertices", vdr, requestID)
		return b.GetAncestorsFailed(vdr, requestID)
	} else if lenVtxs == 0 {
		b.BootstrapConfig.Context.Log.Debug("MultiPut(%s, %d) contains no vertices", vdr, requestID)
		return b.GetAncestorsFailed(vdr, requestID)
	}

	// Make sure this is in response to a request we made
	neededVtxID, needed := b.outstandingRequests.Remove(vdr, requestID)
	if !needed { // this message isn't in response to a request we made
		b.BootstrapConfig.Context.Log.Debug("received unexpected MultiPut from %s with ID %d", vdr, requestID)
		return nil
	}

	neededVtx, err := b.State.ParseVertex(vtxs[0]) // the vertex we requested
	if err != nil {
		b.BootstrapConfig.Context.Log.Debug("Failed to parse requested vertex %s: %w", neededVtxID, err)
		b.BootstrapConfig.Context.Log.Verbo("vertex: %s", formatting.DumpBytes{Bytes: vtxs[0]})
		return b.fetch(neededVtxID)
	} else if actualID := neededVtx.ID(); !actualID.Equals(neededVtxID) {
		b.BootstrapConfig.Context.Log.Debug("expected the first block to be the requested block, %s, but is %s", neededVtxID, actualID)
		return b.fetch(neededVtxID)
	}

	processVertices := make([]avalanche.Vertex, 1, len(vtxs))
	processVertices[0] = neededVtx

	for _, vtxBytes := range vtxs { // Parse/persist all the vertices
		if vtx, err := b.State.ParseVertex(vtxBytes); err != nil { // Persists the vtx
			b.BootstrapConfig.Context.Log.Debug("Failed to parse vertex: %w", err)
			b.BootstrapConfig.Context.Log.Verbo("vertex: %s", formatting.DumpBytes{Bytes: vtxBytes})
		} else {
			processVertices = append(processVertices, vtx)
		}
	}

	return b.process(processVertices...)
}

// GetAncestorsFailed is called when a GetAncestors message we sent fails
func (b *bootstrapper) GetAncestorsFailed(vdr ids.ShortID, requestID uint32) error {
	vtxID, ok := b.outstandingRequests.Remove(vdr, requestID)
	if !ok {
		b.BootstrapConfig.Context.Log.Debug("GetAncestorsFailed(%s, %d) called but there was no outstanding request to this validator with this ID", vdr, requestID)
		return nil
	}
	// Send another request for the vertex
	return b.fetch(vtxID)
}

// ForceAccepted ...
func (b *bootstrapper) ForceAccepted(acceptedContainerIDs ids.Set) error {
	if err := b.VM.Bootstrapping(); err != nil {
		return fmt.Errorf("failed to notify VM that bootstrapping has started: %w",
			err)
	}

	storedVtxs := make([]avalanche.Vertex, 0, acceptedContainerIDs.Len())
	for _, vtxID := range acceptedContainerIDs.List() {
		if vtx, err := b.State.GetVertex(vtxID); err == nil {
			storedVtxs = append(storedVtxs, vtx)
		} else if err := b.fetch(vtxID); err != nil {
			return err
		}
	}
	if err := b.process(storedVtxs...); err != nil {
		return err
	}
	b.processedStartingAcceptedFrontier = true

	if numPending := b.outstandingRequests.Len(); numPending == 0 {
		return b.finish()
	}
	return nil
}

// Finish bootstrapping
func (b *bootstrapper) finish() error {
	if b.finished {
		return nil
	}
	b.BootstrapConfig.Context.Log.Info("finished fetching vertices. executing transaction state transitions...")

	if err := b.executeAll(b.TxBlocked, b.numBSBlockedTx); err != nil {
		return err
	}

	b.BootstrapConfig.Context.Log.Info("executing vertex state transitions...")

	if err := b.executeAll(b.VtxBlocked, b.numBSBlockedVtx); err != nil {
		return err
	}

	if err := b.VM.Bootstrapped(); err != nil {
		return fmt.Errorf("failed to notify VM that bootstrapping has finished: %w",
			err)
	}

	// Start consensus
	if err := b.onFinished(); err != nil {
		return err
	}
	b.finished = true
	return nil
}

func (b *bootstrapper) executeAll(jobs *queue.Jobs, numBlocked prometheus.Gauge) error {
	numExecuted := 0
	for job, err := jobs.Pop(); err == nil; job, err = jobs.Pop() {
		numBlocked.Dec()
		b.BootstrapConfig.Context.Log.Debug("Executing: %s", job.ID())
		if err := jobs.Execute(job); err != nil {
			b.BootstrapConfig.Context.Log.Error("Error executing: %s", err)
			return err
		}
		if err := jobs.Commit(); err != nil {
			return err
		}
		numExecuted++
		if numExecuted%common.StatusUpdateFrequency == 0 { // Periodically print progress
			b.BootstrapConfig.Context.Log.Info("executed %d operations", numExecuted)
		}
	}
	return nil
}
<commit_msg>skip unneccessary parse<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package avalanche

import (
	"fmt"

	"github.com/ava-labs/gecko/cache"
	"github.com/ava-labs/gecko/ids"
	"github.com/ava-labs/gecko/snow/choices"
	"github.com/ava-labs/gecko/snow/consensus/avalanche"
	"github.com/ava-labs/gecko/snow/engine/common"
	"github.com/ava-labs/gecko/snow/engine/common/queue"
	"github.com/ava-labs/gecko/utils/formatting"
	"github.com/prometheus/client_golang/prometheus"
)

const (
	// We cache processed vertices where height = c * stripeDistance for c = {1,2,3...}
	// This forms a "stripe" of cached DAG vertices at height stripeDistance, 2*stripeDistance, etc.
	// This helps to limit the number of repeated DAG traversals performed
	stripeDistance = 2000
	stripeWidth    = 5
	cacheSize      = 100000
)

// BootstrapConfig ...
type BootstrapConfig struct {
	common.Config

	// VtxBlocked tracks operations that are blocked on vertices
	// TxBlocked tracks operations that are blocked on transactions
	VtxBlocked, TxBlocked *queue.Jobs

	State State
	VM    DAGVM
}

type bootstrapper struct {
	BootstrapConfig
	metrics
	common.Bootstrapper

	// true if all of the vertices in the original accepted frontier have been processed
	processedStartingAcceptedFrontier bool

	// number of vertices fetched so far
	numFetched uint32

	// tracks which validators were asked for which containers in which requests
	outstandingRequests common.Requests

	// Contains IDs of vertices that have recently been processed
	processedCache *cache.LRU

	// true if bootstrapping is done
	finished bool

	// Called when bootstrapping is done
	onFinished func() error
}

// Initialize this engine.
func (b *bootstrapper) Initialize(config BootstrapConfig) error {
	b.BootstrapConfig = config
	b.processedCache = &cache.LRU{Size: cacheSize}

	b.VtxBlocked.SetParser(&vtxParser{
		log:         config.Context.Log,
		numAccepted: b.numBSVtx,
		numDropped:  b.numBSDroppedVtx,
		state:       b.State,
	})

	b.TxBlocked.SetParser(&txParser{
		log:         config.Context.Log,
		numAccepted: b.numBSTx,
		numDropped:  b.numBSDroppedTx,
		vm:          b.VM,
	})

	config.Bootstrapable = b
	b.Bootstrapper.Initialize(config.Config)
	return nil
}

// CurrentAcceptedFrontier ...
func (b *bootstrapper) CurrentAcceptedFrontier() ids.Set {
	acceptedFrontier := ids.Set{}
	acceptedFrontier.Add(b.State.Edge()...)
	return acceptedFrontier
}

// FilterAccepted ...
func (b *bootstrapper) FilterAccepted(containerIDs ids.Set) ids.Set {
	acceptedVtxIDs := ids.Set{}
	for _, vtxID := range containerIDs.List() {
		if vtx, err := b.State.GetVertex(vtxID); err == nil && vtx.Status() == choices.Accepted {
			acceptedVtxIDs.Add(vtxID)
		}
	}
	return acceptedVtxIDs
}

// Get vertex [vtxID] and its ancestors
func (b *bootstrapper) fetch(vtxID ids.ID) error {
	// Make sure we haven't already requested this block
	if b.outstandingRequests.Contains(vtxID) {
		return nil
	}

	// Make sure we don't already have this vertex
	if _, err := b.State.GetVertex(vtxID); err == nil {
		if numPending := b.outstandingRequests.Len(); numPending == 0 && b.processedStartingAcceptedFrontier {
			return b.finish()
		}
		return nil
	}

	validators := b.BootstrapConfig.Validators.Sample(1) // validator to send request to
	if len(validators) == 0 {
		return fmt.Errorf("Dropping request for %s as there are no validators", vtxID)
	}
	validatorID := validators[0].ID()
	b.RequestID++

	b.outstandingRequests.Add(validatorID, b.RequestID, vtxID)
	b.BootstrapConfig.Sender.GetAncestors(validatorID, b.RequestID, vtxID) // request vertex and ancestors
	return nil
}

// Process vertices
func (b *bootstrapper) process(vtxs ...avalanche.Vertex) error {
	toProcess := newMaxVertexHeap()
	for _, vtx := range vtxs {
		if _, ok := b.processedCache.Get(vtx.ID()); !ok { // only process if we haven't already
			toProcess.Push(vtx)
		}
	}

	for toProcess.Len() > 0 {
		vtx := toProcess.Pop()
		switch vtx.Status() {
		case choices.Unknown:
			if err := b.fetch(vtx.ID()); err != nil {
				return err
			}
		case choices.Rejected:
			return fmt.Errorf("tried to accept %s even though it was previously rejected", vtx.ID())
		case choices.Processing:
			if err := b.VtxBlocked.Push(&vertexJob{
				log:         b.BootstrapConfig.Context.Log,
				numAccepted: b.numBSVtx,
				numDropped:  b.numBSDroppedVtx,
				vtx:         vtx,
			}); err == nil {
				b.numBSBlockedVtx.Inc()
				b.numFetched++ // Progress tracker
				if b.numFetched%common.StatusUpdateFrequency == 0 {
					b.BootstrapConfig.Context.Log.Info("fetched %d vertices", b.numFetched)
				}
			} else {
				b.BootstrapConfig.Context.Log.Verbo("couldn't push to vtxBlocked: %s", err)
			}
			for _, tx := range vtx.Txs() {
				if err := b.TxBlocked.Push(&txJob{
					log:         b.BootstrapConfig.Context.Log,
					numAccepted: b.numBSTx,
					numDropped:  b.numBSDroppedTx,
					tx:          tx,
				}); err == nil {
					b.numBSBlockedTx.Inc()
				} else {
					b.BootstrapConfig.Context.Log.Verbo("couldn't push to txBlocked: %s", err)
				}
			}
			for _, parent := range vtx.Parents() {
				if _, ok := b.processedCache.Get(parent.ID()); !ok { // already processed this
					toProcess.Push(parent)
				}
			}
			if vtx.Height()%stripeDistance < stripeWidth {
				b.processedCache.Put(vtx.ID(), nil)
			}
		}
	}

	if err := b.VtxBlocked.Commit(); err != nil {
		return err
	}
	if err := b.TxBlocked.Commit(); err != nil {
		return err
	}

	if numPending := b.outstandingRequests.Len(); numPending == 0 && b.processedStartingAcceptedFrontier {
		return b.finish()
	}
	return nil
}

// MultiPut handles the receipt of multiple containers. Should be received in response to a GetAncestors message to [vdr]
// with request ID [requestID]. Expects vtxs[0] to be the vertex requested in the corresponding GetAncestors.
func (b *bootstrapper) MultiPut(vdr ids.ShortID, requestID uint32, vtxs [][]byte) error {
	if lenVtxs := len(vtxs); lenVtxs > common.MaxContainersPerMultiPut {
		b.BootstrapConfig.Context.Log.Debug("MultiPut(%s, %d) contains more than maximum number of vertices", vdr, requestID)
		return b.GetAncestorsFailed(vdr, requestID)
	} else if lenVtxs == 0 {
		b.BootstrapConfig.Context.Log.Debug("MultiPut(%s, %d) contains no vertices", vdr, requestID)
		return b.GetAncestorsFailed(vdr, requestID)
	}

	// Make sure this is in response to a request we made
	neededVtxID, needed := b.outstandingRequests.Remove(vdr, requestID)
	if !needed { // this message isn't in response to a request we made
		b.BootstrapConfig.Context.Log.Debug("received unexpected MultiPut from %s with ID %d", vdr, requestID)
		return nil
	}

	neededVtx, err := b.State.ParseVertex(vtxs[0]) // the vertex we requested
	if err != nil {
		b.BootstrapConfig.Context.Log.Debug("Failed to parse requested vertex %s: %w", neededVtxID, err)
		b.BootstrapConfig.Context.Log.Verbo("vertex: %s", formatting.DumpBytes{Bytes: vtxs[0]})
		return b.fetch(neededVtxID)
	} else if actualID := neededVtx.ID(); !actualID.Equals(neededVtxID) {
		b.BootstrapConfig.Context.Log.Debug("expected the first block to be the requested block, %s, but is %s", neededVtxID, actualID)
		return b.fetch(neededVtxID)
	}

	processVertices := make([]avalanche.Vertex, 1, len(vtxs))
	processVertices[0] = neededVtx

	for _, vtxBytes := range vtxs[1:] { // Parse/persist all the vertices
		if vtx, err := b.State.ParseVertex(vtxBytes); err != nil { // Persists the vtx
			b.BootstrapConfig.Context.Log.Debug("Failed to parse vertex: %w", err)
			b.BootstrapConfig.Context.Log.Verbo("vertex: %s", formatting.DumpBytes{Bytes: vtxBytes})
		} else {
			processVertices = append(processVertices, vtx)
		}
	}

	return b.process(processVertices...)
}

// GetAncestorsFailed is called when a GetAncestors message we sent fails
func (b *bootstrapper) GetAncestorsFailed(vdr ids.ShortID, requestID uint32) error {
	vtxID, ok := b.outstandingRequests.Remove(vdr, requestID)
	if !ok {
		b.BootstrapConfig.Context.Log.Debug("GetAncestorsFailed(%s, %d) called but there was no outstanding request to this validator with this ID", vdr, requestID)
		return nil
	}
	// Send another request for the vertex
	return b.fetch(vtxID)
}

// ForceAccepted ...
func (b *bootstrapper) ForceAccepted(acceptedContainerIDs ids.Set) error {
	if err := b.VM.Bootstrapping(); err != nil {
		return fmt.Errorf("failed to notify VM that bootstrapping has started: %w",
			err)
	}

	storedVtxs := make([]avalanche.Vertex, 0, acceptedContainerIDs.Len())
	for _, vtxID := range acceptedContainerIDs.List() {
		if vtx, err := b.State.GetVertex(vtxID); err == nil {
			storedVtxs = append(storedVtxs, vtx)
		} else if err := b.fetch(vtxID); err != nil {
			return err
		}
	}
	if err := b.process(storedVtxs...); err != nil {
		return err
	}
	b.processedStartingAcceptedFrontier = true

	if numPending := b.outstandingRequests.Len(); numPending == 0 {
		return b.finish()
	}
	return nil
}

// Finish bootstrapping
func (b *bootstrapper) finish() error {
	if b.finished {
		return nil
	}
	b.BootstrapConfig.Context.Log.Info("finished fetching vertices. executing transaction state transitions...")

	if err := b.executeAll(b.TxBlocked, b.numBSBlockedTx); err != nil {
		return err
	}

	b.BootstrapConfig.Context.Log.Info("executing vertex state transitions...")

	if err := b.executeAll(b.VtxBlocked, b.numBSBlockedVtx); err != nil {
		return err
	}

	if err := b.VM.Bootstrapped(); err != nil {
		return fmt.Errorf("failed to notify VM that bootstrapping has finished: %w",
			err)
	}

	// Start consensus
	if err := b.onFinished(); err != nil {
		return err
	}
	b.finished = true
	return nil
}

func (b *bootstrapper) executeAll(jobs *queue.Jobs, numBlocked prometheus.Gauge) error {
	numExecuted := 0
	for job, err := jobs.Pop(); err == nil; job, err = jobs.Pop() {
		numBlocked.Dec()
		b.BootstrapConfig.Context.Log.Debug("Executing: %s", job.ID())
		if err := jobs.Execute(job); err != nil {
			b.BootstrapConfig.Context.Log.Error("Error executing: %s", err)
			return err
		}
		if err := jobs.Commit(); err != nil {
			return err
		}
		numExecuted++
		if numExecuted%common.StatusUpdateFrequency == 0 { // Periodically print progress
			b.BootstrapConfig.Context.Log.Info("executed %d operations", numExecuted)
		}
	}
	return nil
}
<|endoftext|>