<commit_before>// clock counts down to or up from a target time.
package main

import (
	"fmt"
	"os"
	"time"
)

func main() {
	// Exit on Enter key
	go func() {
		buf := make([]byte, 1)
		_, _ = os.Stdin.Read(buf)
		os.Exit(0)
	}()

	const (
		indent          = "\t"
		highlight_start = "\x1b[1;36m"
		highlight_end   = "\x1b[0m"
	)
	fmt.Print(indent, highlight_start, "Just Go", highlight_end, "\n")
	target := time.Date(2016, 2, 29, 0, 0, 0, 0, time.UTC)
	fmt.Print(indent, target.Format(time.UnixDate), "\n")

	var previous time.Time

	for {
		now := time.Now()
		now = now.Add(time.Duration(-now.Nanosecond())) // truncate to second

		if now != previous {
			previous = now
			remaining := target.Sub(now)

			var sign string
			if remaining > 0 {
				sign = "-" // countdown is "T minus..."
			} else {
				sign = "+" // count up is "T plus..."
				remaining = -remaining
			}

			var days int
			if remaining >= 24*time.Hour {
				days = int(remaining / (24 * time.Hour))
				remaining = remaining % (24 * time.Hour)
			}

			fmt.Print(indent, now.Format(time.UnixDate), "  ", sign)
			if days > 0 {
				fmt.Print(days, "d")
			}
			fmt.Print(remaining, "          \r")
		}
		time.Sleep(50 * time.Millisecond)
	}
}
<commit_msg>March<commit_after>// clock counts down to or up from a target time.
package main

import (
	"fmt"
	"os"
	"time"
)

func main() {
	// Exit on Enter key
	go func() {
		buf := make([]byte, 1)
		_, _ = os.Stdin.Read(buf)
		os.Exit(0)
	}()

	const (
		indent          = "\t"
		highlight_start = "\x1b[1;36m"
		highlight_end   = "\x1b[0m"
	)
	fmt.Print(indent, highlight_start, "Just Go", highlight_end, "\n")
	target := time.Date(2016, 3, 1, 0, 0, 0, 0, time.Local)
	fmt.Print(indent, target.Format(time.UnixDate), "\n")

	var previous time.Time

	for {
		now := time.Now()
		now = now.Add(time.Duration(-now.Nanosecond())) // truncate to second

		if now != previous {
			previous = now
			remaining := target.Sub(now)

			var sign string
			if remaining > 0 {
				sign = "-" // countdown is "T minus..."
			} else {
				sign = "+" // count up is "T plus..."
				remaining = -remaining
			}

			var days int
			if remaining >= 24*time.Hour {
				days = int(remaining / (24 * time.Hour))
				remaining = remaining % (24 * time.Hour)
			}

			fmt.Print(indent, now.Format(time.UnixDate), "  ", sign)
			if days > 0 {
				fmt.Print(days, "d")
			}
			fmt.Print(remaining, "          \r")
		}
		time.Sleep(50 * time.Millisecond)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2011 Google Inc.

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

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

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

package localdisk

import (
	"camli/blobref"
	"camli/blobserver"
	"crypto/sha1"
	"fmt"
	"io"
	"os"
	"strings"
	"sync"
	"testing"
)

func cleanUp(ds *diskStorage) {
	os.RemoveAll(ds.root)
}

var (
	epochLock sync.Mutex
	rootEpoch = 0
)

func NewStorage(t *testing.T) *diskStorage {
	epochLock.Lock()
	rootEpoch++
	path := fmt.Sprintf("%s/camli-testroot-%d-%d", os.TempDir(), os.Getpid(), rootEpoch)
	epochLock.Unlock()
	if err := os.Mkdir(path, 0755); err != nil {
		t.Fatalf("Failed to create temp directory %q: %v", path, err)
	}
	ds, err := New(path)
	if err != nil {
		t.Fatalf("Failed to run New: %v", err)
	}
	return ds.(*diskStorage)
}

type testBlob struct {
	val string
}

func (tb *testBlob) BlobRef() *blobref.BlobRef {
	h := sha1.New()
	h.Write([]byte(tb.val))
	return blobref.FromHash("sha1", h)
}

func (tb *testBlob) Size() int64 {
	return int64(len(tb.val))
}

func (tb *testBlob) Reader() io.Reader {
	return strings.NewReader(tb.val)
}

func TestReceiveStat(t *testing.T) {
	ds := NewStorage(t)
	defer cleanUp(ds)

	tb := &testBlob{"Foo"}
	sb, err := ds.ReceiveBlob(tb.BlobRef(), tb.Reader(), nil)
	if err != nil {
		t.Fatalf("ReceiveBlob error: %v", err)
	}
	checkSizedBlob := func() {
		if sb.Size != tb.Size() {
			t.Fatalf("Got size %d; expected %d", sb.Size, tb.Size())
		}
		if sb.BlobRef.String() != tb.BlobRef().String() {
			t.Fatalf("Got blob %q; expected %q", sb.BlobRef.String(), tb.BlobRef())
		}
	}
	checkSizedBlob()

	ch := make(chan *blobref.SizedBlobRef, 0)
	errch := make(chan os.Error, 1)
	go func() {
		errch <- ds.Stat(ch, blobserver.DefaultPartition, []*blobref.BlobRef{tb.BlobRef()}, 0)
		close(ch)
	}()
	got := 0
	for sb = range ch {
		got++
		checkSizedBlob()
	}
	if got != 1 {
		t.Fatalf("Expected %d stat results; got %d", 1, got)
	}
	if err = <-errch; err != nil {
		t.Fatalf("Got error from stat: %v", err)
	}
}

func TestStatWait(t *testing.T) {
	ds := NewStorage(t)
	defer cleanUp(ds)
	// TODO
}
<commit_msg>Writing a (failing) test for Stat waiting.<commit_after>/*
Copyright 2011 Google Inc.

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

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

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

package localdisk

import (
	"camli/blobref"
	"camli/blobserver"
	"crypto/sha1"
	"fmt"
	"io"
	"os"
	"strings"
	"sync"
	"testing"
	"time"
)

func cleanUp(ds *diskStorage) {
	os.RemoveAll(ds.root)
}

var (
	epochLock sync.Mutex
	rootEpoch = 0
)

func NewStorage(t *testing.T) *diskStorage {
	epochLock.Lock()
	rootEpoch++
	path := fmt.Sprintf("%s/camli-testroot-%d-%d", os.TempDir(), os.Getpid(), rootEpoch)
	epochLock.Unlock()
	if err := os.Mkdir(path, 0755); err != nil {
		t.Fatalf("Failed to create temp directory %q: %v", path, err)
	}
	ds, err := New(path)
	if err != nil {
		t.Fatalf("Failed to run New: %v", err)
	}
	return ds.(*diskStorage)
}

type testBlob struct {
	val string
}

func (tb *testBlob) BlobRef() *blobref.BlobRef {
	h := sha1.New()
	h.Write([]byte(tb.val))
	return blobref.FromHash("sha1", h)
}

func (tb *testBlob) BlobRefSlice() []*blobref.BlobRef {
	return []*blobref.BlobRef{tb.BlobRef()}
}

func (tb *testBlob) Size() int64 {
	return int64(len(tb.val))
}

func (tb *testBlob) Reader() io.Reader {
	return strings.NewReader(tb.val)
}

func (tb *testBlob) AssertMatches(t *testing.T, sb *blobref.SizedBlobRef) {
	if sb.Size != tb.Size() {
		t.Fatalf("Got size %d; expected %d", sb.Size, tb.Size())
	}
	if sb.BlobRef.String() != tb.BlobRef().String() {
		t.Fatalf("Got blob %q; expected %q", sb.BlobRef.String(), tb.BlobRef())
	}
}

func (tb *testBlob) ExpectUploadBlob(t *testing.T, ds blobserver.BlobReceiver) {
	sb, err := ds.ReceiveBlob(tb.BlobRef(), tb.Reader(), nil)
	if err != nil {
		t.Fatalf("ReceiveBlob error: %v", err)
	}
	tb.AssertMatches(t, sb)
}

func TestReceiveStat(t *testing.T) {
	ds := NewStorage(t)
	defer cleanUp(ds)

	tb := &testBlob{"Foo"}
	tb.ExpectUploadBlob(t, ds)

	ch := make(chan *blobref.SizedBlobRef, 0)
	errch := make(chan os.Error, 1)
	go func() {
		errch <- ds.Stat(ch, blobserver.DefaultPartition, tb.BlobRefSlice(), 0)
		close(ch)
	}()
	got := 0
	for sb := range ch {
		got++
		tb.AssertMatches(t, sb)
		break
	}
	if got != 1 {
		t.Fatalf("Expected %d stat results; got %d", 1, got)
	}
	if err := <-errch; err != nil {
		t.Fatalf("Got error from stat: %v", err)
	}
}

func TestStatWait(t *testing.T) {
	ds := NewStorage(t)
	defer cleanUp(ds)
	tb := &testBlob{"Foo"}

	// Do a stat before the blob exists, but wait 5 seconds for it to arrive.
	const waitSeconds = 5
	ch := make(chan *blobref.SizedBlobRef, 0)
	errch := make(chan os.Error, 1)
	go func() {
		errch <- ds.Stat(ch, blobserver.DefaultPartition, tb.BlobRefSlice(), waitSeconds)
		close(ch)
	}()

	// Sum and verify the stat results, writing the total number of returned matches
	// to statCountCh (expected: 1)
	statCountCh := make(chan int)
	go func() {
		got := 0
		for sb := range ch {
			got++
			tb.AssertMatches(t, sb)
		}
		statCountCh <- got
	}()

	// Now upload the blob, now that everything else is in-flight.
	// Sleep a bit to make sure the ds.Stat above has had a chance to fail and sleep.
	time.Sleep(1e9 / 5)  // 200ms in nanos
	tb.ExpectUploadBlob(t, ds)

	if got := <- statCountCh; got != 1 {
		t.Fatalf("Expected %d stat results; got %d", 1, got)
	}
}
<|endoftext|>


<commit_before>package device_manager

import (
	"bufio"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/golang/mock/gomock"
	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"

	v1 "kubevirt.io/api/core/v1"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/client-go/kubernetes/fake"

	"kubevirt.io/kubevirt/pkg/testutils"
	virtconfig "kubevirt.io/kubevirt/pkg/virt-config"
)

const (
	fakeMdevNameSelector      = "FAKE 123"
	fakeIntelMdevNameSelector = "i915-GVTg_V5_4"
	fakeMdevResourceName      = "example.org/fake123"
	fakeMdevUUID              = "53764d0e-85a0-42b4-af5c-2046b460b1dc"
	fakeIntelMdevUUID         = "54444d0e-85a0-42b4-af5c-2046b4bbb1aa"
)

var _ = Describe("Mediated Device", func() {
	var mockPCI *MockDeviceHandler
	var fakePermittedHostDevicesConfig string
	var fakePermittedHostDevices v1.PermittedHostDevices
	var ctrl *gomock.Controller
	var fakeSupportedTypesPath string
	var clientTest *fake.Clientset
	resourceNameToTypeName := func(rawName string) string {
		typeNameStr := strings.Replace(string(rawName), " ", "_", -1)
		typeNameStr = strings.TrimSpace(typeNameStr)
		return typeNameStr
	}
	BeforeEach(func() {
		clientTest = fake.NewSimpleClientset()
		By("creating a temporary fake mdev directory tree")
		// create base mdev dir instead of /sys/bus/mdev/devices
		fakeMdevBasePath, err := ioutil.TempDir("/tmp", "mdevs")
		Expect(err).ToNot(HaveOccurred())
		// create an alternative mdev_supported_types dir instead of /sys/class/mdev_bus/[pciAddress]/
		fakeSupportedTypesPath, err = ioutil.TempDir("/tmp", "mdev_supported_types")
		Expect(err).ToNot(HaveOccurred())
		// create a fake path to nvidia mdev type
		fakeNvidiaTypePath := filepath.Join(fakeSupportedTypesPath, "nvidia-222")
		err = os.MkdirAll(fakeNvidiaTypePath, 0700)
		Expect(err).ToNot(HaveOccurred())
		// create a fake path to Intel mdev type
		fakeIntelTypePath := filepath.Join(fakeSupportedTypesPath, fakeIntelMdevNameSelector)
		err = os.MkdirAll(fakeIntelTypePath, 0700)
		Expect(err).ToNot(HaveOccurred())
		mdevBasePath = fakeMdevBasePath
		// create mdev directories and symlinks
		for _, uuid := range []string{fakeMdevUUID, fakeIntelMdevUUID} {
			mdevTypePath := filepath.Join(fakeMdevBasePath, uuid+"real")
			err = os.MkdirAll(mdevTypePath, 0700)
			Expect(err).ToNot(HaveOccurred())
			err = os.Symlink(filepath.Join(fakeMdevBasePath, uuid+"real"), filepath.Join(fakeMdevBasePath, uuid))
			Expect(err).ToNot(HaveOccurred())
		}
		// link nvidia type directory
		err = os.Symlink(fakeNvidiaTypePath, filepath.Join(fakeMdevBasePath, fakeMdevUUID+"real", "mdev_type"))
		Expect(err).ToNot(HaveOccurred())
		err = os.Symlink(fakeIntelTypePath, filepath.Join(fakeMdevBasePath, fakeIntelMdevUUID+"real", "mdev_type"))
		Expect(err).ToNot(HaveOccurred())
		// create a name file in the nvidia type directory
		mdevName, err := os.Create(filepath.Join(fakeNvidiaTypePath, "name"))
		Expect(err).ToNot(HaveOccurred())
		mdevNameWriter := bufio.NewWriter(mdevName)
		n, err := mdevNameWriter.WriteString(fakeMdevNameSelector + "\n")
		Expect(err).ToNot(HaveOccurred())
		Expect(n).To(Equal(len(fakeMdevNameSelector) + 1))
		mdevNameWriter.Flush()

	})

	AfterEach(func() {
		os.RemoveAll(mdevBasePath)
		os.RemoveAll(fakeSupportedTypesPath)
	})
	table.DescribeTable("should get correct file type name", func(namePathExist bool) {
		if namePathExist {
			mdevName, err := getMdevTypeName(fakeMdevUUID)
			Expect(err).ToNot(HaveOccurred())
			Expect(mdevName).To(Equal(resourceNameToTypeName(fakeMdevNameSelector)))
		} else {
			mdevName, err := getMdevTypeName(fakeIntelMdevUUID)
			Expect(err).ToNot(HaveOccurred())
			Expect(mdevName).To(Equal(resourceNameToTypeName(fakeIntelMdevNameSelector)))
		}
	},
		table.Entry("Nvidia name file exist", true),
		table.Entry("Intel name file doesn't exist", false),
	)
	Context("discover devices", func() {
		BeforeEach(func() {
			By("mocking PCI and MDEV functions to simulate an mdev an its parent PCI device")
			ctrl = gomock.NewController(GinkgoT())
			mockPCI = NewMockDeviceHandler(ctrl)
			Handler = mockPCI
			// Force pre-defined returned values and ensure the function only get called exacly once each on 0000:00:00.0
			mockPCI.EXPECT().GetMdevParentPCIAddr(fakeMdevUUID).Return(fakeAddress, nil).Times(1)
			mockPCI.EXPECT().GetDeviceIOMMUGroup(mdevBasePath, fakeMdevUUID).Return(fakeIommuGroup, nil).Times(1)
			mockPCI.EXPECT().GetDeviceNumaNode(pciBasePath, fakeAddress).Return(fakeNumaNode).Times(1)

			By("creating a list of fake device using the yaml decoder")
			fakePermittedHostDevicesConfig = `
    mediatedDevices:
    - mdevNameSelector: "` + fakeMdevNameSelector + `"
      resourceName: "` + fakeMdevResourceName + `"
    `
			err := yaml.NewYAMLOrJSONDecoder(strings.NewReader(fakePermittedHostDevicesConfig), 1024).Decode(&fakePermittedHostDevices)
			Expect(err).ToNot(HaveOccurred())
			Expect(fakePermittedHostDevices.MediatedDevices).To(HaveLen(1))
			Expect(fakePermittedHostDevices.MediatedDevices[0].MDEVNameSelector).To(Equal(fakeMdevNameSelector))
			Expect(fakePermittedHostDevices.MediatedDevices[0].ResourceName).To(Equal(fakeMdevResourceName))
		})
		AfterEach(func() {
			ctrl.Finish()
		})
		It("Should parse the permitted devices and find 1 matching mediated device", func() {
			supportedMdevsMap := make(map[string]string)
			for _, supportedMdev := range fakePermittedHostDevices.MediatedDevices {
				// do not add a device plugin for this resource if it's being provided via an external device plugin
				if !supportedMdev.ExternalResourceProvider {
					selector := removeSelectorSpaces(supportedMdev.MDEVNameSelector)
					supportedMdevsMap[selector] = supportedMdev.ResourceName
				}
			}
			// discoverPermittedHostMediatedDevices() will walk real mdev devices wherever the tests are running
			devices := discoverPermittedHostMediatedDevices(supportedMdevsMap)
			Expect(devices).To(HaveLen(1))
			selector := removeSelectorSpaces(fakeMdevNameSelector)
			Expect(devices[selector]).To(HaveLen(1))
			Expect(devices[selector][0].UUID).To(Equal(fakeMdevUUID))
			Expect(devices[selector][0].typeName).To(Equal(selector))
			Expect(devices[selector][0].parentPciAddress).To(Equal(fakeAddress))
			Expect(devices[selector][0].iommuGroup).To(Equal(fakeIommuGroup))
			Expect(devices[selector][0].numaNode).To(Equal(fakeNumaNode))
		})

		It("Should validate DPI devices", func() {
			iommuToMDEVMap := make(map[string]string)
			supportedMdevsMap := make(map[string]string)
			for _, supportedMdev := range fakePermittedHostDevices.MediatedDevices {
				// do not add a device plugin for this resource if it's being provided via an external device plugin
				if !supportedMdev.ExternalResourceProvider {
					selector := removeSelectorSpaces(supportedMdev.MDEVNameSelector)
					supportedMdevsMap[selector] = supportedMdev.ResourceName
				}
			}
			// discoverPermittedHostMediatedDevices() will walk real mdev devices wherever the tests are running
			mDevices := discoverPermittedHostMediatedDevices(supportedMdevsMap)
			selector := removeSelectorSpaces(fakeMdevNameSelector)
			devs := constructDPIdevicesFromMdev(mDevices[selector], iommuToMDEVMap)
			Expect(devs[0].ID).To(Equal(fakeIommuGroup))
			Expect(devs[0].Topology.Nodes[0].ID).To(Equal(int64(fakeNumaNode)))
		})

		It("Should update the device list according to the configmap", func() {
			By("creating a cluster config")
			kv := &v1.KubeVirt{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "kubevirt",
					Namespace: "kubevirt",
				},
				Spec: v1.KubeVirtSpec{
					Configuration: v1.KubeVirtConfiguration{
						DeveloperConfiguration: &v1.DeveloperConfiguration{},
					},
				},
				Status: v1.KubeVirtStatus{
					Phase: v1.KubeVirtPhaseDeploying,
				},
			}
			fakeClusterConfig, _, kvInformer := testutils.NewFakeClusterConfigUsingKV(kv)

			By("creating an empty device controller")
			var noDevices []Device
			deviceController := NewDeviceController("master", noDevices, fakeClusterConfig, clientTest.CoreV1())

			By("adding a host device to the cluster config")
			kvConfig := kv.DeepCopy()
			kvConfig.Spec.Configuration.DeveloperConfiguration.FeatureGates = []string{virtconfig.HostDevicesGate}
			kvConfig.Spec.Configuration.PermittedHostDevices = &v1.PermittedHostDevices{
				MediatedDevices: []v1.MediatedHostDevice{
					{
						MDEVNameSelector: fakeMdevNameSelector,
						ResourceName:     fakeMdevResourceName,
					},
				},
			}
			testutils.UpdateFakeKubeVirtClusterConfig(kvInformer, kvConfig)
			permittedDevices := fakeClusterConfig.GetPermittedHostDevices()
			Expect(permittedDevices).ToNot(BeNil(), "something went wrong while parsing the configmap(s)")
			Expect(len(permittedDevices.MediatedDevices)).To(Equal(1), "the fake device was not found")

			By("ensuring a device plugin gets created for our fake device")
			enabledDevicePlugins, disabledDevicePlugins := deviceController.splitPermittedDevices(
				deviceController.updatePermittedHostDevicePlugins(),
			)
			Expect(len(enabledDevicePlugins)).To(Equal(1), "a device plugin wasn't created for the fake device")
			Expect(disabledDevicePlugins).To(HaveLen(0))
			Ω(enabledDevicePlugins).Should(HaveKey(fakeMdevResourceName))
			// Manually adding the enabled plugin, since the device controller is not actually running
			deviceController.startedPlugins[fakeMdevResourceName] = controlledDevice{
				devicePlugin: enabledDevicePlugins[fakeMdevResourceName],
			}

			By("deletting the device from the configmap")
			kvConfig.Spec.Configuration.PermittedHostDevices = &v1.PermittedHostDevices{}
			testutils.UpdateFakeKubeVirtClusterConfig(kvInformer, kvConfig)
			permittedDevices = fakeClusterConfig.GetPermittedHostDevices()
			Expect(permittedDevices).ToNot(BeNil(), "something went wrong while parsing the configmap(s)")
			Expect(len(permittedDevices.MediatedDevices)).To(Equal(0), "the fake device was not deleted")

			By("ensuring the device plugin gets stopped")
			enabledDevicePlugins, disabledDevicePlugins = deviceController.splitPermittedDevices(
				deviceController.updatePermittedHostDevicePlugins(),
			)
			Expect(enabledDevicePlugins).To(HaveLen(0))
			Expect(len(disabledDevicePlugins)).To(Equal(1), "the fake device plugin did not get disabled")
			Ω(disabledDevicePlugins).Should(HaveKey(fakeMdevResourceName))
		})
	})
})
<commit_msg>Change HaveLen(0) to BeEmpty()<commit_after>package device_manager

import (
	"bufio"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/golang/mock/gomock"
	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"

	v1 "kubevirt.io/api/core/v1"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/client-go/kubernetes/fake"

	"kubevirt.io/kubevirt/pkg/testutils"
	virtconfig "kubevirt.io/kubevirt/pkg/virt-config"
)

const (
	fakeMdevNameSelector      = "FAKE 123"
	fakeIntelMdevNameSelector = "i915-GVTg_V5_4"
	fakeMdevResourceName      = "example.org/fake123"
	fakeMdevUUID              = "53764d0e-85a0-42b4-af5c-2046b460b1dc"
	fakeIntelMdevUUID         = "54444d0e-85a0-42b4-af5c-2046b4bbb1aa"
)

var _ = Describe("Mediated Device", func() {
	var mockPCI *MockDeviceHandler
	var fakePermittedHostDevicesConfig string
	var fakePermittedHostDevices v1.PermittedHostDevices
	var ctrl *gomock.Controller
	var fakeSupportedTypesPath string
	var clientTest *fake.Clientset
	resourceNameToTypeName := func(rawName string) string {
		typeNameStr := strings.Replace(string(rawName), " ", "_", -1)
		typeNameStr = strings.TrimSpace(typeNameStr)
		return typeNameStr
	}
	BeforeEach(func() {
		clientTest = fake.NewSimpleClientset()
		By("creating a temporary fake mdev directory tree")
		// create base mdev dir instead of /sys/bus/mdev/devices
		fakeMdevBasePath, err := ioutil.TempDir("/tmp", "mdevs")
		Expect(err).ToNot(HaveOccurred())
		// create an alternative mdev_supported_types dir instead of /sys/class/mdev_bus/[pciAddress]/
		fakeSupportedTypesPath, err = ioutil.TempDir("/tmp", "mdev_supported_types")
		Expect(err).ToNot(HaveOccurred())
		// create a fake path to nvidia mdev type
		fakeNvidiaTypePath := filepath.Join(fakeSupportedTypesPath, "nvidia-222")
		err = os.MkdirAll(fakeNvidiaTypePath, 0700)
		Expect(err).ToNot(HaveOccurred())
		// create a fake path to Intel mdev type
		fakeIntelTypePath := filepath.Join(fakeSupportedTypesPath, fakeIntelMdevNameSelector)
		err = os.MkdirAll(fakeIntelTypePath, 0700)
		Expect(err).ToNot(HaveOccurred())
		mdevBasePath = fakeMdevBasePath
		// create mdev directories and symlinks
		for _, uuid := range []string{fakeMdevUUID, fakeIntelMdevUUID} {
			mdevTypePath := filepath.Join(fakeMdevBasePath, uuid+"real")
			err = os.MkdirAll(mdevTypePath, 0700)
			Expect(err).ToNot(HaveOccurred())
			err = os.Symlink(filepath.Join(fakeMdevBasePath, uuid+"real"), filepath.Join(fakeMdevBasePath, uuid))
			Expect(err).ToNot(HaveOccurred())
		}
		// link nvidia type directory
		err = os.Symlink(fakeNvidiaTypePath, filepath.Join(fakeMdevBasePath, fakeMdevUUID+"real", "mdev_type"))
		Expect(err).ToNot(HaveOccurred())
		err = os.Symlink(fakeIntelTypePath, filepath.Join(fakeMdevBasePath, fakeIntelMdevUUID+"real", "mdev_type"))
		Expect(err).ToNot(HaveOccurred())
		// create a name file in the nvidia type directory
		mdevName, err := os.Create(filepath.Join(fakeNvidiaTypePath, "name"))
		Expect(err).ToNot(HaveOccurred())
		mdevNameWriter := bufio.NewWriter(mdevName)
		n, err := mdevNameWriter.WriteString(fakeMdevNameSelector + "\n")
		Expect(err).ToNot(HaveOccurred())
		Expect(n).To(Equal(len(fakeMdevNameSelector) + 1))
		mdevNameWriter.Flush()

	})

	AfterEach(func() {
		os.RemoveAll(mdevBasePath)
		os.RemoveAll(fakeSupportedTypesPath)
	})
	table.DescribeTable("should get correct file type name", func(namePathExist bool) {
		if namePathExist {
			mdevName, err := getMdevTypeName(fakeMdevUUID)
			Expect(err).ToNot(HaveOccurred())
			Expect(mdevName).To(Equal(resourceNameToTypeName(fakeMdevNameSelector)))
		} else {
			mdevName, err := getMdevTypeName(fakeIntelMdevUUID)
			Expect(err).ToNot(HaveOccurred())
			Expect(mdevName).To(Equal(resourceNameToTypeName(fakeIntelMdevNameSelector)))
		}
	},
		table.Entry("Nvidia name file exist", true),
		table.Entry("Intel name file doesn't exist", false),
	)
	Context("discover devices", func() {
		BeforeEach(func() {
			By("mocking PCI and MDEV functions to simulate an mdev an its parent PCI device")
			ctrl = gomock.NewController(GinkgoT())
			mockPCI = NewMockDeviceHandler(ctrl)
			Handler = mockPCI
			// Force pre-defined returned values and ensure the function only get called exacly once each on 0000:00:00.0
			mockPCI.EXPECT().GetMdevParentPCIAddr(fakeMdevUUID).Return(fakeAddress, nil).Times(1)
			mockPCI.EXPECT().GetDeviceIOMMUGroup(mdevBasePath, fakeMdevUUID).Return(fakeIommuGroup, nil).Times(1)
			mockPCI.EXPECT().GetDeviceNumaNode(pciBasePath, fakeAddress).Return(fakeNumaNode).Times(1)

			By("creating a list of fake device using the yaml decoder")
			fakePermittedHostDevicesConfig = `
    mediatedDevices:
    - mdevNameSelector: "` + fakeMdevNameSelector + `"
      resourceName: "` + fakeMdevResourceName + `"
    `
			err := yaml.NewYAMLOrJSONDecoder(strings.NewReader(fakePermittedHostDevicesConfig), 1024).Decode(&fakePermittedHostDevices)
			Expect(err).ToNot(HaveOccurred())
			Expect(fakePermittedHostDevices.MediatedDevices).To(HaveLen(1))
			Expect(fakePermittedHostDevices.MediatedDevices[0].MDEVNameSelector).To(Equal(fakeMdevNameSelector))
			Expect(fakePermittedHostDevices.MediatedDevices[0].ResourceName).To(Equal(fakeMdevResourceName))
		})
		AfterEach(func() {
			ctrl.Finish()
		})
		It("Should parse the permitted devices and find 1 matching mediated device", func() {
			supportedMdevsMap := make(map[string]string)
			for _, supportedMdev := range fakePermittedHostDevices.MediatedDevices {
				// do not add a device plugin for this resource if it's being provided via an external device plugin
				if !supportedMdev.ExternalResourceProvider {
					selector := removeSelectorSpaces(supportedMdev.MDEVNameSelector)
					supportedMdevsMap[selector] = supportedMdev.ResourceName
				}
			}
			// discoverPermittedHostMediatedDevices() will walk real mdev devices wherever the tests are running
			devices := discoverPermittedHostMediatedDevices(supportedMdevsMap)
			Expect(devices).To(HaveLen(1))
			selector := removeSelectorSpaces(fakeMdevNameSelector)
			Expect(devices[selector]).To(HaveLen(1))
			Expect(devices[selector][0].UUID).To(Equal(fakeMdevUUID))
			Expect(devices[selector][0].typeName).To(Equal(selector))
			Expect(devices[selector][0].parentPciAddress).To(Equal(fakeAddress))
			Expect(devices[selector][0].iommuGroup).To(Equal(fakeIommuGroup))
			Expect(devices[selector][0].numaNode).To(Equal(fakeNumaNode))
		})

		It("Should validate DPI devices", func() {
			iommuToMDEVMap := make(map[string]string)
			supportedMdevsMap := make(map[string]string)
			for _, supportedMdev := range fakePermittedHostDevices.MediatedDevices {
				// do not add a device plugin for this resource if it's being provided via an external device plugin
				if !supportedMdev.ExternalResourceProvider {
					selector := removeSelectorSpaces(supportedMdev.MDEVNameSelector)
					supportedMdevsMap[selector] = supportedMdev.ResourceName
				}
			}
			// discoverPermittedHostMediatedDevices() will walk real mdev devices wherever the tests are running
			mDevices := discoverPermittedHostMediatedDevices(supportedMdevsMap)
			selector := removeSelectorSpaces(fakeMdevNameSelector)
			devs := constructDPIdevicesFromMdev(mDevices[selector], iommuToMDEVMap)
			Expect(devs[0].ID).To(Equal(fakeIommuGroup))
			Expect(devs[0].Topology.Nodes[0].ID).To(Equal(int64(fakeNumaNode)))
		})

		It("Should update the device list according to the configmap", func() {
			By("creating a cluster config")
			kv := &v1.KubeVirt{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "kubevirt",
					Namespace: "kubevirt",
				},
				Spec: v1.KubeVirtSpec{
					Configuration: v1.KubeVirtConfiguration{
						DeveloperConfiguration: &v1.DeveloperConfiguration{},
					},
				},
				Status: v1.KubeVirtStatus{
					Phase: v1.KubeVirtPhaseDeploying,
				},
			}
			fakeClusterConfig, _, kvInformer := testutils.NewFakeClusterConfigUsingKV(kv)

			By("creating an empty device controller")
			var noDevices []Device
			deviceController := NewDeviceController("master", noDevices, fakeClusterConfig, clientTest.CoreV1())

			By("adding a host device to the cluster config")
			kvConfig := kv.DeepCopy()
			kvConfig.Spec.Configuration.DeveloperConfiguration.FeatureGates = []string{virtconfig.HostDevicesGate}
			kvConfig.Spec.Configuration.PermittedHostDevices = &v1.PermittedHostDevices{
				MediatedDevices: []v1.MediatedHostDevice{
					{
						MDEVNameSelector: fakeMdevNameSelector,
						ResourceName:     fakeMdevResourceName,
					},
				},
			}
			testutils.UpdateFakeKubeVirtClusterConfig(kvInformer, kvConfig)
			permittedDevices := fakeClusterConfig.GetPermittedHostDevices()
			Expect(permittedDevices).ToNot(BeNil(), "something went wrong while parsing the configmap(s)")
			Expect(len(permittedDevices.MediatedDevices)).To(Equal(1), "the fake device was not found")

			By("ensuring a device plugin gets created for our fake device")
			enabledDevicePlugins, disabledDevicePlugins := deviceController.splitPermittedDevices(
				deviceController.updatePermittedHostDevicePlugins(),
			)
			Expect(len(enabledDevicePlugins)).To(Equal(1), "a device plugin wasn't created for the fake device")
			Expect(disabledDevicePlugins).To(BeEmpty())
			Ω(enabledDevicePlugins).Should(HaveKey(fakeMdevResourceName))
			// Manually adding the enabled plugin, since the device controller is not actually running
			deviceController.startedPlugins[fakeMdevResourceName] = controlledDevice{
				devicePlugin: enabledDevicePlugins[fakeMdevResourceName],
			}

			By("deletting the device from the configmap")
			kvConfig.Spec.Configuration.PermittedHostDevices = &v1.PermittedHostDevices{}
			testutils.UpdateFakeKubeVirtClusterConfig(kvInformer, kvConfig)
			permittedDevices = fakeClusterConfig.GetPermittedHostDevices()
			Expect(permittedDevices).ToNot(BeNil(), "something went wrong while parsing the configmap(s)")
			Expect(len(permittedDevices.MediatedDevices)).To(Equal(0), "the fake device was not deleted")

			By("ensuring the device plugin gets stopped")
			enabledDevicePlugins, disabledDevicePlugins = deviceController.splitPermittedDevices(
				deviceController.updatePermittedHostDevicePlugins(),
			)
			Expect(enabledDevicePlugins).To(BeEmpty())
			Expect(len(disabledDevicePlugins)).To(Equal(1), "the fake device plugin did not get disabled")
			Ω(disabledDevicePlugins).Should(HaveKey(fakeMdevResourceName))
		})
	})
})
<|endoftext|>


<commit_before>// Copyright (c) 2013-2014 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package btcutil

const (
	// SatoshiPerBitcent is the number of satoshi in one bitcoin cent.
	SatoshiPerBitcent = 1e6

	// SatoshiPerBitcoin is the number of satoshi in one bitcoin (1 BTC).
	SatoshiPerBitcoin = 1e8

	// MaxSatoshi is the maximum transaction amount allowed in satoshi.
	MaxSatoshi = 21e6 * SatoshiPerBitcoin
)
<commit_msg>Peercoin constants<commit_after>// Copyright (c) 2013-2014 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package btcutil

const (
	// SatoshiPerBitcent is the number of satoshi in one bitcoin cent.
	SatoshiPerBitcent = 1e4 // ppc: 10000 sunny per peercoin cent

	// SatoshiPerBitcoin is the number of satoshi in one bitcoin (1 BTC).
	SatoshiPerBitcoin = 1e6 // ppc: 1000000 sunny per peercoin

	// MaxSatoshi is the maximum transaction amount allowed in satoshi.
	MaxSatoshi = 21e6 * SatoshiPerBitcoin // ppc: TODO check value for peercoin
)
<|endoftext|>


<commit_before>package stripe

// ISO 3-digit Currency Codes for major currencies (not the full list).
const (
	USD = "usd" // US Dollar ($)
	EUR = "eur" // Euro (€)
	GBP = "gbp" // British Pound Sterling (UK£)
	JPY = "jpy" // Japanese Yen (¥)
	CAD = "cad" // Canadian Dollar (CA$)
	HKD = "hkd" // Hong Kong Dollar (HK$)
	CNY = "cny" // Chinese Yuan (CN¥)
	AUD = "aud" // Australian Dollar (A$)
)

// Subscription Statuses
const (
	SubscriptionTrialing = "trialing"
	SubscriptionActive   = "active"
	SubscriptionPastDue  = "past_due"
	SubscriptionCanceled = "canceled"
	SubscriptionUnpaid   = "unpaid"
)
<commit_msg>moved subscription status constants to substription.go file<commit_after>package stripe

// ISO 3-digit Currency Codes for major currencies (not the full list).
const (
	USD = "usd" // US Dollar ($)
	EUR = "eur" // Euro (€)
	GBP = "gbp" // British Pound Sterling (UK£)
	JPY = "jpy" // Japanese Yen (¥)
	CAD = "cad" // Canadian Dollar (CA$)
	HKD = "hkd" // Hong Kong Dollar (HK$)
	CNY = "cny" // Chinese Yuan (CN¥)
	AUD = "aud" // Australian Dollar (A$)
)
<|endoftext|>


<commit_before>package pass

import (
	"bytes"
	"compress/flate"
	"compress/gzip"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha512"
	"encoding/binary"
	"fmt"
	"io/ioutil"

	"code.google.com/p/go.crypto/scrypt"
)

// the current version of the encrypted format as a byte array
const Version uint32 = 0

// how large our version number is, in bytes. a uint32 should ALWAYS be 4 bytes,
// so we just hard-code this here.
const VersionSize = 4

// the size of the signature appended to signed data
const SignatureSize = sha512.Size

// the size of the random salt in bytes we use during password hashing
const SaltSize = 32

// the size of key to use for encryption. using 32 bytes (256 bits) selects
// AES-256 encryption (see: http://golang.org/pkg/crypto/aes/#NewCipher).
const EncryptionKeySize = 32

// we want our HMAC keys to be the same size as the blocksize (see:
// http://stackoverflow.com/a/12207647 and
// http://en.wikipedia.org/wiki/Hash-based_message_authentication_code#Definition_.28from_RFC_2104.29).
const HMACKeySize = sha512.BlockSize

// the parameters to use when hashing the master password. we shoot for a memory
// requirement of 128Mb (128 * N * r bytes).
const HashN uint32 = 1 << 16 // 2^16
const HashR uint32 = 16
const HashP uint32 = 2

// how large each hash parameter is, in bytes
const HashParamSize = 4

// the minimum size of encrypted content. it must include a version, the
// password salt, the hashing parameters, an initialization vector, and a
// signature - at a minimum!
const minEncryptedLength = (VersionSize + SaltSize + (3 * HashParamSize) +
	aes.BlockSize + SignatureSize)

// compress some data using the GZip algorithm and return it
func compress(data []byte) ([]byte, error) {
	compressed := new(bytes.Buffer)
	writer, err := gzip.NewWriterLevel(compressed, flate.BestCompression)
	if err != nil {
		return nil, err
	}

	// compress our data
	writer.Write(data)
	writer.Close()

	return compressed.Bytes(), nil
}

// decompress some data compressed by the GZip algorithm
func decompress(data []byte) ([]byte, error) {
	b := bytes.NewBuffer(data)
	reader, err := gzip.NewReader(b)
	if err != nil {
		return nil, err
	}

	// decompress our data
	result, err := ioutil.ReadAll(reader)
	if err != nil {
		return nil, err
	}
	reader.Close()

	return result, nil
}

// get the signature of the given data as a byte array using SHA-512. the
// resulting byte array will have a length of SignatureSize.
func sign(data, key []byte) ([]byte, error) {
	// we want the key to be no shorter than the hash algorithm's block size,
	// otherwise it will be zero-padded. longer keys are hashed to obtain a key of
	// the same size as the block size, so there's really no benefit in using a
	// key size that's not equal to the block size of the hash algorithm. it
	// doesn't hurt, however, so we let that case alone.
	if len(key) < HMACKeySize {
		err := fmt.Errorf("Key size is too small (should be %d bytes)",
			HMACKeySize)
		return nil, err
	}

	mac := hmac.New(sha512.New, key)
	mac.Write(data)

	// compute and return the signature
	return mac.Sum(nil), nil
}

// return whether the given signature verifies the given data
func verify(data, suppliedSignature, key []byte) error {
	// make sure the signature is the correct size
	if len(suppliedSignature) != SignatureSize {
		err := fmt.Errorf("Signature must be %d bytes long (got %d)",
			SignatureSize, len(suppliedSignature))
		return err
	}

	// sign the data ourself
	computedSignature, err := sign(data, key)
	if err != nil {
		return err
	}

	// signal an error if the computed signature doesn't match the given one.
	// notice that we securely compare the signatures to avoid timing attacks!
	if !hmac.Equal(suppliedSignature, computedSignature) {
		err := fmt.Errorf(
			"Signatures do not match:\n  supplied: %v\n  computed: %v)",
			suppliedSignature, computedSignature)
		return err
	}

	// return no error since the data authenticated correctly
	return nil
}

// encode the given version number as an array of bytes, then return the array
// and whether there was an error.
func uint32ToBytes(version uint32) ([]byte, error) {
	buf := new(bytes.Buffer)
	if err := binary.Write(buf, binary.BigEndian, version); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// read a version number from an array of bytes and return the version number
// along with an error, if any.
func bytesToUint32(versionBytes []byte) (uint32, error) {
	// make sure we got enough bytes to parse a version out of them
	if len(versionBytes) < VersionSize {
		return 0, fmt.Errorf(
			"Not enough bytes to contain a version (minimum: %d)", VersionSize)
	}

	// read the version from our bytes and return it
	buf := bytes.NewBuffer(versionBytes)
	var version uint32
	if err := binary.Read(buf, binary.BigEndian, &version); err != nil {
		return 0, err
	}

	return version, nil
}

// given a password string and a salt, return two byte arrays. the first should
// be used for encryption, the second for HMAC.
func hashPassword(password string, salt []byte, N, r, p uint32) ([]byte, []byte, error) {
	// ensure that all the encryption paramters meet minimum requirements
	if N <= 1 {
		return nil, nil, fmt.Errorf("N must be larger than one")
	} else if r <= 0 {
		return nil, nil, fmt.Errorf("r must be larger than zero")
	} else if p <= 0 {
		return nil, nil, fmt.Errorf("p must be larger than zero")
	}

	// NOTE: scrypt memory usage is approximately 128 * `N` * `r` bytes. since `p`
	// has little effect on memory usage, it can be used to tune the running time
	// of the algorithm.

	// generate enough bytes for both the encryption and HMAC keys. additionally,
	// since scrypt is checking the sizes of the paramter values for us, we don't
	// need to do it ourselves (see:
	// http://code.google.com/p/go/source/browse/scrypt/scrypt.go?repo=crypto).
	hash, err := scrypt.Key([]byte(password), salt, int(N), int(r), int(p),
		EncryptionKeySize+HMACKeySize)
	if err != nil {
		return nil, nil, err
	}

	// return the keys according to our convention (encryption, then hmac)
	encryptionKey := hash[:EncryptionKeySize]
	hmacKey := hash[EncryptionKeySize:]
	return encryptionKey, hmacKey, nil
}

// encrypt some data using the given password and default scrypt params, then
// return the result.
func Encrypt(plaintext []byte, password string) ([]byte, error) {
	// use the default params to encrypt this text
	return EncryptWithHashParams(plaintext, password, HashN, HashR, HashP)
}

// encrypt some data using the given password and scrypt params, then return the
// result.
func EncryptWithHashParams(plaintext []byte, password string, N, r, p uint32) ([]byte, error) {
	// NOTE: no plaintext padding is needed since we're using CFB mode (see:
	// http://en.wikipedia.org/wiki/Block_cipher_mode_of_operation#Padding).

	// first, compress the plaintext to obfuscate its contents and reduce its size
	compressedPlaintext, err := compress(plaintext)
	if err != nil {
		return nil, err
	}

	// make a blob that conforms to our defined structure
	blob := NewBlob(
		"version", VersionSize,
		"N", HashParamSize,
		"r", HashParamSize,
		"p", HashParamSize,
		"salt", SaltSize,
		"iv", aes.BlockSize,
		"data", len(compressedPlaintext),
		"signature", SignatureSize,
	)

	// get the slices we'll be working with
	version := blob.Get("version")
	salt := blob.Get("salt")
	blobN := blob.Get("N")
	blobR := blob.Get("r")
	blobP := blob.Get("p")
	iv := blob.Get("iv")
	ciphertext := blob.Get("data")
	signature := blob.Get("signature")

	// serialize and store the current version
	versionBytes, err := uint32ToBytes(Version)
	if err != nil {
		return nil, err
	}
	copy(version, versionBytes)

	// randomize the salt and the initialization vector
	if _, err := rand.Read(salt); err != nil {
		return nil, err
	}
	if _, err := rand.Read(iv); err != nil {
		return nil, err
	}

	// serialize and store the hash paramters
	nBytes, err := uint32ToBytes(N)
	if err != nil {
		return nil, err
	}
	copy(blobN, nBytes)

	rBytes, err := uint32ToBytes(r)
	if err != nil {
		return nil, err
	}
	copy(blobR, rBytes)

	pBytes, err := uint32ToBytes(p)
	if err != nil {
		return nil, err
	}
	copy(blobP, pBytes)

	// hash the password into the necessary keys using the salt
	encryptionKey, hmacKey, err := hashPassword(password, salt, N, r, p)
	if err != nil {
		return nil, err
	}

	// encrypt the compressed plaintext
	block, err := aes.NewCipher(encryptionKey)
	if err != nil {
		return nil, err
	}

	// use CFB mode to encrypt the data, so we don't have to pad
	stream := cipher.NewCFBEncrypter(block, iv)
	stream.XORKeyStream(ciphertext, compressedPlaintext)

	// sign our data (everything _but_ the signature space)
	content := blob.To("data")
	signatureData, err := sign(content, hmacKey)
	if err != nil {
		return nil, err
	}

	// store the signature
	copy(signature, signatureData)

	return blob.Bytes(), nil
}

// decrypt some data using the given password
func Decrypt(data []byte, password string) ([]byte, error) {
	// make sure our data is of at least the minimum length
	if len(data) < minEncryptedLength {
		err := fmt.Errorf("Data is too short to be valid (min length: %d)",
			minEncryptedLength)
		return nil, err
	}

	// make a blob that conforms to our defined structure
	blob := NewBlob(
		"version", VersionSize,
		"N", HashParamSize,
		"r", HashParamSize,
		"p", HashParamSize,
		"salt", SaltSize,
		"iv", aes.BlockSize,

		// the ciphertext is everything in the blob _except_ the other fields
		"data", len(data)-(VersionSize+
			SaltSize+
			(3*HashParamSize)+
			aes.BlockSize+
			SignatureSize),

		"signature", SignatureSize,

		// initalize the blob with the encrypted data
		data,
	)

	// make sure we can decrypt this version
	version, err := bytesToUint32(blob.Get("version"))
	if err != nil {
		return nil, err
	}

	// we'll never be able to handle newer versions!
	if version > Version {
		return nil, fmt.Errorf("Latest supported version is %d (got: %d)",
			Version, version)
	}

	// decrypt using a version of the algorithm that matches the given blob
	if version < Version {
		// TODO: add support for older versions once they exist
		panic("No older versions shoud exist at this time!")
	}

	// read the the parts we need from the unverified data
	salt := blob.Get("salt")
	iv := blob.Get("iv")
	ciphertext := blob.Get("data")
	signature := blob.Get("signature")

	// read the hash paramters we need to hash the password
	N, err := bytesToUint32(blob.Get("N"))
	if err != nil {
		return nil, err
	}
	r, err := bytesToUint32(blob.Get("r"))
	if err != nil {
		return nil, err
	}
	p, err := bytesToUint32(blob.Get("p"))
	if err != nil {
		return nil, err
	}

	// hash the password with the supplied salt and paramters to get the keys
	encryptionKey, hmacKey, err := hashPassword(password, salt, N, r, p)
	if err != nil {
		return nil, err
	}

	// verify the integrity of the blob (including the version)
	err = verify(blob.To("data"), signature, hmacKey)
	if err != nil {
		return nil, err
	}

	// decrypt the ciphertext
	block, err := aes.NewCipher(encryptionKey)
	if err != nil {
		return nil, err
	}

	// decrypt directly into the original slice to save creating a new array
	compressedPlaintext := ciphertext[:]
	stream := cipher.NewCFBDecrypter(block, iv)
	stream.XORKeyStream(compressedPlaintext, ciphertext)

	// decompress the compressed plaintext
	plaintext, err := decompress(compressedPlaintext)
	if err != nil {
		return nil, err
	}

	return plaintext, nil
}
<commit_msg>Start implementing new blob format<commit_after>package pass

import (
	"bytes"
	"compress/flate"
	"compress/gzip"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha512"
	"encoding/binary"
	"fmt"
	"io/ioutil"

	"code.google.com/p/go.crypto/scrypt"
)

// the top-level format of the encrypted blob:
//
//   ----
//   magic number (4 bytes, always 0xEA8F73F5)
//   version number as a uint32 (4 bytes)
//   ----
//   metadata size as a uint32 (4 bytes)
//   medadata (same as the given size above)
//   ----
//   payload (the remaining bytes)
//   ---
//
// the version number specifies how the program interprets the metadata and the
// payload. this gives maximum flexibility against future changes.

// the size of our magic number, in bytes
const MagicNumberSize = 4

// the magic number that marks all of our encrypted blobs as belonging to us
const MagicNumber uint32 = 0xEA8F73F5

// the current version of the encrypted blob format, and the one all new blobs
// will be created with.
const Version uint32 = 0

// the size of our version number, in bytes
const VersionSize = 4

// how large the size of the metadata header size block is, in bytes
const MetaDataSizeSize = 4

// the size of the signature added to signed data
const SignatureSize = sha512.Size

// the size of the random salt in bytes we use during password hashing
const SaltSize = 32

// the size of key to use for encryption. using 32 bytes (256 bits) selects
// AES-256 encryption (see: http://golang.org/pkg/crypto/aes/#NewCipher).
const EncryptionKeySize = 32

// we want our HMAC keys to be the same size as the blocksize (see:
// http://stackoverflow.com/a/12207647 and
// http://en.wikipedia.org/wiki/Hash-based_message_authentication_code#Definition_.28from_RFC_2104.29).
const HMACKeySize = sha512.BlockSize

// the parameters to use when hashing the master password. we shoot for a memory
// requirement of 128Mb (128 * N * r bytes).
const HashN uint32 = 1 << 16 // 2^16
const HashR uint32 = 16
const HashP uint32 = 2

// how large each hash parameter is, in bytes
const HashParamSize = 4

// the minimum size of encrypted content. it must include a version, the
// password salt, the hashing parameters, an initialization vector, and a
// signature - at a minimum!
const minEncryptedLength = (VersionSize + SaltSize + (3 * HashParamSize) +
	aes.BlockSize + SignatureSize)

// compress some data using the GZip algorithm and return it
func compress(data []byte) ([]byte, error) {
	compressed := new(bytes.Buffer)
	writer, err := gzip.NewWriterLevel(compressed, flate.BestCompression)
	if err != nil {
		return nil, err
	}

	// compress our data
	writer.Write(data)
	writer.Close()

	return compressed.Bytes(), nil
}

// decompress some data compressed by the GZip algorithm
func decompress(data []byte) ([]byte, error) {
	b := bytes.NewBuffer(data)
	reader, err := gzip.NewReader(b)
	if err != nil {
		return nil, err
	}

	// decompress our data
	result, err := ioutil.ReadAll(reader)
	if err != nil {
		return nil, err
	}
	reader.Close()

	return result, nil
}

// get the signature of the given data as a byte array using SHA-512. the
// resulting byte array will have a length of SignatureSize.
func sign(data, key []byte) ([]byte, error) {
	// we want the key to be no shorter than the hash algorithm's block size,
	// otherwise it will be zero-padded. longer keys are hashed to obtain a key of
	// the same size as the block size, so there's really no benefit in using a
	// key size that's not equal to the block size of the hash algorithm. it
	// doesn't hurt, however, so we let that case alone.
	if len(key) < HMACKeySize {
		err := fmt.Errorf("Key size is too small (should be %d bytes)",
			HMACKeySize)
		return nil, err
	}

	mac := hmac.New(sha512.New, key)
	mac.Write(data)

	// compute and return the signature
	return mac.Sum(nil), nil
}

// return whether the given signature verifies the given data
func verify(data, suppliedSignature, key []byte) error {
	// make sure the signature is the correct size
	if len(suppliedSignature) != SignatureSize {
		err := fmt.Errorf("Signature must be %d bytes long (got %d)",
			SignatureSize, len(suppliedSignature))
		return err
	}

	// sign the data ourself
	computedSignature, err := sign(data, key)
	if err != nil {
		return err
	}

	// signal an error if the computed signature doesn't match the given one.
	// notice that we securely compare the signatures to avoid timing attacks!
	if !hmac.Equal(suppliedSignature, computedSignature) {
		err := fmt.Errorf(
			"Signatures do not match:\n  supplied: %v\n  computed: %v)",
			suppliedSignature, computedSignature)
		return err
	}

	// return no error since the data authenticated correctly
	return nil
}

// encode the given version number as an array of bytes, then return the array
// and whether there was an error.
func uint32ToBytes(version uint32) ([]byte, error) {
	buf := new(bytes.Buffer)
	if err := binary.Write(buf, binary.BigEndian, version); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// read a version number from an array of bytes and return the version number
// along with an error, if any.
func bytesToUint32(versionBytes []byte) (uint32, error) {
	// make sure we got enough bytes to parse a version out of them
	if len(versionBytes) < VersionSize {
		return 0, fmt.Errorf(
			"Not enough bytes to contain a version (minimum: %d)", VersionSize)
	}

	// read the version from our bytes and return it
	buf := bytes.NewBuffer(versionBytes)
	var version uint32
	if err := binary.Read(buf, binary.BigEndian, &version); err != nil {
		return 0, err
	}

	return version, nil
}

// given a password string and a salt, return two byte arrays. the first should
// be used for encryption, the second for HMAC.
func hashPassword(password string, salt []byte, N, r, p uint32) ([]byte, []byte, error) {
	// ensure that all the encryption paramters meet minimum requirements
	if N <= 1 {
		return nil, nil, fmt.Errorf("N must be larger than one")
	} else if r <= 0 {
		return nil, nil, fmt.Errorf("r must be larger than zero")
	} else if p <= 0 {
		return nil, nil, fmt.Errorf("p must be larger than zero")
	}

	// NOTE: scrypt memory usage is approximately 128 * `N` * `r` bytes. since `p`
	// has little effect on memory usage, it can be used to tune the running time
	// of the algorithm.

	// generate enough bytes for both the encryption and HMAC keys. additionally,
	// since scrypt is checking the sizes of the paramter values for us, we don't
	// need to do it ourselves (see:
	// http://code.google.com/p/go/source/browse/scrypt/scrypt.go?repo=crypto).
	hash, err := scrypt.Key([]byte(password), salt, int(N), int(r), int(p),
		EncryptionKeySize+HMACKeySize)
	if err != nil {
		return nil, nil, err
	}

	// return the keys according to our convention (encryption, then hmac)
	encryptionKey := hash[:EncryptionKeySize]
	hmacKey := hash[EncryptionKeySize:]
	return encryptionKey, hmacKey, nil
}

// given a blob, confirms it's of the required format, then returns its
// constituent parts: the version, the metadata, and the payload. returns an
// error if something went wrong.
func loadBlob(blobBytes []byte) (uint32, []byte, []byte, error) {
	// if the blob is too short to be valid, give up
	const minSize = MagicNumberSize + VersionSize + MetaDataSizeSize
	if (len(blobBytes) < minSize) {
		return 0, nil, nil, fmt.Errorf(
			"Blob is too short to be valid (minimum size: %d bytes)", minSize);
	}

	// load as much information as we can
	blob := NewBlob(
		"magic_number", MagicNumberSize,
		"version", VersionSize,
		"meta_size", MetaDataSizeSize,
		blobBytes,
	)

	// validate the magic number to perform a sanity check on our data
	magicNumber, err := bytesToUint32(blob.Get("magic_number"))
	if err != nil {
		return 0, nil, nil, err
	}
	if (magicNumber != MagicNumber) {
		return 0, nil, nil, fmt.Errorf(
			"Unrecognized file format: magic number did not match")
	}

	// parse the version number
	version, err := bytesToUint32(blob.Get("version"))
	if err != nil {
		return 0, nil, nil, err
	}

	// parse the metadata size
	metaSize, err := bytesToUint32(blob.Get("meta_size"))
	if err != nil {
		return 0, nil, nil, err
	}

	// the bytes from the start of the metadata body onwards
	restBytes := blobBytes[minSize:]

	// if the remaining data isn't large enough to cover the given size, return an
	// error.
	if (metaSize > uint32(len(restBytes))) {
		return 0, nil, nil, fmt.Errorf(
			"Got invalid metadata size: %d (remaining data only %d bytes)",
			metaSize,
			len(restBytes),
		)
	}

	// retrieve the metadata and the payload
	metaBytes := restBytes[0:metaSize]
	payloadBytes := restBytes[metaSize:]

	return version, metaBytes, payloadBytes, nil
}

// encrypt some data using the given password and default scrypt params, then
// return the result.
func Encrypt(plaintext []byte, password string) ([]byte, error) {
	// use the default params to encrypt this text
	return EncryptWithHashParams(plaintext, password, HashN, HashR, HashP)
}

// encrypt some data using the given password and scrypt params, then return the
// result.
func EncryptWithHashParams(plaintext []byte, password string, N, r, p uint32) ([]byte, error) {
	// NOTE: no plaintext padding is needed since we're using CFB mode (see:
	// http://en.wikipedia.org/wiki/Block_cipher_mode_of_operation#Padding).

	// first, compress the plaintext to obfuscate its contents and reduce its size
	compressedPlaintext, err := compress(plaintext)
	if err != nil {
		return nil, err
	}

	// make a blob that conforms to our defined structure
	blob := NewBlob(
		"version", VersionSize,
		"N", HashParamSize,
		"r", HashParamSize,
		"p", HashParamSize,
		"salt", SaltSize,
		"iv", aes.BlockSize,
		"data", len(compressedPlaintext),
		"signature", SignatureSize,
	)

	// get the slices we'll be working with
	version := blob.Get("version")
	salt := blob.Get("salt")
	blobN := blob.Get("N")
	blobR := blob.Get("r")
	blobP := blob.Get("p")
	iv := blob.Get("iv")
	ciphertext := blob.Get("data")
	signature := blob.Get("signature")

	// serialize and store the current version
	versionBytes, err := uint32ToBytes(Version)
	if err != nil {
		return nil, err
	}
	copy(version, versionBytes)

	// randomize the salt and the initialization vector
	if _, err := rand.Read(salt); err != nil {
		return nil, err
	}
	if _, err := rand.Read(iv); err != nil {
		return nil, err
	}

	// serialize and store the hash paramters
	nBytes, err := uint32ToBytes(N)
	if err != nil {
		return nil, err
	}
	copy(blobN, nBytes)

	rBytes, err := uint32ToBytes(r)
	if err != nil {
		return nil, err
	}
	copy(blobR, rBytes)

	pBytes, err := uint32ToBytes(p)
	if err != nil {
		return nil, err
	}
	copy(blobP, pBytes)

	// hash the password into the necessary keys using the salt
	encryptionKey, hmacKey, err := hashPassword(password, salt, N, r, p)
	if err != nil {
		return nil, err
	}

	// encrypt the compressed plaintext
	block, err := aes.NewCipher(encryptionKey)
	if err != nil {
		return nil, err
	}

	// use CFB mode to encrypt the data, so we don't have to pad
	stream := cipher.NewCFBEncrypter(block, iv)
	stream.XORKeyStream(ciphertext, compressedPlaintext)

	// sign our data (everything _but_ the signature space)
	content := blob.To("data")
	signatureData, err := sign(content, hmacKey)
	if err != nil {
		return nil, err
	}

	// store the signature
	copy(signature, signatureData)

	return blob.Bytes(), nil
}

// decrypt some data using the given password
func Decrypt(data []byte, password string) ([]byte, error) {
	// make sure our data is of at least the minimum length
	if len(data) < minEncryptedLength {
		err := fmt.Errorf("Data is too short to be valid (min length: %d)",
			minEncryptedLength)
		return nil, err
	}

	// make a blob that conforms to our defined structure
	blob := NewBlob(
		"version", VersionSize,
		"N", HashParamSize,
		"r", HashParamSize,
		"p", HashParamSize,
		"salt", SaltSize,
		"iv", aes.BlockSize,

		// the ciphertext is everything in the blob _except_ the other fields
		"data", len(data)-(VersionSize+
			SaltSize+
			(3*HashParamSize)+
			aes.BlockSize+
			SignatureSize),

		"signature", SignatureSize,

		// initalize the blob with the encrypted data
		data,
	)

	// make sure we can decrypt this version
	version, err := bytesToUint32(blob.Get("version"))
	if err != nil {
		return nil, err
	}

	// we'll never be able to handle newer versions!
	if version > Version {
		return nil, fmt.Errorf("Latest supported version is %d (got: %d)",
			Version, version)
	}

	// decrypt using a version of the algorithm that matches the given blob
	if version < Version {
		// TODO: add support for older versions once they exist
		panic("No older versions shoud exist at this time!")
	}

	// read the the parts we need from the unverified data
	salt := blob.Get("salt")
	iv := blob.Get("iv")
	ciphertext := blob.Get("data")
	signature := blob.Get("signature")

	// read the hash paramters we need to hash the password
	N, err := bytesToUint32(blob.Get("N"))
	if err != nil {
		return nil, err
	}
	r, err := bytesToUint32(blob.Get("r"))
	if err != nil {
		return nil, err
	}
	p, err := bytesToUint32(blob.Get("p"))
	if err != nil {
		return nil, err
	}

	// hash the password with the supplied salt and paramters to get the keys
	encryptionKey, hmacKey, err := hashPassword(password, salt, N, r, p)
	if err != nil {
		return nil, err
	}

	// verify the integrity of the blob (including the version)
	err = verify(blob.To("data"), signature, hmacKey)
	if err != nil {
		return nil, err
	}

	// decrypt the ciphertext
	block, err := aes.NewCipher(encryptionKey)
	if err != nil {
		return nil, err
	}

	// decrypt directly into the original slice to save creating a new array
	compressedPlaintext := ciphertext[:]
	stream := cipher.NewCFBDecrypter(block, iv)
	stream.XORKeyStream(compressedPlaintext, ciphertext)

	// decompress the compressed plaintext
	plaintext, err := decompress(compressedPlaintext)
	if err != nil {
		return nil, err
	}

	return plaintext, nil
}
<|endoftext|>


<commit_before>// Copyright 2016 Cyako Author

// 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 cyako

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
)

type middlewareConfig struct {
	Name  string `json:"name"`
	Key   string `json:"key"`
	Value string `json:"value"`
}

type config struct {
	Middleware []middlewareConfig `json:"middleware"`
}

type middlewareSupport struct {
	Name          string
	AfterReceive  bool
	BeforeProcess bool
	AfterProcess  bool
	BeforeSend    bool
	AfterSend     bool
}

type middleware struct {
	Map               map[string]interface{}
	Support           []middlewareSupport
	AfterReceiveFunc  []func(*Req)
	BeforeProcessFunc []func(*Ctx)
	AfterProcessFunc  []func(*Ctx)
	BeforeSendFunc    []func(*Res)
	AfterSendFunc     []func(*Res)
}

type Cyako struct {
	Config       config
	Middleware   middleware
	ProcessorMap map[string]*Processor
}

func (c *Cyako) loadConfig() {
	var err error
	data, err := ioutil.ReadFile("config.json")
	if err != nil {
		fmt.Println("Read config file error:", err)
	}
	err = json.Unmarshal(data, &c.Config)
	if err != nil {
		fmt.Println("Unmarshal config file content error:", err)
	}
}

func (c *Cyako) PrintLoadInfo() {
	fmt.Println()
	fmt.Println(" Loading...")

	fmt.Printf("\n %-35s %-21s %-10s %-10s\n", "Config", "Name", "Key", "Value")
	for _, config := range c.Config.Middleware {
		fmt.Printf(" %-35s %-21s %-10s %-10s\n", "Middleware", config.Name, config.Key, config.Value)
	}

	fmt.Printf("\n %-35s %-10s %-10s %-10s %-10s %-10s\n", "Middleware", "AR", "BP", "AP", "BS", "AS")
	for _, c := range c.Middleware.Support {
		fmt.Printf(" %-35s %-10v %-10v %-10v %-10v %-10v\n", c.Name, c.AfterReceive, c.BeforeProcess, c.AfterProcess, c.BeforeSend, c.AfterSend)
	}

	fmt.Printf("\n %-35s %-10s %-40s\n", "API", "Module", "Package Path")
	for _, proc := range c.ProcessorMap {
		fmt.Printf(" %-35s %-10s %-40s\n", proc.Module+"."+proc.Name, proc.Module, proc.PkgPath)
	}
	fmt.Println()
}

func (c *Cyako) PrintAPIDoc() {
	type APIDoc struct {
		ParamConfigs ParamConfigs `json:"ParamConfigs"`
	}
	var doc APIDoc
	for _, proc := range c.ProcessorMap {
		var ctx *Ctx
		proc(ctx)
		doc.ParamConfigs = ParamConfigs
	}
	bytes, _ := json.Marshal(doc)
	fmt.Println(bytes)
}

/*
	init
*/

var cyako *Cyako

func init() {
	cyako = &Cyako{
		Middleware: middleware{
			Map: make(map[string]interface{}),
		},
		ProcessorMap: make(map[string]*Processor),
	}
	cyako.loadConfig()
}

/*
	global
*/

// return cyako package's global object: cyako
func Ins() *Cyako {
	return cyako
}

// used in Processor Module package to load itself
func LoadModule(x interface{}) {
	cyako.loadModule(x)
}

// used in Middleware package to load itself
func LoadMiddleware(x interface{}) {
	cyako.loadMiddleware(x)
}
<commit_msg>update<commit_after>// Copyright 2016 Cyako Author

// 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 cyako

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
)

type middlewareConfig struct {
	Name  string `json:"name"`
	Key   string `json:"key"`
	Value string `json:"value"`
}

type config struct {
	Middleware []middlewareConfig `json:"middleware"`
}

type middlewareSupport struct {
	Name          string
	AfterReceive  bool
	BeforeProcess bool
	AfterProcess  bool
	BeforeSend    bool
	AfterSend     bool
}

type middleware struct {
	Map               map[string]interface{}
	Support           []middlewareSupport
	AfterReceiveFunc  []func(*Req)
	BeforeProcessFunc []func(*Ctx)
	AfterProcessFunc  []func(*Ctx)
	BeforeSendFunc    []func(*Res)
	AfterSendFunc     []func(*Res)
}

type Cyako struct {
	Config       config
	Middleware   middleware
	ProcessorMap map[string]*Processor
}

func (c *Cyako) loadConfig() {
	var err error
	data, err := ioutil.ReadFile("config.json")
	if err != nil {
		fmt.Println(" Error config.json:", err)
	} else {
		err = json.Unmarshal(data, &c.Config)
		if err != nil {
			fmt.Println(" Error config.json:", err)
		}
	}
}

func (c *Cyako) PrintLoadInfo() {
	fmt.Println()
	fmt.Println(" Loading...")

	fmt.Printf("\n %-35s %-21s %-10s %-10s\n", "Config", "Name", "Key", "Value")
	for _, config := range c.Config.Middleware {
		fmt.Printf(" %-35s %-21s %-10s %-10s\n", "Middleware", config.Name, config.Key, config.Value)
	}

	fmt.Printf("\n %-35s %-10s %-10s %-10s %-10s %-10s\n", "Middleware", "AR", "BP", "AP", "BS", "AS")
	for _, c := range c.Middleware.Support {
		fmt.Printf(" %-35s %-10v %-10v %-10v %-10v %-10v\n", c.Name, c.AfterReceive, c.BeforeProcess, c.AfterProcess, c.BeforeSend, c.AfterSend)
	}

	fmt.Printf("\n %-35s %-10s %-40s\n", "API", "Module", "Package Path")
	for _, proc := range c.ProcessorMap {
		fmt.Printf(" %-35s %-10s %-40s\n", proc.Module+"."+proc.Name, proc.Module, proc.PkgPath)
	}
	fmt.Println()
}

func (c *Cyako) PrintAPIDoc() {
	type method struct {
		ParamConfigs []*ParamConfig `json:"ParamConfigs"`
		Processor
	}
	type APIDoc struct {
		Methods map[string]method `json:"method"`
	}
	doc := &APIDoc{
		Methods: make(map[string]method),
	}
	for methodName, proc := range c.ProcessorMap {
		req := &Req{}
		req.Init()
		res := &Res{Id: req.Id, Method: req.Method, Temp: req.Temp}
		ctx := &Ctx{res: res, req: req, Method: req.Method, Data: req.Data, Temp: req.Temp}
		res.Init()
		ctx.Init()
		proc.Func(ctx)
		doc.Methods[methodName] = method{
			ParamConfigs: ctx.ParamConfigs,
			Processor:    *proc,
		}
	}
	// bytes, err := json.Marshal(doc)
	// if err != nil {
	// 	fmt.Println(" Error APIDoc:", err)
	// }
	// fmt.Println(string(bytes))
	fmt.Println()
	fmt.Printf("\n %-35s %-10s %-40s\n", "API Detail", "Module", "Package Path")
	for _, proc := range doc.Methods {
		fmt.Printf(" %-35s %-10s %-40s\n", proc.Module+"."+proc.Name, proc.Module, proc.PkgPath)
		for _, cfg := range proc.ParamConfigs {
			fmt.Printf(" -%-10s %+v\n", cfg.Key, *cfg)
		}
	}
	fmt.Println()
}

/*
	init
*/

var cyako *Cyako

func init() {
	cyako = &Cyako{
		Middleware: middleware{
			Map: make(map[string]interface{}),
		},
		ProcessorMap: make(map[string]*Processor),
	}
	cyako.loadConfig()
}

/*
	global
*/

// return cyako package's global object: cyako
func Ins() *Cyako {
	return cyako
}

// used in Processor Module package to load itself
func LoadModule(x interface{}) {
	cyako.loadModule(x)
}

// used in Middleware package to load itself
func LoadMiddleware(x interface{}) {
	cyako.loadMiddleware(x)
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/gcfg"
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"github.com/streadway/amqp"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"time"
)

var (
	debug = flag.Bool("d", false, "turn on debug info")
	cfg   struct {
		DB struct {
			Driver string
			Dsn    string
			Table  string
		}
		RabbitMQ struct {
			Uri      string
			Exchange string
		}
		HTTP struct {
			Host string
			Port string
		}
	}

	db      *sql.DB
	broker  *amqp.Connection
	channel *amqp.Channel
	wakeUp  = make(chan int, 1)
)

type Job struct {
	routingKey string
	body       string
	interval   uint // Seconds
	nextRun    time.Time
	state      string
}

// hadleSchedule is the web server endpoint for path: /schedule
func handleSchedule(w http.ResponseWriter, r *http.Request) {
	routingKey, body, interval_s :=
		r.FormValue("routing_key"), r.FormValue("body"), r.FormValue("interval")
	log.Println("/schedule", routingKey, body)

	interval, err := strconv.ParseInt(interval_s, 10, 64)
	if err != nil {
		panic(err)
	}

	next_run := time.Now().UTC().Add(time.Duration(interval) * time.Second)
	_, err = db.Exec("INSERT INTO "+cfg.DB.Table+" "+
		"(routing_key, body, `interval`, next_run, state) "+
		"VALUES(?, ?, ?, ?, 'WAITING') "+
		"ON DUPLICATE KEY UPDATE `interval`=?",
		routingKey, body, interval, next_run, interval)
	if err != nil {
		panic(err)
	}

	// Wake up the publisher.
	//
	// publisher() may be sleeping for the next job on the queue
	// at the time we schedule a new Job. Let it wake up so it can
	// re-fetch the new Job from the front of the queue.
	//
	// The code below is an idiom for non-blocking send to a channel.
	select {
	case wakeUp <- 1:
		log.Println("Sent wakeup signal")
	default:
		log.Println("Skipped wakeup signal")
	}
}

// handleCancel is the web server endpoint for path: /cancel
func handleCancel(w http.ResponseWriter, r *http.Request) {
	routingKey, body := r.FormValue("routing_key"), r.FormValue("body")
	log.Println("/cancel", routingKey, body)

	_, err := db.Exec("DELETE FROM "+cfg.DB.Table+" "+
		"WHERE routing_key=? AND body=?", routingKey, body)
	if err != nil {
		panic(err)
	}
}

// front returns the first job to be run in the queue.
func front() (*Job, error) {
	j := Job{}
	row := db.QueryRow("SELECT routing_key, body, `interval`, next_run, state " +
		"FROM " + cfg.DB.Table + " " +
		"WHERE next_run = (" +
		"SELECT MIN(next_run) FROM " + cfg.DB.Table + ")")
	err := row.Scan(&j.routingKey, &j.body, &j.interval, &j.nextRun, &j.state)
	if err != nil {
		return nil, err
	}
	return &j, nil
}

// Delete deletes the Job j from the queue.
func (j *Job) Delete() error {
	_, err := db.Exec("DELETE FROM "+cfg.DB.Table+" "+
		"WHERE routing_key=? AND body=?",
		j.routingKey, j.body)
	if err != nil {
		return err
	}
	return nil
}

// Publish sends a message to exchange defined in the config and
// updates the Job's state to RUNNING on the database.
func (j *Job) Publish() error {
	log.Println("publish", *j)

	// Send a message to the broker
	err := channel.Publish(cfg.RabbitMQ.Exchange, j.routingKey, false, false, amqp.Publishing{
		Headers: amqp.Table{
			"interval":     j.interval,
			"published_at": time.Now().UTC(),
		},
		ContentType:     "application/octet-stream",
		ContentEncoding: "",
		Body:            []byte(j.body),
		DeliveryMode:    amqp.Persistent,
		Priority:        0,
	})
	if err != nil {
		return err
	}

	// Update state from database
	_, err = db.Exec("UPDATE "+cfg.DB.Table+" "+
		"SET state=RUNNING "+
		"WHERE routing_key=? AND body=?",
		j.routingKey, j.body)
	if err != nil {
		return err
	}
	return nil
}

// Remaining returns the duration until the job's next scheduled time.
func (j *Job) Remaining() time.Duration {
	return -time.Since(j.nextRun)
}

// publisher runs a loop that reads the next Job from the queue and publishes it.
func publisher() {
	for {
		job, err := front()
		if err != nil {
			fmt.Println(err)
			log.Println("Waiting wakeup signal")
			<-wakeUp
			log.Println("Got wakeup signal")
			continue
		}
		log.Println("Next job:", job, "Remaining:", job.Remaining())

		now := time.Now().UTC()
		if job.nextRun.After(now) {
			// Wait until the next Job time or
			// the webserver's /schedule handler wakes us up
			select {
			case <-time.After(job.Remaining()):
			case <-wakeUp:
				continue
			}
		} else {
			// It's time to publish the Job
			job.Publish()
			err = job.Delete()
			if err != nil {
				log.Println(err)
			}
		}

	}
}

func main() {
	// Setup logging
	flag.Parse()
	if !*debug {
		log.SetOutput(ioutil.Discard)
	}

	// Read config
	err := gcfg.ReadFileInto(&cfg, "dalga.ini")
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Read config: ", cfg)

	// Connect to database
	db, err = sql.Open(cfg.DB.Driver, cfg.DB.Dsn)
	if err != nil {
		log.Fatal(err)
	}
	err = db.Ping()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Connected to DB")

	// Connect to RabbitMQ
	broker, err = amqp.Dial(cfg.RabbitMQ.Uri)
	if err != nil {
		log.Fatal(err)
	}
	channel, err = broker.Channel()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Connected to RabbitMQ")

	// Run publisher
	go publisher()

	// Start HTTP server
	addr := cfg.HTTP.Host + ":" + cfg.HTTP.Port
	http.HandleFunc("/schedule", handleSchedule)
	http.HandleFunc("/cancel", handleCancel)
	http.ListenAndServe(addr, nil)
}
<commit_msg>better debug messages<commit_after>package main

import (
	"code.google.com/p/gcfg"
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"github.com/streadway/amqp"
	"log"
	"net/http"
	"strconv"
	"time"
)

var (
	debugging = flag.Bool("d", false, "turn on debug messages")
	cfg       struct {
		DB struct {
			Driver string
			Dsn    string
			Table  string
		}
		RabbitMQ struct {
			Uri      string
			Exchange string
		}
		HTTP struct {
			Host string
			Port string
		}
	}

	db      *sql.DB
	broker  *amqp.Connection
	channel *amqp.Channel
	wakeUp  = make(chan int, 1)
)

type Job struct {
	routingKey string
	body       string
	interval   uint // Seconds
	nextRun    time.Time
	state      string
}

func debug(args ...interface{}) {
	if *debugging {
		log.Println(args...)
	}
}

// hadleSchedule is the web server endpoint for path: /schedule
func handleSchedule(w http.ResponseWriter, r *http.Request) {
	routingKey, body, interval_s :=
		r.FormValue("routing_key"), r.FormValue("body"), r.FormValue("interval")
	debug("/schedule", routingKey, body)

	interval, err := strconv.ParseInt(interval_s, 10, 64)
	if err != nil {
		panic(err)
	}

	next_run := time.Now().UTC().Add(time.Duration(interval) * time.Second)
	_, err = db.Exec("INSERT INTO "+cfg.DB.Table+" "+
		"(routing_key, body, `interval`, next_run, state) "+
		"VALUES(?, ?, ?, ?, 'WAITING') "+
		"ON DUPLICATE KEY UPDATE `interval`=?",
		routingKey, body, interval, next_run, interval)
	if err != nil {
		panic(err)
	}

	// Wake up the publisher.
	//
	// publisher() may be sleeping for the next job on the queue
	// at the time we schedule a new Job. Let it wake up so it can
	// re-fetch the new Job from the front of the queue.
	//
	// The code below is an idiom for non-blocking send to a channel.
	select {
	case wakeUp <- 1:
		debug("Sent wakeup signal")
	default:
		debug("Skipped wakeup signal")
	}
}

// handleCancel is the web server endpoint for path: /cancel
func handleCancel(w http.ResponseWriter, r *http.Request) {
	routingKey, body := r.FormValue("routing_key"), r.FormValue("body")
	debug("/cancel", routingKey, body)

	_, err := db.Exec("DELETE FROM "+cfg.DB.Table+" "+
		"WHERE routing_key=? AND body=?", routingKey, body)
	if err != nil {
		panic(err)
	}
}

// front returns the first job to be run in the queue.
func front() (*Job, error) {
	j := Job{}
	row := db.QueryRow("SELECT routing_key, body, `interval`, next_run, state " +
		"FROM " + cfg.DB.Table + " " +
		"WHERE next_run = (" +
		"SELECT MIN(next_run) FROM " + cfg.DB.Table + ")")
	err := row.Scan(&j.routingKey, &j.body, &j.interval, &j.nextRun, &j.state)
	if err != nil {
		return nil, err
	}
	return &j, nil
}

// Delete deletes the Job j from the queue.
func (j *Job) Delete() error {
	_, err := db.Exec("DELETE FROM "+cfg.DB.Table+" "+
		"WHERE routing_key=? AND body=?",
		j.routingKey, j.body)
	if err != nil {
		return err
	}
	return nil
}

// Publish sends a message to exchange defined in the config and
// updates the Job's state to RUNNING on the database.
func (j *Job) Publish() error {
	debug("publish", *j)

	// Send a message to the broker
	err := channel.Publish(cfg.RabbitMQ.Exchange, j.routingKey, false, false, amqp.Publishing{
		Headers: amqp.Table{
			"interval":     j.interval,
			"published_at": time.Now().UTC(),
		},
		ContentType:     "application/octet-stream",
		ContentEncoding: "",
		Body:            []byte(j.body),
		DeliveryMode:    amqp.Persistent,
		Priority:        0,
	})
	if err != nil {
		return err
	}

	// Update state from database
	_, err = db.Exec("UPDATE "+cfg.DB.Table+" "+
		"SET state=RUNNING "+
		"WHERE routing_key=? AND body=?",
		j.routingKey, j.body)
	if err != nil {
		return err
	}
	return nil
}

// Remaining returns the duration until the job's next scheduled time.
func (j *Job) Remaining() time.Duration {
	return -time.Since(j.nextRun)
}

// publisher runs a loop that reads the next Job from the queue and publishes it.
func publisher() {
	for {
		job, err := front()
		if err != nil {
			fmt.Println(err)
			debug("Waiting wakeup signal")
			<-wakeUp
			debug("Got wakeup signal")
			continue
		}
		debug("Next job:", job, "Remaining:", job.Remaining())

		now := time.Now().UTC()
		if job.nextRun.After(now) {
			// Wait until the next Job time or
			// the webserver's /schedule handler wakes us up
			select {
			case <-time.After(job.Remaining()):
			case <-wakeUp:
				continue
			}
		} else {
			// It's time to publish the Job
			job.Publish()
			err = job.Delete()
			if err != nil {
				log.Println(err)
			}
		}

	}
}

func main() {
	flag.Parse()

	// Read config
	err := gcfg.ReadFileInto(&cfg, "dalga.ini")
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Read config: ", cfg)

	// Connect to database
	db, err = sql.Open(cfg.DB.Driver, cfg.DB.Dsn)
	if err != nil {
		log.Fatal(err)
	}
	err = db.Ping()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Connected to DB")

	// Connect to RabbitMQ
	broker, err = amqp.Dial(cfg.RabbitMQ.Uri)
	if err != nil {
		log.Fatal(err)
	}
	channel, err = broker.Channel()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Connected to RabbitMQ")

	// Run publisher
	go publisher()

	// Start HTTP server
	addr := cfg.HTTP.Host + ":" + cfg.HTTP.Port
	http.HandleFunc("/schedule", handleSchedule)
	http.HandleFunc("/cancel", handleCancel)
	http.ListenAndServe(addr, nil)
}
<|endoftext|>


<commit_before>// Copyright (C) 2019 The Syncthing Authors.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.

package stun

import (
	"context"
	"net"
	"sync/atomic"
	"time"

	"github.com/AudriusButkevicius/pfilter"
	"github.com/ccding/go-stun/stun"

	"github.com/syncthing/syncthing/lib/config"
	"github.com/syncthing/syncthing/lib/util"
)

const stunRetryInterval = 5 * time.Minute

type Host = stun.Host
type NATType = stun.NATType

// NAT types.

const (
	NATError                = stun.NATError
	NATUnknown              = stun.NATUnknown
	NATNone                 = stun.NATNone
	NATBlocked              = stun.NATBlocked
	NATFull                 = stun.NATFull
	NATSymmetric            = stun.NATSymmetric
	NATRestricted           = stun.NATRestricted
	NATPortRestricted       = stun.NATPortRestricted
	NATSymmetricUDPFirewall = stun.NATSymmetricUDPFirewall
)

type writeTrackingUdpConn struct {
	lastWrite int64 // atomic, must remain 64-bit aligned
	// Needs to be UDPConn not PacketConn, as pfilter checks for WriteMsgUDP/ReadMsgUDP
	// and even if we embed UDPConn here, in place of a PacketConn, seems the interface
	// check fails.
	*net.UDPConn
}

func (c *writeTrackingUdpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
	atomic.StoreInt64(&c.lastWrite, time.Now().Unix())
	return c.UDPConn.WriteTo(p, addr)
}

func (c *writeTrackingUdpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) {
	atomic.StoreInt64(&c.lastWrite, time.Now().Unix())
	return c.UDPConn.WriteMsgUDP(b, oob, addr)
}

func (c *writeTrackingUdpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (int, error) {
	atomic.StoreInt64(&c.lastWrite, time.Now().Unix())
	return c.UDPConn.WriteToUDP(b, addr)
}

func (c *writeTrackingUdpConn) Write(b []byte) (int, error) {
	atomic.StoreInt64(&c.lastWrite, time.Now().Unix())
	return c.UDPConn.Write(b)
}

func (c *writeTrackingUdpConn) getLastWrite() time.Time {
	unix := atomic.LoadInt64(&c.lastWrite)
	return time.Unix(unix, 0)
}

type Subscriber interface {
	OnNATTypeChanged(natType NATType)
	OnExternalAddressChanged(address *Host, via string)
}

type Service struct {
	name       string
	cfg        config.Wrapper
	subscriber Subscriber
	stunConn   net.PacketConn
	client     *stun.Client

	writeTrackingUdpConn *writeTrackingUdpConn

	natType NATType
	addr    *Host
}

func New(cfg config.Wrapper, subscriber Subscriber, conn *net.UDPConn) (*Service, net.PacketConn) {
	// Wrap the original connection to track writes on it
	writeTrackingUdpConn := &writeTrackingUdpConn{lastWrite: 0, UDPConn: conn}

	// Wrap it in a filter and split it up, so that stun packets arrive on stun conn, others arrive on the data conn
	filterConn := pfilter.NewPacketFilter(writeTrackingUdpConn)
	otherDataConn := filterConn.NewConn(otherDataPriority, nil)
	stunConn := filterConn.NewConn(stunFilterPriority, &stunFilter{
		ids: make(map[string]time.Time),
	})

	filterConn.Start()

	// Construct the client to use the stun conn
	client := stun.NewClientWithConnection(stunConn)
	client.SetSoftwareName("") // Explicitly unset this, seems to freak some servers out.

	// Return the service and the other conn to the client
	s := &Service{
		name: "Stun@" + conn.LocalAddr().Network() + "://" + conn.LocalAddr().String(),

		cfg:        cfg,
		subscriber: subscriber,
		stunConn:   stunConn,
		client:     client,

		writeTrackingUdpConn: writeTrackingUdpConn,

		natType: NATUnknown,
		addr:    nil,
	}
	return s, otherDataConn
}

func (s *Service) Serve(ctx context.Context) error {
	defer func() {
		s.setNATType(NATUnknown)
		s.setExternalAddress(nil, "")
	}()

	// Closing s.stunConn unblocks operations that use the connection
	// (Discover, Keepalive) and might otherwise block us from returning.
	go func() {
		<-ctx.Done()
		_ = s.stunConn.Close()
	}()

	timer := time.NewTimer(time.Millisecond)

	for {
	disabled:
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-timer.C:
		}

		if s.cfg.Options().IsStunDisabled() {
			timer.Reset(time.Second)
			continue
		}

		l.Debugf("Starting stun for %s", s)

		for _, addr := range s.cfg.Options().StunServers() {
			// This blocks until we hit an exit condition or there are issues with the STUN server.
			// This returns a boolean signifying if a different STUN server should be tried (oppose to the whole thing
			// shutting down and this winding itself down.
			s.runStunForServer(ctx, addr)

			// Have we been asked to stop?
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
			}

			// Are we disabled?
			if s.cfg.Options().IsStunDisabled() {
				l.Infoln("STUN disabled")
				s.setNATType(NATUnknown)
				s.setExternalAddress(nil, "")
				goto disabled
			}

			// Unpunchable NAT? Chillout for some time.
			if !s.isCurrentNATTypePunchable() {
				break
			}
		}

		// We failed to contact all provided stun servers or the nat is not punchable.
		// Chillout for a while.
		timer.Reset(stunRetryInterval)
	}
}

func (s *Service) runStunForServer(ctx context.Context, addr string) {
	l.Debugf("Running stun for %s via %s", s, addr)

	// Resolve the address, so that in case the server advertises two
	// IPs, we always hit the same one, as otherwise, the mapping might
	// expire as we hit the other address, and cause us to flip flop
	// between servers/external addresses, as a result flooding discovery
	// servers.
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		l.Debugf("%s stun addr resolution on %s: %s", s, addr, err)
		return
	}
	s.client.SetServerAddr(udpAddr.String())

	var natType stun.NATType
	var extAddr *stun.Host
	err = util.CallWithContext(ctx, func() error {
		natType, extAddr, err = s.client.Discover()
		return err
	})
	if err != nil || extAddr == nil {
		l.Debugf("%s stun discovery on %s: %s", s, addr, err)
		return
	}

	// The stun server is most likely borked, try another one.
	if natType == NATError || natType == NATUnknown || natType == NATBlocked {
		l.Debugf("%s stun discovery on %s resolved to %s", s, addr, natType)
		return
	}

	s.setNATType(natType)
	l.Debugf("%s detected NAT type: %s via %s", s, natType, addr)

	// We can't punch through this one, so no point doing keepalives
	// and such, just let the caller check the nat type and work it out themselves.
	if !s.isCurrentNATTypePunchable() {
		l.Debugf("%s cannot punch %s, skipping", s, natType)
		return
	}

	s.setExternalAddress(extAddr, addr)

	s.stunKeepAlive(ctx, addr, extAddr)
}

func (s *Service) stunKeepAlive(ctx context.Context, addr string, extAddr *Host) {
	var err error
	nextSleep := time.Duration(s.cfg.Options().StunKeepaliveStartS) * time.Second

	l.Debugf("%s starting stun keepalive via %s, next sleep %s", s, addr, nextSleep)

	for {
		if areDifferent(s.addr, extAddr) {
			// If the port has changed (addresses are not equal but the hosts are equal),
			// we're probably spending too much time between keepalives, reduce the sleep.
			if s.addr != nil && extAddr != nil && s.addr.IP() == extAddr.IP() {
				nextSleep /= 2
				l.Debugf("%s stun port change (%s to %s), next sleep %s", s, s.addr.TransportAddr(), extAddr.TransportAddr(), nextSleep)
			}

			s.setExternalAddress(extAddr, addr)

			// The stun server is probably stuffed, we've gone beyond min timeout, yet the address keeps changing.
			minSleep := time.Duration(s.cfg.Options().StunKeepaliveMinS) * time.Second
			if nextSleep < minSleep {
				l.Debugf("%s keepalive aborting, sleep below min: %s < %s", s, nextSleep, minSleep)
				return
			}
		}

		// Adjust the keepalives to fire only nextSleep after last write.
		lastWrite := s.writeTrackingUdpConn.getLastWrite()
		minSleep := time.Duration(s.cfg.Options().StunKeepaliveMinS) * time.Second
		if nextSleep < minSleep {
			nextSleep = minSleep
		}
	tryLater:
		sleepFor := nextSleep

		timeUntilNextKeepalive := time.Until(lastWrite.Add(sleepFor))
		if timeUntilNextKeepalive > 0 {
			sleepFor = timeUntilNextKeepalive
		}

		l.Debugf("%s stun sleeping for %s", s, sleepFor)

		select {
		case <-time.After(sleepFor):
		case <-ctx.Done():
			l.Debugf("%s stopping, aborting stun", s)
			return
		}

		if s.cfg.Options().IsStunDisabled() {
			// Disabled, give up
			l.Debugf("%s disabled, aborting stun ", s)
			return
		}

		// Check if any writes happened while we were sleeping, if they did, sleep again
		lastWrite = s.writeTrackingUdpConn.getLastWrite()
		if gap := time.Since(lastWrite); gap < nextSleep {
			l.Debugf("%s stun last write gap less than next sleep: %s < %s. Will try later", s, gap, nextSleep)
			goto tryLater
		}

		l.Debugf("%s stun keepalive", s)

		extAddr, err = s.client.Keepalive()
		if err != nil {
			l.Debugf("%s stun keepalive on %s: %s (%v)", s, addr, err, extAddr)
			return
		}
	}
}

func (s *Service) setNATType(natType NATType) {
	if natType != s.natType {
		l.Debugf("Notifying %s of NAT type change: %s", s.subscriber, natType)
		s.subscriber.OnNATTypeChanged(natType)
	}
	s.natType = natType
}

func (s *Service) setExternalAddress(addr *Host, via string) {
	if areDifferent(s.addr, addr) {
		l.Debugf("Notifying %s of address change: %s via %s", s.subscriber, addr, via)
		s.subscriber.OnExternalAddressChanged(addr, via)
	}
	s.addr = addr
}

func (s *Service) String() string {
	return s.name
}

func (s *Service) isCurrentNATTypePunchable() bool {
	return s.natType == NATNone || s.natType == NATPortRestricted || s.natType == NATRestricted || s.natType == NATFull || s.natType == NATSymmetricUDPFirewall
}

func areDifferent(first, second *Host) bool {
	if (first == nil) != (second == nil) {
		return true
	}
	if first != nil {
		return first.TransportAddr() != second.TransportAddr()
	}
	return false
}
<commit_msg>lib/stun: Prevent nil deref when naming service (#7872)<commit_after>// Copyright (C) 2019 The Syncthing Authors.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.

package stun

import (
	"context"
	"net"
	"sync/atomic"
	"time"

	"github.com/AudriusButkevicius/pfilter"
	"github.com/ccding/go-stun/stun"

	"github.com/syncthing/syncthing/lib/config"
	"github.com/syncthing/syncthing/lib/util"
)

const stunRetryInterval = 5 * time.Minute

type Host = stun.Host
type NATType = stun.NATType

// NAT types.

const (
	NATError                = stun.NATError
	NATUnknown              = stun.NATUnknown
	NATNone                 = stun.NATNone
	NATBlocked              = stun.NATBlocked
	NATFull                 = stun.NATFull
	NATSymmetric            = stun.NATSymmetric
	NATRestricted           = stun.NATRestricted
	NATPortRestricted       = stun.NATPortRestricted
	NATSymmetricUDPFirewall = stun.NATSymmetricUDPFirewall
)

type writeTrackingUdpConn struct {
	lastWrite int64 // atomic, must remain 64-bit aligned
	// Needs to be UDPConn not PacketConn, as pfilter checks for WriteMsgUDP/ReadMsgUDP
	// and even if we embed UDPConn here, in place of a PacketConn, seems the interface
	// check fails.
	*net.UDPConn
}

func (c *writeTrackingUdpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
	atomic.StoreInt64(&c.lastWrite, time.Now().Unix())
	return c.UDPConn.WriteTo(p, addr)
}

func (c *writeTrackingUdpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) {
	atomic.StoreInt64(&c.lastWrite, time.Now().Unix())
	return c.UDPConn.WriteMsgUDP(b, oob, addr)
}

func (c *writeTrackingUdpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (int, error) {
	atomic.StoreInt64(&c.lastWrite, time.Now().Unix())
	return c.UDPConn.WriteToUDP(b, addr)
}

func (c *writeTrackingUdpConn) Write(b []byte) (int, error) {
	atomic.StoreInt64(&c.lastWrite, time.Now().Unix())
	return c.UDPConn.Write(b)
}

func (c *writeTrackingUdpConn) getLastWrite() time.Time {
	unix := atomic.LoadInt64(&c.lastWrite)
	return time.Unix(unix, 0)
}

type Subscriber interface {
	OnNATTypeChanged(natType NATType)
	OnExternalAddressChanged(address *Host, via string)
}

type Service struct {
	name       string
	cfg        config.Wrapper
	subscriber Subscriber
	stunConn   net.PacketConn
	client     *stun.Client

	writeTrackingUdpConn *writeTrackingUdpConn

	natType NATType
	addr    *Host
}

func New(cfg config.Wrapper, subscriber Subscriber, conn *net.UDPConn) (*Service, net.PacketConn) {
	// Wrap the original connection to track writes on it
	writeTrackingUdpConn := &writeTrackingUdpConn{lastWrite: 0, UDPConn: conn}

	// Wrap it in a filter and split it up, so that stun packets arrive on stun conn, others arrive on the data conn
	filterConn := pfilter.NewPacketFilter(writeTrackingUdpConn)
	otherDataConn := filterConn.NewConn(otherDataPriority, nil)
	stunConn := filterConn.NewConn(stunFilterPriority, &stunFilter{
		ids: make(map[string]time.Time),
	})

	filterConn.Start()

	// Construct the client to use the stun conn
	client := stun.NewClientWithConnection(stunConn)
	client.SetSoftwareName("") // Explicitly unset this, seems to freak some servers out.

	// Return the service and the other conn to the client
	name := "Stun@"
	if local := conn.LocalAddr(); local != nil {
		name += local.Network() + "://" + local.String()
	} else {
		name += "unknown"
	}
	s := &Service{
		name: name,

		cfg:        cfg,
		subscriber: subscriber,
		stunConn:   stunConn,
		client:     client,

		writeTrackingUdpConn: writeTrackingUdpConn,

		natType: NATUnknown,
		addr:    nil,
	}
	return s, otherDataConn
}

func (s *Service) Serve(ctx context.Context) error {
	defer func() {
		s.setNATType(NATUnknown)
		s.setExternalAddress(nil, "")
	}()

	// Closing s.stunConn unblocks operations that use the connection
	// (Discover, Keepalive) and might otherwise block us from returning.
	go func() {
		<-ctx.Done()
		_ = s.stunConn.Close()
	}()

	timer := time.NewTimer(time.Millisecond)

	for {
	disabled:
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-timer.C:
		}

		if s.cfg.Options().IsStunDisabled() {
			timer.Reset(time.Second)
			continue
		}

		l.Debugf("Starting stun for %s", s)

		for _, addr := range s.cfg.Options().StunServers() {
			// This blocks until we hit an exit condition or there are issues with the STUN server.
			// This returns a boolean signifying if a different STUN server should be tried (oppose to the whole thing
			// shutting down and this winding itself down.
			s.runStunForServer(ctx, addr)

			// Have we been asked to stop?
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
			}

			// Are we disabled?
			if s.cfg.Options().IsStunDisabled() {
				l.Infoln("STUN disabled")
				s.setNATType(NATUnknown)
				s.setExternalAddress(nil, "")
				goto disabled
			}

			// Unpunchable NAT? Chillout for some time.
			if !s.isCurrentNATTypePunchable() {
				break
			}
		}

		// We failed to contact all provided stun servers or the nat is not punchable.
		// Chillout for a while.
		timer.Reset(stunRetryInterval)
	}
}

func (s *Service) runStunForServer(ctx context.Context, addr string) {
	l.Debugf("Running stun for %s via %s", s, addr)

	// Resolve the address, so that in case the server advertises two
	// IPs, we always hit the same one, as otherwise, the mapping might
	// expire as we hit the other address, and cause us to flip flop
	// between servers/external addresses, as a result flooding discovery
	// servers.
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		l.Debugf("%s stun addr resolution on %s: %s", s, addr, err)
		return
	}
	s.client.SetServerAddr(udpAddr.String())

	var natType stun.NATType
	var extAddr *stun.Host
	err = util.CallWithContext(ctx, func() error {
		natType, extAddr, err = s.client.Discover()
		return err
	})
	if err != nil || extAddr == nil {
		l.Debugf("%s stun discovery on %s: %s", s, addr, err)
		return
	}

	// The stun server is most likely borked, try another one.
	if natType == NATError || natType == NATUnknown || natType == NATBlocked {
		l.Debugf("%s stun discovery on %s resolved to %s", s, addr, natType)
		return
	}

	s.setNATType(natType)
	l.Debugf("%s detected NAT type: %s via %s", s, natType, addr)

	// We can't punch through this one, so no point doing keepalives
	// and such, just let the caller check the nat type and work it out themselves.
	if !s.isCurrentNATTypePunchable() {
		l.Debugf("%s cannot punch %s, skipping", s, natType)
		return
	}

	s.setExternalAddress(extAddr, addr)

	s.stunKeepAlive(ctx, addr, extAddr)
}

func (s *Service) stunKeepAlive(ctx context.Context, addr string, extAddr *Host) {
	var err error
	nextSleep := time.Duration(s.cfg.Options().StunKeepaliveStartS) * time.Second

	l.Debugf("%s starting stun keepalive via %s, next sleep %s", s, addr, nextSleep)

	for {
		if areDifferent(s.addr, extAddr) {
			// If the port has changed (addresses are not equal but the hosts are equal),
			// we're probably spending too much time between keepalives, reduce the sleep.
			if s.addr != nil && extAddr != nil && s.addr.IP() == extAddr.IP() {
				nextSleep /= 2
				l.Debugf("%s stun port change (%s to %s), next sleep %s", s, s.addr.TransportAddr(), extAddr.TransportAddr(), nextSleep)
			}

			s.setExternalAddress(extAddr, addr)

			// The stun server is probably stuffed, we've gone beyond min timeout, yet the address keeps changing.
			minSleep := time.Duration(s.cfg.Options().StunKeepaliveMinS) * time.Second
			if nextSleep < minSleep {
				l.Debugf("%s keepalive aborting, sleep below min: %s < %s", s, nextSleep, minSleep)
				return
			}
		}

		// Adjust the keepalives to fire only nextSleep after last write.
		lastWrite := s.writeTrackingUdpConn.getLastWrite()
		minSleep := time.Duration(s.cfg.Options().StunKeepaliveMinS) * time.Second
		if nextSleep < minSleep {
			nextSleep = minSleep
		}
	tryLater:
		sleepFor := nextSleep

		timeUntilNextKeepalive := time.Until(lastWrite.Add(sleepFor))
		if timeUntilNextKeepalive > 0 {
			sleepFor = timeUntilNextKeepalive
		}

		l.Debugf("%s stun sleeping for %s", s, sleepFor)

		select {
		case <-time.After(sleepFor):
		case <-ctx.Done():
			l.Debugf("%s stopping, aborting stun", s)
			return
		}

		if s.cfg.Options().IsStunDisabled() {
			// Disabled, give up
			l.Debugf("%s disabled, aborting stun ", s)
			return
		}

		// Check if any writes happened while we were sleeping, if they did, sleep again
		lastWrite = s.writeTrackingUdpConn.getLastWrite()
		if gap := time.Since(lastWrite); gap < nextSleep {
			l.Debugf("%s stun last write gap less than next sleep: %s < %s. Will try later", s, gap, nextSleep)
			goto tryLater
		}

		l.Debugf("%s stun keepalive", s)

		extAddr, err = s.client.Keepalive()
		if err != nil {
			l.Debugf("%s stun keepalive on %s: %s (%v)", s, addr, err, extAddr)
			return
		}
	}
}

func (s *Service) setNATType(natType NATType) {
	if natType != s.natType {
		l.Debugf("Notifying %s of NAT type change: %s", s.subscriber, natType)
		s.subscriber.OnNATTypeChanged(natType)
	}
	s.natType = natType
}

func (s *Service) setExternalAddress(addr *Host, via string) {
	if areDifferent(s.addr, addr) {
		l.Debugf("Notifying %s of address change: %s via %s", s.subscriber, addr, via)
		s.subscriber.OnExternalAddressChanged(addr, via)
	}
	s.addr = addr
}

func (s *Service) String() string {
	return s.name
}

func (s *Service) isCurrentNATTypePunchable() bool {
	return s.natType == NATNone || s.natType == NATPortRestricted || s.natType == NATRestricted || s.natType == NATFull || s.natType == NATSymmetricUDPFirewall
}

func areDifferent(first, second *Host) bool {
	if (first == nil) != (second == nil) {
		return true
	}
	if first != nil {
		return first.TransportAddr() != second.TransportAddr()
	}
	return false
}
<|endoftext|>