<commit_before>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package features

import (
	apiextensionsfeatures "k8s.io/apiextensions-apiserver/pkg/features"
	genericfeatures "k8s.io/apiserver/pkg/features"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
)

const (
	// Every feature gate should add method here following this template:
	//
	// // owner: @username
	// // alpha: v1.X
	// MyFeature utilfeature.Feature = "MyFeature"

	// owner: @tallclair
	// beta: v1.4
	AppArmor utilfeature.Feature = "AppArmor"

	// owner: @mtaufen
	// alpha: v1.4
	DynamicKubeletConfig utilfeature.Feature = "DynamicKubeletConfig"

	// owner: @pweil-
	// alpha: v1.5
	//
	// Default userns=host for containers that are using other host namespaces, host mounts, the pod
	// contains a privileged container, or specific non-namespaced capabilities (MKNOD, SYS_MODULE,
	// SYS_TIME). This should only be enabled if user namespace remapping is enabled in the docker daemon.
	ExperimentalHostUserNamespaceDefaultingGate utilfeature.Feature = "ExperimentalHostUserNamespaceDefaulting"

	// owner: @vishh
	// alpha: v1.5
	//
	// Ensures guaranteed scheduling of pods marked with a special pod annotation `scheduler.alpha.kubernetes.io/critical-pod`
	// and also prevents them from being evicted from a node.
	// Note: This feature is not supported for `BestEffort` pods.
	ExperimentalCriticalPodAnnotation utilfeature.Feature = "ExperimentalCriticalPodAnnotation"

	// owner: @vishh
	// alpha: v1.6
	//
	// This is deprecated and will be removed in v1.11. Use DevicePlugins instead.
	//
	// Enables support for GPUs as a schedulable resource.
	// Only Nvidia GPUs are supported as of v1.6.
	// Works only with Docker Container Runtime.
	Accelerators utilfeature.Feature = "Accelerators"

	// owner: @jiayingz
	// alpha: v1.8
	//
	// Enables support for Device Plugins
	// Only Nvidia GPUs are tested as of v1.8.
	DevicePlugins utilfeature.Feature = "DevicePlugins"

	// owner: @gmarek
	// alpha: v1.6
	//
	// Changes the logic behind evicting Pods from not ready Nodes
	// to take advantage of NoExecute Taints and Tolerations.
	TaintBasedEvictions utilfeature.Feature = "TaintBasedEvictions"

	// owner: @jcbsmpsn
	// alpha: v1.7
	//
	// Gets a server certificate for the kubelet from the Certificate Signing
	// Request API instead of generating one self signed and auto rotates the
	// certificate as expiration approaches.
	RotateKubeletServerCertificate utilfeature.Feature = "RotateKubeletServerCertificate"

	// owner: @jcbsmpsn
	// alpha: v1.7
	//
	// Automatically renews the client certificate used for communicating with
	// the API server as the certificate approaches expiration.
	RotateKubeletClientCertificate utilfeature.Feature = "RotateKubeletClientCertificate"

	// owner: @msau42
	// alpha: v1.7
	//
	// A new volume type that supports local disks on a node.
	PersistentLocalVolumes utilfeature.Feature = "PersistentLocalVolumes"

	// owner: @jinxu
	// alpha: v1.7
	//
	// New local storage types to support local storage capacity isolation
	LocalStorageCapacityIsolation utilfeature.Feature = "LocalStorageCapacityIsolation"

	// owner: @gnufied
	// alpha: v1.8
	// Ability to Expand persistent volumes
	ExpandPersistentVolumes utilfeature.Feature = "ExpandPersistentVolumes"

	// owner: @verb
	// alpha: v1.10
	//
	// Allows running a "debug container" in a pod namespaces to troubleshoot a running pod.
	DebugContainers utilfeature.Feature = "DebugContainers"

	// owner: @verb
	// alpha: v1.10
	//
	// Allows all containers in a pod to share a process namespace.
	PodShareProcessNamespace utilfeature.Feature = "PodShareProcessNamespace"

	// owner: @bsalamat
	// alpha: v1.8
	//
	// Add priority to pods. Priority affects scheduling and preemption of pods.
	PodPriority utilfeature.Feature = "PodPriority"

	// owner: @resouer
	// alpha: v1.8
	//
	// Enable equivalence class cache for scheduler.
	EnableEquivalenceClassCache utilfeature.Feature = "EnableEquivalenceClassCache"

	// owner: @k82cn
	// alpha: v1.8
	//
	// Taint nodes based on their condition status for 'NetworkUnavailable',
	// 'MemoryPressure', 'OutOfDisk' and 'DiskPressure'.
	TaintNodesByCondition utilfeature.Feature = "TaintNodesByCondition"

	// owner: @jsafrane
	// alpha: v1.8
	//
	// Enable mount propagation of volumes.
	MountPropagation utilfeature.Feature = "MountPropagation"

	// owner: @ConnorDoyle
	// alpha: v1.8
	//
	// Alternative container-level CPU affinity policies.
	CPUManager utilfeature.Feature = "CPUManager"

	// owner: @derekwaynecarr
	// beta: v1.10
	//
	// Enable pods to consume pre-allocated huge pages of varying page sizes
	HugePages utilfeature.Feature = "HugePages"

	// owner @brendandburns
	// alpha: v1.9
	//
	// Enable nodes to exclude themselves from service load balancers
	ServiceNodeExclusion utilfeature.Feature = "ServiceNodeExclusion"

	// owner @brendandburns
	// deprecated: v1.10
	//
	// Enable the service proxy to contact external IP addresses. Note this feature is present
	// only for backward compatability, it will be removed in the 1.11 release.
	ServiceProxyAllowExternalIPs utilfeature.Feature = "ServiceProxyAllowExternalIPs"

	// owner: @jsafrane
	// alpha: v1.9
	//
	// Enable running mount utilities in containers.
	MountContainers utilfeature.Feature = "MountContainers"

	// owner: @msau42
	// alpha: v1.9
	//
	// Extend the default scheduler to be aware of PV topology and handle PV binding
	// Before moving to beta, resolve Kubernetes issue #56180
	VolumeScheduling utilfeature.Feature = "VolumeScheduling"

	// owner: @vladimirvivien
	// alpha: v1.9
	//
	// Enable mount/attachment of Container Storage Interface (CSI) backed PVs
	CSIPersistentVolume utilfeature.Feature = "CSIPersistentVolume"

	// owner @MrHohn
	// alpha: v1.9
	//
	// Support configurable pod DNS parameters.
	CustomPodDNS utilfeature.Feature = "CustomPodDNS"

	// owner: @screeley44
	// alpha: v1.9
	//
	// Enable Block volume support in containers.
	BlockVolume utilfeature.Feature = "BlockVolume"

	// owner: @pospispa
	//
	// alpha: v1.9
	// Postpone deletion of a persistent volume claim in case it is used by a pod
	PVCProtection utilfeature.Feature = "PVCProtection"

	// owner: @aveshagarwal
	// alpha: v1.9
	//
	// Enable resource limits priority function
	ResourceLimitsPriorityFunction utilfeature.Feature = "ResourceLimitsPriorityFunction"

	// owner: @m1093782566
	// beta: v1.9
	//
	// Implement IPVS-based in-cluster service load balancing
	SupportIPVSProxyMode utilfeature.Feature = "SupportIPVSProxyMode"

	// owner: @dims
	// alpha: v1.10
	//
	// Implement support for limiting pids in pods
	SupportPodPidsLimit utilfeature.Feature = "SupportPodPidsLimit"
)

func init() {
	utilfeature.DefaultFeatureGate.Add(defaultKubernetesFeatureGates)
}

// defaultKubernetesFeatureGates consists of all known Kubernetes-specific feature keys.
// To add a new feature, define a key for it above and add it here. The features will be
// available throughout Kubernetes binaries.
var defaultKubernetesFeatureGates = map[utilfeature.Feature]utilfeature.FeatureSpec{
	AppArmor:                                    {Default: true, PreRelease: utilfeature.Beta},
	DynamicKubeletConfig:                        {Default: false, PreRelease: utilfeature.Alpha},
	ExperimentalHostUserNamespaceDefaultingGate: {Default: false, PreRelease: utilfeature.Beta},
	ExperimentalCriticalPodAnnotation:           {Default: false, PreRelease: utilfeature.Alpha},
	Accelerators:                                {Default: false, PreRelease: utilfeature.Alpha},
	DevicePlugins:                               {Default: false, PreRelease: utilfeature.Alpha},
	TaintBasedEvictions:                         {Default: false, PreRelease: utilfeature.Alpha},
	RotateKubeletServerCertificate:              {Default: false, PreRelease: utilfeature.Alpha},
	RotateKubeletClientCertificate:              {Default: true, PreRelease: utilfeature.Beta},
	PersistentLocalVolumes:                      {Default: false, PreRelease: utilfeature.Alpha},
	LocalStorageCapacityIsolation:               {Default: false, PreRelease: utilfeature.Alpha},
	HugePages:                                   {Default: true, PreRelease: utilfeature.Beta},
	DebugContainers:                             {Default: false, PreRelease: utilfeature.Alpha},
	PodShareProcessNamespace:                    {Default: false, PreRelease: utilfeature.Alpha},
	PodPriority:                                 {Default: false, PreRelease: utilfeature.Alpha},
	EnableEquivalenceClassCache:                 {Default: false, PreRelease: utilfeature.Alpha},
	TaintNodesByCondition:                       {Default: false, PreRelease: utilfeature.Alpha},
	MountPropagation:                            {Default: false, PreRelease: utilfeature.Alpha},
	ExpandPersistentVolumes:                     {Default: false, PreRelease: utilfeature.Alpha},
	CPUManager:                                  {Default: true, PreRelease: utilfeature.Beta},
	ServiceNodeExclusion:                        {Default: false, PreRelease: utilfeature.Alpha},
	MountContainers:                             {Default: false, PreRelease: utilfeature.Alpha},
	VolumeScheduling:                            {Default: false, PreRelease: utilfeature.Alpha},
	CSIPersistentVolume:                         {Default: false, PreRelease: utilfeature.Alpha},
	CustomPodDNS:                                {Default: false, PreRelease: utilfeature.Alpha},
	BlockVolume:                                 {Default: false, PreRelease: utilfeature.Alpha},
	PVCProtection:                               {Default: false, PreRelease: utilfeature.Alpha},
	ResourceLimitsPriorityFunction:              {Default: false, PreRelease: utilfeature.Alpha},
	SupportIPVSProxyMode:                        {Default: false, PreRelease: utilfeature.Beta},
	SupportPodPidsLimit:                         {Default: false, PreRelease: utilfeature.Alpha},

	// inherited features from generic apiserver, relisted here to get a conflict if it is changed
	// unintentionally on either side:
	genericfeatures.StreamingProxyRedirects: {Default: true, PreRelease: utilfeature.Beta},
	genericfeatures.AdvancedAuditing:        {Default: true, PreRelease: utilfeature.Beta},
	genericfeatures.APIResponseCompression:  {Default: false, PreRelease: utilfeature.Alpha},
	genericfeatures.Initializers:            {Default: false, PreRelease: utilfeature.Alpha},
	genericfeatures.APIListChunking:         {Default: true, PreRelease: utilfeature.Beta},

	// inherited features from apiextensions-apiserver, relisted here to get a conflict if it is changed
	// unintentionally on either side:
	apiextensionsfeatures.CustomResourceValidation: {Default: true, PreRelease: utilfeature.Beta},

	// features that enable backwards compatability but are scheduled to be removed
	ServiceProxyAllowExternalIPs: {Default: false, PreRelease: utilfeature.Deprecated},
}
<commit_msg>Change feature gate PreRelease to Beta<commit_after>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package features

import (
	apiextensionsfeatures "k8s.io/apiextensions-apiserver/pkg/features"
	genericfeatures "k8s.io/apiserver/pkg/features"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
)

const (
	// Every feature gate should add method here following this template:
	//
	// // owner: @username
	// // alpha: v1.X
	// MyFeature utilfeature.Feature = "MyFeature"

	// owner: @tallclair
	// beta: v1.4
	AppArmor utilfeature.Feature = "AppArmor"

	// owner: @mtaufen
	// alpha: v1.4
	DynamicKubeletConfig utilfeature.Feature = "DynamicKubeletConfig"

	// owner: @pweil-
	// alpha: v1.5
	//
	// Default userns=host for containers that are using other host namespaces, host mounts, the pod
	// contains a privileged container, or specific non-namespaced capabilities (MKNOD, SYS_MODULE,
	// SYS_TIME). This should only be enabled if user namespace remapping is enabled in the docker daemon.
	ExperimentalHostUserNamespaceDefaultingGate utilfeature.Feature = "ExperimentalHostUserNamespaceDefaulting"

	// owner: @vishh
	// alpha: v1.5
	//
	// Ensures guaranteed scheduling of pods marked with a special pod annotation `scheduler.alpha.kubernetes.io/critical-pod`
	// and also prevents them from being evicted from a node.
	// Note: This feature is not supported for `BestEffort` pods.
	ExperimentalCriticalPodAnnotation utilfeature.Feature = "ExperimentalCriticalPodAnnotation"

	// owner: @vishh
	// alpha: v1.6
	//
	// This is deprecated and will be removed in v1.11. Use DevicePlugins instead.
	//
	// Enables support for GPUs as a schedulable resource.
	// Only Nvidia GPUs are supported as of v1.6.
	// Works only with Docker Container Runtime.
	Accelerators utilfeature.Feature = "Accelerators"

	// owner: @jiayingz
	// alpha: v1.8
	//
	// Enables support for Device Plugins
	// Only Nvidia GPUs are tested as of v1.8.
	DevicePlugins utilfeature.Feature = "DevicePlugins"

	// owner: @gmarek
	// alpha: v1.6
	//
	// Changes the logic behind evicting Pods from not ready Nodes
	// to take advantage of NoExecute Taints and Tolerations.
	TaintBasedEvictions utilfeature.Feature = "TaintBasedEvictions"

	// owner: @jcbsmpsn
	// alpha: v1.7
	//
	// Gets a server certificate for the kubelet from the Certificate Signing
	// Request API instead of generating one self signed and auto rotates the
	// certificate as expiration approaches.
	RotateKubeletServerCertificate utilfeature.Feature = "RotateKubeletServerCertificate"

	// owner: @jcbsmpsn
	// alpha: v1.7
	//
	// Automatically renews the client certificate used for communicating with
	// the API server as the certificate approaches expiration.
	RotateKubeletClientCertificate utilfeature.Feature = "RotateKubeletClientCertificate"

	// owner: @msau42
	// alpha: v1.7
	//
	// A new volume type that supports local disks on a node.
	PersistentLocalVolumes utilfeature.Feature = "PersistentLocalVolumes"

	// owner: @jinxu
	// alpha: v1.7
	//
	// New local storage types to support local storage capacity isolation
	LocalStorageCapacityIsolation utilfeature.Feature = "LocalStorageCapacityIsolation"

	// owner: @gnufied
	// alpha: v1.8
	// Ability to Expand persistent volumes
	ExpandPersistentVolumes utilfeature.Feature = "ExpandPersistentVolumes"

	// owner: @verb
	// alpha: v1.10
	//
	// Allows running a "debug container" in a pod namespaces to troubleshoot a running pod.
	DebugContainers utilfeature.Feature = "DebugContainers"

	// owner: @verb
	// alpha: v1.10
	//
	// Allows all containers in a pod to share a process namespace.
	PodShareProcessNamespace utilfeature.Feature = "PodShareProcessNamespace"

	// owner: @bsalamat
	// alpha: v1.8
	//
	// Add priority to pods. Priority affects scheduling and preemption of pods.
	PodPriority utilfeature.Feature = "PodPriority"

	// owner: @resouer
	// alpha: v1.8
	//
	// Enable equivalence class cache for scheduler.
	EnableEquivalenceClassCache utilfeature.Feature = "EnableEquivalenceClassCache"

	// owner: @k82cn
	// alpha: v1.8
	//
	// Taint nodes based on their condition status for 'NetworkUnavailable',
	// 'MemoryPressure', 'OutOfDisk' and 'DiskPressure'.
	TaintNodesByCondition utilfeature.Feature = "TaintNodesByCondition"

	// owner: @jsafrane
	// alpha: v1.8
	//
	// Enable mount propagation of volumes.
	MountPropagation utilfeature.Feature = "MountPropagation"

	// owner: @ConnorDoyle
	// alpha: v1.8
	//
	// Alternative container-level CPU affinity policies.
	CPUManager utilfeature.Feature = "CPUManager"

	// owner: @derekwaynecarr
	// beta: v1.10
	//
	// Enable pods to consume pre-allocated huge pages of varying page sizes
	HugePages utilfeature.Feature = "HugePages"

	// owner @brendandburns
	// alpha: v1.9
	//
	// Enable nodes to exclude themselves from service load balancers
	ServiceNodeExclusion utilfeature.Feature = "ServiceNodeExclusion"

	// owner @brendandburns
	// deprecated: v1.10
	//
	// Enable the service proxy to contact external IP addresses. Note this feature is present
	// only for backward compatability, it will be removed in the 1.11 release.
	ServiceProxyAllowExternalIPs utilfeature.Feature = "ServiceProxyAllowExternalIPs"

	// owner: @jsafrane
	// alpha: v1.9
	//
	// Enable running mount utilities in containers.
	MountContainers utilfeature.Feature = "MountContainers"

	// owner: @msau42
	// alpha: v1.9
	//
	// Extend the default scheduler to be aware of PV topology and handle PV binding
	// Before moving to beta, resolve Kubernetes issue #56180
	VolumeScheduling utilfeature.Feature = "VolumeScheduling"

	// owner: @vladimirvivien
	// alpha: v1.9
	//
	// Enable mount/attachment of Container Storage Interface (CSI) backed PVs
	CSIPersistentVolume utilfeature.Feature = "CSIPersistentVolume"

	// owner @MrHohn
	// alpha: v1.9
	//
	// Support configurable pod DNS parameters.
	CustomPodDNS utilfeature.Feature = "CustomPodDNS"

	// owner: @screeley44
	// alpha: v1.9
	//
	// Enable Block volume support in containers.
	BlockVolume utilfeature.Feature = "BlockVolume"

	// owner: @pospispa
	//
	// alpha: v1.9
	// Postpone deletion of a persistent volume claim in case it is used by a pod
	PVCProtection utilfeature.Feature = "PVCProtection"

	// owner: @aveshagarwal
	// alpha: v1.9
	//
	// Enable resource limits priority function
	ResourceLimitsPriorityFunction utilfeature.Feature = "ResourceLimitsPriorityFunction"

	// owner: @m1093782566
	// beta: v1.9
	//
	// Implement IPVS-based in-cluster service load balancing
	SupportIPVSProxyMode utilfeature.Feature = "SupportIPVSProxyMode"

	// owner: @dims
	// alpha: v1.10
	//
	// Implement support for limiting pids in pods
	SupportPodPidsLimit utilfeature.Feature = "SupportPodPidsLimit"
)

func init() {
	utilfeature.DefaultFeatureGate.Add(defaultKubernetesFeatureGates)
}

// defaultKubernetesFeatureGates consists of all known Kubernetes-specific feature keys.
// To add a new feature, define a key for it above and add it here. The features will be
// available throughout Kubernetes binaries.
var defaultKubernetesFeatureGates = map[utilfeature.Feature]utilfeature.FeatureSpec{
	AppArmor:                                    {Default: true, PreRelease: utilfeature.Beta},
	DynamicKubeletConfig:                        {Default: false, PreRelease: utilfeature.Alpha},
	ExperimentalHostUserNamespaceDefaultingGate: {Default: false, PreRelease: utilfeature.Beta},
	ExperimentalCriticalPodAnnotation:           {Default: false, PreRelease: utilfeature.Alpha},
	Accelerators:                                {Default: false, PreRelease: utilfeature.Alpha},
	DevicePlugins:                               {Default: false, PreRelease: utilfeature.Alpha},
	TaintBasedEvictions:                         {Default: false, PreRelease: utilfeature.Alpha},
	RotateKubeletServerCertificate:              {Default: false, PreRelease: utilfeature.Alpha},
	RotateKubeletClientCertificate:              {Default: true, PreRelease: utilfeature.Beta},
	PersistentLocalVolumes:                      {Default: false, PreRelease: utilfeature.Alpha},
	LocalStorageCapacityIsolation:               {Default: false, PreRelease: utilfeature.Alpha},
	HugePages:                                   {Default: true, PreRelease: utilfeature.Beta},
	DebugContainers:                             {Default: false, PreRelease: utilfeature.Alpha},
	PodShareProcessNamespace:                    {Default: false, PreRelease: utilfeature.Alpha},
	PodPriority:                                 {Default: false, PreRelease: utilfeature.Alpha},
	EnableEquivalenceClassCache:                 {Default: false, PreRelease: utilfeature.Alpha},
	TaintNodesByCondition:                       {Default: false, PreRelease: utilfeature.Alpha},
	MountPropagation:                            {Default: false, PreRelease: utilfeature.Alpha},
	ExpandPersistentVolumes:                     {Default: false, PreRelease: utilfeature.Alpha},
	CPUManager:                                  {Default: true, PreRelease: utilfeature.Beta},
	ServiceNodeExclusion:                        {Default: false, PreRelease: utilfeature.Alpha},
	MountContainers:                             {Default: false, PreRelease: utilfeature.Alpha},
	VolumeScheduling:                            {Default: false, PreRelease: utilfeature.Alpha},
	CSIPersistentVolume:                         {Default: true, PreRelease: utilfeature.Beta},
	CustomPodDNS:                                {Default: false, PreRelease: utilfeature.Alpha},
	BlockVolume:                                 {Default: false, PreRelease: utilfeature.Alpha},
	PVCProtection:                               {Default: false, PreRelease: utilfeature.Alpha},
	ResourceLimitsPriorityFunction:              {Default: false, PreRelease: utilfeature.Alpha},
	SupportIPVSProxyMode:                        {Default: false, PreRelease: utilfeature.Beta},
	SupportPodPidsLimit:                         {Default: false, PreRelease: utilfeature.Alpha},

	// inherited features from generic apiserver, relisted here to get a conflict if it is changed
	// unintentionally on either side:
	genericfeatures.StreamingProxyRedirects: {Default: true, PreRelease: utilfeature.Beta},
	genericfeatures.AdvancedAuditing:        {Default: true, PreRelease: utilfeature.Beta},
	genericfeatures.APIResponseCompression:  {Default: false, PreRelease: utilfeature.Alpha},
	genericfeatures.Initializers:            {Default: false, PreRelease: utilfeature.Alpha},
	genericfeatures.APIListChunking:         {Default: true, PreRelease: utilfeature.Beta},

	// inherited features from apiextensions-apiserver, relisted here to get a conflict if it is changed
	// unintentionally on either side:
	apiextensionsfeatures.CustomResourceValidation: {Default: true, PreRelease: utilfeature.Beta},

	// features that enable backwards compatability but are scheduled to be removed
	ServiceProxyAllowExternalIPs: {Default: false, PreRelease: utilfeature.Deprecated},
}
<|endoftext|>


<commit_before><commit_msg>增加时间范围生成代码的支持<commit_after><|endoftext|>


<commit_before>package bitswap

import (
	"context"
	"sync"
	"time"

	engine "github.com/ipfs/go-ipfs/exchange/bitswap/decision"
	bsmsg "github.com/ipfs/go-ipfs/exchange/bitswap/message"
	bsnet "github.com/ipfs/go-ipfs/exchange/bitswap/network"
	wantlist "github.com/ipfs/go-ipfs/exchange/bitswap/wantlist"

	metrics "gx/ipfs/QmRg1gKTHzc3CZXSKzem8aR4E3TubFhbgXwfVuWnSK5CC5/go-metrics-interface"
	cid "gx/ipfs/Qma4RJSuh7mMeJQYCqMbKzekn6EwBo7HEs5AQYjVRMQATB/go-cid"
	peer "gx/ipfs/QmdS9KpbDyPrieswibZhkod1oXqRwZJrUPzxCofAMWpFGq/go-libp2p-peer"
)

type WantManager struct {
	// sync channels for Run loop
	incoming   chan []*bsmsg.Entry
	connect    chan peer.ID        // notification channel for new peers connecting
	disconnect chan peer.ID        // notification channel for peers disconnecting
	peerReqs   chan chan []peer.ID // channel to request connected peers on

	// synchronized by Run loop, only touch inside there
	peers map[peer.ID]*msgQueue
	wl    *wantlist.ThreadSafe

	network bsnet.BitSwapNetwork
	ctx     context.Context
	cancel  func()

	wantlistGauge metrics.Gauge
	sentHistogram metrics.Histogram
}

func NewWantManager(ctx context.Context, network bsnet.BitSwapNetwork) *WantManager {
	ctx, cancel := context.WithCancel(ctx)
	wantlistGauge := metrics.NewCtx(ctx, "wantlist_total",
		"Number of items in wantlist.").Gauge()
	sentHistogram := metrics.NewCtx(ctx, "sent_all_blocks_bytes", "Histogram of blocks sent by"+
		" this bitswap").Histogram(metricsBuckets)
	return &WantManager{
		incoming:      make(chan []*bsmsg.Entry, 10),
		connect:       make(chan peer.ID, 10),
		disconnect:    make(chan peer.ID, 10),
		peerReqs:      make(chan chan []peer.ID),
		peers:         make(map[peer.ID]*msgQueue),
		wl:            wantlist.NewThreadSafe(),
		network:       network,
		ctx:           ctx,
		cancel:        cancel,
		wantlistGauge: wantlistGauge,
		sentHistogram: sentHistogram,
	}
}

type msgQueue struct {
	p peer.ID

	outlk   sync.Mutex
	out     bsmsg.BitSwapMessage
	network bsnet.BitSwapNetwork

	sender bsnet.MessageSender

	refcnt int

	work chan struct{}
	done chan struct{}
}

func (pm *WantManager) WantBlocks(ctx context.Context, ks []*cid.Cid) {
	log.Infof("want blocks: %s", ks)
	pm.addEntries(ctx, ks, false)
}

func (pm *WantManager) CancelWants(ks []*cid.Cid) {
	log.Infof("cancel wants: %s", ks)
	pm.addEntries(context.TODO(), ks, true)
}

func (pm *WantManager) addEntries(ctx context.Context, ks []*cid.Cid, cancel bool) {
	var entries []*bsmsg.Entry
	for i, k := range ks {
		entries = append(entries, &bsmsg.Entry{
			Cancel: cancel,
			Entry: &wantlist.Entry{
				Cid:      k,
				Priority: kMaxPriority - i,
				RefCnt:   1,
			},
		})
	}
	select {
	case pm.incoming <- entries:
	case <-pm.ctx.Done():
	case <-ctx.Done():
	}
}

func (pm *WantManager) ConnectedPeers() []peer.ID {
	resp := make(chan []peer.ID)
	pm.peerReqs <- resp
	return <-resp
}

func (pm *WantManager) SendBlock(ctx context.Context, env *engine.Envelope) {
	// Blocks need to be sent synchronously to maintain proper backpressure
	// throughout the network stack
	defer env.Sent()

	pm.sentHistogram.Observe(float64(len(env.Block.RawData())))

	msg := bsmsg.New(false)
	msg.AddBlock(env.Block)
	log.Infof("Sending block %s to %s", env.Block, env.Peer)
	err := pm.network.SendMessage(ctx, env.Peer, msg)
	if err != nil {
		log.Infof("sendblock error: %s", err)
	}
}

func (pm *WantManager) startPeerHandler(p peer.ID) *msgQueue {
	mq, ok := pm.peers[p]
	if ok {
		mq.refcnt++
		return nil
	}

	mq = pm.newMsgQueue(p)

	// new peer, we will want to give them our full wantlist
	fullwantlist := bsmsg.New(true)
	for _, e := range pm.wl.Entries() {
		fullwantlist.AddEntry(e.Cid, e.Priority)
	}
	mq.out = fullwantlist
	mq.work <- struct{}{}

	pm.peers[p] = mq
	go mq.runQueue(pm.ctx)
	return mq
}

func (pm *WantManager) stopPeerHandler(p peer.ID) {
	pq, ok := pm.peers[p]
	if !ok {
		// TODO: log error?
		return
	}

	pq.refcnt--
	if pq.refcnt > 0 {
		return
	}

	close(pq.done)
	delete(pm.peers, p)
}

func (mq *msgQueue) runQueue(ctx context.Context) {
	defer func() {
		if mq.sender != nil {
			mq.sender.Close()
		}
	}()
	for {
		select {
		case <-mq.work: // there is work to be done
			mq.doWork(ctx)
		case <-mq.done:
			return
		case <-ctx.Done():
			return
		}
	}
}

func (mq *msgQueue) doWork(ctx context.Context) {
	if mq.sender == nil {
		err := mq.openSender(ctx)
		if err != nil {
			log.Infof("cant open message sender to peer %s: %s", mq.p, err)
			// TODO: cant connect, what now?
			return
		}
	}

	// grab outgoing message
	mq.outlk.Lock()
	wlm := mq.out
	if wlm == nil || wlm.Empty() {
		mq.outlk.Unlock()
		return
	}
	mq.out = nil
	mq.outlk.Unlock()

	// send wantlist updates
	for { // try to send this message until we fail.
		err := mq.sender.SendMsg(ctx, wlm)
		if err == nil {
			return
		}

		log.Infof("bitswap send error: %s", err)
		mq.sender.Close()
		mq.sender = nil

		select {
		case <-mq.done:
			return
		case <-ctx.Done():
			return
		case <-time.After(time.Millisecond * 100):
			// wait 100ms in case disconnect notifications are still propogating
			log.Warning("SendMsg errored but neither 'done' nor context.Done() were set")
		}

		err = mq.openSender(ctx)
		if err != nil {
			log.Errorf("couldnt open sender again after SendMsg(%s) failed: %s", mq.p, err)
			// TODO(why): what do we do now?
			// I think the *right* answer is to probably put the message we're
			// trying to send back, and then return to waiting for new work or
			// a disconnect.
			return
		}

		// TODO: Is this the same instance for the remote peer?
		// If its not, we should resend our entire wantlist to them
		/*
			if mq.sender.InstanceID() != mq.lastSeenInstanceID {
				wlm = mq.getFullWantlistMessage()
			}
		*/
	}
}

func (mq *msgQueue) openSender(ctx context.Context) error {
	// allow ten minutes for connections this includes looking them up in the
	// dht dialing them, and handshaking
	conctx, cancel := context.WithTimeout(ctx, time.Minute*10)
	defer cancel()

	err := mq.network.ConnectTo(conctx, mq.p)
	if err != nil {
		return err
	}

	nsender, err := mq.network.NewMessageSender(ctx, mq.p)
	if err != nil {
		return err
	}

	mq.sender = nsender
	return nil
}

func (pm *WantManager) Connected(p peer.ID) {
	select {
	case pm.connect <- p:
	case <-pm.ctx.Done():
	}
}

func (pm *WantManager) Disconnected(p peer.ID) {
	select {
	case pm.disconnect <- p:
	case <-pm.ctx.Done():
	}
}

// TODO: use goprocess here once i trust it
func (pm *WantManager) Run() {
	tock := time.NewTicker(rebroadcastDelay.Get())
	defer tock.Stop()
	for {
		select {
		case entries := <-pm.incoming:

			// add changes to our wantlist
			var filtered []*bsmsg.Entry
			for _, e := range entries {
				if e.Cancel {
					if pm.wl.Remove(e.Cid) {
						pm.wantlistGauge.Dec()
						filtered = append(filtered, e)
					}
				} else {
					if pm.wl.AddEntry(e.Entry) {
						pm.wantlistGauge.Inc()
						filtered = append(filtered, e)
					}
				}
			}

			// broadcast those wantlist changes
			for _, p := range pm.peers {
				p.addMessage(filtered)
			}

		case <-tock.C:
			// resend entire wantlist every so often (REALLY SHOULDNT BE NECESSARY)
			var es []*bsmsg.Entry
			for _, e := range pm.wl.Entries() {
				es = append(es, &bsmsg.Entry{Entry: e})
			}

			for _, p := range pm.peers {
				p.outlk.Lock()
				p.out = bsmsg.New(true)
				p.outlk.Unlock()

				p.addMessage(es)
			}
		case p := <-pm.connect:
			pm.startPeerHandler(p)
		case p := <-pm.disconnect:
			pm.stopPeerHandler(p)
		case req := <-pm.peerReqs:
			var peers []peer.ID
			for p := range pm.peers {
				peers = append(peers, p)
			}
			req <- peers
		case <-pm.ctx.Done():
			return
		}
	}
}

func (wm *WantManager) newMsgQueue(p peer.ID) *msgQueue {
	return &msgQueue{
		done:    make(chan struct{}),
		work:    make(chan struct{}, 1),
		network: wm.network,
		p:       p,
		refcnt:  1,
	}
}

func (mq *msgQueue) addMessage(entries []*bsmsg.Entry) {
	mq.outlk.Lock()
	defer func() {
		mq.outlk.Unlock()
		select {
		case mq.work <- struct{}{}:
		default:
		}
	}()

	// if we have no message held allocate a new one
	if mq.out == nil {
		mq.out = bsmsg.New(false)
	}

	// TODO: add a msg.Combine(...) method
	// otherwise, combine the one we are holding with the
	// one passed in
	for _, e := range entries {
		if e.Cancel {
			mq.out.Cancel(e.Cid)
		} else {
			mq.out.AddEntry(e.Cid, e.Priority)
		}
	}
}
<commit_msg>track wantlists sent to peers individually<commit_after>package bitswap

import (
	"context"
	"sync"
	"time"

	engine "github.com/ipfs/go-ipfs/exchange/bitswap/decision"
	bsmsg "github.com/ipfs/go-ipfs/exchange/bitswap/message"
	bsnet "github.com/ipfs/go-ipfs/exchange/bitswap/network"
	wantlist "github.com/ipfs/go-ipfs/exchange/bitswap/wantlist"

	metrics "gx/ipfs/QmRg1gKTHzc3CZXSKzem8aR4E3TubFhbgXwfVuWnSK5CC5/go-metrics-interface"
	cid "gx/ipfs/Qma4RJSuh7mMeJQYCqMbKzekn6EwBo7HEs5AQYjVRMQATB/go-cid"
	peer "gx/ipfs/QmdS9KpbDyPrieswibZhkod1oXqRwZJrUPzxCofAMWpFGq/go-libp2p-peer"
)

type WantManager struct {
	// sync channels for Run loop
	incoming   chan *wantSet
	connect    chan peer.ID        // notification channel for new peers connecting
	disconnect chan peer.ID        // notification channel for peers disconnecting
	peerReqs   chan chan []peer.ID // channel to request connected peers on

	// synchronized by Run loop, only touch inside there
	peers map[peer.ID]*msgQueue
	wl    *wantlist.ThreadSafe

	network bsnet.BitSwapNetwork
	ctx     context.Context
	cancel  func()

	wantlistGauge metrics.Gauge
	sentHistogram metrics.Histogram
}

func NewWantManager(ctx context.Context, network bsnet.BitSwapNetwork) *WantManager {
	ctx, cancel := context.WithCancel(ctx)
	wantlistGauge := metrics.NewCtx(ctx, "wantlist_total",
		"Number of items in wantlist.").Gauge()
	sentHistogram := metrics.NewCtx(ctx, "sent_all_blocks_bytes", "Histogram of blocks sent by"+
		" this bitswap").Histogram(metricsBuckets)
	return &WantManager{
		incoming:      make(chan *wantSet, 10),
		connect:       make(chan peer.ID, 10),
		disconnect:    make(chan peer.ID, 10),
		peerReqs:      make(chan chan []peer.ID),
		peers:         make(map[peer.ID]*msgQueue),
		wl:            wantlist.NewThreadSafe(),
		network:       network,
		ctx:           ctx,
		cancel:        cancel,
		wantlistGauge: wantlistGauge,
		sentHistogram: sentHistogram,
	}
}

type msgQueue struct {
	p peer.ID

	outlk   sync.Mutex
	out     bsmsg.BitSwapMessage
	network bsnet.BitSwapNetwork
	wl      *wantlist.Wantlist

	sender bsnet.MessageSender

	refcnt int

	work chan struct{}
	done chan struct{}
}

func (pm *WantManager) WantBlocks(ctx context.Context, ks []*cid.Cid) {
	log.Infof("want blocks: %s", ks)
	pm.addEntries(ctx, ks, false)
}

func (pm *WantManager) CancelWants(ks []*cid.Cid) {
	pm.addEntries(context.Background(), ks, true)
}

type wantSet struct {
	entries []*bsmsg.Entry
	targets []peer.ID
}

func (pm *WantManager) addEntries(ctx context.Context, ks []*cid.Cid, cancel bool) {
	var entries []*bsmsg.Entry
	for i, k := range ks {
		entries = append(entries, &bsmsg.Entry{
			Cancel: cancel,
			Entry: &wantlist.Entry{
				Cid:      k,
				Priority: kMaxPriority - i,
				RefCnt:   1,
			},
		})
	}
	select {
	case pm.incoming <- &wantSet{entries: entries}:
	case <-pm.ctx.Done():
	case <-ctx.Done():
	}
}

func (pm *WantManager) ConnectedPeers() []peer.ID {
	resp := make(chan []peer.ID)
	pm.peerReqs <- resp
	return <-resp
}

func (pm *WantManager) SendBlock(ctx context.Context, env *engine.Envelope) {
	// Blocks need to be sent synchronously to maintain proper backpressure
	// throughout the network stack
	defer env.Sent()

	pm.sentHistogram.Observe(float64(len(env.Block.RawData())))

	msg := bsmsg.New(false)
	msg.AddBlock(env.Block)
	log.Infof("Sending block %s to %s", env.Block, env.Peer)
	err := pm.network.SendMessage(ctx, env.Peer, msg)
	if err != nil {
		log.Infof("sendblock error: %s", err)
	}
}

func (pm *WantManager) startPeerHandler(p peer.ID) *msgQueue {
	mq, ok := pm.peers[p]
	if ok {
		mq.refcnt++
		return nil
	}

	mq = pm.newMsgQueue(p)

	// new peer, we will want to give them our full wantlist
	fullwantlist := bsmsg.New(true)
	for _, e := range pm.wl.Entries() {
		ne := *e
		mq.wl.AddEntry(&ne)
		fullwantlist.AddEntry(e.Cid, e.Priority)
	}
	mq.out = fullwantlist
	mq.work <- struct{}{}

	pm.peers[p] = mq
	go mq.runQueue(pm.ctx)
	return mq
}

func (pm *WantManager) stopPeerHandler(p peer.ID) {
	pq, ok := pm.peers[p]
	if !ok {
		// TODO: log error?
		return
	}

	pq.refcnt--
	if pq.refcnt > 0 {
		return
	}

	close(pq.done)
	delete(pm.peers, p)
}

func (mq *msgQueue) runQueue(ctx context.Context) {
	defer func() {
		if mq.sender != nil {
			mq.sender.Close()
		}
	}()
	for {
		select {
		case <-mq.work: // there is work to be done
			mq.doWork(ctx)
		case <-mq.done:
			return
		case <-ctx.Done():
			return
		}
	}
}

func (mq *msgQueue) doWork(ctx context.Context) {
	if mq.sender == nil {
		err := mq.openSender(ctx)
		if err != nil {
			log.Infof("cant open message sender to peer %s: %s", mq.p, err)
			// TODO: cant connect, what now?
			return
		}
	}

	// grab outgoing message
	mq.outlk.Lock()
	wlm := mq.out
	if wlm == nil || wlm.Empty() {
		mq.outlk.Unlock()
		return
	}
	mq.out = nil
	mq.outlk.Unlock()

	// send wantlist updates
	for { // try to send this message until we fail.
		err := mq.sender.SendMsg(ctx, wlm)
		if err == nil {
			return
		}

		log.Infof("bitswap send error: %s", err)
		mq.sender.Close()
		mq.sender = nil

		select {
		case <-mq.done:
			return
		case <-ctx.Done():
			return
		case <-time.After(time.Millisecond * 100):
			// wait 100ms in case disconnect notifications are still propogating
			log.Warning("SendMsg errored but neither 'done' nor context.Done() were set")
		}

		err = mq.openSender(ctx)
		if err != nil {
			log.Errorf("couldnt open sender again after SendMsg(%s) failed: %s", mq.p, err)
			// TODO(why): what do we do now?
			// I think the *right* answer is to probably put the message we're
			// trying to send back, and then return to waiting for new work or
			// a disconnect.
			return
		}

		// TODO: Is this the same instance for the remote peer?
		// If its not, we should resend our entire wantlist to them
		/*
			if mq.sender.InstanceID() != mq.lastSeenInstanceID {
				wlm = mq.getFullWantlistMessage()
			}
		*/
	}
}

func (mq *msgQueue) openSender(ctx context.Context) error {
	// allow ten minutes for connections this includes looking them up in the
	// dht dialing them, and handshaking
	conctx, cancel := context.WithTimeout(ctx, time.Minute*10)
	defer cancel()

	err := mq.network.ConnectTo(conctx, mq.p)
	if err != nil {
		return err
	}

	nsender, err := mq.network.NewMessageSender(ctx, mq.p)
	if err != nil {
		return err
	}

	mq.sender = nsender
	return nil
}

func (pm *WantManager) Connected(p peer.ID) {
	select {
	case pm.connect <- p:
	case <-pm.ctx.Done():
	}
}

func (pm *WantManager) Disconnected(p peer.ID) {
	select {
	case pm.disconnect <- p:
	case <-pm.ctx.Done():
	}
}

// TODO: use goprocess here once i trust it
func (pm *WantManager) Run() {
	tock := time.NewTicker(rebroadcastDelay.Get())
	defer tock.Stop()
	for {
		select {
		case ws := <-pm.incoming:

			// add changes to our wantlist
			for _, e := range ws.entries {
				if e.Cancel {
					if pm.wl.Remove(e.Cid) {
						pm.wantlistGauge.Dec()
					}
				} else {
					if pm.wl.AddEntry(e.Entry) {
						pm.wantlistGauge.Inc()
					}
				}
			}

			// broadcast those wantlist changes
			if len(ws.targets) == 0 {
				for _, p := range pm.peers {
					p.addMessage(ws.entries)
				}
			} else {
				for _, t := range ws.targets {
					p, ok := pm.peers[t]
					if !ok {
						log.Warning("tried sending wantlist change to non-partner peer")
						continue
					}
					p.addMessage(ws.entries)
				}
			}

		case <-tock.C:
			// resend entire wantlist every so often (REALLY SHOULDNT BE NECESSARY)
			var es []*bsmsg.Entry
			for _, e := range pm.wl.Entries() {
				es = append(es, &bsmsg.Entry{Entry: e})
			}

			for _, p := range pm.peers {
				p.outlk.Lock()
				p.out = bsmsg.New(true)
				p.outlk.Unlock()

				p.addMessage(es)
			}
		case p := <-pm.connect:
			pm.startPeerHandler(p)
		case p := <-pm.disconnect:
			pm.stopPeerHandler(p)
		case req := <-pm.peerReqs:
			var peers []peer.ID
			for p := range pm.peers {
				peers = append(peers, p)
			}
			req <- peers
		case <-pm.ctx.Done():
			return
		}
	}
}

func (wm *WantManager) newMsgQueue(p peer.ID) *msgQueue {
	return &msgQueue{
		done:    make(chan struct{}),
		work:    make(chan struct{}, 1),
		wl:      wantlist.New(),
		network: wm.network,
		p:       p,
		refcnt:  1,
	}
}

func (mq *msgQueue) addMessage(entries []*bsmsg.Entry) {
	var work bool
	mq.outlk.Lock()
	defer func() {
		mq.outlk.Unlock()
		if !work {
			return
		}
		select {
		case mq.work <- struct{}{}:
		default:
		}
	}()

	// if we have no message held allocate a new one
	if mq.out == nil {
		mq.out = bsmsg.New(false)
	}

	// TODO: add a msg.Combine(...) method
	// otherwise, combine the one we are holding with the
	// one passed in
	for _, e := range entries {
		if e.Cancel {
			if mq.wl.Remove(e.Cid) {
				work = true
				mq.out.Cancel(e.Cid)
			}
		} else {
			if mq.wl.Add(e.Cid, e.Priority) {
				work = true
				mq.out.AddEntry(e.Cid, e.Priority)
			}
		}
	}
}
<|endoftext|>


<commit_before>/*
The MIT License (MIT)

Copyright (c) 2014 Boldly Go Ventures

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

package predicate

import (
	"encoding/json"
	"testing"
)

var js = [][]byte{
	[]byte(`"42"`),
	[]byte(`42`),
	[]byte(`true`),
	[]byte(`null`),
	[]byte(`{"x":42}`),
	[]byte(`[42,true]`),
}

func TestUnmarshalJSON(t *testing.T) {
	var (
		data []byte
		//		v    And //interface{}
	)

	data = []byte(`[{"and":[{"x":1,"y":2}]},{"or":null},{"not":null},{"xor":null},{"xyz":"e"},{"acb":["f","g"]}]`)
	data = []byte(`{"and":{"x":1}}`)
	var v Set //interface{}
	if err := json.Unmarshal(data, &v); err != nil {
		t.Error(err)
	} else {
		t.Logf("%#v\n", v)
	}
}
<commit_msg>Added several negative test cases<commit_after>/*
The MIT License (MIT)

Copyright (c) 2014 Boldly Go Ventures

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

package predicate

import (
	"encoding/json"
	"testing"
)

var js = [][]byte{
	[]byte(`"42"`),
	[]byte(`42`),
	[]byte(`true`),
	[]byte(`null`),
	[]byte(`{"x":42}`),
	[]byte(`[42,true]`),
}

func TestUnmarshalJSON(t *testing.T) {
	data := []byte(`[{"and":[{"x":"a"}]}]`)
	var v Set
	if err := json.Unmarshal(data, &v); err != nil {
		t.Error(err)
	} else {
		//		t.Logf("%#v\n", v)
	}

	if And(v).P(map[string]interface{}{"x": "a"}) {
		//		t.Log(true)
	} else {
		//		t.Log(false)
		t.Fail()
	}
}

func TestUnmarsalJSON_Invalid(t *testing.T) {
	var (
		data []byte
		v    Set
	)

	data = nil
	if err := json.Unmarshal(data, &v); err == nil {
		t.Fail()
	}

	data = make([]byte, 0)
	if err := json.Unmarshal(data, &v); err == nil {
		t.Fail()
	}

	data = []byte(``)
	if err := json.Unmarshal(data, &v); err == nil {
		t.Fail()
	}

	data = []byte(`"`)
	if err := json.Unmarshal(data, &v); err == nil {
		t.Fail()
	}

	data = []byte(`{`)
	if err := json.Unmarshal(data, &v); err == nil {
		t.Fail()
	}

	data = []byte(`[`)
	if err := json.Unmarshal(data, &v); err == nil {
		t.Fail()
	}

	data = []byte(`{null}`)
	if err := json.Unmarshal(data, &v); err == nil {
		t.Fail()
	}

	data = []byte(`[null]`)
	if err := json.Unmarshal(data, &v); err == nil {
		t.Fail()
	}

	data = []byte(`[{"and":null}]`)
	if err := json.Unmarshal(data, &v); err == nil {
		t.Fail()
	}
}
<|endoftext|>


<commit_before><commit_msg>Make account ID optional for accounts followers|following|statuses<commit_after><|endoftext|>


<commit_before>package mlxscraper

import (
	"log"
	"sort"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"

	"github.com/gbl08ma/disturbancesmlx/dataobjects"
	"github.com/mmcdole/gofeed"
)

// RSSScraper is an announcement scraper for the Metro de Lisboa website
// It reads the RSS feed from the official website
type RSSScraper struct {
	ticker         *time.Ticker
	stopChan       chan struct{}
	log            *log.Logger
	newAnnCallback func(announcement *dataobjects.Announcement)
	firstUpdate    bool
	fp             *gofeed.Parser
	announcements  []*dataobjects.Announcement

	URL     string
	Network *dataobjects.Network
	Period  time.Duration
}

// Begin starts the scraper
func (sc *RSSScraper) Begin(log *log.Logger,
	newAnnCallback func(announcement *dataobjects.Announcement)) {
	sc.stopChan = make(chan struct{})
	sc.ticker = time.NewTicker(sc.Period)
	sc.log = log
	sc.newAnnCallback = newAnnCallback
	sc.firstUpdate = true
	sc.fp = gofeed.NewParser()

	sc.log.Println("RSSScraper starting")
	sc.update()
	sc.firstUpdate = false
	sc.log.Println("RSSScraper completed first fetch")
	go sc.scrape()
}

// End stops the scraper
func (sc *RSSScraper) End() {
	sc.ticker.Stop()
	close(sc.stopChan)
}

func (sc *RSSScraper) copyAnnouncements() []*dataobjects.Announcement {
	c := make([]*dataobjects.Announcement, len(sc.announcements))
	for i, annPointer := range sc.announcements {
		ann := *annPointer
		c[i] = &ann
	}
	return c
}

func (sc *RSSScraper) scrape() {
	sc.update()
	sc.log.Println("RSSScraper completed second fetch")
	for {
		select {
		case <-sc.ticker.C:
			sc.update()
			sc.log.Println("RSSScraper fetch complete")
		case <-sc.stopChan:
			return
		}
	}
}

func (sc *RSSScraper) update() {
	feed, _ := sc.fp.ParseURL(sc.URL)

	announcements := []*dataobjects.Announcement{}
	for _, item := range feed.Items {
		ann := dataobjects.Announcement{
			Time:    *item.PublishedParsed,
			Network: sc.Network,
			Title:   item.Title,
			Body:    sc.adaptPostBody(item.Description),
			URL:     item.Link,
			Source:  "pt-ml-rss",
		}
		announcements = append(announcements, &ann)
	}

	sort.SliceStable(announcements, func(i, j int) bool {
		return announcements[i].Time.Before(announcements[j].Time)
	})

	if !sc.firstUpdate && len(announcements) > 0 {
		isNew := false
		curLast := announcements[len(announcements)-1]
		if len(sc.announcements) == 0 {
			isNew = true
		} else {
			// decide if an announcement is new by looking only into the published date
			prevLast := sc.announcements[len(sc.announcements)-1]
			if curLast.Time.After(prevLast.Time) {
				isNew = true
			}
		}
		if isNew {
			sc.newAnnCallback(curLast)
		}
	}
	sc.announcements = announcements
}

func (sc *RSSScraper) adaptPostBody(original string) string {
	doc, err := goquery.NewDocumentFromReader(strings.NewReader(original))
	if err != nil {
		sc.log.Println(err)
		return ""
	}
	return doc.Find("p").First().Text()
}

// Networks returns the networks monitored by this scraper
func (sc *RSSScraper) Networks() []*dataobjects.Network {
	return []*dataobjects.Network{sc.Network}
}

// Sources returns the sources provided by this scraper
func (sc *RSSScraper) Sources() []string {
	return []string{"pt-ml-rss"}
}

// Announcements returns the announcements read by this scraper
func (sc *RSSScraper) Announcements(source string) []*dataobjects.Announcement {
	return sc.copyAnnouncements()
}
<commit_msg>Fix error handling in RSS feed scraper<commit_after>package mlxscraper

import (
	"log"
	"sort"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"

	"github.com/gbl08ma/disturbancesmlx/dataobjects"
	"github.com/mmcdole/gofeed"
)

// RSSScraper is an announcement scraper for the Metro de Lisboa website
// It reads the RSS feed from the official website
type RSSScraper struct {
	ticker         *time.Ticker
	stopChan       chan struct{}
	log            *log.Logger
	newAnnCallback func(announcement *dataobjects.Announcement)
	firstUpdate    bool
	fp             *gofeed.Parser
	announcements  []*dataobjects.Announcement

	URL     string
	Network *dataobjects.Network
	Period  time.Duration
}

// Begin starts the scraper
func (sc *RSSScraper) Begin(log *log.Logger,
	newAnnCallback func(announcement *dataobjects.Announcement)) {
	sc.stopChan = make(chan struct{})
	sc.ticker = time.NewTicker(sc.Period)
	sc.log = log
	sc.newAnnCallback = newAnnCallback
	sc.firstUpdate = true
	sc.fp = gofeed.NewParser()

	sc.log.Println("RSSScraper starting")
	sc.update()
	sc.firstUpdate = false
	sc.log.Println("RSSScraper completed first fetch")
	go sc.scrape()
}

// End stops the scraper
func (sc *RSSScraper) End() {
	sc.ticker.Stop()
	close(sc.stopChan)
}

func (sc *RSSScraper) copyAnnouncements() []*dataobjects.Announcement {
	c := make([]*dataobjects.Announcement, len(sc.announcements))
	for i, annPointer := range sc.announcements {
		ann := *annPointer
		c[i] = &ann
	}
	return c
}

func (sc *RSSScraper) scrape() {
	sc.update()
	sc.log.Println("RSSScraper completed second fetch")
	for {
		select {
		case <-sc.ticker.C:
			sc.update()
			sc.log.Println("RSSScraper fetch complete")
		case <-sc.stopChan:
			return
		}
	}
}

func (sc *RSSScraper) update() {
	feed, err := sc.fp.ParseURL(sc.URL)
	if err != nil {
		sc.log.Println(err)
		return
	}

	announcements := []*dataobjects.Announcement{}
	for _, item := range feed.Items {
		ann := dataobjects.Announcement{
			Time:    *item.PublishedParsed,
			Network: sc.Network,
			Title:   item.Title,
			Body:    sc.adaptPostBody(item.Description),
			URL:     item.Link,
			Source:  "pt-ml-rss",
		}
		announcements = append(announcements, &ann)
	}

	sort.SliceStable(announcements, func(i, j int) bool {
		return announcements[i].Time.Before(announcements[j].Time)
	})

	if !sc.firstUpdate && len(announcements) > 0 {
		isNew := false
		curLast := announcements[len(announcements)-1]
		if len(sc.announcements) == 0 {
			isNew = true
		} else {
			// decide if an announcement is new by looking only into the published date
			prevLast := sc.announcements[len(sc.announcements)-1]
			if curLast.Time.After(prevLast.Time) {
				isNew = true
			}
		}
		if isNew {
			sc.newAnnCallback(curLast)
		}
	}
	sc.announcements = announcements
}

func (sc *RSSScraper) adaptPostBody(original string) string {
	doc, err := goquery.NewDocumentFromReader(strings.NewReader(original))
	if err != nil {
		sc.log.Println(err)
		return ""
	}
	return doc.Find("p").First().Text()
}

// Networks returns the networks monitored by this scraper
func (sc *RSSScraper) Networks() []*dataobjects.Network {
	return []*dataobjects.Network{sc.Network}
}

// Sources returns the sources provided by this scraper
func (sc *RSSScraper) Sources() []string {
	return []string{"pt-ml-rss"}
}

// Announcements returns the announcements read by this scraper
func (sc *RSSScraper) Announcements(source string) []*dataobjects.Announcement {
	return sc.copyAnnouncements()
}
<|endoftext|>


<commit_before>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied.  See the License for the
specific language governing permissions and limitations
under the License.
*/

package crypto

import (
	obcca "github.com/openblockchain/obc-peer/obc-ca/protos"

	"bytes"
	"crypto/ecdsa"
	"crypto/hmac"
	"errors"
	"github.com/golang/protobuf/proto"
	"github.com/openblockchain/obc-peer/openchain/crypto/utils"
	"golang.org/x/net/context"
	"google/protobuf"
	"math/big"
	"time"
)

func (client *clientImpl) initTCertEngine() (err error) {
	// load TCertOwnerKDFKey
	if err = client.loadTCertOwnerKDFKey(); err != nil {
		return
	}

	// open tCertPoolChannel and start the writer
	client.tCertPoolChannel = make(chan tCert, client.node.conf.getTCertBathSize()*2)
	go client.tCertPoolWriter()

	return
}

func (client *clientImpl) storeTCertOwnerKDFKey() error {
	if err := client.node.ks.storeKey(client.node.conf.getTCertOwnerKDFKeyFilename(), client.tCertOwnerKDFKey); err != nil {
		client.node.log.Error("Failed storing TCertOwnerKDFKey [%s].", err.Error())

		return err
	}
	return nil
}

func (client *clientImpl) loadTCertOwnerKDFKey() error {
	// Load TCertOwnerKDFKey
	client.node.log.Debug("Loading TCertOwnerKDFKey...")

	if !client.node.ks.isAliasSet(client.node.conf.getTCertOwnerKDFKeyFilename()) {
		client.node.log.Debug("Failed loading TCertOwnerKDFKey. Key is missing.")

		return nil
	}

	tCertOwnerKDFKey, err := client.node.ks.loadKey(client.node.conf.getTCertOwnerKDFKeyFilename())
	if err != nil {
		client.node.log.Error("Failed parsing TCertOwnerKDFKey [%s].", err.Error())

		return err
	}
	client.tCertOwnerKDFKey = tCertOwnerKDFKey

	client.node.log.Debug("Loading TCertOwnerKDFKey...done!")

	return nil
}

func (client *clientImpl) getNextTCert() (tCert tCert, err error) {
	client.node.log.Debug("Getting next TCert...")
	//rawCert, err := client.node.ks.getNextTCert(client.getTCertsFromTCA)
	//if err != nil {
	//	client.node.log.Error("Failed accessing db [%s].", err.Error())
	//
	//	return nil, err
	//}
	for {
		select {
		case tCert = <-client.tCertPoolChannel:
			break
		case <-time.After(30 * time.Second):
			client.node.log.Error("Failed getting a new TCert. Buffer is empty!")

			return nil, errors.New("Failed getting a new TCert. Buffer is empty!")
		}
		if tCert != nil {
			break
		}
	}

	client.node.log.Debug("Cert [% x].", tCert.GetCertificate().Raw)
	client.node.log.Debug("Getting next TCert...done!")

	return
}

func (client *clientImpl) getTCertFromDER(der []byte) (tCert, error) {
	client.node.log.Debug("Validating TCert [% x]", der)

	// DER to x509
	x509Cert, err := utils.DERToX509Certificate(der)
	if err != nil {
		client.node.log.Debug("Failed parsing certificate: [%s].", err)

		return nil, err
	}

	// Handle Critical Extension TCertEncTCertIndex
	if _, err = utils.GetCriticalExtension(x509Cert, utils.TCertEncTCertIndex); err != nil {
		client.node.log.Error("Failed getting extension TCERT_ENC_TCERTINDEX [%s].", err.Error())

		return nil, err
	}

	// Verify certificate against root
	if _, err := utils.CheckCertAgainRoot(x509Cert, client.node.tcaCertPool); err != nil {
		client.node.log.Warning("Warning verifing certificate [%s].", err.Error())

		return nil, err
	}

	return &tCertImpl{client, x509Cert, nil}, nil
}

func (client *clientImpl) tCertPoolWriter() {
	// Fill the channel
	for {
		if len(client.tCertPoolChannel) < 50 {
			client.node.log.Debug("Refill TCert Pool. Current size [%d].", len(client.tCertPoolChannel))
			client.getTCertsFromTCA(
				cap(client.tCertPoolChannel) - len(client.tCertPoolChannel),
			)
		}
		// TODO: configure this
		time.Sleep(1 * time.Second)
	}
}

func (client *clientImpl) getTCertsFromTCA(num int) error {
	client.node.log.Debug("Get [%d] certificates from the TCA...", num)

	// Contact the TCA
	TCertOwnerKDFKey, certDERs, err := client.callTCACreateCertificateSet(num)
	if err != nil {
		client.node.log.Debug("Failed contacting TCA [%s].", err.Error())

		return err
	}

	//	client.node.log.Debug("TCertOwnerKDFKey [%s].", utils.EncodeBase64(TCertOwnerKDFKey))

	// Store TCertOwnerKDFKey and checks that every time it is always the same key
	if client.tCertOwnerKDFKey != nil {
		// Check that the keys are the same
		equal := bytes.Equal(client.tCertOwnerKDFKey, TCertOwnerKDFKey)
		if !equal {
			return errors.New("Failed reciving kdf key from TCA. The keys are different.")
		}
	} else {
		client.tCertOwnerKDFKey = TCertOwnerKDFKey

		// TODO: handle this situation more carefully
		if err := client.storeTCertOwnerKDFKey(); err != nil {
			client.node.log.Error("Failed storing TCertOwnerKDFKey [%s].", err.Error())

			return err
		}
	}

	// Validate the Certificates obtained

	TCertOwnerEncryptKey := utils.HMACTruncated(TCertOwnerKDFKey, []byte{1}, utils.AESKeyLength)
	ExpansionKey := utils.HMAC(TCertOwnerKDFKey, []byte{2})

	j := 0
	for i := 0; i < num; i++ {
		client.node.log.Debug("Validating certificate [%d], [% x]", i, certDERs[i])

		// DER to x509
		x509Cert, err := utils.DERToX509Certificate(certDERs[i])
		if err != nil {
			client.node.log.Debug("Failed parsing certificate: [%s].", err)

			continue
		}

		// Handle Critical Extenstion TCertEncTCertIndex
		tCertIndexCT, err := utils.GetCriticalExtension(x509Cert, utils.TCertEncTCertIndex)
		if err != nil {
			client.node.log.Error("Failed getting extension TCERT_ENC_TCERTINDEX [%s].", err.Error())

			continue
		}

		// Verify certificate against root
		if _, err := utils.CheckCertAgainRoot(x509Cert, client.node.tcaCertPool); err != nil {
			client.node.log.Warning("Warning verifing certificate [%s].", err.Error())

			continue
		}

		// Verify public key

		// 384-bit ExpansionValue = HMAC(Expansion_Key, TCertIndex)
		// Let TCertIndex = Timestamp, RandValue, 1,2,…
		// Timestamp assigned, RandValue assigned and counter reinitialized to 1 per batch

		// Decrypt ct to TCertIndex (TODO: || EnrollPub_Key || EnrollID ?)
		pt, err := utils.CBCPKCS7Decrypt(TCertOwnerEncryptKey, tCertIndexCT)
		if err != nil {
			client.node.log.Error("Failed decrypting extension TCERT_ENC_TCERTINDEX [%s].", err.Error())

			continue
		}

		// Compute ExpansionValue based on TCertIndex
		TCertIndex := pt
		//		TCertIndex := []byte(strconv.Itoa(i))

		client.node.log.Debug("TCertIndex: [% x].", TCertIndex)
		mac := hmac.New(utils.NewHash, ExpansionKey)
		mac.Write(TCertIndex)
		ExpansionValue := mac.Sum(nil)

		// Derive tpk and tsk accordingly to ExapansionValue from enrollment pk,sk
		// Computable by TCA / Auditor: TCertPub_Key = EnrollPub_Key + ExpansionValue G
		// using elliptic curve point addition per NIST FIPS PUB 186-4- specified P-384

		// Compute temporary secret key
		tempSK := &ecdsa.PrivateKey{
			PublicKey: ecdsa.PublicKey{
				Curve: client.node.enrollPrivKey.Curve,
				X:     new(big.Int),
				Y:     new(big.Int),
			},
			D: new(big.Int),
		}

		var k = new(big.Int).SetBytes(ExpansionValue)
		var one = new(big.Int).SetInt64(1)
		n := new(big.Int).Sub(client.node.enrollPrivKey.Params().N, one)
		k.Mod(k, n)
		k.Add(k, one)

		tempSK.D.Add(client.node.enrollPrivKey.D, k)
		tempSK.D.Mod(tempSK.D, client.node.enrollPrivKey.PublicKey.Params().N)

		// Compute temporary public key
		tempX, tempY := client.node.enrollPrivKey.PublicKey.ScalarBaseMult(k.Bytes())
		tempSK.PublicKey.X, tempSK.PublicKey.Y =
			tempSK.PublicKey.Add(
				client.node.enrollPrivKey.PublicKey.X, client.node.enrollPrivKey.PublicKey.Y,
				tempX, tempY,
			)

		// Verify temporary public key is a valid point on the reference curve
		isOn := tempSK.Curve.IsOnCurve(tempSK.PublicKey.X, tempSK.PublicKey.Y)
		if !isOn {
			client.node.log.Error("Failed temporary public key IsOnCurve check.")

			continue
		}

		// Check that the derived public key is the same as the one in the certificate
		certPK := x509Cert.PublicKey.(*ecdsa.PublicKey)

		if cmp := certPK.X.Cmp(tempSK.PublicKey.X); cmp != 0 {
			client.node.log.Error("Derived public key is different on X")

			continue
		}

		if cmp = certPK.Y.Cmp(tempSK.PublicKey.Y); cmp != 0 {
			client.node.log.Error("Derived public key is different on Y")

			continue
		}

		// Verify the signing capability of tempSK
		err = utils.VerifySignCapability(tempSK, x509Cert.PublicKey)
		if err != nil {
			client.node.log.Error("Failed verifing signing capability [%s].", err.Error())

			continue
		}

		// Marshall certificate and secret key to be stored in the database
		if err != nil {
			client.node.log.Error("Failed marshalling private key [%s].", err.Error())

			continue
		}

		if err := utils.CheckCertPKAgainstSK(x509Cert, interface{}(tempSK)); err != nil {
			client.node.log.Error("Failed checking TCA cert PK against private key [%s].", err.Error())

			continue
		}

		client.node.log.Debug("Sub index [%d]", j)
		j++
		client.node.log.Debug("Certificate [%d] validated.", i)

		client.tCertPoolChannel <- &tCertImpl{client, x509Cert, tempSK}
	}

	if j == 0 {
		client.node.log.Error("No valid TCert was sent")

		return errors.New("No valid TCert was sent.")
	}

	return nil
}

func (client *clientImpl) callTCACreateCertificateSet(num int) ([]byte, [][]byte, error) {
	// Get a TCA Client
	sock, tcaP, err := client.node.getTCAClient()
	defer sock.Close()

	// Execute the protocol
	now := time.Now()
	timestamp := google_protobuf.Timestamp{int64(now.Second()), int32(now.Nanosecond())}
	req := &obcca.TCertCreateSetReq{
		&timestamp,
		&obcca.Identity{Id: client.node.enrollID},
		uint32(num),
		nil,
	}
	rawReq, err := proto.Marshal(req)
	if err != nil {
		client.node.log.Error("Failed marshaling request [%s] [%s].", err.Error())
		return nil, nil, err
	}

	// 2. Sign rawReq
	client.node.log.Debug("Signing req [% x]", rawReq)
	r, s, err := client.node.ecdsaSignWithEnrollmentKey(rawReq)
	if err != nil {
		client.node.log.Error("Failed creating signature [%s] [%s].", err.Error())
		return nil, nil, err
	}

	R, _ := r.MarshalText()
	S, _ := s.MarshalText()

	// 3. Append the signature
	req.Sig = &obcca.Signature{obcca.CryptoType_ECDSA, R, S}

	// 4. Send request
	certSet, err := tcaP.CreateCertificateSet(context.Background(), req)
	if err != nil {
		client.node.log.Error("Failed requesting tca create certificate set [%s].", err.Error())

		return nil, nil, err
	}

	return certSet.Certs.Key, certSet.Certs.Certs, nil
}
<commit_msg>retry getting a new tcert only three time<commit_after>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied.  See the License for the
specific language governing permissions and limitations
under the License.
*/

package crypto

import (
	obcca "github.com/openblockchain/obc-peer/obc-ca/protos"

	"bytes"
	"crypto/ecdsa"
	"crypto/hmac"
	"errors"
	"github.com/golang/protobuf/proto"
	"github.com/openblockchain/obc-peer/openchain/crypto/utils"
	"golang.org/x/net/context"
	"google/protobuf"
	"math/big"
	"time"
)

func (client *clientImpl) initTCertEngine() (err error) {
	// load TCertOwnerKDFKey
	if err = client.loadTCertOwnerKDFKey(); err != nil {
		return
	}

	// open tCertPoolChannel and start the writer
	client.tCertPoolChannel = make(chan tCert, client.node.conf.getTCertBathSize()*2)
	go client.tCertPoolWriter()

	return
}

func (client *clientImpl) storeTCertOwnerKDFKey() error {
	if err := client.node.ks.storeKey(client.node.conf.getTCertOwnerKDFKeyFilename(), client.tCertOwnerKDFKey); err != nil {
		client.node.log.Error("Failed storing TCertOwnerKDFKey [%s].", err.Error())

		return err
	}
	return nil
}

func (client *clientImpl) loadTCertOwnerKDFKey() error {
	// Load TCertOwnerKDFKey
	client.node.log.Debug("Loading TCertOwnerKDFKey...")

	if !client.node.ks.isAliasSet(client.node.conf.getTCertOwnerKDFKeyFilename()) {
		client.node.log.Debug("Failed loading TCertOwnerKDFKey. Key is missing.")

		return nil
	}

	tCertOwnerKDFKey, err := client.node.ks.loadKey(client.node.conf.getTCertOwnerKDFKeyFilename())
	if err != nil {
		client.node.log.Error("Failed parsing TCertOwnerKDFKey [%s].", err.Error())

		return err
	}
	client.tCertOwnerKDFKey = tCertOwnerKDFKey

	client.node.log.Debug("Loading TCertOwnerKDFKey...done!")

	return nil
}

func (client *clientImpl) getNextTCert() (tCert tCert, err error) {
	client.node.log.Debug("Getting next TCert...")
	//rawCert, err := client.node.ks.getNextTCert(client.getTCertsFromTCA)
	//if err != nil {
	//	client.node.log.Error("Failed accessing db [%s].", err.Error())
	//
	//	return nil, err
	//}
	for i := 0; i < 3; i++ {
		select {
		case tCert = <-client.tCertPoolChannel:
			break
		case <-time.After(30 * time.Second):
			client.node.log.Error("Failed getting a new TCert. Buffer is empty!")

			return nil, errors.New("Failed getting a new TCert. Buffer is empty!")
		}
		if tCert != nil {
			break
		}
	}

	client.node.log.Debug("Cert [% x].", tCert.GetCertificate().Raw)
	client.node.log.Debug("Getting next TCert...done!")

	return
}

func (client *clientImpl) getTCertFromDER(der []byte) (tCert, error) {
	client.node.log.Debug("Validating TCert [% x]", der)

	// DER to x509
	x509Cert, err := utils.DERToX509Certificate(der)
	if err != nil {
		client.node.log.Debug("Failed parsing certificate: [%s].", err)

		return nil, err
	}

	// Handle Critical Extension TCertEncTCertIndex
	if _, err = utils.GetCriticalExtension(x509Cert, utils.TCertEncTCertIndex); err != nil {
		client.node.log.Error("Failed getting extension TCERT_ENC_TCERTINDEX [%s].", err.Error())

		return nil, err
	}

	// Verify certificate against root
	if _, err := utils.CheckCertAgainRoot(x509Cert, client.node.tcaCertPool); err != nil {
		client.node.log.Warning("Warning verifing certificate [%s].", err.Error())

		return nil, err
	}

	return &tCertImpl{client, x509Cert, nil}, nil
}

func (client *clientImpl) tCertPoolWriter() {
	// Fill the channel
	for {
		if len(client.tCertPoolChannel) < 50 {
			client.node.log.Debug("Refill TCert Pool. Current size [%d].", len(client.tCertPoolChannel))
			client.getTCertsFromTCA(
				cap(client.tCertPoolChannel) - len(client.tCertPoolChannel),
			)
		}
		// TODO: configure this
		time.Sleep(1 * time.Second)
	}
}

func (client *clientImpl) getTCertsFromTCA(num int) error {
	client.node.log.Debug("Get [%d] certificates from the TCA...", num)

	// Contact the TCA
	TCertOwnerKDFKey, certDERs, err := client.callTCACreateCertificateSet(num)
	if err != nil {
		client.node.log.Debug("Failed contacting TCA [%s].", err.Error())

		return err
	}

	//	client.node.log.Debug("TCertOwnerKDFKey [%s].", utils.EncodeBase64(TCertOwnerKDFKey))

	// Store TCertOwnerKDFKey and checks that every time it is always the same key
	if client.tCertOwnerKDFKey != nil {
		// Check that the keys are the same
		equal := bytes.Equal(client.tCertOwnerKDFKey, TCertOwnerKDFKey)
		if !equal {
			return errors.New("Failed reciving kdf key from TCA. The keys are different.")
		}
	} else {
		client.tCertOwnerKDFKey = TCertOwnerKDFKey

		// TODO: handle this situation more carefully
		if err := client.storeTCertOwnerKDFKey(); err != nil {
			client.node.log.Error("Failed storing TCertOwnerKDFKey [%s].", err.Error())

			return err
		}
	}

	// Validate the Certificates obtained

	TCertOwnerEncryptKey := utils.HMACTruncated(TCertOwnerKDFKey, []byte{1}, utils.AESKeyLength)
	ExpansionKey := utils.HMAC(TCertOwnerKDFKey, []byte{2})

	j := 0
	for i := 0; i < num; i++ {
		client.node.log.Debug("Validating certificate [%d], [% x]", i, certDERs[i])

		// DER to x509
		x509Cert, err := utils.DERToX509Certificate(certDERs[i])
		if err != nil {
			client.node.log.Debug("Failed parsing certificate: [%s].", err)

			continue
		}

		// Handle Critical Extenstion TCertEncTCertIndex
		tCertIndexCT, err := utils.GetCriticalExtension(x509Cert, utils.TCertEncTCertIndex)
		if err != nil {
			client.node.log.Error("Failed getting extension TCERT_ENC_TCERTINDEX [%s].", err.Error())

			continue
		}

		// Verify certificate against root
		if _, err := utils.CheckCertAgainRoot(x509Cert, client.node.tcaCertPool); err != nil {
			client.node.log.Warning("Warning verifing certificate [%s].", err.Error())

			continue
		}

		// Verify public key

		// 384-bit ExpansionValue = HMAC(Expansion_Key, TCertIndex)
		// Let TCertIndex = Timestamp, RandValue, 1,2,…
		// Timestamp assigned, RandValue assigned and counter reinitialized to 1 per batch

		// Decrypt ct to TCertIndex (TODO: || EnrollPub_Key || EnrollID ?)
		pt, err := utils.CBCPKCS7Decrypt(TCertOwnerEncryptKey, tCertIndexCT)
		if err != nil {
			client.node.log.Error("Failed decrypting extension TCERT_ENC_TCERTINDEX [%s].", err.Error())

			continue
		}

		// Compute ExpansionValue based on TCertIndex
		TCertIndex := pt
		//		TCertIndex := []byte(strconv.Itoa(i))

		client.node.log.Debug("TCertIndex: [% x].", TCertIndex)
		mac := hmac.New(utils.NewHash, ExpansionKey)
		mac.Write(TCertIndex)
		ExpansionValue := mac.Sum(nil)

		// Derive tpk and tsk accordingly to ExapansionValue from enrollment pk,sk
		// Computable by TCA / Auditor: TCertPub_Key = EnrollPub_Key + ExpansionValue G
		// using elliptic curve point addition per NIST FIPS PUB 186-4- specified P-384

		// Compute temporary secret key
		tempSK := &ecdsa.PrivateKey{
			PublicKey: ecdsa.PublicKey{
				Curve: client.node.enrollPrivKey.Curve,
				X:     new(big.Int),
				Y:     new(big.Int),
			},
			D: new(big.Int),
		}

		var k = new(big.Int).SetBytes(ExpansionValue)
		var one = new(big.Int).SetInt64(1)
		n := new(big.Int).Sub(client.node.enrollPrivKey.Params().N, one)
		k.Mod(k, n)
		k.Add(k, one)

		tempSK.D.Add(client.node.enrollPrivKey.D, k)
		tempSK.D.Mod(tempSK.D, client.node.enrollPrivKey.PublicKey.Params().N)

		// Compute temporary public key
		tempX, tempY := client.node.enrollPrivKey.PublicKey.ScalarBaseMult(k.Bytes())
		tempSK.PublicKey.X, tempSK.PublicKey.Y =
			tempSK.PublicKey.Add(
				client.node.enrollPrivKey.PublicKey.X, client.node.enrollPrivKey.PublicKey.Y,
				tempX, tempY,
			)

		// Verify temporary public key is a valid point on the reference curve
		isOn := tempSK.Curve.IsOnCurve(tempSK.PublicKey.X, tempSK.PublicKey.Y)
		if !isOn {
			client.node.log.Error("Failed temporary public key IsOnCurve check.")

			continue
		}

		// Check that the derived public key is the same as the one in the certificate
		certPK := x509Cert.PublicKey.(*ecdsa.PublicKey)

		if certPK.X.Cmp(tempSK.PublicKey.X) != 0 {
			client.node.log.Error("Derived public key is different on X")

			continue
		}

		if certPK.Y.Cmp(tempSK.PublicKey.Y) != 0 {
			client.node.log.Error("Derived public key is different on Y")

			continue
		}

		// Verify the signing capability of tempSK
		err = utils.VerifySignCapability(tempSK, x509Cert.PublicKey)
		if err != nil {
			client.node.log.Error("Failed verifing signing capability [%s].", err.Error())

			continue
		}

		// Marshall certificate and secret key to be stored in the database
		if err != nil {
			client.node.log.Error("Failed marshalling private key [%s].", err.Error())

			continue
		}

		if err := utils.CheckCertPKAgainstSK(x509Cert, interface{}(tempSK)); err != nil {
			client.node.log.Error("Failed checking TCA cert PK against private key [%s].", err.Error())

			continue
		}

		client.node.log.Debug("Sub index [%d]", j)
		j++
		client.node.log.Debug("Certificate [%d] validated.", i)

		client.tCertPoolChannel <- &tCertImpl{client, x509Cert, tempSK}
	}

	if j == 0 {
		client.node.log.Error("No valid TCert was sent")

		return errors.New("No valid TCert was sent.")
	}

	return nil
}

func (client *clientImpl) callTCACreateCertificateSet(num int) ([]byte, [][]byte, error) {
	// Get a TCA Client
	sock, tcaP, err := client.node.getTCAClient()
	defer sock.Close()

	// Execute the protocol
	now := time.Now()
	timestamp := google_protobuf.Timestamp{int64(now.Second()), int32(now.Nanosecond())}
	req := &obcca.TCertCreateSetReq{
		&timestamp,
		&obcca.Identity{Id: client.node.enrollID},
		uint32(num),
		nil,
	}
	rawReq, err := proto.Marshal(req)
	if err != nil {
		client.node.log.Error("Failed marshaling request [%s] [%s].", err.Error())
		return nil, nil, err
	}

	// 2. Sign rawReq
	client.node.log.Debug("Signing req [% x]", rawReq)
	r, s, err := client.node.ecdsaSignWithEnrollmentKey(rawReq)
	if err != nil {
		client.node.log.Error("Failed creating signature [%s] [%s].", err.Error())
		return nil, nil, err
	}

	R, _ := r.MarshalText()
	S, _ := s.MarshalText()

	// 3. Append the signature
	req.Sig = &obcca.Signature{obcca.CryptoType_ECDSA, R, S}

	// 4. Send request
	certSet, err := tcaP.CreateCertificateSet(context.Background(), req)
	if err != nil {
		client.node.log.Error("Failed requesting tca create certificate set [%s].", err.Error())

		return nil, nil, err
	}

	return certSet.Certs.Key, certSet.Certs.Certs, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package featuretests

import (
	"github.com/juju/names"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/crossmodel"
	jujutesting "github.com/juju/juju/juju/testing"
	"github.com/juju/juju/testing"
)

type cmdCrossModelSuite struct {
	jujutesting.RepoSuite
}

func (s *cmdCrossModelSuite) SetUpTest(c *gc.C) {
	s.RepoSuite.SetUpTest(c)
	crossmodel.TempPlaceholder = make(map[names.ServiceTag]crossmodel.Offer)
}

func runOffer(c *gc.C, expectedError string, args ...string) {
	cmdArgs := append([]string{"offer"}, args...)
	context, err := runJujuCommand(c, cmdArgs...)
	if expectedError == "" {
		c.Assert(err, jc.ErrorIsNil)
	} else {
		c.Assert(err, gc.NotNil)
		c.Assert(testing.Stderr(context), jc.Contains, expectedError)
	}
}

func (s *cmdCrossModelSuite) TestOfferEmpty(c *gc.C) {
	runOffer(c, "error: an offer must at least specify service endpoint\n")
}

func (s *cmdCrossModelSuite) TestOfferInvalidEndpoints(c *gc.C) {
	runOffer(c, `error: endpoints must conform to format "<service-name>:<endpoint-name>[,...]" `, "fluff")
}

func (s *cmdCrossModelSuite) TestOfferAndShow(c *gc.C) {
	_, err := runJujuCommand(c, "offer", "test:db", "local:/u/fred/prod/hosted-db2")
	c.Assert(err, jc.ErrorIsNil)

	context, err := runJujuCommand(c, "show-endpoints", "local:/u/fred/prod/hosted-db2")
	c.Assert(err, jc.ErrorIsNil)

	expected := `
SERVICE  INTERFACES  DESCRIPTION
test     db

`[1:]
	c.Assert(testing.Stdout(context), gc.Matches, expected)
}
<commit_msg>Removed non-feature test...<commit_after><|endoftext|>


<commit_before>package mesos_test

import (
	"testing"

	"github.com/mesos/mesos-go"
)

func TestOpUnreserve(t *testing.T) {
	var (
		reservedCPU = resources(
			resource(name("cpus"),
				valueScalar(1),
				role("role"),
				reservation(reservedBy("principal"))))
		reservedMem = resources(
			resource(name("mem"),
				valueScalar(512),
				role("role"),
				reservation(reservedBy("principal"))))
		reserved = reservedCPU.Plus(reservedMem...)
	)

	// test case 1: unreserve some amount of CPU that's already been reserved
	unreservedCPU := reservedCPU.Flatten("", nil)
	wantsUnreserved := reservedMem.Plus(unreservedCPU...)
	actualUnreserved, err := unreserve(reservedCPU).Apply(reserved)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !wantsUnreserved.Equivalent(actualUnreserved) {
		t.Errorf("expected resources %+v instead of %+v", wantsUnreserved, actualUnreserved)
	}

	// test case 2: unreserve some amount of CPU greater than that which already been reserved
	reservedCPU2 := resources(
		resource(name("cpus"),
			valueScalar(2),
			role("role"),
			reservation(reservedBy("principal"))))
	_, err = unreserve(reservedCPU2).Apply(reserved)
	if err == nil {
		t.Fatalf("expected reservation error")
	}
}

func TestOpReserve(t *testing.T) {
	// func opReserve(operation mesos.Offer_Operation, resources mesos.Resources) (mesos.Resources, error)
	var (
		unreservedCPU = resources(resource(name("cpus"), valueScalar(1)))
		unreservedMem = resources(resource(name("mem"), valueScalar(512)))
		unreserved    = unreservedCPU.Plus(unreservedMem...)
		reservedCPU1  = unreservedCPU.Flatten("role", reservedBy("principal"))
	)

	// test case 1: reserve an amount of CPU that's available
	wantsReserved := unreservedMem.Plus(reservedCPU1...)
	actualReserved, err := reserve(reservedCPU1).Apply(unreserved)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !wantsReserved.Equivalent(actualReserved) {
		t.Errorf("expected resources %+v instead of %+v", wantsReserved, actualReserved)
	}

	// test case 2: reserve an amount of CPU that's NOT available
	reservedCPU2 := resources(
		resource(name("cpus"),
			valueScalar(2),
			role("role"),
			reservation(reservedBy("principal"))))
	_, err = reserve(reservedCPU2).Apply(unreserved)
	if err == nil {
		t.Fatalf("expected reservation error")
	}
}

func reservedBy(principal string) *mesos.Resource_ReservationInfo {
	return &mesos.Resource_ReservationInfo{
		Principal: principal,
	}
}

func reserve(r mesos.Resources) *mesos.Offer_Operation {
	return &mesos.Offer_Operation{
		Type: mesos.RESERVE.Enum(),
		Reserve: &mesos.Offer_Operation_Reserve{
			Resources: r,
		},
	}
}

func unreserve(r mesos.Resources) *mesos.Offer_Operation {
	return &mesos.Offer_Operation{
		Type: mesos.UNRESERVE.Enum(),
		Unreserve: &mesos.Offer_Operation_Unreserve{
			Resources: r,
		},
	}
}

func create(r mesos.Resources) *mesos.Offer_Operation {
	return &mesos.Offer_Operation{
		Type: mesos.CREATE.Enum(),
		Create: &mesos.Offer_Operation_Create{
			Volumes: r,
		},
	}
}

func reservation(ri *mesos.Resource_ReservationInfo) resourceOpt {
	return func(r *mesos.Resource) {
		r.Reservation = ri
	}
}

func disk(persistenceID, containerPath string) resourceOpt {
	return func(r *mesos.Resource) {
		r.Disk = &mesos.Resource_DiskInfo{}
		if containerPath != "" {
			r.Disk.Volume = &mesos.Volume{ContainerPath: containerPath}
		}
		if persistenceID != "" {
			r.Disk.Persistence = &mesos.Resource_DiskInfo_Persistence{ID: persistenceID}
		}
	}
}
<commit_msg>create volume test<commit_after>package mesos_test

import (
	"testing"

	"github.com/mesos/mesos-go"
)

func TestOpCreate(t *testing.T) {
	var (
		total = resources(
			resource(name("cpus"), valueScalar(1)),
			resource(name("mem"), valueScalar(512)),
			resource(name("disk"), valueScalar(1000), role("role")),
		)
		volume1 = resource(name("disk"), valueScalar(200), role("role"), disk("1", "path"))
		volume2 = resource(name("disk"), valueScalar(2000), role("role"), disk("1", "path"))
	)
	op := create(resources(volume1))
	rs, err := op.Apply(total)
	if err != nil {
		t.Fatalf("unexpected error: %+v", err)
	}
	expected := resources(
		resource(name("cpus"), valueScalar(1)),
		resource(name("mem"), valueScalar(512)),
		resource(name("disk"), valueScalar(800), role("role")),
		volume1,
	)
	if !expected.Equivalent(rs) {
		t.Fatalf("expected %v instead of %v", expected, rs)
	}

	// check the case of insufficient disk resources
	op = create(resources(volume2))
	_, err = op.Apply(total)
	if err == nil {
		t.Fatalf("expected an error due to insufficient disk resources")
	}
}

func TestOpUnreserve(t *testing.T) {
	var (
		reservedCPU = resources(
			resource(name("cpus"),
				valueScalar(1),
				role("role"),
				reservation(reservedBy("principal"))))
		reservedMem = resources(
			resource(name("mem"),
				valueScalar(512),
				role("role"),
				reservation(reservedBy("principal"))))
		reserved = reservedCPU.Plus(reservedMem...)
	)

	// test case 1: unreserve some amount of CPU that's already been reserved
	unreservedCPU := reservedCPU.Flatten("", nil)
	wantsUnreserved := reservedMem.Plus(unreservedCPU...)
	actualUnreserved, err := unreserve(reservedCPU).Apply(reserved)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !wantsUnreserved.Equivalent(actualUnreserved) {
		t.Errorf("expected resources %+v instead of %+v", wantsUnreserved, actualUnreserved)
	}

	// test case 2: unreserve some amount of CPU greater than that which already been reserved
	reservedCPU2 := resources(
		resource(name("cpus"),
			valueScalar(2),
			role("role"),
			reservation(reservedBy("principal"))))
	_, err = unreserve(reservedCPU2).Apply(reserved)
	if err == nil {
		t.Fatalf("expected reservation error")
	}
}

func TestOpReserve(t *testing.T) {
	// func opReserve(operation mesos.Offer_Operation, resources mesos.Resources) (mesos.Resources, error)
	var (
		unreservedCPU = resources(resource(name("cpus"), valueScalar(1)))
		unreservedMem = resources(resource(name("mem"), valueScalar(512)))
		unreserved    = unreservedCPU.Plus(unreservedMem...)
		reservedCPU1  = unreservedCPU.Flatten("role", reservedBy("principal"))
	)

	// test case 1: reserve an amount of CPU that's available
	wantsReserved := unreservedMem.Plus(reservedCPU1...)
	actualReserved, err := reserve(reservedCPU1).Apply(unreserved)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !wantsReserved.Equivalent(actualReserved) {
		t.Errorf("expected resources %+v instead of %+v", wantsReserved, actualReserved)
	}

	// test case 2: reserve an amount of CPU that's NOT available
	reservedCPU2 := resources(
		resource(name("cpus"),
			valueScalar(2),
			role("role"),
			reservation(reservedBy("principal"))))
	_, err = reserve(reservedCPU2).Apply(unreserved)
	if err == nil {
		t.Fatalf("expected reservation error")
	}
}

func reservedBy(principal string) *mesos.Resource_ReservationInfo {
	return &mesos.Resource_ReservationInfo{
		Principal: principal,
	}
}

func reserve(r mesos.Resources) *mesos.Offer_Operation {
	return &mesos.Offer_Operation{
		Type: mesos.RESERVE.Enum(),
		Reserve: &mesos.Offer_Operation_Reserve{
			Resources: r,
		},
	}
}

func unreserve(r mesos.Resources) *mesos.Offer_Operation {
	return &mesos.Offer_Operation{
		Type: mesos.UNRESERVE.Enum(),
		Unreserve: &mesos.Offer_Operation_Unreserve{
			Resources: r,
		},
	}
}

func create(r mesos.Resources) *mesos.Offer_Operation {
	return &mesos.Offer_Operation{
		Type: mesos.CREATE.Enum(),
		Create: &mesos.Offer_Operation_Create{
			Volumes: r,
		},
	}
}

func reservation(ri *mesos.Resource_ReservationInfo) resourceOpt {
	return func(r *mesos.Resource) {
		r.Reservation = ri
	}
}

func disk(persistenceID, containerPath string) resourceOpt {
	return func(r *mesos.Resource) {
		r.Disk = &mesos.Resource_DiskInfo{}
		if containerPath != "" {
			r.Disk.Volume = &mesos.Volume{ContainerPath: containerPath}
		}
		if persistenceID != "" {
			r.Disk.Persistence = &mesos.Resource_DiskInfo_Persistence{ID: persistenceID}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Eric Wollesen <ericw at xmtp dot net>
//
// 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 discord

import (
	"flag"
	"os"
	"strings"
	"sync"
	"time"

	redis "gopkg.in/redis.v5"

	"github.com/ewollesen/discordgo"
	"github.com/spacemonkeygo/errors"
	"github.com/spacemonkeygo/spacelog"

	"github.com/ewollesen/zenbot/cache"
	memorycache "github.com/ewollesen/zenbot/cache/memory"
	"github.com/ewollesen/zenbot/cache/rediscache"
	"github.com/ewollesen/zenbot/commands"
	"github.com/ewollesen/zenbot/httpapi"
	"github.com/ewollesen/zenbot/overwatch"
	"github.com/ewollesen/zenbot/overwatch/lootbox"
	"github.com/ewollesen/zenbot/queue"
	memoryqueue "github.com/ewollesen/zenbot/queue/memory"
	"github.com/ewollesen/zenbot/queue/redisqueue"
)

var (
	Error  = errors.NewClass("discord")
	logger = spacelog.GetLogger()

	commandPrefix = flag.String("discord.command_prefix", "!",
		"String that prefixes bot commands")
	discordToken  = flag.String("discord.token", "", "Discord bot API token")
	game          = flag.String("discord.game", "!queue help", "Game being played")
	redisKeySpace = flag.String("discord.redis_keyspace", "discord",
		"redis keyspace prefix")
)

type bot struct {
	command_handlers_mu sync.Mutex
	command_handlers    map[string]DiscordHandler

	handler_callbacks []func()
	user_id           string

	session_cache cache.Cache

	oauth_mu     sync.Mutex
	oauth_states map[string]string
}

func New(redis_client *redis.Client) *bot {
	b := &bot{
		command_handlers: make(map[string]DiscordHandler),
		oauth_states:     make(map[string]string),
	}

	b.RegisterCommand("ping", &discordHandler{commands.Pong})

	var q queue.Queue
	var c, owc cache.Cache
	if redis_client != nil {
		logger.Infof("using redis queue and cache")
		q = redisqueue.New(redis_client, *redisKeySpace+".queues.scrimmages")
		c = rediscache.New(redis_client, *redisKeySpace+".caches.battletags", 0)
		owc = rediscache.New(redis_client, *redisKeySpace+".caches.overwatch", time.Hour*2)
	} else {
		logger.Infof("using memory queue and cache")
		q = memoryqueue.New()
		c = memorycache.New()
		owc = memorycache.New()
	}

	b.session_cache = memorycache.New()

	btq := newBattleTagQueue(q)
	btc := NewBattleTagCache(c)
	qh := newQueueHandler(btq, btc)
	b.RegisterCommand("dequeue", qh)
	b.RegisterCommand("enqueue", qh)
	b.RegisterCommand("queue", qh)

	dh := newDebugHandler(btq, btc)
	b.RegisterCommand("debug", dh)

	cow := overwatch.NewCaching(lootbox.New(), owc)
	srh := newSkillRankHandler(cow)
	b.RegisterCommand("sr", srh)

	return b
}

func (b *bot) Run(quit chan os.Signal) error {
	session, err := b.logIn()
	if err != nil {
		return err
	}

	logger.Info("online")

	if *game != "" {
		logger.Warne(session.UpdateStatus(0, *game))
	}

	<-quit
	logger.Info("shutting down")
	b.logOut(session)

	return nil
}

func (b *bot) logIn() (*discordgo.Session, error) {
	token, err := getToken()
	if err != nil {
		return nil, err
	}

	session, err := discordgo.New(token)
	if err != nil {
		return nil, Error.Wrap(err)
	}
	if session == nil {
		return nil, Error.New("log in failed")
	}

	if err = b.addHandlers(session); err != nil {
		return nil, Error.Wrap(err)
	}

	if err = session.Open(); err != nil {
		return nil, Error.Wrap(err)
	}

	return session, nil
}

func getToken() (token string, err error) {
	if token = os.Getenv("DISCORD_TOKEN"); token != "" {
		return token, nil
	}

	if *discordToken != "" {
		return *discordToken, nil
	}

	return "", Error.New("no discord token specified")
}

func (b *bot) addHandlers(session *discordgo.Session) error {
	b.handler_callbacks = append(b.handler_callbacks,
		session.AddHandler(b.messageHandler),
		session.AddHandler(b.presenceHandler))

	return nil
}

func (b *bot) logOut(session *discordgo.Session) {
	logger.Errore(session.Close())
	logger.Errore(b.removeHandlers())
	logger.Info("offline")
}

func (b *bot) removeHandlers() error {
	for _, cb := range b.handler_callbacks {
		if cb != nil {
			cb()
		}
	}

	return nil
}

func (b *bot) messageHandler(ds *discordgo.Session,
	m *discordgo.MessageCreate) {

	s := newCachingSession(ds, b.session_cache)

	if m.Author.ID == b.myUserId(s) {
		logger.Debugf("-> %s: %s", m.Author.Username, m.Content)
		return
	}

	private := isPrivateMessage(s, m)
	if !private && !strings.HasPrefix(m.Content, *commandPrefix) {
		return
	}
	logger.Debugf("<-%t %s: %s", private, m.Author.Username, m.Content)

	if err := b.handleCommand(s, m); err != nil {
		logger.Warne(err)
		return
	}
}

func (b *bot) myUserId(s Session) string {
	if b.user_id != "" {
		return b.user_id
	}

	user, err := s.User("@me")
	if err != nil {
		logger.Warnf("unable to find my user id: %v", err)
		return ""
	}
	b.user_id = user.ID

	return user.ID
}

func (b *bot) presenceHandler(s *discordgo.Session, m *discordgo.PresenceUpdate) {
	logger.Debugf("presenceHandler isn't implemented yet")
	return
}

func isPrivateMessage(s Session, m *discordgo.MessageCreate) bool {
	ch, err := s.Channel(m.ChannelID)
	if err != nil {
		return false
	}
	return ch.IsPrivate
}

func (b *bot) ReceiveRouter(router httpapi.Router) {
	router.HandleFunc("/", b.handleHTTP)
	router.HandleFunc("/oauth/redirect", b.oauthRedirect)
}
<commit_msg>enhanced exit logging<commit_after>// Copyright 2016 Eric Wollesen <ericw at xmtp dot net>
//
// 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 discord

import (
	"flag"
	"os"
	"strings"
	"sync"
	"time"

	redis "gopkg.in/redis.v5"

	"github.com/ewollesen/discordgo"
	"github.com/spacemonkeygo/errors"
	"github.com/spacemonkeygo/spacelog"

	"github.com/ewollesen/zenbot/cache"
	memorycache "github.com/ewollesen/zenbot/cache/memory"
	"github.com/ewollesen/zenbot/cache/rediscache"
	"github.com/ewollesen/zenbot/commands"
	"github.com/ewollesen/zenbot/httpapi"
	"github.com/ewollesen/zenbot/overwatch"
	"github.com/ewollesen/zenbot/overwatch/lootbox"
	"github.com/ewollesen/zenbot/queue"
	memoryqueue "github.com/ewollesen/zenbot/queue/memory"
	"github.com/ewollesen/zenbot/queue/redisqueue"
)

var (
	Error  = errors.NewClass("discord")
	logger = spacelog.GetLogger()

	commandPrefix = flag.String("discord.command_prefix", "!",
		"String that prefixes bot commands")
	discordToken  = flag.String("discord.token", "", "Discord bot API token")
	game          = flag.String("discord.game", "!queue help", "Game being played")
	redisKeySpace = flag.String("discord.redis_keyspace", "discord",
		"redis keyspace prefix")
)

type bot struct {
	command_handlers_mu sync.Mutex
	command_handlers    map[string]DiscordHandler

	handler_callbacks []func()
	user_id           string

	session_cache cache.Cache

	oauth_mu     sync.Mutex
	oauth_states map[string]string
}

func New(redis_client *redis.Client) *bot {
	b := &bot{
		command_handlers: make(map[string]DiscordHandler),
		oauth_states:     make(map[string]string),
	}

	b.RegisterCommand("ping", &discordHandler{commands.Pong})

	var q queue.Queue
	var c, owc cache.Cache
	if redis_client != nil {
		logger.Infof("using redis queue and cache")
		q = redisqueue.New(redis_client, *redisKeySpace+".queues.scrimmages")
		c = rediscache.New(redis_client, *redisKeySpace+".caches.battletags", 0)
		owc = rediscache.New(redis_client, *redisKeySpace+".caches.overwatch", time.Hour*2)
	} else {
		logger.Infof("using memory queue and cache")
		q = memoryqueue.New()
		c = memorycache.New()
		owc = memorycache.New()
	}

	b.session_cache = memorycache.New()

	btq := newBattleTagQueue(q)
	btc := NewBattleTagCache(c)
	qh := newQueueHandler(btq, btc)
	b.RegisterCommand("dequeue", qh)
	b.RegisterCommand("enqueue", qh)
	b.RegisterCommand("queue", qh)

	dh := newDebugHandler(btq, btc)
	b.RegisterCommand("debug", dh)

	cow := overwatch.NewCaching(lootbox.New(), owc)
	srh := newSkillRankHandler(cow)
	b.RegisterCommand("sr", srh)

	return b
}

func (b *bot) Run(quit chan os.Signal) error {
	session, err := b.logIn()
	if err != nil {
		return err
	}

	logger.Info("online")

	if *game != "" {
		logger.Warne(session.UpdateStatus(0, *game))
	}

	signal, closed := <-quit
	if closed {
		logger.Debugf("quit channel closed, shutting down")
	}
	if signal != nil {
		logger.Debugf("quit channel received signal %v, shutting down", signal)
	}

	logger.Info("shutting down")
	b.logOut(session)

	return nil
}

func (b *bot) logIn() (*discordgo.Session, error) {
	token, err := getToken()
	if err != nil {
		return nil, err
	}

	session, err := discordgo.New(token)
	if err != nil {
		return nil, Error.Wrap(err)
	}
	if session == nil {
		return nil, Error.New("log in failed")
	}

	if err = b.addHandlers(session); err != nil {
		return nil, Error.Wrap(err)
	}

	if err = session.Open(); err != nil {
		return nil, Error.Wrap(err)
	}

	return session, nil
}

func getToken() (token string, err error) {
	if token = os.Getenv("DISCORD_TOKEN"); token != "" {
		return token, nil
	}

	if *discordToken != "" {
		return *discordToken, nil
	}

	return "", Error.New("no discord token specified")
}

func (b *bot) addHandlers(session *discordgo.Session) error {
	b.handler_callbacks = append(b.handler_callbacks,
		session.AddHandler(b.messageHandler),
		session.AddHandler(b.presenceHandler))

	return nil
}

func (b *bot) logOut(session *discordgo.Session) {
	logger.Errore(session.Close())
	logger.Errore(b.removeHandlers())
	logger.Info("offline")
}

func (b *bot) removeHandlers() error {
	for _, cb := range b.handler_callbacks {
		if cb != nil {
			cb()
		}
	}

	return nil
}

func (b *bot) messageHandler(ds *discordgo.Session,
	m *discordgo.MessageCreate) {

	s := newCachingSession(ds, b.session_cache)

	if m.Author.ID == b.myUserId(s) {
		logger.Debugf("-> %s: %s", m.Author.Username, m.Content)
		return
	}

	private := isPrivateMessage(s, m)
	if !private && !strings.HasPrefix(m.Content, *commandPrefix) {
		return
	}
	logger.Debugf("<-%t %s: %s", private, m.Author.Username, m.Content)

	if err := b.handleCommand(s, m); err != nil {
		logger.Warne(err)
		return
	}
}

func (b *bot) myUserId(s Session) string {
	if b.user_id != "" {
		return b.user_id
	}

	user, err := s.User("@me")
	if err != nil {
		logger.Warnf("unable to find my user id: %v", err)
		return ""
	}
	b.user_id = user.ID

	return user.ID
}

func (b *bot) presenceHandler(s *discordgo.Session, m *discordgo.PresenceUpdate) {
	logger.Debugf("presenceHandler isn't implemented yet")
	return
}

func isPrivateMessage(s Session, m *discordgo.MessageCreate) bool {
	ch, err := s.Channel(m.ChannelID)
	if err != nil {
		return false
	}
	return ch.IsPrivate
}

func (b *bot) ReceiveRouter(router httpapi.Router) {
	router.HandleFunc("/", b.handleHTTP)
	router.HandleFunc("/oauth/redirect", b.oauthRedirect)
}
<|endoftext|>


<commit_before>package admin

import (
	"diskette/collections"
	"diskette/util"
	"net/http"
	"time"

	"github.com/labstack/echo"
	"labix.org/v2/mgo/bson"
)

// http DELETE localhost:5025/admin/remove-unconfirmed-users X-Diskette-Session-Token:<session_token>
func (service *serviceImpl) RemoveUnconfirmedUsers(c *echo.Context) error {

	d, err := time.ParseDuration("12h")
	if err != nil {
		return err
	}

	objectIds := []bson.ObjectId{}
	iter := service.userCollection.Find(nil).Iter()
	var userDoc collections.UserDocument
	for iter.Next(&userDoc) {
		isAccountOld := userDoc.CreatedAt.Before(time.Now().Add(-1 * d))
		isAccountNotConfirmed := userDoc.ConfirmedAt.Before(userDoc.CreatedAt)
		if isAccountOld && isAccountNotConfirmed {
			objectIds = append(objectIds, userDoc.Id)
		}
	}
	if err := iter.Close(); err != nil {
		return err
	}

	info, err := service.userCollection.RemoveAll(bson.M{
		"_id": bson.M{
			"$in": objectIds,
		},
	})
	if err != nil {
		return c.JSON(http.StatusInternalServerError, util.CreateErrResponse(err))
	}

	return c.JSON(http.StatusOK, util.CreateOkResponse(info.Updated))
}
<commit_msg>small refactoring<commit_after>package admin

import (
	"diskette/collections"
	"diskette/util"
	"net/http"
	"time"

	"github.com/labstack/echo"
	"labix.org/v2/mgo/bson"
)

// http DELETE localhost:5025/admin/remove-unconfirmed-users X-Diskette-Session-Token:<session_token>
func (service *serviceImpl) RemoveUnconfirmedUsers(c *echo.Context) error {

	d, err := time.ParseDuration("12h")
	if err != nil {
		return err
	}

	objectIds := []bson.ObjectId{}
	iter := service.userCollection.Find(nil).Iter()
	var userDoc collections.UserDocument
	for iter.Next(&userDoc) {
		isAccountOld := userDoc.CreatedAt.Before(time.Now().Add(-1 * d))
		isAccountConfirmed := userDoc.ConfirmedAt.After(userDoc.CreatedAt)
		if isAccountOld && !isAccountConfirmed {
			objectIds = append(objectIds, userDoc.Id)
		}
	}
	if err := iter.Close(); err != nil {
		return err
	}

	info, err := service.userCollection.RemoveAll(bson.M{
		"_id": bson.M{
			"$in": objectIds,
		},
	})
	if err != nil {
		return c.JSON(http.StatusInternalServerError, util.CreateErrResponse(err))
	}

	return c.JSON(http.StatusOK, util.CreateOkResponse(info.Updated))
}
<|endoftext|>


<commit_before><commit_msg>Add short for separator flag.<commit_after><|endoftext|>