<commit_before>package sortutil

import (
	"sort"
	"testing"
	"time"
)

const (
	day = 24 * time.Hour
)

type Item struct {
	Id    int64
	Name  string
	Date  time.Time
	Valid bool
}

type SortableItems []Item

func (s SortableItems) Len() int {
	return len(s)
}

func (s SortableItems) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s SortableItems) Less(i, j int) bool {
	return s[i].Id > s[j].Id
}

func names() []string {
	return []string{"A", "C", "a", "b", "d", "g", "h", "y", "z"}
}

func namesInsensitive() []string {
	return []string{"A", "a", "b", "C", "d", "g", "h", "y", "z"}
}

var now = time.Now()

func dates() []time.Time {
	return []time.Time{
		now.Add(-3 * day),
		now.Add(-2 * day),
		now.Add(-1 * day),
		now,
		now.Add(1 * day),
		now.Add(2 * day),
		now.Add(4 * day),
		now.Add(5 * day),
		now.Add(7 * day),
	}
}

func items() []Item {
	n := names()
	d := dates()
	is := []Item{
		{6, n[4], d[0], true},
		{1, n[3], d[5], true},
		{9, n[1], d[6], true},
		{3, n[8], d[2], false},
		{7, n[7], d[8], true},
		{2, n[2], d[4], false},
		{8, n[0], d[1], false},
		{5, n[5], d[7], false},
		{4, n[6], d[3], true},
	}
	return is
}

func nestedIntSlice() [][]int {
	return [][]int{
		{4, 5, 1},
		{2, 1, 7},
		{9, 3, 3},
		{1, 6, 2},
	}
}

func TestSortReverse(t *testing.T) {
	is := items()
	SortReverse(SortableItems(is))
	for i, v := range is {
		if v.Id != int64(i+1) {
			t.Errorf("is[%d].Id is not %d, but %d", i, i+1, v.Id)
		}
	}
}

// func TestSortNoGetterReverse(t *testing.T) {

func TestSortByStringFieldAscending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Name"), Ascending)
	c := names()
	for i, v := range is {
		if v.Name != c[i] {
			t.Errorf("is[%d].Name is not %s, but %s", i, c[i], v.Name)
		}
	}
}

func TestSortByStringFieldDescending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Name"), Descending)
	c := names()
	Reverse(sort.StringSlice(c))
	for i, v := range is {
		if v.Name != c[i] {
			t.Errorf("is[%d].Name is not %s, but %s", i, c[i], v.Name)
		}
	}
}

func TestSortByStringFieldCaseInsensitiveAscending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Name"), CaseInsensitiveAscending)
	c := namesInsensitive()
	for i, v := range is {
		if v.Name != c[i] {
			t.Errorf("is[%d].Name is not %s, but %s", i, c[i], v.Name)
		}
	}
}

func TestSortByStringFieldCaseInsensitiveDescending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Name"), CaseInsensitiveDescending)
	c := namesInsensitive()
	Reverse(sort.StringSlice(c))
	for i, v := range is {
		if v.Name != c[i] {
			t.Errorf("is[%d].Name is not %s, but %s", i, c[i], v.Name)
		}
	}
}

func TestSortByInt64FieldAscending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Id"), Ascending)
	for i, v := range is {
		if v.Id != int64(i+1) {
			t.Errorf("is[%d].Id is not %d, but %d", i, i+1, v.Id)
		}
	}
}

func TestSortByInt64FieldDescending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Id"), Descending)
	l := len(is)
	for i, v := range is {
		if v.Id != int64(l-i) {
			t.Errorf("is[%d].Id is not %d, but %d", i, l-i, v.Id)
		}
	}
}

func TestSortByIntIndexAscending(t *testing.T) {
	is := nestedIntSlice()
	Sort(is, IndexGetter(2), Ascending)
	if !sort.IntsAreSorted([]int{is[0][2], is[1][2], is[2][2], is[3][2]}) {
		t.Errorf("Nested int slice is not sorted by index 2 in child slices: %v", is)
	}
}

func TestSortByTimeFieldAscending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Date"), Ascending)
	c := dates()
	for i, v := range is {
		if !v.Date.Equal(c[i]) {
			t.Errorf("is[%d].Date is not %v, but %v", i, c[i], v.Date)
		}
	}
}

func TestSortByTimeFieldDescending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Date"), Descending)
	c := dates()
	l := len(is)
	for i, v := range is {
		if !v.Date.Equal(c[l-i-1]) {
			t.Errorf("is[%d].Date is not %v, but %v", i, c[l-i], v.Date)
		}
	}
}

type TestStruct struct {
	TimePtr    *time.Time
	Invalid    InvalidType
	unexported int
}

type InvalidType struct {
	Foo string
	Bar int
}

func testStructs() []TestStruct {
	return []TestStruct{
		{
			TimePtr:    &now,
			Invalid:    InvalidType{"foo", 123},
			unexported: 5,
		},
	}
}

func TestSortInvalidType(t *testing.T) {
	// Sorting an invalid type should cause a panic
	defer func() {
		if x := recover(); x == nil {
			t.Fatal("Sorting an unrecognized type didn't cause a panic")
		}
	}()
	s := testStructs()
	Sort(s, FieldGetter("Invalid"), Ascending)
}

func TestSortUnexportedType(t *testing.T) {
	// Sorting an unexported type should cause a panic
	// TODO: This should test on a field outside the package
	return // TEMP
	defer func() {
		if x := recover(); x == nil {
			t.Fatal("Sorting an unexported type didn't cause a panic")
		}
	}()
	s := testStructs()
	Sort(s, FieldGetter("unexported"), Ascending)
}

func TestSortPointerType(t *testing.T) {
	// Sorting a pointer type shouldn't cause a panic
	s := testStructs()
	Sort(s, FieldGetter("TimePtr"), Ascending)
}

func BenchmarkSortStructByInt64(b *testing.B) {
	for i := 0; i < b.N; i++ {
		is := items()
		sort.Sort(SortableItems(is))
	}
}

func BenchmarkSortReverseStructByInt64(b *testing.B) {
	for i := 0; i < b.N; i++ {
		is := items()
		SortReverse(SortableItems(is))
	}
}
<commit_msg>Remove Old reference<commit_after>package sortutil

import (
	"sort"
	"testing"
	"time"
)

const (
	day = 24 * time.Hour
)

type Item struct {
	Id    int64
	Name  string
	Date  time.Time
	Valid bool
}

type SortableItems []Item

func (s SortableItems) Len() int {
	return len(s)
}

func (s SortableItems) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s SortableItems) Less(i, j int) bool {
	return s[i].Id > s[j].Id
}

func names() []string {
	return []string{"A", "C", "a", "b", "d", "g", "h", "y", "z"}
}

func namesInsensitive() []string {
	return []string{"A", "a", "b", "C", "d", "g", "h", "y", "z"}
}

var now = time.Now()

func dates() []time.Time {
	return []time.Time{
		now.Add(-3 * day),
		now.Add(-2 * day),
		now.Add(-1 * day),
		now,
		now.Add(1 * day),
		now.Add(2 * day),
		now.Add(4 * day),
		now.Add(5 * day),
		now.Add(7 * day),
	}
}

func items() []Item {
	n := names()
	d := dates()
	is := []Item{
		{6, n[4], d[0], true},
		{1, n[3], d[5], true},
		{9, n[1], d[6], true},
		{3, n[8], d[2], false},
		{7, n[7], d[8], true},
		{2, n[2], d[4], false},
		{8, n[0], d[1], false},
		{5, n[5], d[7], false},
		{4, n[6], d[3], true},
	}
	return is
}

func nestedIntSlice() [][]int {
	return [][]int{
		{4, 5, 1},
		{2, 1, 7},
		{9, 3, 3},
		{1, 6, 2},
	}
}

func TestSortReverse(t *testing.T) {
	is := items()
	SortReverse(SortableItems(is))
	for i, v := range is {
		if v.Id != int64(i+1) {
			t.Errorf("is[%d].Id is not %d, but %d", i, i+1, v.Id)
		}
	}
}

func TestSortByStringFieldAscending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Name"), Ascending)
	c := names()
	for i, v := range is {
		if v.Name != c[i] {
			t.Errorf("is[%d].Name is not %s, but %s", i, c[i], v.Name)
		}
	}
}

func TestSortByStringFieldDescending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Name"), Descending)
	c := names()
	Reverse(sort.StringSlice(c))
	for i, v := range is {
		if v.Name != c[i] {
			t.Errorf("is[%d].Name is not %s, but %s", i, c[i], v.Name)
		}
	}
}

func TestSortByStringFieldCaseInsensitiveAscending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Name"), CaseInsensitiveAscending)
	c := namesInsensitive()
	for i, v := range is {
		if v.Name != c[i] {
			t.Errorf("is[%d].Name is not %s, but %s", i, c[i], v.Name)
		}
	}
}

func TestSortByStringFieldCaseInsensitiveDescending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Name"), CaseInsensitiveDescending)
	c := namesInsensitive()
	Reverse(sort.StringSlice(c))
	for i, v := range is {
		if v.Name != c[i] {
			t.Errorf("is[%d].Name is not %s, but %s", i, c[i], v.Name)
		}
	}
}

func TestSortByInt64FieldAscending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Id"), Ascending)
	for i, v := range is {
		if v.Id != int64(i+1) {
			t.Errorf("is[%d].Id is not %d, but %d", i, i+1, v.Id)
		}
	}
}

func TestSortByInt64FieldDescending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Id"), Descending)
	l := len(is)
	for i, v := range is {
		if v.Id != int64(l-i) {
			t.Errorf("is[%d].Id is not %d, but %d", i, l-i, v.Id)
		}
	}
}

func TestSortByIntIndexAscending(t *testing.T) {
	is := nestedIntSlice()
	Sort(is, IndexGetter(2), Ascending)
	if !sort.IntsAreSorted([]int{is[0][2], is[1][2], is[2][2], is[3][2]}) {
		t.Errorf("Nested int slice is not sorted by index 2 in child slices: %v", is)
	}
}

func TestSortByTimeFieldAscending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Date"), Ascending)
	c := dates()
	for i, v := range is {
		if !v.Date.Equal(c[i]) {
			t.Errorf("is[%d].Date is not %v, but %v", i, c[i], v.Date)
		}
	}
}

func TestSortByTimeFieldDescending(t *testing.T) {
	is := items()
	Sort(is, FieldGetter("Date"), Descending)
	c := dates()
	l := len(is)
	for i, v := range is {
		if !v.Date.Equal(c[l-i-1]) {
			t.Errorf("is[%d].Date is not %v, but %v", i, c[l-i], v.Date)
		}
	}
}

type TestStruct struct {
	TimePtr    *time.Time
	Invalid    InvalidType
	unexported int
}

type InvalidType struct {
	Foo string
	Bar int
}

func testStructs() []TestStruct {
	return []TestStruct{
		{
			TimePtr:    &now,
			Invalid:    InvalidType{"foo", 123},
			unexported: 5,
		},
	}
}

func TestSortInvalidType(t *testing.T) {
	// Sorting an invalid type should cause a panic
	defer func() {
		if x := recover(); x == nil {
			t.Fatal("Sorting an unrecognized type didn't cause a panic")
		}
	}()
	s := testStructs()
	Sort(s, FieldGetter("Invalid"), Ascending)
}

func TestSortUnexportedType(t *testing.T) {
	// Sorting an unexported type should cause a panic
	// TODO: This should test on a field outside the package
	return // TEMP
	defer func() {
		if x := recover(); x == nil {
			t.Fatal("Sorting an unexported type didn't cause a panic")
		}
	}()
	s := testStructs()
	Sort(s, FieldGetter("unexported"), Ascending)
}

func TestSortPointerType(t *testing.T) {
	// Sorting a pointer type shouldn't cause a panic
	s := testStructs()
	Sort(s, FieldGetter("TimePtr"), Ascending)
}

func BenchmarkSortStructByInt64(b *testing.B) {
	for i := 0; i < b.N; i++ {
		is := items()
		sort.Sort(SortableItems(is))
	}
}

func BenchmarkSortReverseStructByInt64(b *testing.B) {
	for i := 0; i < b.N; i++ {
		is := items()
		SortReverse(SortableItems(is))
	}
}
<|endoftext|>


<commit_before>// +build darwin

package cpu

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	common "github.com/shirou/gopsutil/common"
)

var HELPER_PATH = filepath.Join(os.Getenv("HOME"), ".gopsutil_cpu_helper")

// enable cpu helper. It may become security problem.
// This env valiable approach will be changed.
const HELPER_ENABLE_ENV = "ALLLOW_INSECURE_CPU_HELPER"

var ClocksPerSec = float64(100)

func init() {
	out, err := exec.Command("/usr/bin/getconf", "CLK_TCK").Output()
	// ignore errors
	if err == nil {
		i, err := strconv.ParseFloat(strings.TrimSpace(string(out)), 64)
		if err == nil {
			ClocksPerSec = float64(i)
		}
	}

	// adhoc compile on the host. Errors will be ignored.
	// gcc is required to compile.
	if !common.PathExists(HELPER_PATH) && os.Getenv(HELPER_ENABLE_ENV) == "yes" {
		cmd := exec.Command("gcc", "-o", HELPER_PATH, "-x", "c", "-")
		stdin, err := cmd.StdinPipe()
		if err != nil {
			return
		}
		io.WriteString(stdin, cpu_helper_src)
		stdin.Close()
		cmd.Output()
	}
}

func CPUTimes(percpu bool) ([]CPUTimesStat, error) {
	var ret []CPUTimesStat
	if !common.PathExists(HELPER_PATH) {
		return nil, fmt.Errorf("could not get cpu time")
	}

	out, err := exec.Command(HELPER_PATH).Output()
	if err != nil {
		return ret, err
	}

	for _, line := range strings.Split(string(out), "\n") {
		f := strings.Split(string(line), ",")
		if len(f) != 5 {
			continue
		}
		cpu, err := strconv.ParseFloat(f[0], 64)
		if err != nil {
			return ret, err
		}
		// cpu:99 means total, so just ignore if percpu
		if (percpu && cpu == 99) || (!percpu && cpu != 99) {
			continue
		}
		user, err := strconv.ParseFloat(f[1], 64)
		if err != nil {
			return ret, err
		}
		sys, err := strconv.ParseFloat(f[2], 64)
		if err != nil {
			return ret, err
		}
		idle, err := strconv.ParseFloat(f[3], 64)
		if err != nil {
			return ret, err
		}
		nice, err := strconv.ParseFloat(f[4], 64)
		if err != nil {
			return ret, err
		}
		c := CPUTimesStat{
			User:   float64(user / ClocksPerSec),
			Nice:   float64(nice / ClocksPerSec),
			System: float64(sys / ClocksPerSec),
			Idle:   float64(idle / ClocksPerSec),
		}
		if !percpu {
			c.CPU = "cpu-total"
		} else {
			c.CPU = fmt.Sprintf("cpu%d", uint16(cpu))
		}
		ret = append(ret, c)
	}
	return ret, nil
}

// Returns only one CPUInfoStat on FreeBSD
func CPUInfo() ([]CPUInfoStat, error) {
	var ret []CPUInfoStat

	out, err := exec.Command("/usr/sbin/sysctl", "machdep.cpu").Output()
	if err != nil {
		return ret, err
	}

	c := CPUInfoStat{}
	for _, line := range strings.Split(string(out), "\n") {
		values := strings.Fields(line)
		if len(values) < 1 {
			continue
		}

		t, err := strconv.ParseInt(values[1], 10, 64)
		// err is not checked here because some value is string.
		if strings.HasPrefix(line, "machdep.cpu.brand_string") {
			c.ModelName = strings.Join(values[1:], " ")
		} else if strings.HasPrefix(line, "machdep.cpu.family") {
			c.Family = values[1]
		} else if strings.HasPrefix(line, "machdep.cpu.model") {
			c.Model = values[1]
		} else if strings.HasPrefix(line, "machdep.cpu.stepping") {
			if err != nil {
				return ret, err
			}
			c.Stepping = int32(t)
		} else if strings.HasPrefix(line, "machdep.cpu.features") {
			for _, v := range values[1:] {
				c.Flags = append(c.Flags, strings.ToLower(v))
			}
		} else if strings.HasPrefix(line, "machdep.cpu.leaf7_features") {
			for _, v := range values[1:] {
				c.Flags = append(c.Flags, strings.ToLower(v))
			}
		} else if strings.HasPrefix(line, "machdep.cpu.extfeatures") {
			for _, v := range values[1:] {
				c.Flags = append(c.Flags, strings.ToLower(v))
			}
		} else if strings.HasPrefix(line, "machdep.cpu.core_count") {
			if err != nil {
				return ret, err
			}
			c.Cores = int32(t)
		} else if strings.HasPrefix(line, "machdep.cpu.cache.size") {
			if err != nil {
				return ret, err
			}
			c.CacheSize = int32(t)
		} else if strings.HasPrefix(line, "machdep.cpu.vendor") {
			c.VendorID = values[1]
		}

		// TODO:
		// c.Mhz = mustParseFloat64(values[1])
	}

	return append(ret, c), nil
}

const cpu_helper_src = `
#include <mach/mach.h>
#include <mach/mach_error.h>
#include <stdio.h>

int main() {
  	natural_t cpuCount;
	processor_info_array_t ia;
	mach_msg_type_number_t ic;

	kern_return_t error = host_processor_info(mach_host_self(),
		PROCESSOR_CPU_LOAD_INFO, &cpuCount, &ia, &ic);
	if (error) {
		return error;
	}

	processor_cpu_load_info_data_t* cpuLoadInfo =
		(processor_cpu_load_info_data_t*) ia;

	unsigned int all[4];
	// cpu_no,user,system,idle,nice
	for (int cpu=0; cpu<cpuCount; cpu++){
	  printf("%d,", cpu);
	  for (int i=0; i<4; i++){
		printf("%d", cpuLoadInfo[cpu].cpu_ticks[i]);
		all[i] = all[i] + cpuLoadInfo[cpu].cpu_ticks[i];
		if (i != 3){
		  printf(",");
		}else{
		  printf("\n");
		}
	  }
	}
	printf("99,");
	for (int i=0; i<4; i++){
	  printf("%d", all[i]);
	  if (i != 3){
		printf(",");
	  }else{
		printf("\n");
	  }
	}

	vm_deallocate(mach_task_self(), (vm_address_t)ia, ic);
	return error;
}
`
<commit_msg>Read Darwin CPU frequency from sysctl hw.cpufrequency<commit_after>// +build darwin

package cpu

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	common "github.com/shirou/gopsutil/common"
)

var HELPER_PATH = filepath.Join(os.Getenv("HOME"), ".gopsutil_cpu_helper")

// enable cpu helper. It may become security problem.
// This env valiable approach will be changed.
const HELPER_ENABLE_ENV = "ALLLOW_INSECURE_CPU_HELPER"

var ClocksPerSec = float64(100)

func init() {
	out, err := exec.Command("/usr/bin/getconf", "CLK_TCK").Output()
	// ignore errors
	if err == nil {
		i, err := strconv.ParseFloat(strings.TrimSpace(string(out)), 64)
		if err == nil {
			ClocksPerSec = float64(i)
		}
	}

	// adhoc compile on the host. Errors will be ignored.
	// gcc is required to compile.
	if !common.PathExists(HELPER_PATH) && os.Getenv(HELPER_ENABLE_ENV) == "yes" {
		cmd := exec.Command("gcc", "-o", HELPER_PATH, "-x", "c", "-")
		stdin, err := cmd.StdinPipe()
		if err != nil {
			return
		}
		io.WriteString(stdin, cpu_helper_src)
		stdin.Close()
		cmd.Output()
	}
}

func CPUTimes(percpu bool) ([]CPUTimesStat, error) {
	var ret []CPUTimesStat
	if !common.PathExists(HELPER_PATH) {
		return nil, fmt.Errorf("could not get cpu time")
	}

	out, err := exec.Command(HELPER_PATH).Output()
	if err != nil {
		return ret, err
	}

	for _, line := range strings.Split(string(out), "\n") {
		f := strings.Split(string(line), ",")
		if len(f) != 5 {
			continue
		}
		cpu, err := strconv.ParseFloat(f[0], 64)
		if err != nil {
			return ret, err
		}
		// cpu:99 means total, so just ignore if percpu
		if (percpu && cpu == 99) || (!percpu && cpu != 99) {
			continue
		}
		user, err := strconv.ParseFloat(f[1], 64)
		if err != nil {
			return ret, err
		}
		sys, err := strconv.ParseFloat(f[2], 64)
		if err != nil {
			return ret, err
		}
		idle, err := strconv.ParseFloat(f[3], 64)
		if err != nil {
			return ret, err
		}
		nice, err := strconv.ParseFloat(f[4], 64)
		if err != nil {
			return ret, err
		}
		c := CPUTimesStat{
			User:   float64(user / ClocksPerSec),
			Nice:   float64(nice / ClocksPerSec),
			System: float64(sys / ClocksPerSec),
			Idle:   float64(idle / ClocksPerSec),
		}
		if !percpu {
			c.CPU = "cpu-total"
		} else {
			c.CPU = fmt.Sprintf("cpu%d", uint16(cpu))
		}
		ret = append(ret, c)
	}
	return ret, nil
}

// Returns only one CPUInfoStat on FreeBSD
func CPUInfo() ([]CPUInfoStat, error) {
	var ret []CPUInfoStat

	out, err := exec.Command("/usr/sbin/sysctl", "machdep.cpu").Output()
	if err != nil {
		return ret, err
	}

	c := CPUInfoStat{}
	for _, line := range strings.Split(string(out), "\n") {
		values := strings.Fields(line)
		if len(values) < 1 {
			continue
		}

		t, err := strconv.ParseInt(values[1], 10, 64)
		// err is not checked here because some value is string.
		if strings.HasPrefix(line, "machdep.cpu.brand_string") {
			c.ModelName = strings.Join(values[1:], " ")
		} else if strings.HasPrefix(line, "machdep.cpu.family") {
			c.Family = values[1]
		} else if strings.HasPrefix(line, "machdep.cpu.model") {
			c.Model = values[1]
		} else if strings.HasPrefix(line, "machdep.cpu.stepping") {
			if err != nil {
				return ret, err
			}
			c.Stepping = int32(t)
		} else if strings.HasPrefix(line, "machdep.cpu.features") {
			for _, v := range values[1:] {
				c.Flags = append(c.Flags, strings.ToLower(v))
			}
		} else if strings.HasPrefix(line, "machdep.cpu.leaf7_features") {
			for _, v := range values[1:] {
				c.Flags = append(c.Flags, strings.ToLower(v))
			}
		} else if strings.HasPrefix(line, "machdep.cpu.extfeatures") {
			for _, v := range values[1:] {
				c.Flags = append(c.Flags, strings.ToLower(v))
			}
		} else if strings.HasPrefix(line, "machdep.cpu.core_count") {
			if err != nil {
				return ret, err
			}
			c.Cores = int32(t)
		} else if strings.HasPrefix(line, "machdep.cpu.cache.size") {
			if err != nil {
				return ret, err
			}
			c.CacheSize = int32(t)
		} else if strings.HasPrefix(line, "machdep.cpu.vendor") {
			c.VendorID = values[1]
		}
	}

	// Use the rated frequency of the CPU. This is a static value and does not
	// account for low power or Turbo Boost modes.
	out, err = exec.Command("/usr/sbin/sysctl", "hw.cpufrequency").Output()
	if err != nil {
		return ret, err
	}

	values := strings.Fields(string(out))
	c.Mhz, err = strconv.ParseFloat(values[1], 64)
	if err != nil {
		return ret, err
	}

	return append(ret, c), nil
}

const cpu_helper_src = `
#include <mach/mach.h>
#include <mach/mach_error.h>
#include <stdio.h>

int main() {
  	natural_t cpuCount;
	processor_info_array_t ia;
	mach_msg_type_number_t ic;

	kern_return_t error = host_processor_info(mach_host_self(),
		PROCESSOR_CPU_LOAD_INFO, &cpuCount, &ia, &ic);
	if (error) {
		return error;
	}

	processor_cpu_load_info_data_t* cpuLoadInfo =
		(processor_cpu_load_info_data_t*) ia;

	unsigned int all[4];
	// cpu_no,user,system,idle,nice
	for (int cpu=0; cpu<cpuCount; cpu++){
	  printf("%d,", cpu);
	  for (int i=0; i<4; i++){
		printf("%d", cpuLoadInfo[cpu].cpu_ticks[i]);
		all[i] = all[i] + cpuLoadInfo[cpu].cpu_ticks[i];
		if (i != 3){
		  printf(",");
		}else{
		  printf("\n");
		}
	  }
	}
	printf("99,");
	for (int i=0; i<4; i++){
	  printf("%d", all[i]);
	  if (i != 3){
		printf(",");
	  }else{
		printf("\n");
	  }
	}

	vm_deallocate(mach_task_self(), (vm_address_t)ia, ic);
	return error;
}
`
<|endoftext|>


<commit_before>package expiration_test

import (
	"errors"
	"time"

	"code.cloudfoundry.org/clock/fakeclock"
	"code.cloudfoundry.org/lager/lagertest"
	"code.cloudfoundry.org/locket/db"
	"code.cloudfoundry.org/locket/db/dbfakes"
	"code.cloudfoundry.org/locket/expiration"
	"code.cloudfoundry.org/locket/models"

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

var _ = Describe("LockPick", func() {
	var (
		lockPick expiration.LockPick

		logger     *lagertest.TestLogger
		fakeLockDB *dbfakes.FakeLockDB
		fakeClock  *fakeclock.FakeClock

		ttl time.Duration

		lock *db.Lock
	)

	BeforeEach(func() {
		lock = &db.Lock{
			Resource: &models.Resource{
				Key:   "funky",
				Owner: "town",
				Value: "won't you take me to",
			},
			TtlInSeconds:  25,
			ModifiedIndex: 6,
		}

		ttl = time.Duration(lock.TtlInSeconds) * time.Second

		fakeClock = fakeclock.NewFakeClock(time.Now())
		logger = lagertest.NewTestLogger("lock-pick")
		fakeLockDB = &dbfakes.FakeLockDB{}

		lockPick = expiration.NewLockPick(fakeLockDB, fakeClock)
	})

	Context("RegisterTTL", func() {
		BeforeEach(func() {
			fakeLockDB.FetchReturns(lock, nil)
		})

		It("checks that the lock expires after the ttl", func() {
			lockPick.RegisterTTL(logger, lock)

			fakeClock.WaitForWatcherAndIncrement(ttl)

			Eventually(fakeLockDB.FetchCallCount).Should(Equal(1))
			_, key := fakeLockDB.FetchArgsForCall(0)
			Expect(key).To(Equal(lock.Key))

			Eventually(fakeLockDB.ReleaseCallCount).Should(Equal(1))
			_, resource := fakeLockDB.ReleaseArgsForCall(0)
			Expect(resource).To(Equal(lock.Resource))
		})

		Context("when the modified index has been incremented", func() {
			var returnedLock *db.Lock
			BeforeEach(func() {
				returnedLock = &db.Lock{
					Resource: &models.Resource{
						Key:   "funky",
						Owner: "town",
						Value: "won't you take me to",
					},
					TtlInSeconds:  25,
					ModifiedIndex: 7,
				}

				fakeLockDB.FetchReturns(returnedLock, nil)
			})

			It("does not release the lock", func() {
				lockPick.RegisterTTL(logger, lock)

				fakeClock.WaitForWatcherAndIncrement(ttl)

				Eventually(fakeLockDB.FetchCallCount).Should(Equal(1))
				_, key := fakeLockDB.FetchArgsForCall(0)
				Expect(key).To(Equal(lock.Key))

				Consistently(fakeLockDB.ReleaseCallCount).Should(Equal(0))
			})
		})

		Context("when fetching the lock fails", func() {
			BeforeEach(func() {
				fakeLockDB.FetchReturns(nil, errors.New("failed-to-fetch-lock"))
			})

			It("logs the error and does not release the lock", func() {
				lockPick.RegisterTTL(logger, lock)

				fakeClock.WaitForWatcherAndIncrement(ttl)

				Eventually(fakeLockDB.FetchCallCount).Should(Equal(1))

				Eventually(logger).Should(gbytes.Say("failed-to-fetch-lock"))
				Consistently(fakeLockDB.ReleaseCallCount).Should(Equal(0))
			})
		})

		Context("when releasing the lock fails", func() {
			BeforeEach(func() {
				fakeLockDB.ReleaseReturns(errors.New("failed-to-release-lock"))
			})

			It("logs the error", func() {
				lockPick.RegisterTTL(logger, lock)

				fakeClock.WaitForWatcherAndIncrement(ttl)

				Eventually(fakeLockDB.ReleaseCallCount).Should(Equal(1))
				Eventually(logger).Should(gbytes.Say("failed-to-release-lock"))
			})
		})

		Context("when there is already a check process running", func() {
			BeforeEach(func() {
				lockPick.RegisterTTL(logger, lock)
			})

			It("does not add a new check process", func() {
				lockPick.RegisterTTL(logger, lock)

				Eventually(fakeClock.WatcherCount).Should(Equal(1))
				Consistently(fakeClock.WatcherCount).Should(Equal(1))
				fakeClock.WaitForWatcherAndIncrement(ttl)

				Eventually(fakeLockDB.FetchCallCount).Should(Equal(1))
				_, key := fakeLockDB.FetchArgsForCall(0)
				Expect(key).To(Equal(lock.Key))

				Consistently(fakeLockDB.ReleaseCallCount).Should(Equal(1))
			})

			Context("and the check process finishes", func() {
				BeforeEach(func() {
					fakeClock.WaitForWatcherAndIncrement(ttl)

					Eventually(fakeLockDB.FetchCallCount).Should(Equal(1))
					_, key := fakeLockDB.FetchArgsForCall(0)
					Expect(key).To(Equal(lock.Key))

					Consistently(fakeLockDB.ReleaseCallCount).Should(Equal(1))
				})

				It("performs the expiration check", func() {
					lockPick.RegisterTTL(logger, lock)

					Eventually(fakeClock.WatcherCount).Should(Equal(1))
					fakeClock.WaitForWatcherAndIncrement(ttl)

					Eventually(fakeLockDB.FetchCallCount).Should(Equal(2))
					_, key := fakeLockDB.FetchArgsForCall(0)
					Expect(key).To(Equal(lock.Key))

					Consistently(fakeLockDB.ReleaseCallCount).Should(Equal(2))
				})
			})
		})
	})
})
<commit_msg>fix locket test suite flaky test<commit_after>package expiration_test

import (
	"errors"
	"time"

	"code.cloudfoundry.org/clock/fakeclock"
	"code.cloudfoundry.org/lager/lagertest"
	"code.cloudfoundry.org/locket/db"
	"code.cloudfoundry.org/locket/db/dbfakes"
	"code.cloudfoundry.org/locket/expiration"
	"code.cloudfoundry.org/locket/models"

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

var _ = Describe("LockPick", func() {
	var (
		lockPick expiration.LockPick

		logger     *lagertest.TestLogger
		fakeLockDB *dbfakes.FakeLockDB
		fakeClock  *fakeclock.FakeClock

		ttl time.Duration

		lock *db.Lock
	)

	BeforeEach(func() {
		lock = &db.Lock{
			Resource: &models.Resource{
				Key:   "funky",
				Owner: "town",
				Value: "won't you take me to",
			},
			TtlInSeconds:  25,
			ModifiedIndex: 6,
		}

		ttl = time.Duration(lock.TtlInSeconds) * time.Second

		fakeClock = fakeclock.NewFakeClock(time.Now())
		logger = lagertest.NewTestLogger("lock-pick")
		fakeLockDB = &dbfakes.FakeLockDB{}

		lockPick = expiration.NewLockPick(fakeLockDB, fakeClock)
	})

	Context("RegisterTTL", func() {
		BeforeEach(func() {
			fakeLockDB.FetchReturns(lock, nil)
		})

		It("checks that the lock expires after the ttl", func() {
			lockPick.RegisterTTL(logger, lock)

			fakeClock.WaitForWatcherAndIncrement(ttl)

			Eventually(fakeLockDB.FetchCallCount).Should(Equal(1))
			_, key := fakeLockDB.FetchArgsForCall(0)
			Expect(key).To(Equal(lock.Key))

			Eventually(fakeLockDB.ReleaseCallCount).Should(Equal(1))
			_, resource := fakeLockDB.ReleaseArgsForCall(0)
			Expect(resource).To(Equal(lock.Resource))
		})

		Context("when the modified index has been incremented", func() {
			var returnedLock *db.Lock
			BeforeEach(func() {
				returnedLock = &db.Lock{
					Resource: &models.Resource{
						Key:   "funky",
						Owner: "town",
						Value: "won't you take me to",
					},
					TtlInSeconds:  25,
					ModifiedIndex: 7,
				}

				fakeLockDB.FetchReturns(returnedLock, nil)
			})

			It("does not release the lock", func() {
				lockPick.RegisterTTL(logger, lock)

				fakeClock.WaitForWatcherAndIncrement(ttl)

				Eventually(fakeLockDB.FetchCallCount).Should(Equal(1))
				_, key := fakeLockDB.FetchArgsForCall(0)
				Expect(key).To(Equal(lock.Key))

				Consistently(fakeLockDB.ReleaseCallCount).Should(Equal(0))
			})
		})

		Context("when fetching the lock fails", func() {
			BeforeEach(func() {
				fakeLockDB.FetchReturns(nil, errors.New("failed-to-fetch-lock"))
			})

			It("logs the error and does not release the lock", func() {
				lockPick.RegisterTTL(logger, lock)

				fakeClock.WaitForWatcherAndIncrement(ttl)

				Eventually(fakeLockDB.FetchCallCount).Should(Equal(1))

				Eventually(logger).Should(gbytes.Say("failed-to-fetch-lock"))
				Consistently(fakeLockDB.ReleaseCallCount).Should(Equal(0))
			})
		})

		Context("when releasing the lock fails", func() {
			BeforeEach(func() {
				fakeLockDB.ReleaseReturns(errors.New("failed-to-release-lock"))
			})

			It("logs the error", func() {
				lockPick.RegisterTTL(logger, lock)

				fakeClock.WaitForWatcherAndIncrement(ttl)

				Eventually(fakeLockDB.ReleaseCallCount).Should(Equal(1))
				Eventually(logger).Should(gbytes.Say("failed-to-release-lock"))
			})
		})

		Context("when there is already a check process running", func() {
			BeforeEach(func() {
				lockPick.RegisterTTL(logger, lock)
			})

			It("does not add a new check process", func() {
				lockPick.RegisterTTL(logger, lock)

				Eventually(fakeClock.WatcherCount).Should(Equal(1))
				Consistently(fakeClock.WatcherCount).Should(Equal(1))
				fakeClock.WaitForWatcherAndIncrement(ttl)

				Eventually(fakeLockDB.FetchCallCount).Should(Equal(1))
				_, key := fakeLockDB.FetchArgsForCall(0)
				Expect(key).To(Equal(lock.Key))

				Eventually(fakeLockDB.ReleaseCallCount).Should(Equal(1))
				Consistently(fakeLockDB.ReleaseCallCount).Should(Equal(1))
			})

			Context("and the check process finishes", func() {
				BeforeEach(func() {
					fakeClock.WaitForWatcherAndIncrement(ttl)

					Eventually(fakeLockDB.FetchCallCount).Should(Equal(1))
					_, key := fakeLockDB.FetchArgsForCall(0)
					Expect(key).To(Equal(lock.Key))

					Eventually(fakeLockDB.ReleaseCallCount).Should(Equal(1))
					Consistently(fakeLockDB.ReleaseCallCount).Should(Equal(1))
				})

				It("performs the expiration check", func() {
					lockPick.RegisterTTL(logger, lock)

					Eventually(fakeClock.WatcherCount).Should(Equal(1))
					fakeClock.WaitForWatcherAndIncrement(ttl)

					Eventually(fakeLockDB.FetchCallCount).Should(Equal(2))
					_, key := fakeLockDB.FetchArgsForCall(0)
					Expect(key).To(Equal(lock.Key))

					Eventually(fakeLockDB.ReleaseCallCount).Should(Equal(2))
					Consistently(fakeLockDB.ReleaseCallCount).Should(Equal(2))
				})
			})
		})
	})
})
<|endoftext|>


<commit_before>package polygon

import (
	"encoding/json"
	"errors"
	"net/http"
	"strconv"
)

const polygonHost = "https://api.polygon.io"

type Client struct {
	client http.Client
	apiKey string
}

func NewClient(httpClient *http.Client, apiKey string) *Client {
	return &Client{
		client: *httpClient,
		apiKey: apiKey,
	}
}

type TickersResponse struct {
	Count   int      `json:"count"`
	Tickers []Ticker `json:"tickers"`
}

type Ticker struct {
	Ticker          string `json:"ticker"`
	Name            string `json:"name"`
	Market          string `json:"market"`
	Locale          string `json:"locale"`
	Currency        string `json:"currency"`
	Active          bool   `json:"active"`
	PrimaryExchange string `json:"primaryExch"`
}

type TickersQueryParams struct {
	PerPage int
	Page    int
	Active  bool
	Type    string
	Market  string
	Locale  string
	Search  string
	Sort    string
}

func (c *Client) GetTickers(params TickersQueryParams) (result TickersResponse, err error) {
	req, err := http.NewRequest("GET", polygonHost+"/v2/reference/tickers", nil)
	if err != nil {
		return result, err
	}
	queryParams := req.URL.Query()
	queryParams.Add("apiKey", c.apiKey)
	queryParams.Add("sort", params.Sort)
	queryParams.Add("type", params.Type)
	queryParams.Add("market", params.Market)
	queryParams.Add("locale", params.Locale)
	queryParams.Add("search", params.Search)
	queryParams.Add("perpage", strconv.Itoa(params.PerPage))
	queryParams.Add("page", strconv.Itoa(params.Page))
	queryParams.Add("active", strconv.FormatBool(params.Active))
	req.URL.RawQuery = queryParams.Encode()

	resp, err := c.client.Do(req)
	if err != nil {
		return result, err
	}
	if resp.StatusCode != http.StatusOK {
		return result, errors.New("Response failed with status code: " + resp.Status)
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return result, err
	}
	return result, nil
}
<commit_msg>feat: Get ticker news from Polygon<commit_after>package polygon

import (
	"encoding/json"
	"errors"
	"net/http"
	"strconv"
	"time"
)

const polygonHost = "https://api.polygon.io"

type Client struct {
	client http.Client
	apiKey string
}

func NewClient(httpClient *http.Client, apiKey string) *Client {
	return &Client{
		client: *httpClient,
		apiKey: apiKey,
	}
}

type TickersResponse struct {
	Count   int      `json:"count"`
	Tickers []Ticker `json:"tickers"`
}

type Ticker struct {
	Ticker          string `json:"ticker"`
	Name            string `json:"name"`
	Market          string `json:"market"`
	Locale          string `json:"locale"`
	Currency        string `json:"currency"`
	Active          bool   `json:"active"`
	PrimaryExchange string `json:"primaryExch"`
}

type TickersQueryParams struct {
	PerPage int
	Page    int
	Active  bool
	Type    string
	Market  string
	Locale  string
	Search  string
	Sort    string
}

type Article struct {
	Title     string    `json:"title"`
	URL       string    `json:"url"`
	Source    string    `json:"source"`
	Summary   string    `json:"summary"`
	Timestamp time.Time `json:"timestamp"`
}

func (c *Client) GetTickers(params TickersQueryParams) (result TickersResponse, err error) {
	req, err := http.NewRequest("GET", polygonHost+"/v2/reference/tickers", nil)
	if err != nil {
		return result, err
	}
	queryParams := req.URL.Query()
	queryParams.Add("apiKey", c.apiKey)
	queryParams.Add("sort", params.Sort)
	queryParams.Add("type", params.Type)
	queryParams.Add("market", params.Market)
	queryParams.Add("locale", params.Locale)
	queryParams.Add("search", params.Search)
	queryParams.Add("perpage", strconv.Itoa(params.PerPage))
	queryParams.Add("page", strconv.Itoa(params.Page))
	queryParams.Add("active", strconv.FormatBool(params.Active))
	req.URL.RawQuery = queryParams.Encode()

	resp, err := c.client.Do(req)
	if err != nil {
		return result, err
	}
	if resp.StatusCode != http.StatusOK {
		return result, errors.New("Response failed with status code: " + resp.Status)
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return result, err
	}
	return result, nil
}

func (c *Client) GetTickerNews(ticker string, perPage int, page int) (articles []Article, err error) {
	url := polygonHost + "/v1/meta/symbols/" + ticker + "/news"
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return articles, err
	}
	queryParams := req.URL.Query()
	queryParams.Add("apiKey", c.apiKey)
	queryParams.Add("perpage", strconv.Itoa(perPage))
	queryParams.Add("page", strconv.Itoa(page))
	req.URL.RawQuery = queryParams.Encode()

	resp, err := c.client.Do(req)
	if err != nil {
		return articles, err
	}
	if resp.StatusCode != http.StatusOK {
		return articles, errors.New("Response failed with status code: " + resp.Status)
	}
	if err := json.NewDecoder(resp.Body).Decode(&articles); err != nil {
		return articles, err
	}
	return articles, nil
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 Nippon Telegraph and Telephone Corporation.
//
// 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 server

import (
	"encoding/json"
	log "github.com/Sirupsen/logrus"
	"github.com/osrg/gobgp/api"
	"github.com/osrg/gobgp/config"
	"github.com/osrg/gobgp/packet"
	"github.com/osrg/gobgp/table"
	"gopkg.in/tomb.v2"
	"net"
	"time"
)

type peerMsgType int

const (
	_ peerMsgType = iota
	PEER_MSG_PATH
	PEER_MSG_PEER_DOWN
)

type peerMsg struct {
	msgType peerMsgType
	msgData interface{}
}

type Peer struct {
	t              tomb.Tomb
	globalConfig   config.GlobalType
	peerConfig     config.NeighborType
	acceptedConnCh chan *net.TCPConn
	incoming       chan *fsmMsg
	outgoing       chan *bgp.BGPMessage
	serverMsgCh    chan *serverMsg
	peerMsgCh      chan *peerMsg
	fsm            *FSM
	adjRib         *table.AdjRib
	// peer and rib are always not one-to-one so should not be
	// here but it's the simplest and works our first target.
	rib *table.TableManager
	// for now we support only the same afi as transport
	rf       bgp.RouteFamily
	capMap   map[bgp.BGPCapabilityCode]bgp.ParameterCapabilityInterface
	peerInfo *table.PeerInfo
	siblings map[string]*serverMsgDataPeer
}

func NewPeer(g config.GlobalType, peer config.NeighborType, serverMsgCh chan *serverMsg, peerMsgCh chan *peerMsg, peerList []*serverMsgDataPeer) *Peer {
	p := &Peer{
		globalConfig:   g,
		peerConfig:     peer,
		acceptedConnCh: make(chan *net.TCPConn),
		incoming:       make(chan *fsmMsg, 4096),
		outgoing:       make(chan *bgp.BGPMessage, 4096),
		serverMsgCh:    serverMsgCh,
		peerMsgCh:      peerMsgCh,
		capMap:         make(map[bgp.BGPCapabilityCode]bgp.ParameterCapabilityInterface),
	}
	p.siblings = make(map[string]*serverMsgDataPeer)
	for _, s := range peerList {
		p.siblings[s.address.String()] = s
	}
	p.fsm = NewFSM(&g, &peer, p.acceptedConnCh, p.incoming, p.outgoing)
	peer.BgpNeighborCommonState.State = uint32(bgp.BGP_FSM_IDLE)
	if peer.NeighborAddress.To4() != nil {
		p.rf = bgp.RF_IPv4_UC
	} else {
		p.rf = bgp.RF_IPv6_UC
	}
	p.peerInfo = &table.PeerInfo{
		AS:         peer.PeerAs,
		VersionNum: 1,
		LocalID:    g.RouterId,
		RF:         p.rf,
		Address:    peer.NeighborAddress,
	}
	p.adjRib = table.NewAdjRib()
	p.rib = table.NewTableManager()
	p.t.Go(p.loop)
	return p
}

func (peer *Peer) handleBGPmessage(m *bgp.BGPMessage) {
	j, _ := json.Marshal(m)
	log.Debug(string(j))

	switch m.Header.Type {
	case bgp.BGP_MSG_OPEN:
		body := m.Body.(*bgp.BGPOpen)
		peer.peerInfo.ID = m.Body.(*bgp.BGPOpen).ID
		for _, p := range body.OptParams {
			paramCap, y := p.(*bgp.OptionParameterCapability)
			if !y {
				continue
			}
			for _, c := range paramCap.Capability {
				peer.capMap[c.Code()] = c
			}
		}

	case bgp.BGP_MSG_ROUTE_REFRESH:
		pathList := peer.adjRib.GetOutPathList(peer.rf)
		peer.sendMessages(table.CreateUpdateMsgFromPaths(pathList))
	case bgp.BGP_MSG_UPDATE:
		peer.peerConfig.BgpNeighborCommonState.UpdateRecvTime = time.Now()
		body := m.Body.(*bgp.BGPUpdate)
		table.UpdatePathAttrs4ByteAs(body)
		msg := table.NewProcessMessage(m, peer.peerInfo)
		pathList := msg.ToPathList()
		if len(pathList) == 0 {
			return
		}
		peer.adjRib.UpdateIn(pathList)
		pm := &peerMsg{
			msgType: PEER_MSG_PATH,
			msgData: pathList,
		}
		for _, s := range peer.siblings {
			// TODO: check rf
			s.peerMsgCh <- pm
		}
	}
}

func (peer *Peer) sendMessages(msgs []*bgp.BGPMessage) {
	for _, m := range msgs {
		if peer.peerConfig.BgpNeighborCommonState.State != uint32(bgp.BGP_FSM_ESTABLISHED) {
			continue
		}

		if m.Header.Type != bgp.BGP_MSG_UPDATE {
			log.Fatal("not update message ", m.Header.Type)
		}

		_, y := peer.capMap[bgp.BGP_CAP_FOUR_OCTET_AS_NUMBER]
		if !y {
			log.Debug("update BGPUpdate for 2byte AS peer, ", peer.peerConfig.NeighborAddress.String())
			table.UpdatePathAttrs2ByteAs(m.Body.(*bgp.BGPUpdate))
		}

		peer.outgoing <- m
	}
}

func (peer *Peer) handleREST(restReq *api.RestRequest) {
	result := &api.RestResponse{}
	j, _ := json.Marshal(peer.rib.Tables[peer.rf])
	result.Data = j
	restReq.ResponseCh <- result
	close(restReq.ResponseCh)
}

func (peer *Peer) handlePeerMsg(m *peerMsg) {
	sendpath := func(pList []table.Path, wList []table.Path) {
		pathList := append([]table.Path(nil), pList...)
		pathList = append(pathList, wList...)

		for _, p := range wList {
			if !p.IsWithdraw() {
				log.Fatal("withdraw pathlist has non withdraw path")
			}
		}
		peer.adjRib.UpdateOut(pathList)
		peer.sendMessages(table.CreateUpdateMsgFromPaths(pathList))
	}

	switch m.msgType {
	case PEER_MSG_PATH:
		pList, wList, _ := peer.rib.ProcessPaths(m.msgData.([]table.Path))
		sendpath(pList, wList)
	case PEER_MSG_PEER_DOWN:
		pList, wList, _ := peer.rib.DeletePathsforPeer(m.msgData.(*table.PeerInfo))
		sendpath(pList, wList)
	}
}

func (peer *Peer) handleServerMsg(m *serverMsg) {
	switch m.msgType {
	case SRV_MSG_PEER_ADDED:
		d := m.msgData.(*serverMsgDataPeer)
		peer.siblings[d.address.String()] = d
		pathList := peer.adjRib.GetInPathList(d.rf)
		if len(pathList) == 0 {
			return
		}
		pm := &peerMsg{
			msgType: PEER_MSG_PATH,
			msgData: pathList,
		}
		for _, s := range peer.siblings {
			// TODO: check rf
			s.peerMsgCh <- pm
		}
	case SRV_MSG_PEER_DELETED:
		d := m.msgData.(*table.PeerInfo)
		_, found := peer.siblings[d.Address.String()]
		if found {
			delete(peer.siblings, d.Address.String())
			// TODO: do the same that PEER_MSG_PEER_DOWN handler
		} else {
			log.Warning("can not find peer: ", d.Address.String())
		}
	case SRV_MSG_API:
		peer.handleREST(m.msgData.(*api.RestRequest))
	default:
		log.Fatal("unknown server msg type ", m.msgType)
	}
}

// this goroutine handles routing table operations
func (peer *Peer) loop() error {
	for {
		h := NewFSMHandler(peer.fsm)
		sameState := true
		for sameState {
			select {
			case <-peer.t.Dying():
				close(peer.acceptedConnCh)
				h.Stop()
				close(peer.incoming)
				close(peer.outgoing)
				return nil
			case e := <-peer.incoming:
				switch e.MsgType {
				case FSM_MSG_STATE_CHANGE:
					nextState := e.MsgData.(bgp.FSMState)
					// waits for all goroutines created for the current state
					h.Wait()
					oldState := bgp.FSMState(peer.peerConfig.BgpNeighborCommonState.State)
					peer.peerConfig.BgpNeighborCommonState.State = uint32(nextState)
					peer.fsm.StateChange(nextState)
					sameState = false
					if nextState == bgp.BGP_FSM_ESTABLISHED {
						pathList := peer.adjRib.GetOutPathList(peer.rf)
						peer.sendMessages(table.CreateUpdateMsgFromPaths(pathList))
						peer.fsm.peerConfig.BgpNeighborCommonState.Uptime = time.Now()
						peer.fsm.peerConfig.BgpNeighborCommonState.EstablishedCount++
						if oldState >= bgp.BGP_FSM_OPENSENT {
							peer.peerInfo.VersionNum++
						}
					}
					if oldState == bgp.BGP_FSM_ESTABLISHED {
						peer.fsm.peerConfig.BgpNeighborCommonState.Uptime = time.Time{}
						pm := &peerMsg{
							msgType: PEER_MSG_PEER_DOWN,
							msgData: peer.peerInfo,
						}
						for _, s := range peer.siblings {
							s.peerMsgCh <- pm
						}
					}
				case FSM_MSG_BGP_MESSAGE:
					peer.handleBGPmessage(e.MsgData.(*bgp.BGPMessage))
				}
			case m := <-peer.serverMsgCh:
				peer.handleServerMsg(m)
			case m := <-peer.peerMsgCh:
				peer.handlePeerMsg(m)
			}
		}
	}
}

func (peer *Peer) Stop() error {
	peer.t.Kill(nil)
	return peer.t.Wait()
}

func (peer *Peer) PassConn(conn *net.TCPConn) {
	peer.acceptedConnCh <- conn
}

func (peer *Peer) MarshalJSON() ([]byte, error) {

	f := peer.fsm
	c := f.peerConfig

	p := make(map[string]interface{})

	p["conf"] = struct {
		RemoteIP string `json:"remote_ip"`
		Id       string `json:"id"`
		//Description        string `json:"description"`
		RemoteAS uint32 `json:"remote_as"`
		//LocalAddress       string `json:"local_address"`
		//LocalPort          int    `json:"local_port"`
		CapRefresh         bool `json:"cap_refresh"`
		CapEnhancedRefresh bool `json:"cap_enhanced_refresh"`
	}{
		RemoteIP: c.NeighborAddress.String(),
		Id:       peer.peerInfo.ID.To4().String(),
		//Description: "",
		RemoteAS: c.PeerAs,
		//LocalAddress:       f.passiveConn.LocalAddr().String(),
		//LocalPort:          f.passiveConn.LocalAddr().(*net.TCPAddr).Port,
		CapRefresh:         false,
		CapEnhancedRefresh: false,
	}

	s := c.BgpNeighborCommonState

	uptime := float64(0)
	if !s.Uptime.IsZero() {
		uptime = time.Now().Sub(s.Uptime).Seconds()
	}
	p["info"] = struct {
		BgpState                  string  `json:"bgp_state"`
		FsmEstablishedTransitions uint32  `json:"fsm_established_transitions"`
		TotalMessageOut           uint32  `json:"total_message_out"`
		TotalMessageIn            uint32  `json:"total_message_in"`
		UpdateMessageOut          uint32  `json:"update_message_out"`
		UpdateMessageIn           uint32  `json:"update_message_in"`
		KeepAliveMessageOut       uint32  `json:"keepalive_message_out"`
		KeepAliveMessageIn        uint32  `json:"keepalive_message_in"`
		OpenMessageOut            uint32  `json:"open_message_out"`
		OpenMessageIn             uint32  `json:"open_message_in"`
		NotificationOut           uint32  `json:"notification_out"`
		NotificationIn            uint32  `json:"notification_in"`
		RefreshMessageOut         uint32  `json:"refresh_message_out"`
		RefreshMessageIn          uint32  `json:"refresh_message_in"`
		Uptime                    float64 `json:"uptime"`
		LastError                 string  `json:"last_error"`
	}{

		BgpState:                  f.state.String(),
		FsmEstablishedTransitions: s.EstablishedCount,
		TotalMessageOut:           s.TotalOut,
		TotalMessageIn:            s.TotalIn,
		UpdateMessageOut:          s.UpdateOut,
		UpdateMessageIn:           s.UpdateIn,
		KeepAliveMessageOut:       s.KeepaliveOut,
		KeepAliveMessageIn:        s.KeepaliveIn,
		OpenMessageOut:            s.OpenOut,
		OpenMessageIn:             s.OpenIn,
		NotificationOut:           s.NotifyOut,
		NotificationIn:            s.NotifyIn,
		RefreshMessageOut:         s.RefreshOut,
		RefreshMessageIn:          s.RefreshIn,
		Uptime:                    uptime,
	}

	return json.Marshal(p)
}
<commit_msg>server: fix PEER_MSG_PATH about rf<commit_after>// Copyright (C) 2014 Nippon Telegraph and Telephone Corporation.
//
// 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 server

import (
	"encoding/json"
	log "github.com/Sirupsen/logrus"
	"github.com/osrg/gobgp/api"
	"github.com/osrg/gobgp/config"
	"github.com/osrg/gobgp/packet"
	"github.com/osrg/gobgp/table"
	"gopkg.in/tomb.v2"
	"net"
	"time"
)

type peerMsgType int

const (
	_ peerMsgType = iota
	PEER_MSG_PATH
	PEER_MSG_PEER_DOWN
)

type peerMsg struct {
	msgType peerMsgType
	msgData interface{}
}

type Peer struct {
	t              tomb.Tomb
	globalConfig   config.GlobalType
	peerConfig     config.NeighborType
	acceptedConnCh chan *net.TCPConn
	incoming       chan *fsmMsg
	outgoing       chan *bgp.BGPMessage
	serverMsgCh    chan *serverMsg
	peerMsgCh      chan *peerMsg
	fsm            *FSM
	adjRib         *table.AdjRib
	// peer and rib are always not one-to-one so should not be
	// here but it's the simplest and works our first target.
	rib *table.TableManager
	// for now we support only the same afi as transport
	rf       bgp.RouteFamily
	capMap   map[bgp.BGPCapabilityCode]bgp.ParameterCapabilityInterface
	peerInfo *table.PeerInfo
	siblings map[string]*serverMsgDataPeer
}

func NewPeer(g config.GlobalType, peer config.NeighborType, serverMsgCh chan *serverMsg, peerMsgCh chan *peerMsg, peerList []*serverMsgDataPeer) *Peer {
	p := &Peer{
		globalConfig:   g,
		peerConfig:     peer,
		acceptedConnCh: make(chan *net.TCPConn),
		incoming:       make(chan *fsmMsg, 4096),
		outgoing:       make(chan *bgp.BGPMessage, 4096),
		serverMsgCh:    serverMsgCh,
		peerMsgCh:      peerMsgCh,
		capMap:         make(map[bgp.BGPCapabilityCode]bgp.ParameterCapabilityInterface),
	}
	p.siblings = make(map[string]*serverMsgDataPeer)
	for _, s := range peerList {
		p.siblings[s.address.String()] = s
	}
	p.fsm = NewFSM(&g, &peer, p.acceptedConnCh, p.incoming, p.outgoing)
	peer.BgpNeighborCommonState.State = uint32(bgp.BGP_FSM_IDLE)
	if peer.NeighborAddress.To4() != nil {
		p.rf = bgp.RF_IPv4_UC
	} else {
		p.rf = bgp.RF_IPv6_UC
	}
	p.peerInfo = &table.PeerInfo{
		AS:         peer.PeerAs,
		VersionNum: 1,
		LocalID:    g.RouterId,
		RF:         p.rf,
		Address:    peer.NeighborAddress,
	}
	p.adjRib = table.NewAdjRib()
	p.rib = table.NewTableManager()
	p.t.Go(p.loop)
	return p
}

func (peer *Peer) handleBGPmessage(m *bgp.BGPMessage) {
	j, _ := json.Marshal(m)
	log.Debug(string(j))

	switch m.Header.Type {
	case bgp.BGP_MSG_OPEN:
		body := m.Body.(*bgp.BGPOpen)
		peer.peerInfo.ID = m.Body.(*bgp.BGPOpen).ID
		for _, p := range body.OptParams {
			paramCap, y := p.(*bgp.OptionParameterCapability)
			if !y {
				continue
			}
			for _, c := range paramCap.Capability {
				peer.capMap[c.Code()] = c
			}
		}

	case bgp.BGP_MSG_ROUTE_REFRESH:
		pathList := peer.adjRib.GetOutPathList(peer.rf)
		peer.sendMessages(table.CreateUpdateMsgFromPaths(pathList))
	case bgp.BGP_MSG_UPDATE:
		peer.peerConfig.BgpNeighborCommonState.UpdateRecvTime = time.Now()
		body := m.Body.(*bgp.BGPUpdate)
		table.UpdatePathAttrs4ByteAs(body)
		msg := table.NewProcessMessage(m, peer.peerInfo)
		pathList := msg.ToPathList()
		if len(pathList) == 0 {
			return
		}
		peer.adjRib.UpdateIn(pathList)
		pm := &peerMsg{
			msgType: PEER_MSG_PATH,
			msgData: pathList,
		}
		for _, s := range peer.siblings {
			if s.rf != peer.rf {
				continue
			}
			s.peerMsgCh <- pm
		}
	}
}

func (peer *Peer) sendMessages(msgs []*bgp.BGPMessage) {
	for _, m := range msgs {
		if peer.peerConfig.BgpNeighborCommonState.State != uint32(bgp.BGP_FSM_ESTABLISHED) {
			continue
		}

		if m.Header.Type != bgp.BGP_MSG_UPDATE {
			log.Fatal("not update message ", m.Header.Type)
		}

		_, y := peer.capMap[bgp.BGP_CAP_FOUR_OCTET_AS_NUMBER]
		if !y {
			log.Debug("update BGPUpdate for 2byte AS peer, ", peer.peerConfig.NeighborAddress.String())
			table.UpdatePathAttrs2ByteAs(m.Body.(*bgp.BGPUpdate))
		}

		peer.outgoing <- m
	}
}

func (peer *Peer) handleREST(restReq *api.RestRequest) {
	result := &api.RestResponse{}
	j, _ := json.Marshal(peer.rib.Tables[peer.rf])
	result.Data = j
	restReq.ResponseCh <- result
	close(restReq.ResponseCh)
}

func (peer *Peer) sendUpdateMsgFromPaths(pList []table.Path, wList []table.Path) {
	pathList := append([]table.Path(nil), pList...)
	pathList = append(pathList, wList...)

	for _, p := range wList {
		if !p.IsWithdraw() {
			log.Fatal("withdraw pathlist has non withdraw path")
		}
	}
	peer.adjRib.UpdateOut(pathList)
	peer.sendMessages(table.CreateUpdateMsgFromPaths(pathList))
}

func (peer *Peer) handlePeerMsg(m *peerMsg) {
	switch m.msgType {
	case PEER_MSG_PATH:
		pList, wList, _ := peer.rib.ProcessPaths(m.msgData.([]table.Path))
		peer.sendUpdateMsgFromPaths(pList, wList)
	case PEER_MSG_PEER_DOWN:
		pList, wList, _ := peer.rib.DeletePathsforPeer(m.msgData.(*table.PeerInfo))
		peer.sendUpdateMsgFromPaths(pList, wList)
	}
}

func (peer *Peer) handleServerMsg(m *serverMsg) {
	switch m.msgType {
	case SRV_MSG_PEER_ADDED:
		d := m.msgData.(*serverMsgDataPeer)
		peer.siblings[d.address.String()] = d
		pathList := peer.adjRib.GetInPathList(d.rf)
		if len(pathList) == 0 {
			return
		}
		pm := &peerMsg{
			msgType: PEER_MSG_PATH,
			msgData: pathList,
		}
		for _, s := range peer.siblings {
			if s.rf != peer.rf {
				continue
			}
			s.peerMsgCh <- pm
		}
	case SRV_MSG_PEER_DELETED:
		d := m.msgData.(*table.PeerInfo)
		_, found := peer.siblings[d.Address.String()]
		if found {
			delete(peer.siblings, d.Address.String())
			pList, wList, _ := peer.rib.DeletePathsforPeer(d)
			peer.sendUpdateMsgFromPaths(pList, wList)
		} else {
			log.Warning("can not find peer: ", d.Address.String())
		}
	case SRV_MSG_API:
		peer.handleREST(m.msgData.(*api.RestRequest))
	default:
		log.Fatal("unknown server msg type ", m.msgType)
	}
}

// this goroutine handles routing table operations
func (peer *Peer) loop() error {
	for {
		h := NewFSMHandler(peer.fsm)
		sameState := true
		for sameState {
			select {
			case <-peer.t.Dying():
				close(peer.acceptedConnCh)
				h.Stop()
				close(peer.incoming)
				close(peer.outgoing)
				return nil
			case e := <-peer.incoming:
				switch e.MsgType {
				case FSM_MSG_STATE_CHANGE:
					nextState := e.MsgData.(bgp.FSMState)
					// waits for all goroutines created for the current state
					h.Wait()
					oldState := bgp.FSMState(peer.peerConfig.BgpNeighborCommonState.State)
					peer.peerConfig.BgpNeighborCommonState.State = uint32(nextState)
					peer.fsm.StateChange(nextState)
					sameState = false
					if nextState == bgp.BGP_FSM_ESTABLISHED {
						pathList := peer.adjRib.GetOutPathList(peer.rf)
						peer.sendMessages(table.CreateUpdateMsgFromPaths(pathList))
						peer.fsm.peerConfig.BgpNeighborCommonState.Uptime = time.Now()
						peer.fsm.peerConfig.BgpNeighborCommonState.EstablishedCount++
						if oldState >= bgp.BGP_FSM_OPENSENT {
							peer.peerInfo.VersionNum++
						}
					}
					if oldState == bgp.BGP_FSM_ESTABLISHED {
						peer.fsm.peerConfig.BgpNeighborCommonState.Uptime = time.Time{}
						pm := &peerMsg{
							msgType: PEER_MSG_PEER_DOWN,
							msgData: peer.peerInfo,
						}
						for _, s := range peer.siblings {
							s.peerMsgCh <- pm
						}
					}
				case FSM_MSG_BGP_MESSAGE:
					peer.handleBGPmessage(e.MsgData.(*bgp.BGPMessage))
				}
			case m := <-peer.serverMsgCh:
				peer.handleServerMsg(m)
			case m := <-peer.peerMsgCh:
				peer.handlePeerMsg(m)
			}
		}
	}
}

func (peer *Peer) Stop() error {
	peer.t.Kill(nil)
	return peer.t.Wait()
}

func (peer *Peer) PassConn(conn *net.TCPConn) {
	peer.acceptedConnCh <- conn
}

func (peer *Peer) MarshalJSON() ([]byte, error) {

	f := peer.fsm
	c := f.peerConfig

	p := make(map[string]interface{})

	p["conf"] = struct {
		RemoteIP string `json:"remote_ip"`
		Id       string `json:"id"`
		//Description        string `json:"description"`
		RemoteAS uint32 `json:"remote_as"`
		//LocalAddress       string `json:"local_address"`
		//LocalPort          int    `json:"local_port"`
		CapRefresh         bool `json:"cap_refresh"`
		CapEnhancedRefresh bool `json:"cap_enhanced_refresh"`
	}{
		RemoteIP: c.NeighborAddress.String(),
		Id:       peer.peerInfo.ID.To4().String(),
		//Description: "",
		RemoteAS: c.PeerAs,
		//LocalAddress:       f.passiveConn.LocalAddr().String(),
		//LocalPort:          f.passiveConn.LocalAddr().(*net.TCPAddr).Port,
		CapRefresh:         false,
		CapEnhancedRefresh: false,
	}

	s := c.BgpNeighborCommonState

	uptime := float64(0)
	if !s.Uptime.IsZero() {
		uptime = time.Now().Sub(s.Uptime).Seconds()
	}
	p["info"] = struct {
		BgpState                  string  `json:"bgp_state"`
		FsmEstablishedTransitions uint32  `json:"fsm_established_transitions"`
		TotalMessageOut           uint32  `json:"total_message_out"`
		TotalMessageIn            uint32  `json:"total_message_in"`
		UpdateMessageOut          uint32  `json:"update_message_out"`
		UpdateMessageIn           uint32  `json:"update_message_in"`
		KeepAliveMessageOut       uint32  `json:"keepalive_message_out"`
		KeepAliveMessageIn        uint32  `json:"keepalive_message_in"`
		OpenMessageOut            uint32  `json:"open_message_out"`
		OpenMessageIn             uint32  `json:"open_message_in"`
		NotificationOut           uint32  `json:"notification_out"`
		NotificationIn            uint32  `json:"notification_in"`
		RefreshMessageOut         uint32  `json:"refresh_message_out"`
		RefreshMessageIn          uint32  `json:"refresh_message_in"`
		Uptime                    float64 `json:"uptime"`
		LastError                 string  `json:"last_error"`
	}{

		BgpState:                  f.state.String(),
		FsmEstablishedTransitions: s.EstablishedCount,
		TotalMessageOut:           s.TotalOut,
		TotalMessageIn:            s.TotalIn,
		UpdateMessageOut:          s.UpdateOut,
		UpdateMessageIn:           s.UpdateIn,
		KeepAliveMessageOut:       s.KeepaliveOut,
		KeepAliveMessageIn:        s.KeepaliveIn,
		OpenMessageOut:            s.OpenOut,
		OpenMessageIn:             s.OpenIn,
		NotificationOut:           s.NotifyOut,
		NotificationIn:            s.NotifyIn,
		RefreshMessageOut:         s.RefreshOut,
		RefreshMessageIn:          s.RefreshIn,
		Uptime:                    uptime,
	}

	return json.Marshal(p)
}
<|endoftext|>


<commit_before>package dht

import (
	"context"
	"crypto/ed25519"
	"fmt"
	"net"
	"testing"

	"github.com/anacrolix/dht/v2/bep44"
	"github.com/anacrolix/log"
	"github.com/stretchr/testify/require"
)

func TestPutGet(t *testing.T) {
	require := require.New(t)

	s1 := newServerFromPort(4580)
	s2 := newServerFromPort(4581)

	s2Addr := NewAddr(s2.Addr())

	immuItem, err := bep44.NewItem("Hello World! immu", nil, 1, 1, nil)
	require.NoError(err)

	// send get request to s2, we need a write token to put data
	qr := s1.Get(context.TODO(), s2Addr, immuItem.Target(), immuItem.Seq, QueryRateLimiting{})
	require.NoError(qr.ToError())
	require.NotNil(qr.Reply.R)
	require.NotNil(qr.Reply.R.Token)

	// send put request to s2
	qr = s1.Put(context.TODO(), s2Addr, immuItem, *qr.Reply.R.Token, QueryRateLimiting{})
	require.NoError(qr.ToError())

	qr = s1.Get(context.TODO(), s2Addr, immuItem.Target(), immuItem.Seq, QueryRateLimiting{})
	require.NoError(qr.ToError())
	require.Equal("Hello World! immu", qr.Reply.R.V)

	_, priv, err := ed25519.GenerateKey(nil)
	require.NoError(err)

	mutItem, err := bep44.NewItem("Hello World!", []byte("s1"), 1, 1, priv)
	require.NoError(err)

	// send get request to s2, we need a write token to put data
	qr = s1.Get(context.TODO(), s2Addr, mutItem.Target(), mutItem.Seq, QueryRateLimiting{})
	require.NoError(qr.ToError())
	require.NotNil(qr.Reply.R)

	mutToken := qr.Reply.R.Token
	require.NotNil(mutToken)

	// send put request to s2
	qr = s1.Put(context.TODO(), s2Addr, mutItem, *mutToken, QueryRateLimiting{})
	require.NoError(qr.ToError())

	qr = s1.Get(context.TODO(), s2Addr, mutItem.Target(), mutItem.Seq, QueryRateLimiting{})
	require.NoError(qr.ToError())
	require.Equal("Hello World!", qr.Reply.R.V)

	ii, err := s2.store.Get(immuItem.Target())
	require.NoError(err)
	require.Equal("Hello World! immu", ii.V)

	mi, err := s2.store.Get(mutItem.Target())
	require.NoError(err)
	require.Equal("Hello World!", mi.V)

	//change mutable item
	ok := mutItem.Modify("Bye World!", priv)
	require.True(ok)
	qr = s1.Put(context.TODO(), s2Addr, mutItem, *mutToken, QueryRateLimiting{})
	require.NoError(qr.ToError())

	mi, err = s2.store.Get(mutItem.Target())
	require.NoError(err)
	require.Equal("Bye World!", mi.V)

	qr = s1.Get(context.TODO(), s2Addr, mutItem.Target(), mutItem.Seq, QueryRateLimiting{})
	require.NoError(qr.ToError())
	require.Equal("Bye World!", qr.Reply.R.V)

	qr = s1.Get(context.TODO(), s2Addr, mutItem.Target(), 3, QueryRateLimiting{})
	require.NoError(qr.ToError())
	require.Nil(qr.Reply.R.V)
	require.Equal(int64(2), qr.Reply.R.Seq)
}

func newServerFromPort(port int) *Server {
	cfg := NewDefaultServerConfig()
	conn1, err := net.ListenPacket("udp", fmt.Sprintf("localhost:%d", port))
	if err != nil {
		panic(err)
	}

	cfg.Conn = conn1
	cfg.Logger = log.Default.FilterLevel(log.Debug)
	s, err := NewServer(cfg)
	if err != nil {
		panic(err)
	}

	return s
}
<commit_msg>Fix tests on circle CI. (#49)<commit_after>package dht

import (
	"context"
	"crypto/ed25519"
	"net"
	"testing"

	"github.com/anacrolix/dht/v2/bep44"
	"github.com/anacrolix/log"
	"github.com/stretchr/testify/require"
)

func TestPutGet(t *testing.T) {
	require := require.New(t)

	s1 := newServer(t)
	s2 := newServer(t)

	s2Addr := NewAddr(s2.Addr())

	immuItem, err := bep44.NewItem("Hello World! immu", nil, 1, 1, nil)
	require.NoError(err)

	// send get request to s2, we need a write token to put data
	qr := s1.Get(context.TODO(), s2Addr, immuItem.Target(), immuItem.Seq, QueryRateLimiting{})
	require.NoError(qr.ToError())
	require.NotNil(qr.Reply.R)
	require.NotNil(qr.Reply.R.Token)

	// send put request to s2
	qr = s1.Put(context.TODO(), s2Addr, immuItem, *qr.Reply.R.Token, QueryRateLimiting{})
	require.NoError(qr.ToError())

	qr = s1.Get(context.TODO(), s2Addr, immuItem.Target(), immuItem.Seq, QueryRateLimiting{})
	require.NoError(qr.ToError())
	require.Equal("Hello World! immu", qr.Reply.R.V)

	_, priv, err := ed25519.GenerateKey(nil)
	require.NoError(err)

	mutItem, err := bep44.NewItem("Hello World!", []byte("s1"), 1, 1, priv)
	require.NoError(err)

	// send get request to s2, we need a write token to put data
	qr = s1.Get(context.TODO(), s2Addr, mutItem.Target(), mutItem.Seq, QueryRateLimiting{})
	require.NoError(qr.ToError())
	require.NotNil(qr.Reply.R)

	mutToken := qr.Reply.R.Token
	require.NotNil(mutToken)

	// send put request to s2
	qr = s1.Put(context.TODO(), s2Addr, mutItem, *mutToken, QueryRateLimiting{})
	require.NoError(qr.ToError())

	qr = s1.Get(context.TODO(), s2Addr, mutItem.Target(), mutItem.Seq, QueryRateLimiting{})
	require.NoError(qr.ToError())
	require.Equal("Hello World!", qr.Reply.R.V)

	ii, err := s2.store.Get(immuItem.Target())
	require.NoError(err)
	require.Equal("Hello World! immu", ii.V)

	mi, err := s2.store.Get(mutItem.Target())
	require.NoError(err)
	require.Equal("Hello World!", mi.V)

	//change mutable item
	ok := mutItem.Modify("Bye World!", priv)
	require.True(ok)
	qr = s1.Put(context.TODO(), s2Addr, mutItem, *mutToken, QueryRateLimiting{})
	require.NoError(qr.ToError())

	mi, err = s2.store.Get(mutItem.Target())
	require.NoError(err)
	require.Equal("Bye World!", mi.V)

	qr = s1.Get(context.TODO(), s2Addr, mutItem.Target(), mutItem.Seq, QueryRateLimiting{})
	require.NoError(qr.ToError())
	require.Equal("Bye World!", qr.Reply.R.V)

	qr = s1.Get(context.TODO(), s2Addr, mutItem.Target(), 3, QueryRateLimiting{})
	require.NoError(qr.ToError())
	require.Nil(qr.Reply.R.V)
	require.Equal(int64(2), qr.Reply.R.Seq)
}

func newServer(t *testing.T) *Server {
	cfg := NewDefaultServerConfig()
	conn1, err := net.ListenPacket("udp", "localhost:0")
	if err != nil {
		panic(err)
	}

	cfg.Conn = conn1
	cfg.Logger = log.Default.FilterLevel(log.Debug)
	s, err := NewServer(cfg)
	if err != nil {
		panic(err)
	}

	t.Cleanup(func() {
		s.Close()
	})

	return s
}
<|endoftext|>


<commit_before>package sftpd

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"io"
	"net"
	"os"
	"strings"
	"testing"

	client "github.com/pkg/sftp"
	"github.com/taruti/sshutil"
	"golang.org/x/crypto/ssh"
)

const alnum = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

var testUser = "test"
var testPass = func() []byte {
	bs := make([]byte, 16)
	_, e := io.ReadFull(rand.Reader, bs)
	if e != nil {
		panic(e)
	}
	for i, b := range bs {
		bs[i] = alnum[int(b)%len(alnum)]
	}
	return bs
}()

func createKeyPEM() ([]byte, error) {
	k, e := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if e != nil {
		return nil, e
	}
	b, e := x509.MarshalECPrivateKey(k)
	if e != nil {
		return nil, e
	}
	return pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: b}), nil

}

func failOnErr(t *testing.T, err error, reason string) {
	if err != nil {
		t.Fatalf("%s: %v", reason, err)
	}
}

func PrintDiscardRequests(in <-chan *ssh.Request) {
	for req := range in {
		debug("PRINTDISC", req.Type, *req)
		if req.WantReply {
			req.Reply(false, nil)
		}
	}
}

func TestServer(t *testing.T) {
	debug = t.Log
	debugf = t.Logf
	debugf("Listening on port 2022 user %s pass %s\n", testUser, testPass)

	config := &ssh.ServerConfig{
		PasswordCallback: sshutil.CreatePasswordCheck(testUser, testPass),
	}

	hkey, e := sshutil.KeyLoader{Flags: sshutil.Create}.Load()
	failOnErr(t, e, "Failed to parse host key")
	debugf("Public key: %s\n", sshutil.PublicKeyHash(hkey.PublicKey()))

	config.AddHostKey(hkey)

	listener, e := net.Listen("tcp", "127.0.0.1:2022")
	failOnErr(t, e, "Failed to listen")

	go ClientDo()

	//	for {
	nConn, e := listener.Accept()
	failOnErr(t, e, "Failed to accept")
	handleConn(nConn, config, t, EmptyFS{})
	//	}

	go ClientDo()

	//	for {
	nConn, e = listener.Accept()
	failOnErr(t, e, "Failed to accept")
	os.Mkdir("/tmp/test-sftpd", 0700)
	handleConn(nConn, config, t, rfs{})
	//	}
}

func handleConn(conn net.Conn, config *ssh.ServerConfig, t *testing.T, fs FileSystem) {
	sc, chans, reqs, e := ssh.NewServerConn(conn, config)
	failOnErr(t, e, "Failed to initiate new connection")

	debug("sc", sc)

	// The incoming Request channel must be serviced.
	go PrintDiscardRequests(reqs)

	// Service the incoming Channel channel.
	for newChannel := range chans {
		debug("NEWCHANNEL", newChannel, newChannel.ChannelType())
		if newChannel.ChannelType() != "session" {
			newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}
		channel, requests, err := newChannel.Accept()
		if err != nil {
			panic("could not accept channel.")
		}

		go func(in <-chan *ssh.Request) {
			for req := range in {
				debug("REQUEST:", *req)
				ok := false
				switch {
				case IsSftpRequest(req):
					ok = true
					go func() { debug(ServeChannel(channel, fs)) }()
				}
				req.Reply(ok, nil)
			}
		}(requests)

	}
}

func ClientDo() {
	e := clientDo()
	if e != nil {
		debug("CLIENT ERROR", e)
	}
}
func clientDo() error {
	var cc ssh.ClientConfig
	cc.User = string(testUser)
	cc.Auth = append(cc.Auth, ssh.Password(string(testPass)))
	conn, e := ssh.Dial("tcp4", "127.0.0.1:2022", &cc)
	if e != nil {
		return e
	}
	defer conn.Close()
	cl, e := client.NewClient(conn)
	if e != nil {
		return e
	}
	rs, e := cl.ReadDir("/")
	if e != nil {
		return e
	}
	debug(rs)
	return nil
}

func TestRandomInput(t *testing.T) {
	fs := EmptyFS{}
	rd := &fakeRandChannel{}
	for i := 0; i < 10000; i++ {
		rd.rem = 5
		ServeChannel(rd, fs)
	}
	for i := 0; i < 257; i++ {
		for j := 0; j < 100; j++ {
			rd.rem = i
			ServeChannel(rd, fs)
		}
	}
}

type fakeRandChannel struct{ rem int }

func (fr *fakeRandChannel) Read(data []byte) (int, error) {
	if len(data) > fr.rem {
		data = data[0:fr.rem]
	}
	n, e := rand.Read(data)
	fr.rem -= n
	if fr.rem <= 0 {
		fr.rem = 0
		e = io.EOF
	}
	if len(data) >= 4 && data[0]&1 == 0 {
		data[0], data[1], data[2] = 0, 0, 0
		for i := 5; len(data) >= 4+i; i += 4 {
			data[i], data[i+1], data[i+2] = 0, 0, 0
		}
	}
	return n, e
}
func (*fakeRandChannel) Write(bs []byte) (int, error) {
	tmp := make([]byte, 1)
	rand.Read(tmp)
	if tmp[0]&7 == 7 {
		return 0, errors.New("Random write error")
	}
	return len(bs), nil
}
func (*fakeRandChannel) Close() error      { return nil }
func (*fakeRandChannel) CloseWrite() error { return nil }
func (*fakeRandChannel) SendRequest(name string, wantReply bool, payload []byte) (bool, error) {
	return true, nil
}
func (fr *fakeRandChannel) Stderr() io.ReadWriter { return fr }

type rfile struct {
	EmptyFile
	f *os.File
}

func (rf rfile) Close() error                             { return rf.f.Close() }
func (rf rfile) ReadAt(bs []byte, pos int64) (int, error) { return rf.f.ReadAt(bs, pos) }

type rfs struct {
	EmptyFS
}

type rdir struct {
	d *os.File
}

func (d rdir) Readdir(count int) ([]NamedAttr, error) {
	fis, e := d.d.Readdir(count)
	if e != nil {
		return nil, e
	}
	nas := make([]NamedAttr, len(fis))
	for i, fi := range fis {
		nas[i].Name = fi.Name()
		nas[i].FillFrom(fi)
	}
	return nas, nil
}
func (d rdir) Close() error {
	return d.d.Close()
}

// Warning:
// Use your own path mangling functionality in production code.
// This can be quite non-trivial depending on the operating system.
// The code below is not sufficient for production servers.
func rfsMangle(path string) (string, error) {
	if strings.Contains(path, "..") {
		return "<invalid>", errors.New("Invalid path")
	}
	if len(path) > 0 && path[0] == '/' {
		path = path[1:]
	}
	path = "/tmp/test-sftpd/" + path
	debug("MANGLE -> " + path)
	return path, nil
}

func (fs rfs) OpenDir(path string) (Dir, error) {
	p, e := rfsMangle(path)
	if e != nil {
		return nil, e
	}
	f, e := os.Open(p)
	if e != nil {
		return nil, e
	}
	return rdir{f}, nil
}

func (fs rfs) OpenFile(path string, mode uint32, a *Attr) (File, error) {
	p, e := rfsMangle(path)
	if e != nil {
		return nil, e
	}
	f, e := os.Open(p)
	if e != nil {
		return nil, e
	}
	return rfile{f: f}, nil
}

func (fs rfs) Stat(name string, islstat bool) (*Attr, error) {
	p, e := rfsMangle(name)
	if e != nil {
		return nil, e
	}
	var fi os.FileInfo
	if islstat {
		fi, e = os.Lstat(p)
	} else {
		fi, e = os.Stat(p)
	}
	if e != nil {
		return nil, e
	}
	var a Attr
	e = a.FillFrom(fi)

	return &a, e
}
<commit_msg>More food for coverage<commit_after>package sftpd

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"io"
	"net"
	"os"
	"strings"
	"testing"

	client "github.com/pkg/sftp"
	"github.com/taruti/sshutil"
	"golang.org/x/crypto/ssh"
)

const alnum = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

var testUser = "test"
var testPass = func() []byte {
	bs := make([]byte, 16)
	_, e := io.ReadFull(rand.Reader, bs)
	if e != nil {
		panic(e)
	}
	for i, b := range bs {
		bs[i] = alnum[int(b)%len(alnum)]
	}
	return bs
}()

func createKeyPEM() ([]byte, error) {
	k, e := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if e != nil {
		return nil, e
	}
	b, e := x509.MarshalECPrivateKey(k)
	if e != nil {
		return nil, e
	}
	return pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: b}), nil

}

func failOnErr(t *testing.T, err error, reason string) {
	if err != nil {
		t.Fatalf("%s: %v", reason, err)
	}
}

func PrintDiscardRequests(in <-chan *ssh.Request) {
	for req := range in {
		debug("PRINTDISC", req.Type, *req)
		if req.WantReply {
			req.Reply(false, nil)
		}
	}
}

func TestServer(t *testing.T) {
	debug = t.Log
	debugf = t.Logf
	debugf("Listening on port 2022 user %s pass %s\n", testUser, testPass)

	config := &ssh.ServerConfig{
		PasswordCallback: sshutil.CreatePasswordCheck(testUser, testPass),
	}

	hkey, e := sshutil.KeyLoader{Flags: sshutil.Create}.Load()
	failOnErr(t, e, "Failed to parse host key")
	debugf("Public key: %s\n", sshutil.PublicKeyHash(hkey.PublicKey()))

	config.AddHostKey(hkey)

	listener, e := net.Listen("tcp", "127.0.0.1:2022")
	failOnErr(t, e, "Failed to listen")

	go ClientDo()

	//	for {
	nConn, e := listener.Accept()
	failOnErr(t, e, "Failed to accept")
	handleConn(nConn, config, t, EmptyFS{})
	//	}

	go ClientDo()

	//	for {
	nConn, e = listener.Accept()
	failOnErr(t, e, "Failed to accept")
	os.Mkdir("/tmp/test-sftpd", 0700)
	handleConn(nConn, config, t, rfs{})
	//	}
}

func handleConn(conn net.Conn, config *ssh.ServerConfig, t *testing.T, fs FileSystem) {
	sc, chans, reqs, e := ssh.NewServerConn(conn, config)
	failOnErr(t, e, "Failed to initiate new connection")

	debug("sc", sc)

	// The incoming Request channel must be serviced.
	go PrintDiscardRequests(reqs)

	// Service the incoming Channel channel.
	for newChannel := range chans {
		debug("NEWCHANNEL", newChannel, newChannel.ChannelType())
		if newChannel.ChannelType() != "session" {
			newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}
		channel, requests, err := newChannel.Accept()
		if err != nil {
			panic("could not accept channel.")
		}

		go func(in <-chan *ssh.Request) {
			for req := range in {
				debug("REQUEST:", *req)
				ok := false
				switch {
				case IsSftpRequest(req):
					ok = true
					go func() { debug(ServeChannel(channel, fs)) }()
				}
				req.Reply(ok, nil)
			}
		}(requests)

	}
}

func ClientDo() {
	e := clientDo()
	if e != nil {
		debug("CLIENT ERROR", e)
	}
}
func clientDo() error {
	var cc ssh.ClientConfig
	cc.User = string(testUser)
	cc.Auth = append(cc.Auth, ssh.Password(string(testPass)))
	conn, e := ssh.Dial("tcp4", "127.0.0.1:2022", &cc)
	if e != nil {
		return e
	}
	defer conn.Close()
	cl, e := client.NewClient(conn)
	if e != nil {
		return e
	}
	cl.Mkdir("/D")
	cl.Chmod("/D", 0700)
	cl.Remove("/D")
	rs, e := cl.ReadDir("/")
	if e != nil {
		return e
	}
	debug(rs)
	return nil
}

func TestRandomInput(t *testing.T) {
	fs := EmptyFS{}
	rd := &fakeRandChannel{}
	for i := 0; i < 10000; i++ {
		rd.rem = 5
		ServeChannel(rd, fs)
	}
	for i := 0; i < 257; i++ {
		for j := 0; j < 1000; j++ {
			rd.rem = i
			ServeChannel(rd, fs)
		}
	}
}

type fakeRandChannel struct{ rem int }

func (fr *fakeRandChannel) Read(data []byte) (int, error) {
	if len(data) > fr.rem {
		data = data[0:fr.rem]
	}
	n, e := rand.Read(data)
	fr.rem -= n
	if fr.rem <= 0 {
		fr.rem = 0
		e = io.EOF
	}
	if len(data) >= 4 && data[0]&1 == 0 {
		data[0], data[1], data[2] = 0, 0, 0
		for i := 5; len(data) >= 4+i; i += 4 {
			data[i], data[i+1], data[i+2] = 0, 0, 0
		}
	}
	return n, e
}
func (*fakeRandChannel) Write(bs []byte) (int, error) {
	tmp := make([]byte, 1)
	rand.Read(tmp)
	if tmp[0]&7 == 7 {
		return 0, errors.New("Random write error")
	}
	return len(bs), nil
}
func (*fakeRandChannel) Close() error      { return nil }
func (*fakeRandChannel) CloseWrite() error { return nil }
func (*fakeRandChannel) SendRequest(name string, wantReply bool, payload []byte) (bool, error) {
	return true, nil
}
func (fr *fakeRandChannel) Stderr() io.ReadWriter { return fr }

type rfile struct {
	EmptyFile
	f *os.File
}

func (rf rfile) Close() error                             { return rf.f.Close() }
func (rf rfile) ReadAt(bs []byte, pos int64) (int, error) { return rf.f.ReadAt(bs, pos) }

type rfs struct {
	EmptyFS
}

type rdir struct {
	d *os.File
}

func (d rdir) Readdir(count int) ([]NamedAttr, error) {
	fis, e := d.d.Readdir(count)
	if e != nil {
		return nil, e
	}
	nas := make([]NamedAttr, len(fis))
	for i, fi := range fis {
		nas[i].Name = fi.Name()
		nas[i].FillFrom(fi)
	}
	return nas, nil
}
func (d rdir) Close() error {
	return d.d.Close()
}

// Warning:
// Use your own path mangling functionality in production code.
// This can be quite non-trivial depending on the operating system.
// The code below is not sufficient for production servers.
func rfsMangle(path string) (string, error) {
	if strings.Contains(path, "..") {
		return "<invalid>", errors.New("Invalid path")
	}
	if len(path) > 0 && path[0] == '/' {
		path = path[1:]
	}
	path = "/tmp/test-sftpd/" + path
	debug("MANGLE -> " + path)
	return path, nil
}

func (fs rfs) OpenDir(path string) (Dir, error) {
	p, e := rfsMangle(path)
	if e != nil {
		return nil, e
	}
	f, e := os.Open(p)
	if e != nil {
		return nil, e
	}
	return rdir{f}, nil
}

func (fs rfs) OpenFile(path string, mode uint32, a *Attr) (File, error) {
	p, e := rfsMangle(path)
	if e != nil {
		return nil, e
	}
	f, e := os.Open(p)
	if e != nil {
		return nil, e
	}
	return rfile{f: f}, nil
}

func (fs rfs) Stat(name string, islstat bool) (*Attr, error) {
	p, e := rfsMangle(name)
	if e != nil {
		return nil, e
	}
	var fi os.FileInfo
	if islstat {
		fi, e = os.Lstat(p)
	} else {
		fi, e = os.Stat(p)
	}
	if e != nil {
		return nil, e
	}
	var a Attr
	e = a.FillFrom(fi)

	return &a, e
}
<|endoftext|>


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

package xorm

import "database/sql"

// Count counts the records. bean's non-empty fields
// are conditions.
func (session *Session) Count(bean interface{}) (int64, error) {
	defer session.resetStatement()
	if session.IsAutoClose {
		defer session.Close()
	}

	var sqlStr string
	var args []interface{}
	if session.Statement.RawSQL == "" {
		sqlStr, args = session.Statement.genCountSQL(bean)
	} else {
		sqlStr = session.Statement.RawSQL
		args = session.Statement.RawParams
	}

	session.queryPreprocess(&sqlStr, args...)

	var err error
	var total int64
	if session.IsAutoCommit {
		err = session.DB().QueryRow(sqlStr, args...).Scan(&total)
	} else {
		err = session.Tx.QueryRow(sqlStr, args...).Scan(&total)
	}

	if err == sql.ErrNoRows || err == nil {
		return total, nil
	}

	return 0, err
}

// Sum call sum some column. bean's non-empty fields are conditions.
func (session *Session) Sum(bean interface{}, columnName string) (float64, error) {
	defer session.resetStatement()
	if session.IsAutoClose {
		defer session.Close()
	}

	var sqlStr string
	var args []interface{}
	if len(session.Statement.RawSQL) == 0 {
		sqlStr, args = session.Statement.genSumSQL(bean, columnName)
	} else {
		sqlStr = session.Statement.RawSQL
		args = session.Statement.RawParams
	}

	session.queryPreprocess(&sqlStr, args...)

	var err error
	var res float64
	if session.IsAutoCommit {
		err = session.DB().QueryRow(sqlStr, args...).Scan(&res)
	} else {
		err = session.Tx.QueryRow(sqlStr, args...).Scan(&res)
	}

	if err == sql.ErrNoRows || err == nil {
		return res, nil
	}
	return 0, err
}

// Sums call sum some columns. bean's non-empty fields are conditions.
func (session *Session) Sums(bean interface{}, columnNames ...string) ([]float64, error) {
	defer session.resetStatement()
	if session.IsAutoClose {
		defer session.Close()
	}

	var sqlStr string
	var args []interface{}
	if len(session.Statement.RawSQL) == 0 {
		sqlStr, args = session.Statement.genSumSQL(bean, columnNames...)
	} else {
		sqlStr = session.Statement.RawSQL
		args = session.Statement.RawParams
	}

	session.queryPreprocess(&sqlStr, args...)

	var err error
	var res = make([]float64, len(columnNames), len(columnNames))
	if session.IsAutoCommit {
		err = session.DB().QueryRow(sqlStr, args...).ScanSlice(&res)
	} else {
		err = session.Tx.QueryRow(sqlStr, args...).ScanSlice(&res)
	}

	if err == sql.ErrNoRows || err == nil {
		return res, nil
	}
	return nil, err
}

// SumsInt sum specify columns and return as []int64 instead of []float64
func (session *Session) SumsInt(bean interface{}, columnNames ...string) ([]int64, error) {
	defer session.resetStatement()
	if session.IsAutoClose {
		defer session.Close()
	}

	var sqlStr string
	var args []interface{}
	if len(session.Statement.RawSQL) == 0 {
		sqlStr, args = session.Statement.genSumSQL(bean, columnNames...)
	} else {
		sqlStr = session.Statement.RawSQL
		args = session.Statement.RawParams
	}

	session.queryPreprocess(&sqlStr, args...)

	var err error
	var res = make([]int64, 0, len(columnNames))
	if session.IsAutoCommit {
		err = session.DB().QueryRow(sqlStr, args...).ScanSlice(&res)
	} else {
		err = session.Tx.QueryRow(sqlStr, args...).ScanSlice(&res)
	}

	if err == sql.ErrNoRows || err == nil {
		return res, nil
	}
	return nil, err
}
<commit_msg>bug fix for SumsInt return empty slice (#542)<commit_after>// Copyright 2016 The Xorm Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package xorm

import "database/sql"

// Count counts the records. bean's non-empty fields
// are conditions.
func (session *Session) Count(bean interface{}) (int64, error) {
	defer session.resetStatement()
	if session.IsAutoClose {
		defer session.Close()
	}

	var sqlStr string
	var args []interface{}
	if session.Statement.RawSQL == "" {
		sqlStr, args = session.Statement.genCountSQL(bean)
	} else {
		sqlStr = session.Statement.RawSQL
		args = session.Statement.RawParams
	}

	session.queryPreprocess(&sqlStr, args...)

	var err error
	var total int64
	if session.IsAutoCommit {
		err = session.DB().QueryRow(sqlStr, args...).Scan(&total)
	} else {
		err = session.Tx.QueryRow(sqlStr, args...).Scan(&total)
	}

	if err == sql.ErrNoRows || err == nil {
		return total, nil
	}

	return 0, err
}

// Sum call sum some column. bean's non-empty fields are conditions.
func (session *Session) Sum(bean interface{}, columnName string) (float64, error) {
	defer session.resetStatement()
	if session.IsAutoClose {
		defer session.Close()
	}

	var sqlStr string
	var args []interface{}
	if len(session.Statement.RawSQL) == 0 {
		sqlStr, args = session.Statement.genSumSQL(bean, columnName)
	} else {
		sqlStr = session.Statement.RawSQL
		args = session.Statement.RawParams
	}

	session.queryPreprocess(&sqlStr, args...)

	var err error
	var res float64
	if session.IsAutoCommit {
		err = session.DB().QueryRow(sqlStr, args...).Scan(&res)
	} else {
		err = session.Tx.QueryRow(sqlStr, args...).Scan(&res)
	}

	if err == sql.ErrNoRows || err == nil {
		return res, nil
	}
	return 0, err
}

// Sums call sum some columns. bean's non-empty fields are conditions.
func (session *Session) Sums(bean interface{}, columnNames ...string) ([]float64, error) {
	defer session.resetStatement()
	if session.IsAutoClose {
		defer session.Close()
	}

	var sqlStr string
	var args []interface{}
	if len(session.Statement.RawSQL) == 0 {
		sqlStr, args = session.Statement.genSumSQL(bean, columnNames...)
	} else {
		sqlStr = session.Statement.RawSQL
		args = session.Statement.RawParams
	}

	session.queryPreprocess(&sqlStr, args...)

	var err error
	var res = make([]float64, len(columnNames), len(columnNames))
	if session.IsAutoCommit {
		err = session.DB().QueryRow(sqlStr, args...).ScanSlice(&res)
	} else {
		err = session.Tx.QueryRow(sqlStr, args...).ScanSlice(&res)
	}

	if err == sql.ErrNoRows || err == nil {
		return res, nil
	}
	return nil, err
}

// SumsInt sum specify columns and return as []int64 instead of []float64
func (session *Session) SumsInt(bean interface{}, columnNames ...string) ([]int64, error) {
	defer session.resetStatement()
	if session.IsAutoClose {
		defer session.Close()
	}

	var sqlStr string
	var args []interface{}
	if len(session.Statement.RawSQL) == 0 {
		sqlStr, args = session.Statement.genSumSQL(bean, columnNames...)
	} else {
		sqlStr = session.Statement.RawSQL
		args = session.Statement.RawParams
	}

	session.queryPreprocess(&sqlStr, args...)

	var err error
	var res = make([]int64, len(columnNames), len(columnNames))
	if session.IsAutoCommit {
		err = session.DB().QueryRow(sqlStr, args...).ScanSlice(&res)
	} else {
		err = session.Tx.QueryRow(sqlStr, args...).ScanSlice(&res)
	}

	if err == sql.ErrNoRows || err == nil {
		return res, nil
	}
	return nil, err
}
<|endoftext|>


<commit_before>package crypto

// encrypt.go contains functions for encrypting and decrypting data byte slices
// and readers.

import (
	"crypto/cipher"
	"crypto/rand"
	"encoding/json"
	"errors"
	"io"

	"golang.org/x/crypto/twofish"
)

const (
	TwofishOverhead = 28 // number of bytes added by EncryptBytes
)

var (
	ErrInsufficientLen = errors.New("supplied ciphertext is not long enough to contain a nonce")
)

type (
	Ciphertext []byte
	TwofishKey [EntropySize]byte
)

// GenerateEncryptionKey produces a key that can be used for encrypting and
// decrypting files.
func GenerateTwofishKey() (key TwofishKey, err error) {
	_, err = rand.Read(key[:])
	return key, err
}

// NewCipher creates a new Twofish cipher from the key.
func (key TwofishKey) NewCipher() cipher.Block {
	// NOTE: NewCipher only returns an error if len(key) != 16, 24, or 32.
	//       Assumption tested in TestTwofishNewCipherAssumption.
	cipher, _ := twofish.NewCipher(key[:])
	return cipher
}

// EncryptBytes encrypts a []byte using the key. EncryptBytes uses GCM and
// prepends the nonce (12 bytes) to the ciphertext.
func (key TwofishKey) EncryptBytes(plaintext []byte) (Ciphertext, error) {
	// Create the cipher.
	// NOTE: NewGCM only returns an error if twofishCipher.BlockSize != 16.
	//       Assumption tested in TestCipherNewGCMAssumption.
	aead, _ := cipher.NewGCM(key.NewCipher())

	// Create the nonce.
	nonce, err := RandBytes(aead.NonceSize())
	if err != nil {
		return nil, err
	}

	// Encrypt the data. No authenticated data is provided, as EncryptBytes is
	// meant for file encryption.
	return aead.Seal(nonce, nonce, plaintext, nil), nil
}

// DecryptBytes decrypts the ciphertext created by EncryptBytes. The nonce is
// expected to be the first 12 bytes of the ciphertext.
func (key TwofishKey) DecryptBytes(ct Ciphertext) ([]byte, error) {
	// Create the cipher.
	// NOTE: NewGCM only returns an error if twofishCipher.BlockSize != 16.
	//       Assumption tested in TestCipherNewGCMAssumption.
	aead, _ := cipher.NewGCM(key.NewCipher())

	// Check for a nonce.
	if len(ct) < aead.NonceSize() {
		return nil, ErrInsufficientLen
	}

	// Decrypt the data.
	return aead.Open(nil, ct[:aead.NonceSize()], ct[aead.NonceSize():], nil)
}

// NewWriter returns a writer that encrypts or decrypts its input stream.
func (key TwofishKey) NewWriter(w io.Writer) io.Writer {
	// OK to use a zero IV if the key is unique for each ciphertext.
	iv := make([]byte, twofish.BlockSize)
	stream := cipher.NewOFB(key.NewCipher(), iv)

	return &cipher.StreamWriter{S: stream, W: w}
}

// NewReader returns a reader that encrypts or decrypts its input stream.
func (key TwofishKey) NewReader(r io.Reader) io.Reader {
	// OK to use a zero IV if the key is unique for each ciphertext.
	iv := make([]byte, twofish.BlockSize)
	stream := cipher.NewOFB(key.NewCipher(), iv)

	return &cipher.StreamReader{S: stream, R: r}
}

func (c Ciphertext) MarshalJSON() ([]byte, error) {
	return json.Marshal([]byte(c))
}

func (c *Ciphertext) UnmarshalJSON(b []byte) error {
	var umarB []byte
	err := json.Unmarshal(b, &umarB)
	if err != nil {
		return err
	}
	*c = Ciphertext(umarB)
	return nil
}
<commit_msg>Remove unnecessary comments<commit_after>package crypto

// encrypt.go contains functions for encrypting and decrypting data byte slices
// and readers.

import (
	"crypto/cipher"
	"crypto/rand"
	"encoding/json"
	"errors"
	"io"

	"golang.org/x/crypto/twofish"
)

const (
	TwofishOverhead = 28 // number of bytes added by EncryptBytes
)

var (
	ErrInsufficientLen = errors.New("supplied ciphertext is not long enough to contain a nonce")
)

type (
	Ciphertext []byte
	TwofishKey [EntropySize]byte
)

// GenerateEncryptionKey produces a key that can be used for encrypting and
// decrypting files.
func GenerateTwofishKey() (key TwofishKey, err error) {
	_, err = rand.Read(key[:])
	return key, err
}

// NewCipher creates a new Twofish cipher from the key.
func (key TwofishKey) NewCipher() cipher.Block {
	// NOTE: NewCipher only returns an error if len(key) != 16, 24, or 32.
	cipher, _ := twofish.NewCipher(key[:])
	return cipher
}

// EncryptBytes encrypts a []byte using the key. EncryptBytes uses GCM and
// prepends the nonce (12 bytes) to the ciphertext.
func (key TwofishKey) EncryptBytes(plaintext []byte) (Ciphertext, error) {
	// Create the cipher.
	// NOTE: NewGCM only returns an error if twofishCipher.BlockSize != 16.
	aead, _ := cipher.NewGCM(key.NewCipher())

	// Create the nonce.
	nonce, err := RandBytes(aead.NonceSize())
	if err != nil {
		return nil, err
	}

	// Encrypt the data. No authenticated data is provided, as EncryptBytes is
	// meant for file encryption.
	return aead.Seal(nonce, nonce, plaintext, nil), nil
}

// DecryptBytes decrypts the ciphertext created by EncryptBytes. The nonce is
// expected to be the first 12 bytes of the ciphertext.
func (key TwofishKey) DecryptBytes(ct Ciphertext) ([]byte, error) {
	// Create the cipher.
	// NOTE: NewGCM only returns an error if twofishCipher.BlockSize != 16.
	aead, _ := cipher.NewGCM(key.NewCipher())

	// Check for a nonce.
	if len(ct) < aead.NonceSize() {
		return nil, ErrInsufficientLen
	}

	// Decrypt the data.
	return aead.Open(nil, ct[:aead.NonceSize()], ct[aead.NonceSize():], nil)
}

// NewWriter returns a writer that encrypts or decrypts its input stream.
func (key TwofishKey) NewWriter(w io.Writer) io.Writer {
	// OK to use a zero IV if the key is unique for each ciphertext.
	iv := make([]byte, twofish.BlockSize)
	stream := cipher.NewOFB(key.NewCipher(), iv)

	return &cipher.StreamWriter{S: stream, W: w}
}

// NewReader returns a reader that encrypts or decrypts its input stream.
func (key TwofishKey) NewReader(r io.Reader) io.Reader {
	// OK to use a zero IV if the key is unique for each ciphertext.
	iv := make([]byte, twofish.BlockSize)
	stream := cipher.NewOFB(key.NewCipher(), iv)

	return &cipher.StreamReader{S: stream, R: r}
}

func (c Ciphertext) MarshalJSON() ([]byte, error) {
	return json.Marshal([]byte(c))
}

func (c *Ciphertext) UnmarshalJSON(b []byte) error {
	var umarB []byte
	err := json.Unmarshal(b, &umarB)
	if err != nil {
		return err
	}
	*c = Ciphertext(umarB)
	return nil
}
<|endoftext|>


<commit_before>package csiplugin

import (
	"fmt"
	"path"
	"reflect"
	"sync"

	"golang.org/x/net/context"

	"google.golang.org/grpc"

	"code.cloudfoundry.org/goshims/grpcshim"
	"code.cloudfoundry.org/goshims/osshim"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/volman"
	"code.cloudfoundry.org/csishim"
	"github.com/container-storage-interface/spec/lib/go/csi"
)

func csiVersion() *csi.Version {
	return &csi.Version{
		Major: 0,
		Minor: 1,
		Patch: 0,
	}
}

type nodeWrapper struct {
	Impl            interface{}
	Spec            volman.PluginSpec
	grpcShim        grpcshim.Grpc
	csiShim         csishim.Csi
	osShim          osshim.Os
	volumes         map[string]int
	volumesMutex    sync.RWMutex
	csiMountRootDir string
}

func (dw *nodeWrapper) Unmount(logger lager.Logger, volumeId string) error {
	logger = logger.Session("unmount")
	logger.Info("start")
	defer logger.Info("end")

	conn, err := dw.grpcShim.Dial(dw.Spec.Address, grpc.WithInsecure())
	defer conn.Close()
	if err != nil {
		logger.Error("grpc-dial", err, lager.Data{"address": dw.Spec.Address})
		return err
	}

	mountPath := path.Join(dw.csiMountRootDir, dw.Spec.Name)
	targetPath := path.Join(mountPath, volumeId)

	nodePlugin := dw.csiShim.NewNodeClient(conn)
	nodeResponse, err := nodePlugin.NodeUnpublishVolume(context.TODO(), &csi.NodeUnpublishVolumeRequest{
		Version:    csiVersion(),
		VolumeId:   volumeId,
		TargetPath: targetPath,
	})

	logger.Debug(fmt.Sprintf("nodeResponse: %#v", nodeResponse))

	if err != nil {
		logger.Error("node-response-error", err)
		return err
	}

	dw.volumesMutex.Lock()
	defer dw.volumesMutex.Unlock()

	if count, ok := dw.volumes[volumeId]; ok {
		if count > 1 {
			dw.volumes[volumeId] = count - 1
		} else {
			delete(dw.volumes, volumeId)
		}
	}
	logger.Debug(fmt.Sprintf("reference count: %#v", dw.volumes))

	return nil
}

func (dw *nodeWrapper) Mount(logger lager.Logger, volumeId string, config map[string]interface{}) (volman.MountResponse, error) {
	logger = logger.Session("mount")
	logger.Info("start")
	defer logger.Info("end")

	conn, err := dw.grpcShim.Dial(dw.Spec.Address, grpc.WithInsecure())
	defer conn.Close()
	if err != nil {
		logger.Error("grpc-dial", err, lager.Data{"address": dw.Spec.Address})
		return volman.MountResponse{}, err
	}

	mountPath := path.Join(dw.csiMountRootDir, dw.Spec.Name)

	targetPath := path.Join(mountPath, volumeId)

	nodePlugin := dw.csiShim.NewNodeClient(conn)
	nodeResponse, err := nodePlugin.NodePublishVolume(context.TODO(), &csi.NodePublishVolumeRequest{
		Version:    csiVersion(),
		VolumeId:   volumeId,
		TargetPath: targetPath,
		VolumeCapability: &csi.VolumeCapability{
			AccessType: &csi.VolumeCapability_Mount{
				Mount: &csi.VolumeCapability_MountVolume{},
			},
			AccessMode: &csi.VolumeCapability_AccessMode{
				csi.VolumeCapability_AccessMode_MULTI_NODE_MULTI_WRITER,
			},
		},
	})

	logger.Debug(fmt.Sprintf("nodeResponse: %#v", nodeResponse))

	if err != nil {
		logger.Error("node-response-error", err)
		return volman.MountResponse{}, err
	}

	dw.volumesMutex.Lock()
	defer dw.volumesMutex.Unlock()

	if count, ok := dw.volumes[volumeId]; ok {
		dw.volumes[volumeId] = count + 1
	} else {
		dw.volumes[volumeId] = 1
	}

	logger.Debug(fmt.Sprintf("reference count: %#v", dw.volumes))
	return volman.MountResponse{Path: targetPath}, nil
}

func (dw *nodeWrapper) ListVolumes(logger lager.Logger) ([]string, error) {
	logger = logger.Session("listvolumes")
	logger.Info("start")
	defer logger.Info("end")

	logger.Debug(fmt.Sprintf("reference count: %#v", dw.volumes))

	dw.volumesMutex.RLock()
	defer dw.volumesMutex.RUnlock()

	volumes := make([]string, len(dw.volumes))
	i := 0
	for volume := range dw.volumes {
		volumes[i] = volume
		i++
	}
	return volumes, nil
}

func (dw *nodeWrapper) Matches(logger lager.Logger, otherSpec volman.PluginSpec) bool {
	logger = logger.Session("matches")
	logger.Info("start")
	defer logger.Info("end")

	matches := reflect.DeepEqual(dw.Spec, otherSpec)
	logger.Info("matches", lager.Data{"matches": matches})
	return matches
}

func NewCsiPlugin(plugin csi.NodeClient, pluginSpec volman.PluginSpec, grpcShim grpcshim.Grpc, csiShim csishim.Csi, osShim osshim.Os, csiMountRootDir string) volman.Plugin {
	return &nodeWrapper{
		Impl:            plugin,
		Spec:            pluginSpec,
		grpcShim:        grpcShim,
		csiShim:         csiShim,
		osShim:          osShim,
		volumes:         map[string]int{},
		csiMountRootDir: csiMountRootDir,
	}
}
<commit_msg>clean up lager logging formats<commit_after>package csiplugin

import (
	"path"
	"reflect"
	"sync"

	"golang.org/x/net/context"

	"google.golang.org/grpc"

	"code.cloudfoundry.org/goshims/grpcshim"
	"code.cloudfoundry.org/goshims/osshim"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/volman"
	"code.cloudfoundry.org/csishim"
	"github.com/container-storage-interface/spec/lib/go/csi"
)

func csiVersion() *csi.Version {
	return &csi.Version{
		Major: 0,
		Minor: 1,
		Patch: 0,
	}
}

type nodeWrapper struct {
	Impl            interface{}
	Spec            volman.PluginSpec
	grpcShim        grpcshim.Grpc
	csiShim         csishim.Csi
	osShim          osshim.Os
	volumes         map[string]int
	volumesMutex    sync.RWMutex
	csiMountRootDir string
}

func (dw *nodeWrapper) Unmount(logger lager.Logger, volumeId string) error {
	logger = logger.Session("unmount")
	logger.Info("start")
	defer logger.Info("end")

	conn, err := dw.grpcShim.Dial(dw.Spec.Address, grpc.WithInsecure())
	defer conn.Close()
	if err != nil {
		logger.Error("grpc-dial", err, lager.Data{"address": dw.Spec.Address})
		return err
	}

	mountPath := path.Join(dw.csiMountRootDir, dw.Spec.Name)
	targetPath := path.Join(mountPath, volumeId)

	nodePlugin := dw.csiShim.NewNodeClient(conn)
	nodeResponse, err := nodePlugin.NodeUnpublishVolume(context.TODO(), &csi.NodeUnpublishVolumeRequest{
		Version:    csiVersion(),
		VolumeId:   volumeId,
		TargetPath: targetPath,
	})

	logger.Debug("node-response", lager.Data{"nodeResponse": nodeResponse})

	if err != nil {
		logger.Error("node-response-error", err)
		return err
	}

	dw.volumesMutex.Lock()
	defer dw.volumesMutex.Unlock()

	if count, ok := dw.volumes[volumeId]; ok {
		if count > 1 {
			dw.volumes[volumeId] = count - 1
		} else {
			delete(dw.volumes, volumeId)
		}
	}

	return nil
}

func (dw *nodeWrapper) Mount(logger lager.Logger, volumeId string, config map[string]interface{}) (volman.MountResponse, error) {
	logger = logger.Session("mount")
	logger.Info("start")
	defer logger.Info("end")

	conn, err := dw.grpcShim.Dial(dw.Spec.Address, grpc.WithInsecure())
	defer conn.Close()
	if err != nil {
		logger.Error("grpc-dial", err, lager.Data{"address": dw.Spec.Address})
		return volman.MountResponse{}, err
	}

	mountPath := path.Join(dw.csiMountRootDir, dw.Spec.Name)

	targetPath := path.Join(mountPath, volumeId)

	nodePlugin := dw.csiShim.NewNodeClient(conn)
	nodeResponse, err := nodePlugin.NodePublishVolume(context.TODO(), &csi.NodePublishVolumeRequest{
		Version:    csiVersion(),
		VolumeId:   volumeId,
		TargetPath: targetPath,
		VolumeCapability: &csi.VolumeCapability{
			AccessType: &csi.VolumeCapability_Mount{
				Mount: &csi.VolumeCapability_MountVolume{},
			},
			AccessMode: &csi.VolumeCapability_AccessMode{
				csi.VolumeCapability_AccessMode_MULTI_NODE_MULTI_WRITER,
			},
		},
	})

	logger.Debug("node-response", lager.Data{"nodeResponse": nodeResponse})

	if err != nil {
		logger.Error("node-response-error", err)
		return volman.MountResponse{}, err
	}

	dw.volumesMutex.Lock()
	defer dw.volumesMutex.Unlock()

	if count, ok := dw.volumes[volumeId]; ok {
		dw.volumes[volumeId] = count + 1
	} else {
		dw.volumes[volumeId] = 1
	}

	return volman.MountResponse{Path: targetPath}, nil
}

func (dw *nodeWrapper) ListVolumes(logger lager.Logger) ([]string, error) {
	logger = logger.Session("listvolumes")
	logger.Info("start")
	defer logger.Info("end")

	dw.volumesMutex.RLock()
	defer dw.volumesMutex.RUnlock()

	volumes := make([]string, len(dw.volumes))
	i := 0
	for volume := range dw.volumes {
		volumes[i] = volume
		i++
	}
	return volumes, nil
}

func (dw *nodeWrapper) Matches(logger lager.Logger, otherSpec volman.PluginSpec) bool {
	logger = logger.Session("matches")
	logger.Info("start")
	defer logger.Info("end")

	matches := reflect.DeepEqual(dw.Spec, otherSpec)
	logger.Info("matches", lager.Data{"matches": matches})
	return matches
}

func NewCsiPlugin(plugin csi.NodeClient, pluginSpec volman.PluginSpec, grpcShim grpcshim.Grpc, csiShim csishim.Csi, osShim osshim.Os, csiMountRootDir string) volman.Plugin {
	return &nodeWrapper{
		Impl:            plugin,
		Spec:            pluginSpec,
		grpcShim:        grpcShim,
		csiShim:         csiShim,
		osShim:          osShim,
		volumes:         map[string]int{},
		csiMountRootDir: csiMountRootDir,
	}
}
<|endoftext|>