<commit_before>package deamon

import (
	"fmt"
	"github.com/ekalinin/gogalert/api"
	"time"
)

type StatServer struct {
	// Period to refresh data from ganglia
	refreshRate time.Duration
	// Source of the XML data
	source *api.DataSource
	// threshholds
	threshholds *api.MetricFilter
	// config file path
	configPath string
	// config details
	config *Config
}

func NewStatServer(source *api.DataSource, threshholds *api.MetricFilter, configPath string) *StatServer {
	return &StatServer{2 * time.Minute, source, threshholds, configPath, nil}
}

func (this *StatServer) Start() error {
	fmt.Println(" Server: started, config: ", *this)
	this.config = ParseConfig(this.configPath)
	fmt.Println(" Server: parsed config: ", *this.config)
	for {
		go this.CollectStats()
		time.Sleep(this.refreshRate)
	}

	return nil
}

func (this *StatServer) CollectStats() {
	fmt.Println(" Server: collecting data ...")
	for _, metric := range api.Parse(this.source).Find(this.threshholds) {
		fmt.Println(" >> Alert: ", metric.GMetric)
	}
}
<commit_msg>better deamon logic (added handlers for several system signals)<commit_after>package deamon

import (
	"fmt"
	"github.com/ekalinin/gogalert/api"
	"os"
	"os/signal"
	"syscall"
	"time"
)

type StatServer struct {
	// Period to refresh data from ganglia
	refreshRate int
	// Source of the XML data
	source *api.DataSource
	// threshholds
	threshholds *api.MetricFilter
	// config file path
	configPath string
	// config details
	config *Config
}

func NewStatServer(source *api.DataSource, threshholds *api.MetricFilter, configPath string) *StatServer {
	return &StatServer{60, source, threshholds, configPath, nil}
}

func (this *StatServer) Start() error {
	fmt.Println(" Server: started, config: ", *this)
	this.config = ParseConfig(this.configPath)
	this.refreshRate = this.config.Interval
	fmt.Println(" Server: parsed config: ", *this.config)

	go this.CollectStats()

	signalChan := make(chan os.Signal, 1)
	// http://adampresley.com/2015/02/16/waiting-for-goroutines-to-finish-running-before-exiting.html
	// https://gist.github.com/reiki4040/be3705f307d3cd136e85
	signal.Notify(signalChan,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)

	needStop := false
	for {
		select {
		case <-time.After(time.Duration(this.refreshRate) * time.Second):
			go this.CollectStats()
		case <-signalChan:
			fmt.Println("Stopping...")
			needStop = true
		}

		if needStop {
			break
		}
	}

	return nil
}

func (this *StatServer) CollectStats() {
	fmt.Println(" Server: collecting data ...")
	for _, metric := range api.Parse(this.source).Find(this.threshholds) {
		fmt.Println(" >> Alert: ", metric.GMetric)
	}
}
<|endoftext|>


<commit_before>package vcf

import (
	"fmt"
	"log"
	"os"
)

// Channels should be initialized and passed to the ToChannel function. Notice closing the channels should
// not be done by the client. This will happen inside ToChannel, when the input is exhausted.
func Example() {
	validVariants := make(chan *Variant, 100)      // buffered channel for correctly parsed variants
	invalidVariants := make(chan InvalidLine, 100) // buffered channel for variants that fail to parse

	filename := "example_vcfs/test.vcf"

	vcfFile, err := os.Open(filename)
	if err != nil {
		log.Fatalln("can't open file", filename)
	}
	defer vcfFile.Close()

	go func() {
		err := ToChannel(vcfFile, validVariants, invalidVariants)
		if err != nil {
			log.Fatalln(err)
		}
	}()

	go func() {
		// consume invalid variants channel asynchronously
		for invalid := range invalidVariants {
			log.Println("failed to parse line", invalid.Line, "with error", invalid.Err)
		}
	}()

	for variant := range validVariants {
		fmt.Println(variant)
		if variant.Qual != nil {
			fmt.Println("Quality:", *variant.Qual)
		}
		fmt.Println("Filter:", variant.Filter)
		fmt.Println("Allele Count:", *variant.AlleleCount)
		fmt.Println("Allele Frequency:", *variant.AlleleFrequency)
		fmt.Println("Total Alleles:", *variant.TotalAlleles)
		fmt.Println("Depth:", *variant.Depth)
		fmt.Println("Mapping Quality:", *variant.MappingQuality)
		fmt.Println("MAPQ0 Reads:", *variant.MAPQ0Reads)

		rawInfo := variant.Info
		vqslod := rawInfo["VQSLOD"]
		fmt.Println("VQSLOD:", vqslod)
	}

	// output:
	// Chromosome: 1 Position: 762588 Reference: G Alternative: C
	// Quality: 40
	// Filter: PASS
	// Allele Count: 2
	// Allele Frequency: 1
	// Total Alleles: 2
	// Depth: 5
	// Mapping Quality: 43.32
	// MAPQ0 Reads: 0
	// VQSLOD: 1.18
}

func ExampleSampleIDs() {
	filename := "example_vcfs/testsamples.vcf"
	vcfFile, err := os.Open(filename)
	if err != nil {
		log.Fatalln("can't open file", filename)
	}
	defer vcfFile.Close()

	sampleIDs, err := SampleIDs(vcfFile)
	if err == nil && sampleIDs != nil {
		for i, sample := range sampleIDs {
			fmt.Printf("sample %d: %s\n", i, sample)
		}
	}
	// output:
	// sample 0: 111222
}
<commit_msg>Example reworded<commit_after>package vcf

import (
	"fmt"
	"log"
	"os"
)

// Channels should be initialized and passed to the ToChannel function. The client should not close the channels
// This will happen inside ToChannel, when the input is exhausted.
func Example() {
	validVariants := make(chan *Variant, 100)      // buffered channel for correctly parsed variants
	invalidVariants := make(chan InvalidLine, 100) // buffered channel for variants that fail to parse

	filename := "example_vcfs/test.vcf"

	vcfFile, err := os.Open(filename)
	if err != nil {
		log.Fatalln("can't open file", filename)
	}
	defer vcfFile.Close()

	go func() {
		err := ToChannel(vcfFile, validVariants, invalidVariants)
		if err != nil {
			log.Fatalln(err)
		}
	}()

	go func() {
		// consume invalid variants channel asynchronously
		for invalid := range invalidVariants {
			log.Println("failed to parse line", invalid.Line, "with error", invalid.Err)
		}
	}()

	for variant := range validVariants {
		fmt.Println(variant)
		if variant.Qual != nil {
			fmt.Println("Quality:", *variant.Qual)
		}
		fmt.Println("Filter:", variant.Filter)
		fmt.Println("Allele Count:", *variant.AlleleCount)
		fmt.Println("Allele Frequency:", *variant.AlleleFrequency)
		fmt.Println("Total Alleles:", *variant.TotalAlleles)
		fmt.Println("Depth:", *variant.Depth)
		fmt.Println("Mapping Quality:", *variant.MappingQuality)
		fmt.Println("MAPQ0 Reads:", *variant.MAPQ0Reads)

		rawInfo := variant.Info
		vqslod := rawInfo["VQSLOD"]
		fmt.Println("VQSLOD:", vqslod)
	}

	// output:
	// Chromosome: 1 Position: 762588 Reference: G Alternative: C
	// Quality: 40
	// Filter: PASS
	// Allele Count: 2
	// Allele Frequency: 1
	// Total Alleles: 2
	// Depth: 5
	// Mapping Quality: 43.32
	// MAPQ0 Reads: 0
	// VQSLOD: 1.18
}

func ExampleSampleIDs() {
	filename := "example_vcfs/testsamples.vcf"
	vcfFile, err := os.Open(filename)
	if err != nil {
		log.Fatalln("can't open file", filename)
	}
	defer vcfFile.Close()

	sampleIDs, err := SampleIDs(vcfFile)
	if err == nil && sampleIDs != nil {
		for i, sample := range sampleIDs {
			fmt.Printf("sample %d: %s\n", i, sample)
		}
	}
	// output:
	// sample 0: 111222
}
<|endoftext|>


<commit_before>package jsoniter_test

import (
	"fmt"
	"github.com/json-iterator/go"
	"os"
)

func ExampleMarshal() {
	type ColorGroup struct {
		ID     int
		Name   string
		Colors []string
	}
	group := ColorGroup{
		ID:     1,
		Name:   "Reds",
		Colors: []string{"Crimson", "Red", "Ruby", "Maroon"},
	}
	b, err := jsoniter.Marshal(group)
	if err != nil {
		fmt.Println("error:", err)
	}
	os.Stdout.Write(b)
	// Output:
	// {"ID":1,"Name":"Reds","Colors":["Crimson","Red","Ruby","Maroon"]}
}
<commit_msg>update test<commit_after>package jsoniter_test

import (
	"fmt"
	"github.com/json-iterator/go"
	"os"
	"testing"
)

func Test_ExampleMarshal(t *testing.T) {
	type ColorGroup struct {
		ID     int
		Name   string
		Colors []string
	}
	group := ColorGroup{
		ID:     1,
		Name:   "Reds",
		Colors: []string{"Crimson", "Red", "Ruby", "Maroon"},
	}
	b, err := jsoniter.Marshal(group)
	if err != nil {
		fmt.Println("error:", err)
	}
	os.Stdout.Write(b)
	// Output:
	// {"ID":1,"Name":"Reds","Colors":["Crimson","Red","Ruby","Maroon"]}
}
<|endoftext|>


<commit_before>// Copyright 2022 Google LLC
//
// 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 v1alpha1 contains API Schema definitions for the v1alpha1 API group
//+kubebuilder:object:generate=true
//+groupName=config.porch.kpt.dev
package v1alpha1

import (
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
)

//go:generate go run sigs.k8s.io/controller-tools/cmd/controller-gen@v0.8.0 object object:headerFile="../../../../../hack/boilerplate.go.txt" crd:crdVersions=v1 output:crd:artifacts:config=. paths=./...

var (
	// GroupVersion is group version used to register these objects
	GroupVersion = schema.GroupVersion{Group: "config.porch.kpt.dev", Version: "v1alpha1"}

	// We removed SchemeBuilder to keep our dependencies small

	KindRepository = KindInfo{
		Resource: GroupVersion.WithResource("repositories"),
		objects:  []runtime.Object{&Repository{}, &RepositoryList{}},
	}

	AllKinds = []KindInfo{KindRepository}
)

//+kubebuilder:object:generate=false

// KindInfo holds type meta-information
type KindInfo struct {
	Resource schema.GroupVersionResource
	objects  []runtime.Object
}

// GroupResource returns the GroupResource for the kind
func (k *KindInfo) GroupResource() schema.GroupResource {
	return k.Resource.GroupResource()
}

func AddToScheme(scheme *runtime.Scheme) error {
	for _, kind := range AllKinds {
		scheme.AddKnownTypes(GroupVersion, kind.objects...)
	}
	metav1.AddToGroupVersion(scheme, GroupVersion)
	return nil
}
<commit_msg>Fix controller generation (#2969)<commit_after>// Copyright 2022 Google LLC
//
// 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 v1alpha1 contains API Schema definitions for the v1alpha1 API group
//+kubebuilder:object:generate=true
//+groupName=config.porch.kpt.dev
package v1alpha1

import (
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
)

//go:generate go run sigs.k8s.io/controller-tools/cmd/controller-gen@v0.8.0 object object:headerFile="../../../hack/boilerplate.go.txt" crd:crdVersions=v1 output:crd:artifacts:config=. paths=./...

var (
	// GroupVersion is group version used to register these objects
	GroupVersion = schema.GroupVersion{Group: "config.porch.kpt.dev", Version: "v1alpha1"}

	// We removed SchemeBuilder to keep our dependencies small

	KindRepository = KindInfo{
		Resource: GroupVersion.WithResource("repositories"),
		objects:  []runtime.Object{&Repository{}, &RepositoryList{}},
	}

	AllKinds = []KindInfo{KindRepository}
)

//+kubebuilder:object:generate=false

// KindInfo holds type meta-information
type KindInfo struct {
	Resource schema.GroupVersionResource
	objects  []runtime.Object
}

// GroupResource returns the GroupResource for the kind
func (k *KindInfo) GroupResource() schema.GroupResource {
	return k.Resource.GroupResource()
}

func AddToScheme(scheme *runtime.Scheme) error {
	for _, kind := range AllKinds {
		scheme.AddKnownTypes(GroupVersion, kind.objects...)
	}
	metav1.AddToGroupVersion(scheme, GroupVersion)
	return nil
}
<|endoftext|>


<commit_before>package vagrantcloud

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"

	"github.com/hashicorp/packer/packer"
	"github.com/stretchr/testify/assert"
)

type tarFiles []struct {
	Name, Body string
}

func testGoodConfig() map[string]interface{} {
	return map[string]interface{}{
		"access_token":        "foo",
		"version_description": "bar",
		"box_tag":             "hashicorp/precise64",
		"version":             "0.5",
	}
}

func testBadConfig() map[string]interface{} {
	return map[string]interface{}{
		"access_token":        "foo",
		"box_tag":             "baz",
		"version_description": "bar",
	}
}

func newSecureServer(token string, handler http.HandlerFunc) *httptest.Server {
	token = fmt.Sprintf("Bearer %s", token)
	return httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		if req.Header.Get("authorization") != token {
			http.Error(rw, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
			return
		}
		if handler != nil {
			handler(rw, req)
		}
	}))
}

func newSelfSignedSslServer(token string, handler http.HandlerFunc) *httptest.Server {
	token = fmt.Sprintf("Bearer %s", token)
	return httptest.NewTLSServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		if req.Header.Get("authorization") != token {
			http.Error(rw, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
			return
		}
		if handler != nil {
			handler(rw, req)
		}
	}))
}

func TestPostProcessor_Insecure_Ssl(t *testing.T) {
	var p PostProcessor
	server := newSelfSignedSslServer("foo", nil)
	defer server.Close()

	config := testGoodConfig()
	config["vagrant_cloud_url"] = server.URL
	config["insecure_skip_tls_verify"] = true
	if err := p.Configure(config); err != nil {
		t.Fatalf("Expected TLS to skip certificate validation: %s", err)
	}
}

func TestPostProcessor_Configure_fromVagrantEnv(t *testing.T) {
	var p PostProcessor
	config := testGoodConfig()
	server := newSecureServer("bar", nil)
	defer server.Close()
	config["vagrant_cloud_url"] = server.URL
	config["access_token"] = ""
	os.Setenv("VAGRANT_CLOUD_TOKEN", "bar")
	defer func() {
		os.Setenv("VAGRANT_CLOUD_TOKEN", "")
	}()

	if err := p.Configure(config); err != nil {
		t.Fatalf("err: %s", err)
	}

	if p.config.AccessToken != "bar" {
		t.Fatalf("Expected to get token from VAGRANT_CLOUD_TOKEN env var. Got '%s' instead",
			p.config.AccessToken)
	}
}

func TestPostProcessor_Configure_fromAtlasEnv(t *testing.T) {
	var p PostProcessor
	config := testGoodConfig()
	config["access_token"] = ""
	server := newSecureServer("foo", nil)
	defer server.Close()
	config["vagrant_cloud_url"] = server.URL
	os.Setenv("ATLAS_TOKEN", "foo")
	defer func() {
		os.Setenv("ATLAS_TOKEN", "")
	}()

	if err := p.Configure(config); err != nil {
		t.Fatalf("err: %s", err)
	}

	if p.config.AccessToken != "foo" {
		t.Fatalf("Expected to get token from ATLAS_TOKEN env var. Got '%s' instead",
			p.config.AccessToken)
	}

	if !p.warnAtlasToken {
		t.Fatal("Expected warn flag to be set when getting token from atlas env var.")
	}
}

func TestPostProcessor_Configure_Good(t *testing.T) {
	config := testGoodConfig()
	server := newSecureServer("foo", nil)
	defer server.Close()
	config["vagrant_cloud_url"] = server.URL
	var p PostProcessor
	if err := p.Configure(config); err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestPostProcessor_Configure_Bad(t *testing.T) {
	config := testBadConfig()
	server := newSecureServer("foo", nil)
	defer server.Close()
	config["vagrant_cloud_url"] = server.URL
	var p PostProcessor
	if err := p.Configure(config); err == nil {
		t.Fatalf("should have err")
	}
}

func TestPostProcessor_PostProcess_checkArtifactType(t *testing.T) {
	artifact := &packer.MockArtifact{
		BuilderIdValue: "invalid.builder",
	}

	config := testGoodConfig()
	server := newSecureServer("foo", nil)
	defer server.Close()
	config["vagrant_cloud_url"] = server.URL
	var p PostProcessor

	p.Configure(config)
	_, _, _, err := p.PostProcess(context.Background(), testUi(), artifact)
	if !strings.Contains(err.Error(), "Unknown artifact type") {
		t.Fatalf("Should error with message 'Unknown artifact type...' with BuilderId: %s", artifact.BuilderIdValue)
	}
}

func TestPostProcessor_PostProcess_checkArtifactFileIsBox(t *testing.T) {
	artifact := &packer.MockArtifact{
		BuilderIdValue: "mitchellh.post-processor.vagrant", // good
		FilesValue:     []string{"invalid.boxfile"},        // should have .box extension
	}

	config := testGoodConfig()
	server := newSecureServer("foo", nil)
	defer server.Close()
	config["vagrant_cloud_url"] = server.URL
	var p PostProcessor

	p.Configure(config)
	_, _, _, err := p.PostProcess(context.Background(), testUi(), artifact)
	if !strings.Contains(err.Error(), "Unknown files in artifact") {
		t.Fatalf("Should error with message 'Unknown files in artifact...' with artifact file: %s",
			artifact.FilesValue[0])
	}
}

func testUi() *packer.BasicUi {
	return &packer.BasicUi{
		Reader: new(bytes.Buffer),
		Writer: new(bytes.Buffer),
	}
}

func TestPostProcessor_ImplementsPostProcessor(t *testing.T) {
	var _ packer.PostProcessor = new(PostProcessor)
}

func TestProviderFromBuilderName(t *testing.T) {
	if providerFromBuilderName("foobar") != "foobar" {
		t.Fatal("should copy unknown provider")
	}

	if providerFromBuilderName("vmware") != "vmware_desktop" {
		t.Fatal("should convert provider")
	}
}

func TestProviderFromVagrantBox_missing_box(t *testing.T) {
	// Bad: Box does not exist
	boxfile := "i_dont_exist.box"
	_, err := providerFromVagrantBox(boxfile)
	if err == nil {
		t.Fatal("Should have error as box file does not exist")
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_empty_box(t *testing.T) {
	// Bad: Empty box file
	boxfile, err := newBoxFile()
	if err != nil {
		t.Fatalf("%s", err)
	}
	defer os.Remove(boxfile.Name())

	_, err = providerFromVagrantBox(boxfile.Name())
	if err == nil {
		t.Fatal("Should have error as box file is empty")
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_gzip_only_box(t *testing.T) {
	boxfile, err := newBoxFile()
	if err != nil {
		t.Fatalf("%s", err)
	}
	defer os.Remove(boxfile.Name())

	// Bad: Box is just a plain gzip file
	aw := gzip.NewWriter(boxfile)
	_, err = aw.Write([]byte("foo content"))
	if err != nil {
		t.Fatal("Error zipping test box file")
	}
	aw.Close() // Flush the gzipped contents to file

	_, err = providerFromVagrantBox(boxfile.Name())
	if err == nil {
		t.Fatalf("Should have error as box file is a plain gzip file: %s", err)
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_no_files_in_archive(t *testing.T) {
	// Bad: Box contains no files
	boxfile, err := createBox(tarFiles{})
	if err != nil {
		t.Fatalf("Error creating test box: %s", err)
	}
	defer os.Remove(boxfile.Name())

	_, err = providerFromVagrantBox(boxfile.Name())
	if err == nil {
		t.Fatalf("Should have error as box file has no contents")
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_no_metadata(t *testing.T) {
	// Bad: Box contains no metadata/metadata.json file
	files := tarFiles{
		{"foo.txt", "This is a foo file"},
		{"bar.txt", "This is a bar file"},
	}
	boxfile, err := createBox(files)
	if err != nil {
		t.Fatalf("Error creating test box: %s", err)
	}
	defer os.Remove(boxfile.Name())

	_, err = providerFromVagrantBox(boxfile.Name())
	if err == nil {
		t.Fatalf("Should have error as box file does not include metadata.json file")
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_metadata_empty(t *testing.T) {
	// Bad: Create a box with an empty metadata.json file
	files := tarFiles{
		{"foo.txt", "This is a foo file"},
		{"bar.txt", "This is a bar file"},
		{"metadata.json", ""},
	}
	boxfile, err := createBox(files)
	if err != nil {
		t.Fatalf("Error creating test box: %s", err)
	}
	defer os.Remove(boxfile.Name())

	_, err = providerFromVagrantBox(boxfile.Name())
	if err == nil {
		t.Fatalf("Should have error as box files metadata.json file is empty")
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_metadata_provider_value_empty(t *testing.T) {
	// Bad: The boxes metadata.json file 'provider' key has an empty value
	files := tarFiles{
		{"foo.txt", "This is a foo file"},
		{"bar.txt", "This is a bar file"},
		{"metadata.json", `{"provider":""}`},
	}
	boxfile, err := createBox(files)
	if err != nil {
		t.Fatalf("Error creating test box: %s", err)
	}
	defer os.Remove(boxfile.Name())

	_, err = providerFromVagrantBox(boxfile.Name())
	if err == nil {
		t.Fatalf("Should have error as boxes metadata.json file 'provider' key is empty")
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_metadata_ok(t *testing.T) {
	// Good: The boxes metadata.json file has the 'provider' key/value pair
	expectedProvider := "virtualbox"
	files := tarFiles{
		{"foo.txt", "This is a foo file"},
		{"bar.txt", "This is a bar file"},
		{"metadata.json", `{"provider":"` + expectedProvider + `"}`},
	}
	boxfile, err := createBox(files)
	if err != nil {
		t.Fatalf("Error creating test box: %s", err)
	}
	defer os.Remove(boxfile.Name())

	provider, err := providerFromVagrantBox(boxfile.Name())
	assert.Equal(t, expectedProvider, provider, "Error: Expected provider: '%s'. Got '%s'", expectedProvider, provider)
	t.Logf("Expected provider '%s'. Got provider '%s'", expectedProvider, provider)
}

func newBoxFile() (boxfile *os.File, err error) {
	boxfile, err = ioutil.TempFile(os.TempDir(), "test*.box")
	if err != nil {
		return boxfile, fmt.Errorf("Error creating test box file: %s", err)
	}
	return boxfile, nil
}

func createBox(files tarFiles) (boxfile *os.File, err error) {
	boxfile, err = newBoxFile()
	if err != nil {
		return boxfile, err
	}

	// Box files are gzipped tar archives
	aw := gzip.NewWriter(boxfile)
	tw := tar.NewWriter(aw)

	// Add each file to the box
	for _, file := range files {
		// Create and write the tar file header
		hdr := &tar.Header{
			Name: file.Name,
			Mode: 0644,
			Size: int64(len(file.Body)),
		}
		err = tw.WriteHeader(hdr)
		if err != nil {
			return boxfile, fmt.Errorf("Error writing box tar file header: %s", err)
		}
		// Write the file contents
		_, err = tw.Write([]byte(file.Body))
		if err != nil {
			return boxfile, fmt.Errorf("Error writing box tar file contents: %s", err)
		}
	}
	// Flush and close each writer
	err = tw.Close()
	if err != nil {
		return boxfile, fmt.Errorf("Error flushing tar file contents: %s", err)
	}
	err = aw.Close()
	if err != nil {
		return boxfile, fmt.Errorf("Error flushing gzip file contents: %s", err)
	}

	return boxfile, nil
}
<commit_msg>Should return an error if the metadata file has badly formatted JSON<commit_after>package vagrantcloud

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"

	"github.com/hashicorp/packer/packer"
	"github.com/stretchr/testify/assert"
)

type tarFiles []struct {
	Name, Body string
}

func testGoodConfig() map[string]interface{} {
	return map[string]interface{}{
		"access_token":        "foo",
		"version_description": "bar",
		"box_tag":             "hashicorp/precise64",
		"version":             "0.5",
	}
}

func testBadConfig() map[string]interface{} {
	return map[string]interface{}{
		"access_token":        "foo",
		"box_tag":             "baz",
		"version_description": "bar",
	}
}

func newSecureServer(token string, handler http.HandlerFunc) *httptest.Server {
	token = fmt.Sprintf("Bearer %s", token)
	return httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		if req.Header.Get("authorization") != token {
			http.Error(rw, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
			return
		}
		if handler != nil {
			handler(rw, req)
		}
	}))
}

func newSelfSignedSslServer(token string, handler http.HandlerFunc) *httptest.Server {
	token = fmt.Sprintf("Bearer %s", token)
	return httptest.NewTLSServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		if req.Header.Get("authorization") != token {
			http.Error(rw, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
			return
		}
		if handler != nil {
			handler(rw, req)
		}
	}))
}

func TestPostProcessor_Insecure_Ssl(t *testing.T) {
	var p PostProcessor
	server := newSelfSignedSslServer("foo", nil)
	defer server.Close()

	config := testGoodConfig()
	config["vagrant_cloud_url"] = server.URL
	config["insecure_skip_tls_verify"] = true
	if err := p.Configure(config); err != nil {
		t.Fatalf("Expected TLS to skip certificate validation: %s", err)
	}
}

func TestPostProcessor_Configure_fromVagrantEnv(t *testing.T) {
	var p PostProcessor
	config := testGoodConfig()
	server := newSecureServer("bar", nil)
	defer server.Close()
	config["vagrant_cloud_url"] = server.URL
	config["access_token"] = ""
	os.Setenv("VAGRANT_CLOUD_TOKEN", "bar")
	defer func() {
		os.Setenv("VAGRANT_CLOUD_TOKEN", "")
	}()

	if err := p.Configure(config); err != nil {
		t.Fatalf("err: %s", err)
	}

	if p.config.AccessToken != "bar" {
		t.Fatalf("Expected to get token from VAGRANT_CLOUD_TOKEN env var. Got '%s' instead",
			p.config.AccessToken)
	}
}

func TestPostProcessor_Configure_fromAtlasEnv(t *testing.T) {
	var p PostProcessor
	config := testGoodConfig()
	config["access_token"] = ""
	server := newSecureServer("foo", nil)
	defer server.Close()
	config["vagrant_cloud_url"] = server.URL
	os.Setenv("ATLAS_TOKEN", "foo")
	defer func() {
		os.Setenv("ATLAS_TOKEN", "")
	}()

	if err := p.Configure(config); err != nil {
		t.Fatalf("err: %s", err)
	}

	if p.config.AccessToken != "foo" {
		t.Fatalf("Expected to get token from ATLAS_TOKEN env var. Got '%s' instead",
			p.config.AccessToken)
	}

	if !p.warnAtlasToken {
		t.Fatal("Expected warn flag to be set when getting token from atlas env var.")
	}
}

func TestPostProcessor_Configure_Good(t *testing.T) {
	config := testGoodConfig()
	server := newSecureServer("foo", nil)
	defer server.Close()
	config["vagrant_cloud_url"] = server.URL
	var p PostProcessor
	if err := p.Configure(config); err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestPostProcessor_Configure_Bad(t *testing.T) {
	config := testBadConfig()
	server := newSecureServer("foo", nil)
	defer server.Close()
	config["vagrant_cloud_url"] = server.URL
	var p PostProcessor
	if err := p.Configure(config); err == nil {
		t.Fatalf("should have err")
	}
}

func TestPostProcessor_PostProcess_checkArtifactType(t *testing.T) {
	artifact := &packer.MockArtifact{
		BuilderIdValue: "invalid.builder",
	}

	config := testGoodConfig()
	server := newSecureServer("foo", nil)
	defer server.Close()
	config["vagrant_cloud_url"] = server.URL
	var p PostProcessor

	p.Configure(config)
	_, _, _, err := p.PostProcess(context.Background(), testUi(), artifact)
	if !strings.Contains(err.Error(), "Unknown artifact type") {
		t.Fatalf("Should error with message 'Unknown artifact type...' with BuilderId: %s", artifact.BuilderIdValue)
	}
}

func TestPostProcessor_PostProcess_checkArtifactFileIsBox(t *testing.T) {
	artifact := &packer.MockArtifact{
		BuilderIdValue: "mitchellh.post-processor.vagrant", // good
		FilesValue:     []string{"invalid.boxfile"},        // should have .box extension
	}

	config := testGoodConfig()
	server := newSecureServer("foo", nil)
	defer server.Close()
	config["vagrant_cloud_url"] = server.URL
	var p PostProcessor

	p.Configure(config)
	_, _, _, err := p.PostProcess(context.Background(), testUi(), artifact)
	if !strings.Contains(err.Error(), "Unknown files in artifact") {
		t.Fatalf("Should error with message 'Unknown files in artifact...' with artifact file: %s",
			artifact.FilesValue[0])
	}
}

func testUi() *packer.BasicUi {
	return &packer.BasicUi{
		Reader: new(bytes.Buffer),
		Writer: new(bytes.Buffer),
	}
}

func TestPostProcessor_ImplementsPostProcessor(t *testing.T) {
	var _ packer.PostProcessor = new(PostProcessor)
}

func TestProviderFromBuilderName(t *testing.T) {
	if providerFromBuilderName("foobar") != "foobar" {
		t.Fatal("should copy unknown provider")
	}

	if providerFromBuilderName("vmware") != "vmware_desktop" {
		t.Fatal("should convert provider")
	}
}

func TestProviderFromVagrantBox_missing_box(t *testing.T) {
	// Bad: Box does not exist
	boxfile := "i_dont_exist.box"
	_, err := providerFromVagrantBox(boxfile)
	if err == nil {
		t.Fatal("Should have error as box file does not exist")
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_empty_box(t *testing.T) {
	// Bad: Empty box file
	boxfile, err := newBoxFile()
	if err != nil {
		t.Fatalf("%s", err)
	}
	defer os.Remove(boxfile.Name())

	_, err = providerFromVagrantBox(boxfile.Name())
	if err == nil {
		t.Fatal("Should have error as box file is empty")
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_gzip_only_box(t *testing.T) {
	boxfile, err := newBoxFile()
	if err != nil {
		t.Fatalf("%s", err)
	}
	defer os.Remove(boxfile.Name())

	// Bad: Box is just a plain gzip file
	aw := gzip.NewWriter(boxfile)
	_, err = aw.Write([]byte("foo content"))
	if err != nil {
		t.Fatal("Error zipping test box file")
	}
	aw.Close() // Flush the gzipped contents to file

	_, err = providerFromVagrantBox(boxfile.Name())
	if err == nil {
		t.Fatalf("Should have error as box file is a plain gzip file: %s", err)
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_no_files_in_archive(t *testing.T) {
	// Bad: Box contains no files
	boxfile, err := createBox(tarFiles{})
	if err != nil {
		t.Fatalf("Error creating test box: %s", err)
	}
	defer os.Remove(boxfile.Name())

	_, err = providerFromVagrantBox(boxfile.Name())
	if err == nil {
		t.Fatalf("Should have error as box file has no contents")
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_no_metadata(t *testing.T) {
	// Bad: Box contains no metadata/metadata.json file
	files := tarFiles{
		{"foo.txt", "This is a foo file"},
		{"bar.txt", "This is a bar file"},
	}
	boxfile, err := createBox(files)
	if err != nil {
		t.Fatalf("Error creating test box: %s", err)
	}
	defer os.Remove(boxfile.Name())

	_, err = providerFromVagrantBox(boxfile.Name())
	if err == nil {
		t.Fatalf("Should have error as box file does not include metadata.json file")
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_metadata_empty(t *testing.T) {
	// Bad: Create a box with an empty metadata.json file
	files := tarFiles{
		{"foo.txt", "This is a foo file"},
		{"bar.txt", "This is a bar file"},
		{"metadata.json", ""},
	}
	boxfile, err := createBox(files)
	if err != nil {
		t.Fatalf("Error creating test box: %s", err)
	}
	defer os.Remove(boxfile.Name())

	_, err = providerFromVagrantBox(boxfile.Name())
	if err == nil {
		t.Fatalf("Should have error as box files metadata.json file is empty")
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_metadata_bad_json(t *testing.T) {
	// Bad: Create a box with bad JSON in the metadata.json file
	files := tarFiles{
		{"foo.txt", "This is a foo file"},
		{"bar.txt", "This is a bar file"},
		{"metadata.json", "{provider: badjson}"},
	}
	boxfile, err := createBox(files)
	if err != nil {
		t.Fatalf("Error creating test box: %s", err)
	}
	defer os.Remove(boxfile.Name())

	_, err = providerFromVagrantBox(boxfile.Name())
	if err == nil {
		t.Fatalf("Should have error as box files metadata.json file is empty")
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_metadata_provider_value_empty(t *testing.T) {
	// Bad: The boxes metadata.json file 'provider' key has an empty value
	files := tarFiles{
		{"foo.txt", "This is a foo file"},
		{"bar.txt", "This is a bar file"},
		{"metadata.json", `{"provider":""}`},
	}
	boxfile, err := createBox(files)
	if err != nil {
		t.Fatalf("Error creating test box: %s", err)
	}
	defer os.Remove(boxfile.Name())

	_, err = providerFromVagrantBox(boxfile.Name())
	if err == nil {
		t.Fatalf("Should have error as boxes metadata.json file 'provider' key is empty")
	}
	t.Logf("%s", err)
}

func TestProviderFromVagrantBox_metadata_ok(t *testing.T) {
	// Good: The boxes metadata.json file has the 'provider' key/value pair
	expectedProvider := "virtualbox"
	files := tarFiles{
		{"foo.txt", "This is a foo file"},
		{"bar.txt", "This is a bar file"},
		{"metadata.json", `{"provider":"` + expectedProvider + `"}`},
	}
	boxfile, err := createBox(files)
	if err != nil {
		t.Fatalf("Error creating test box: %s", err)
	}
	defer os.Remove(boxfile.Name())

	provider, err := providerFromVagrantBox(boxfile.Name())
	assert.Equal(t, expectedProvider, provider, "Error: Expected provider: '%s'. Got '%s'", expectedProvider, provider)
	t.Logf("Expected provider '%s'. Got provider '%s'", expectedProvider, provider)
}

func newBoxFile() (boxfile *os.File, err error) {
	boxfile, err = ioutil.TempFile(os.TempDir(), "test*.box")
	if err != nil {
		return boxfile, fmt.Errorf("Error creating test box file: %s", err)
	}
	return boxfile, nil
}

func createBox(files tarFiles) (boxfile *os.File, err error) {
	boxfile, err = newBoxFile()
	if err != nil {
		return boxfile, err
	}

	// Box files are gzipped tar archives
	aw := gzip.NewWriter(boxfile)
	tw := tar.NewWriter(aw)

	// Add each file to the box
	for _, file := range files {
		// Create and write the tar file header
		hdr := &tar.Header{
			Name: file.Name,
			Mode: 0644,
			Size: int64(len(file.Body)),
		}
		err = tw.WriteHeader(hdr)
		if err != nil {
			return boxfile, fmt.Errorf("Error writing box tar file header: %s", err)
		}
		// Write the file contents
		_, err = tw.Write([]byte(file.Body))
		if err != nil {
			return boxfile, fmt.Errorf("Error writing box tar file contents: %s", err)
		}
	}
	// Flush and close each writer
	err = tw.Close()
	if err != nil {
		return boxfile, fmt.Errorf("Error flushing tar file contents: %s", err)
	}
	err = aw.Close()
	if err != nil {
		return boxfile, fmt.Errorf("Error flushing gzip file contents: %s", err)
	}

	return boxfile, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"io/ioutil"
	"net/http"

	"encoding/base64"
	"encoding/json"
	"io"
	"net/url"

	"github.com/pborman/uuid"
	"golang.org/x/net/websocket"
)

type PodWatch struct {
	Object Object `json:"object"`
}

type Object struct {
	Meta   Metadata `json:"metadata"`
	Status Status   `json:"status"`
}

type Metadata struct {
	Name string `json:"name"`
}

type Status struct {
	Statuses []XContainerStatus `json:"containerStatuses"`
}

type XContainerStatus struct {
	Name  string `json:"name"`
	Ready bool   `json:"ready"`
	State State  `json:"state"`
}

type State struct {
	Terminated Terminated `json:"terminated"`
}

type Terminated struct {
	ContainerID string `json:"containerID"`
	ExitCode    int    `json:"exitCode"`
}

// TODO distinguish between pushes and branch creation.  Github has a header value that allows these to be differentitated.
// https://developer.github.com/webhooks/#delivery-headers
// https://gist.githubusercontent.com/ae6rt/53a25e726ac00b4cb535/raw/e3f412f6e7f408a56d0d691a1ec8b7658a495124/gh-create.json
// https://gist.githubusercontent.com/ae6rt/2be93f7d5edef8030b52/raw/29f591eb8ecc5555c55f1878b545613c1f9839b7/gh-push.json
type BuildEvent interface {
	Team() string
	Library() string
	ProjectKey() string
	Branches() []string
}

type DefaultDecap struct {
	MasterURL       string
	UserName        string
	Password        string
	AWSAccessKeyID  string
	AWSAccessSecret string
	AWSRegion       string
	Locker          Locker

	apiToken  string
	apiClient *http.Client
}

type RepoManagerCredential struct {
	User     string
	Password string
}

type StorageService interface {
	GetBuildsByProject(project Project, sinceUnixTime uint64, limit uint64) ([]Build, error)
	GetArtifacts(buildID string) ([]byte, error)
	GetConsoleLog(buildID string) ([]byte, error)
}

type Decap interface {
	LaunchBuild(buildEvent BuildEvent) error
	DeletePod(podName string) error
}

func NewDefaultDecap(apiServerURL, username, password, awsKey, awsSecret, awsRegion string, locker Locker) DefaultDecap {
	// todo when running in cluster, provide root certificate via /var/run/secrets/kubernetes.io/serviceaccount/ca.crt
	apiClient := &http.Client{Transport: &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}}

	data, _ := ioutil.ReadFile("/var/run/secrets/kubernetes.io/serviceaccount/token")

	return DefaultDecap{
		MasterURL:       apiServerURL,
		apiToken:        string(data),
		UserName:        username,
		Password:        password,
		Locker:          locker,
		AWSAccessKeyID:  awsKey,
		AWSAccessSecret: awsSecret,
		AWSRegion:       awsRegion,
		apiClient:       apiClient,
	}
}

func (k8s DefaultDecap) LaunchBuild(buildEvent BuildEvent) error {
	projectKey := buildEvent.ProjectKey()

	projs := getProjects()

	for _, branch := range buildEvent.Branches() {
		key := k8s.Locker.Key(projectKey, branch)
		buildID := uuid.NewRandom().String()

		containers := make([]Container, 1+len(projs[projectKey].Sidecars))

		baseContainer := Container{
			Name:  "build-server",
			Image: projs[projectKey].Descriptor.Image,
			VolumeMounts: []VolumeMount{
				VolumeMount{
					Name:      "build-scripts",
					MountPath: "/home/decap/buildscripts",
				},
				VolumeMount{
					Name:      "decap-credentials",
					MountPath: "/etc/secrets",
				},
			},
			Env: []EnvVar{
				EnvVar{
					Name:  "BUILD_ID",
					Value: buildID,
				},
				EnvVar{
					Name:  "PROJECT_KEY",
					Value: projectKey,
				},
				EnvVar{
					Name:  "BRANCH_TO_BUILD",
					Value: branch,
				},
				EnvVar{
					Name:  "BUILD_LOCK_KEY",
					Value: key,
				},
				EnvVar{
					Name:  "AWS_ACCESS_KEY_ID",
					Value: k8s.AWSAccessKeyID,
				},
				EnvVar{
					Name:  "AWS_SECRET_ACCESS_KEY",
					Value: k8s.AWSAccessSecret,
				},
				EnvVar{
					Name:  "AWS_DEFAULT_REGION",
					Value: k8s.AWSRegion,
				},
			},
		}

		containers[0] = baseContainer
		for i, v := range projs[projectKey].Sidecars {
			var c Container
			err := json.Unmarshal([]byte(v), &c)
			if err != nil {
				Log.Println(err)
				continue
			}
			containers[i+1] = c
		}

		pod := Pod{
			TypeMeta: TypeMeta{
				Kind:       "Pod",
				APIVersion: "v1",
			},
			ObjectMeta: ObjectMeta{
				Name:      buildID,
				Namespace: "decap",
				Labels: map[string]string{
					"type":    "decap-build",
					"team":    buildEvent.Team(),
					"library": buildEvent.Library(),
					"branch":  branch,
				},
			},
			Spec: PodSpec{
				Volumes: []Volume{
					Volume{
						Name: "build-scripts",
						VolumeSource: VolumeSource{
							GitRepo: &GitRepoVolumeSource{
								Repository: *buildScriptsRepo,
								Revision:   *buildScriptsRepoBranch,
							},
						},
					},
					Volume{
						Name: "decap-credentials",
						VolumeSource: VolumeSource{
							Secret: &SecretVolumeSource{
								SecretName: "decap-credentials",
							},
						},
					},
				},
				Containers: containers,
			},
		}

		podBytes, err := json.Marshal(&pod)
		if err != nil {
			Log.Println(err)
			continue
		}

		fullyQualifiedKey := "/buildlocks/" + key
		resp, err := k8s.Locker.Lock(fullyQualifiedKey, buildID)
		if err != nil {
			Log.Printf("Failed to acquire lock %s on build %s: %v\n", key, buildID, err)
			continue
		}

		if resp.Node.Value == buildID {
			Log.Printf("Acquired lock on build %s with key %s\n", buildID, key)
			if podError := k8s.CreatePod(podBytes); podError != nil {
				Log.Println(podError)
				if _, err := k8s.Locker.Unlock(fullyQualifiedKey, buildID); err != nil {
					Log.Println(err)
				} else {
					Log.Printf("Released lock on build %s with key %s because of pod creation error %v\n", buildID, key, podError)
				}
			}
			Log.Printf("Created pod=%s\n", buildID)
		} else {
			Log.Printf("Failed to acquire lock %s on build %s\n", key, buildID)
		}
	}
	return nil
}

func (base DefaultDecap) CreatePod(pod []byte) error {
	Log.Printf("spec pod:\n%s\n", string(pod))

	req, err := http.NewRequest("POST", fmt.Sprintf("%s/api/v1/namespaces/decap/pods", base.MasterURL), bytes.NewReader(pod))
	if err != nil {
		Log.Println(err)
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	if base.apiToken != "" {
		req.Header.Set("Authorization", "Bearer "+base.apiToken)
	} else {
		req.SetBasicAuth(base.UserName, base.Password)
	}

	resp, err := base.apiClient.Do(req)
	if err != nil {
		return err
	}
	defer func() {
		resp.Body.Close()
	}()

	if resp.StatusCode != 201 {
		if data, err := ioutil.ReadAll(resp.Body); err != nil {
			Log.Printf("Error reading non-201 response body: %v\n", err)
			return err
		} else {
			Log.Printf("%s\n", string(data))
			return nil
		}
	}
	return nil
}

func (base DefaultDecap) DeletePod(podName string) error {
	req, err := http.NewRequest("DELETE", fmt.Sprintf("%s/api/v1/namespaces/decap/pods/%s", base.MasterURL, podName), nil)
	if err != nil {
		Log.Println(err)
		return err
	}
	if base.apiToken != "" {
		req.Header.Set("Authorization", "Bearer "+base.apiToken)
	} else {
		req.SetBasicAuth(base.UserName, base.Password)
	}

	resp, err := base.apiClient.Do(req)
	if err != nil {
		return err
	}
	defer func() {
		resp.Body.Close()
	}()

	if resp.StatusCode != 200 {
		if data, err := ioutil.ReadAll(resp.Body); err != nil {
			Log.Printf("Error reading non-200 response body: %v\n", err)
			return err
		} else {
			Log.Printf("%s\n", string(data))
			return nil
		}
	}
	return nil
}

func (decap DefaultDecap) Websock() {
	originURL, err := url.Parse(decap.MasterURL + "/api/v1/watch/namespaces/decap/pods?watch=true&labelSelector=type=decap-build")
	if err != nil {
		Log.Printf("Error parsing websocket origin URL.  Will be unable to reap exited pods.: %v\n", err)
		return
	}
	serviceURL, err := url.Parse("wss://" + originURL.Host + "/api/v1/watch/namespaces/decap/pods?watch=true&labelSelector=type=decap-build")
	if err != nil {
		Log.Printf("Error parsing websocket service URL.  Will be unable to reap exited pods.: %v\n", err)
		return
	}

	var hdrs http.Header
	if decap.apiToken != "" {
		hdrs = map[string][]string{"Authorization": []string{"Bearer " + decap.apiToken}}
	} else {
		hdrs = map[string][]string{"Authorization": []string{"Basic " + base64.StdEncoding.EncodeToString([]byte(decap.UserName+":"+decap.Password))}}
	}

	cfg := websocket.Config{
		Location:  serviceURL,
		Origin:    originURL,
		TlsConfig: &tls.Config{InsecureSkipVerify: true},
		Header:    hdrs,
		Version:   websocket.ProtocolVersionHybi13,
	}

	conn, err := websocket.DialConfig(&cfg)
	if err != nil {
		Log.Printf("Error opening websocket connection.  Will be unable to reap exited pods.: %v\n", err)
		return
	}
	Log.Print("Watching pods on websocket")

	var msg string
	for {
		err := websocket.Message.Receive(conn, &msg)
		if err != nil {
			if err == io.EOF {
				break
			}
			Log.Println("Couldn't receive msg " + err.Error())
		}
		var pod PodWatch
		if err := json.Unmarshal([]byte(msg), &pod); err != nil {
			Log.Println(err)
			continue
		}
		var deletePod bool
		for _, status := range pod.Object.Status.Statuses {
			if status.Name == "build-server" && status.State.Terminated.ContainerID != "" {
				deletePod = true
				break
			}
		}
		if deletePod {
			_, err := decap.Locker.Lock("/pods/"+pod.Object.Meta.Name, "anyvalue")
			if err == nil {
				if err := decap.DeletePod(pod.Object.Meta.Name); err != nil {
					Log.Print(err)
				} else {
					Log.Printf("Pod deleted: %s\n", pod.Object.Meta.Name)
				}
			}
		}
	}
}

func kubeSecret(file string, defaultValue string) string {
	if v, err := ioutil.ReadFile(file); err != nil {
		Log.Printf("Secret %s not found in the filesystem.  Using default.\n", file)
		return defaultValue
	} else {
		Log.Printf("Successfully read secret %s from the filesystem\n", file)
		return string(v)
	}
}
<commit_msg>set restart policy<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"io/ioutil"
	"net/http"

	"encoding/base64"
	"encoding/json"
	"io"
	"net/url"

	"github.com/pborman/uuid"
	"golang.org/x/net/websocket"
)

type PodWatch struct {
	Object Object `json:"object"`
}

type Object struct {
	Meta   Metadata `json:"metadata"`
	Status Status   `json:"status"`
}

type Metadata struct {
	Name string `json:"name"`
}

type Status struct {
	Statuses []XContainerStatus `json:"containerStatuses"`
}

type XContainerStatus struct {
	Name  string `json:"name"`
	Ready bool   `json:"ready"`
	State State  `json:"state"`
}

type State struct {
	Terminated Terminated `json:"terminated"`
}

type Terminated struct {
	ContainerID string `json:"containerID"`
	ExitCode    int    `json:"exitCode"`
}

// TODO distinguish between pushes and branch creation.  Github has a header value that allows these to be differentitated.
// https://developer.github.com/webhooks/#delivery-headers
// https://gist.githubusercontent.com/ae6rt/53a25e726ac00b4cb535/raw/e3f412f6e7f408a56d0d691a1ec8b7658a495124/gh-create.json
// https://gist.githubusercontent.com/ae6rt/2be93f7d5edef8030b52/raw/29f591eb8ecc5555c55f1878b545613c1f9839b7/gh-push.json
type BuildEvent interface {
	Team() string
	Library() string
	ProjectKey() string
	Branches() []string
}

type DefaultDecap struct {
	MasterURL       string
	UserName        string
	Password        string
	AWSAccessKeyID  string
	AWSAccessSecret string
	AWSRegion       string
	Locker          Locker

	apiToken  string
	apiClient *http.Client
}

type RepoManagerCredential struct {
	User     string
	Password string
}

type StorageService interface {
	GetBuildsByProject(project Project, sinceUnixTime uint64, limit uint64) ([]Build, error)
	GetArtifacts(buildID string) ([]byte, error)
	GetConsoleLog(buildID string) ([]byte, error)
}

type Decap interface {
	LaunchBuild(buildEvent BuildEvent) error
	DeletePod(podName string) error
}

func NewDefaultDecap(apiServerURL, username, password, awsKey, awsSecret, awsRegion string, locker Locker) DefaultDecap {
	// todo when running in cluster, provide root certificate via /var/run/secrets/kubernetes.io/serviceaccount/ca.crt
	apiClient := &http.Client{Transport: &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}}

	data, _ := ioutil.ReadFile("/var/run/secrets/kubernetes.io/serviceaccount/token")

	return DefaultDecap{
		MasterURL:       apiServerURL,
		apiToken:        string(data),
		UserName:        username,
		Password:        password,
		Locker:          locker,
		AWSAccessKeyID:  awsKey,
		AWSAccessSecret: awsSecret,
		AWSRegion:       awsRegion,
		apiClient:       apiClient,
	}
}

func (k8s DefaultDecap) LaunchBuild(buildEvent BuildEvent) error {
	projectKey := buildEvent.ProjectKey()

	projs := getProjects()

	for _, branch := range buildEvent.Branches() {
		key := k8s.Locker.Key(projectKey, branch)
		buildID := uuid.NewRandom().String()

		containers := make([]Container, 1+len(projs[projectKey].Sidecars))

		baseContainer := Container{
			Name:  "build-server",
			Image: projs[projectKey].Descriptor.Image,
			VolumeMounts: []VolumeMount{
				VolumeMount{
					Name:      "build-scripts",
					MountPath: "/home/decap/buildscripts",
				},
				VolumeMount{
					Name:      "decap-credentials",
					MountPath: "/etc/secrets",
				},
			},
			Env: []EnvVar{
				EnvVar{
					Name:  "BUILD_ID",
					Value: buildID,
				},
				EnvVar{
					Name:  "PROJECT_KEY",
					Value: projectKey,
				},
				EnvVar{
					Name:  "BRANCH_TO_BUILD",
					Value: branch,
				},
				EnvVar{
					Name:  "BUILD_LOCK_KEY",
					Value: key,
				},
				EnvVar{
					Name:  "AWS_ACCESS_KEY_ID",
					Value: k8s.AWSAccessKeyID,
				},
				EnvVar{
					Name:  "AWS_SECRET_ACCESS_KEY",
					Value: k8s.AWSAccessSecret,
				},
				EnvVar{
					Name:  "AWS_DEFAULT_REGION",
					Value: k8s.AWSRegion,
				},
			},
		}

		containers[0] = baseContainer
		for i, v := range projs[projectKey].Sidecars {
			var c Container
			err := json.Unmarshal([]byte(v), &c)
			if err != nil {
				Log.Println(err)
				continue
			}
			containers[i+1] = c
		}

		pod := Pod{
			TypeMeta: TypeMeta{
				Kind:       "Pod",
				APIVersion: "v1",
			},
			ObjectMeta: ObjectMeta{
				Name:      buildID,
				Namespace: "decap",
				Labels: map[string]string{
					"type":    "decap-build",
					"team":    buildEvent.Team(),
					"library": buildEvent.Library(),
					"branch":  branch,
				},
			},
			Spec: PodSpec{
				Volumes: []Volume{
					Volume{
						Name: "build-scripts",
						VolumeSource: VolumeSource{
							GitRepo: &GitRepoVolumeSource{
								Repository: *buildScriptsRepo,
								Revision:   *buildScriptsRepoBranch,
							},
						},
					},
					Volume{
						Name: "decap-credentials",
						VolumeSource: VolumeSource{
							Secret: &SecretVolumeSource{
								SecretName: "decap-credentials",
							},
						},
					},
				},
				Containers:    containers,
				RestartPolicy: "Never",
			},
		}

		podBytes, err := json.Marshal(&pod)
		if err != nil {
			Log.Println(err)
			continue
		}

		fullyQualifiedKey := "/buildlocks/" + key
		resp, err := k8s.Locker.Lock(fullyQualifiedKey, buildID)
		if err != nil {
			Log.Printf("Failed to acquire lock %s on build %s: %v\n", key, buildID, err)
			continue
		}

		if resp.Node.Value == buildID {
			Log.Printf("Acquired lock on build %s with key %s\n", buildID, key)
			if podError := k8s.CreatePod(podBytes); podError != nil {
				Log.Println(podError)
				if _, err := k8s.Locker.Unlock(fullyQualifiedKey, buildID); err != nil {
					Log.Println(err)
				} else {
					Log.Printf("Released lock on build %s with key %s because of pod creation error %v\n", buildID, key, podError)
				}
			}
			Log.Printf("Created pod=%s\n", buildID)
		} else {
			Log.Printf("Failed to acquire lock %s on build %s\n", key, buildID)
		}
	}
	return nil
}

func (base DefaultDecap) CreatePod(pod []byte) error {
	Log.Printf("spec pod:\n%s\n", string(pod))

	req, err := http.NewRequest("POST", fmt.Sprintf("%s/api/v1/namespaces/decap/pods", base.MasterURL), bytes.NewReader(pod))
	if err != nil {
		Log.Println(err)
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	if base.apiToken != "" {
		req.Header.Set("Authorization", "Bearer "+base.apiToken)
	} else {
		req.SetBasicAuth(base.UserName, base.Password)
	}

	resp, err := base.apiClient.Do(req)
	if err != nil {
		return err
	}
	defer func() {
		resp.Body.Close()
	}()

	if resp.StatusCode != 201 {
		if data, err := ioutil.ReadAll(resp.Body); err != nil {
			Log.Printf("Error reading non-201 response body: %v\n", err)
			return err
		} else {
			Log.Printf("%s\n", string(data))
			return nil
		}
	}
	return nil
}

func (base DefaultDecap) DeletePod(podName string) error {
	req, err := http.NewRequest("DELETE", fmt.Sprintf("%s/api/v1/namespaces/decap/pods/%s", base.MasterURL, podName), nil)
	if err != nil {
		Log.Println(err)
		return err
	}
	if base.apiToken != "" {
		req.Header.Set("Authorization", "Bearer "+base.apiToken)
	} else {
		req.SetBasicAuth(base.UserName, base.Password)
	}

	resp, err := base.apiClient.Do(req)
	if err != nil {
		return err
	}
	defer func() {
		resp.Body.Close()
	}()

	if resp.StatusCode != 200 {
		if data, err := ioutil.ReadAll(resp.Body); err != nil {
			Log.Printf("Error reading non-200 response body: %v\n", err)
			return err
		} else {
			Log.Printf("%s\n", string(data))
			return nil
		}
	}
	return nil
}

func (decap DefaultDecap) Websock() {
	originURL, err := url.Parse(decap.MasterURL + "/api/v1/watch/namespaces/decap/pods?watch=true&labelSelector=type=decap-build")
	if err != nil {
		Log.Printf("Error parsing websocket origin URL.  Will be unable to reap exited pods.: %v\n", err)
		return
	}
	serviceURL, err := url.Parse("wss://" + originURL.Host + "/api/v1/watch/namespaces/decap/pods?watch=true&labelSelector=type=decap-build")
	if err != nil {
		Log.Printf("Error parsing websocket service URL.  Will be unable to reap exited pods.: %v\n", err)
		return
	}

	var hdrs http.Header
	if decap.apiToken != "" {
		hdrs = map[string][]string{"Authorization": []string{"Bearer " + decap.apiToken}}
	} else {
		hdrs = map[string][]string{"Authorization": []string{"Basic " + base64.StdEncoding.EncodeToString([]byte(decap.UserName+":"+decap.Password))}}
	}

	cfg := websocket.Config{
		Location:  serviceURL,
		Origin:    originURL,
		TlsConfig: &tls.Config{InsecureSkipVerify: true},
		Header:    hdrs,
		Version:   websocket.ProtocolVersionHybi13,
	}

	conn, err := websocket.DialConfig(&cfg)
	if err != nil {
		Log.Printf("Error opening websocket connection.  Will be unable to reap exited pods.: %v\n", err)
		return
	}
	Log.Print("Watching pods on websocket")

	var msg string
	for {
		err := websocket.Message.Receive(conn, &msg)
		if err != nil {
			if err == io.EOF {
				break
			}
			Log.Println("Couldn't receive msg " + err.Error())
		}
		var pod PodWatch
		if err := json.Unmarshal([]byte(msg), &pod); err != nil {
			Log.Println(err)
			continue
		}
		var deletePod bool
		for _, status := range pod.Object.Status.Statuses {
			if status.Name == "build-server" && status.State.Terminated.ContainerID != "" {
				deletePod = true
				break
			}
		}
		if deletePod {
			_, err := decap.Locker.Lock("/pods/"+pod.Object.Meta.Name, "anyvalue")
			if err == nil {
				if err := decap.DeletePod(pod.Object.Meta.Name); err != nil {
					Log.Print(err)
				} else {
					Log.Printf("Pod deleted: %s\n", pod.Object.Meta.Name)
				}
			}
		}
	}
}

func kubeSecret(file string, defaultValue string) string {
	if v, err := ioutil.ReadFile(file); err != nil {
		Log.Printf("Secret %s not found in the filesystem.  Using default.\n", file)
		return defaultValue
	} else {
		Log.Printf("Successfully read secret %s from the filesystem\n", file)
		return string(v)
	}
}
<|endoftext|>


<commit_before>package device

import (
	"sync"

	"github.com/tarm/serial"
)

// Manager manages devices that are plugged into the system. It supports auto
// detection of devices.
//
// Serial ports are opened each for a device, and a clean API for communicating
// is provided via Read, Write and Flush methods.
//
// The devices are monitored via udev, and any changes that requires reloading
// of the  ports are handled by reloading the ports to the devices.
//
// This is safe to use concurrently in multiple goroutines
type Manager struct {
	devices map[string]serial.Config
	conn    []*Conn
	mu      sync.RWMutex
}

// AddDevice adds device name to the manager
func (m *Manager) AddDevice(name string) error {
	cfg := serial.Config{Name: name}
	m.mu.Lock()
	m.devices[name] = cfg
	m.mu.Unlock()
	return nil
}

// RemoveDevice removes device name from the manager
func (m *Manager) RemoveDevice(name string) error {
	m.mu.RLock()
	delete(m.devices, name)
	m.mu.RUnlock()
	return nil
}

type Conn struct {
	device serial.Config
	port   *serial.Port
	isOpen bool
}

func (c *Conn) Open() error {
	p, err := serial.OpenPort(&c.device)
	if err != nil {
		return nil
	}
	c.port = p
	c.isOpen = true
	return nil
}

// Close closes the port helt by *Conn.
func (c *Conn) Close() error {
	if c.isOpen {
		return c.port.Close()
	}
	return nil
}
<commit_msg>Add documentation<commit_after>package device

import (
	"sync"

	"github.com/tarm/serial"
)

// Manager manages devices that are plugged into the system. It supports auto
// detection of devices.
//
// Serial ports are opened each for a device, and a clean API for communicating
// is provided via Read, Write and Flush methods.
//
// The devices are monitored via udev, and any changes that requires reloading
// of the  ports are handled by reloading the ports to the devices.
//
// This is safe to use concurrently in multiple goroutines
type Manager struct {
	devices map[string]serial.Config
	conn    []*Conn
	mu      sync.RWMutex
}

// AddDevice adds device name to the manager
func (m *Manager) AddDevice(name string) error {
	cfg := serial.Config{Name: name}
	m.mu.Lock()
	m.devices[name] = cfg
	m.mu.Unlock()
	return nil
}

// RemoveDevice removes device name from the manager
func (m *Manager) RemoveDevice(name string) error {
	m.mu.RLock()
	delete(m.devices, name)
	m.mu.RUnlock()
	return nil
}

// Conn is a device serial connection
type Conn struct {
	device serial.Config
	port   *serial.Port
	isOpen bool
}

// Open opens a serial port to the undelying device
func (c *Conn) Open() error {
	p, err := serial.OpenPort(&c.device)
	if err != nil {
		return nil
	}
	c.port = p
	c.isOpen = true
	return nil
}

// Close closes the port helt by *Conn.
func (c *Conn) Close() error {
	if c.isOpen {
		return c.port.Close()
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 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 dhcp4

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"sort"
)

// Option is a DHCP option.
type Option byte

// Some of the more commonly seen DHCP options. Refer to
// http://www.iana.org/assignments/bootp-dhcp-parameters/bootp-dhcp-parameters.xhtml
// for the full authoritative list.
const (
	OptSubnetMask         Option = 1  // IPMask
	OptTimeOffset                = 2  // int32
	OptRouters                   = 3  // IPs
	OptDNSServers                = 6  // IPs
	OptHostname                  = 12 // string
	OptBootFileSize              = 13 // uint16
	OptDomainName                = 15 // string
	OptBroadcastAddr             = 28 // IP
	OptNTPServers                = 42 // IP
	OptVendorSpecific            = 43 // []byte
	OptRequestedIP               = 50 // IP
	OptLeaseTime                 = 51 // uint32
	OptOverload                  = 52 // byte
	OptServerIdentifier          = 54 // IP
	OptRequestedOptions          = 55 // []byte
	OptMessage                   = 56 // string
	OptMaximumMessageSize        = 57 // uint16
	OptRenewalTime               = 58 // uint32
	OptRebindingTime             = 59 // uint32
	OptVendorIdentifier          = 60 // string
	OptClientIdentifier          = 61 // string
	OptFQDN                      = 81 // string

	// You shouldn't need to use the following directly. Instead,
	// refer to the fields in the Packet struct, and Marshal/Unmarshal
	// will handle encoding for you.

	OptTFTPServer      = 66 // string
	OptBootFile        = 67 // string
	OptDHCPMessageType = 53 // byte
)

// Options stores DHCP options.
type Options map[Option][]byte

var (
	errOptionNotPresent = errors.New("option not present in Options")
	errOptionWrongSize  = errors.New("option value is the wrong size")
)

// Unmarshal parses DHCP options into o.
func (o Options) Unmarshal(bs []byte) error {
	for len(bs) > 0 {
		opt := Option(bs[0])
		switch opt {
		case 0:
			// Padding byte
			bs = bs[1:]
		case 255:
			// End of options
			return nil
		default:
			// In theory, DHCP permits multiple instances of the same
			// option in a packet, as a way to have option values >255
			// bytes. Unfortunately, this is very loosely specified as
			// "up to individual options", and AFAICT, isn't at all
			// used in the wild.
			//
			// So, for now, seeing the same option twice in a packet
			// is going to be an error, until I get a bug report about
			// something that actually does it.
			if _, ok := o[opt]; ok {
				// Okay fine option 56 can be duped.
				if opt != 56 {
					return fmt.Errorf("packet has duplicate option %d (please file a bug with a pcap!)", opt)
				}
			}
			if len(bs) < 2 {
				return fmt.Errorf("option %d has no length byte", opt)
			}
			l := int(bs[1])
			if len(bs[2:]) < l {
				return fmt.Errorf("option %d claims to have %d bytes of payload, but only has %d bytes", opt, l, len(bs[2:]))
			}
			o[opt] = bs[2 : 2+l]
			bs = bs[2+l:]
		}
	}

	return errors.New("options are not terminated by a 255 byte")
}

// Marshal returns the wire encoding of o.
func (o Options) Marshal() ([]byte, error) {
	var ret bytes.Buffer
	opts, err := o.marshalLimited(&ret, 0, false)
	if err != nil {
		return nil, err
	}
	if len(opts) > 0 {
		return nil, errors.New("some options not written, but no limit was given (please file a bug)")
	}

	return ret.Bytes(), nil
}

// Copy returns a shallow copy of o.
func (o Options) Copy() Options {
	ret := make(Options, len(o))
	for k, v := range o {
		ret[k] = v
	}
	return ret
}

// marshalLimited serializes o into w. If nBytes > 0, as many options
// as possible are packed into that many bytes, inserting padding as
// needed, and the remaining unwritten options are returned.
func (o Options) marshalLimited(w io.Writer, nBytes int, skip52 bool) (Options, error) {
	ks := make([]int, 0, len(o))
	for n := range o {
		if n <= 0 || n >= 255 {
			return nil, fmt.Errorf("invalid DHCP option number %d", n)
		}
		ks = append(ks, int(n))
	}
	sort.Ints(ks)

	ret := make(Options)
	for _, n := range ks {
		opt := o[Option(n)]
		if len(opt) > 255 {
			return nil, fmt.Errorf("DHCP option %d has value >255 bytes", n)
		}

		// If space is limited, verify that we can fit the option plus
		// the final end-of-options marker.
		if nBytes > 0 && ((skip52 && n == 52) || len(opt)+3 > nBytes) {
			ret[Option(n)] = opt
			continue
		}

		w.Write([]byte{byte(n), byte(len(opt))})
		w.Write(opt)
		nBytes -= len(opt) + 2
	}

	w.Write([]byte{255})
	nBytes--
	if nBytes > 0 {
		w.Write(make([]byte, nBytes))
	}

	return ret, nil
}

// Bytes returns the value of option n as a byte slice.
func (o Options) Bytes(n Option) ([]byte, error) {
	bs := o[n]
	if bs == nil {
		return nil, errOptionNotPresent
	}
	return bs, nil
}

// String returns the value of option n as a string.
func (o Options) String(n Option) (string, error) {
	bs, err := o.Bytes(n)
	if err != nil {
		return "", err
	}
	return string(bs), err
}

// Byte returns the value of option n as a byte.
func (o Options) Byte(n Option) (byte, error) {
	bs, err := o.Bytes(n)
	if err != nil {
		return 0, err
	}
	if len(bs) != 1 {
		return 0, errOptionWrongSize
	}
	return bs[0], nil
}

// Uint16 returns the value of option n as a uint16.
func (o Options) Uint16(n Option) (uint16, error) {
	bs, err := o.Bytes(n)
	if err != nil {
		return 0, err
	}
	if len(bs) != 2 {
		return 0, errOptionWrongSize
	}
	return binary.BigEndian.Uint16(bs), nil
}

// Uint32 returns the value of option n as a uint32.
func (o Options) Uint32(n Option) (uint32, error) {
	bs, err := o.Bytes(n)
	if err != nil {
		return 0, err
	}
	if len(bs) != 4 {
		return 0, errOptionWrongSize
	}
	return binary.BigEndian.Uint32(bs), nil
}

// Int32 returns the value of option n as an int32.
func (o Options) Int32(n Option) (int32, error) {
	bs, err := o.Bytes(n)
	if err != nil {
		return 0, err
	}
	if len(bs) != 4 {
		return 0, errOptionWrongSize
	}
	return int32(binary.BigEndian.Uint32(bs)), nil
}

// IPs returns the value of option n as a list of IPv4 addresses.
func (o Options) IPs(n Option) ([]net.IP, error) {
	bs, err := o.Bytes(n)
	if err != nil {
		return nil, err
	}
	if len(bs) < 4 || len(bs)%4 != 0 {
		return nil, errOptionWrongSize
	}
	ret := make([]net.IP, 0, len(bs)/4)
	for i := 0; i < len(bs); i += 4 {
		ret = append(ret, net.IP(bs[i:i+4]))
	}
	return ret, nil
}

// IP returns the value of option n as an IPv4 address.
func (o Options) IP(n Option) (net.IP, error) {
	ips, err := o.IPs(n)
	if err != nil {
		return nil, err
	}
	if len(ips) != 1 {
		return nil, errOptionWrongSize
	}
	return ips[0], nil
}

// IPMask returns the value of option n as a net.IPMask.
func (o Options) IPMask(n Option) (net.IPMask, error) {
	bs := o[n]
	if bs == nil {
		return nil, fmt.Errorf("option %d not found", n)
	}
	if len(bs) != 4 {
		return nil, fmt.Errorf("option %d is the wrong size for an IPMask", n)
	}
	return net.IPMask(bs), nil
}
<commit_msg>dhcp4: make all options use type Option<commit_after>// Copyright 2016 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 dhcp4

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"sort"
)

// Option is a DHCP option.
type Option byte

// Some of the more commonly seen DHCP options. Refer to
// http://www.iana.org/assignments/bootp-dhcp-parameters/bootp-dhcp-parameters.xhtml
// for the full authoritative list.
const (
	OptSubnetMask         Option = 1  // IPMask
	OptTimeOffset         Option = 2  // int32
	OptRouters            Option = 3  // IPs
	OptDNSServers         Option = 6  // IPs
	OptHostname           Option = 12 // string
	OptBootFileSize       Option = 13 // uint16
	OptDomainName         Option = 15 // string
	OptBroadcastAddr      Option = 28 // IP
	OptNTPServers         Option = 42 // IP
	OptVendorSpecific     Option = 43 // []byte
	OptRequestedIP        Option = 50 // IP
	OptLeaseTime          Option = 51 // uint32
	OptOverload           Option = 52 // byte
	OptServerIdentifier   Option = 54 // IP
	OptRequestedOptions   Option = 55 // []byte
	OptMessage            Option = 56 // string
	OptMaximumMessageSize Option = 57 // uint16
	OptRenewalTime        Option = 58 // uint32
	OptRebindingTime      Option = 59 // uint32
	OptVendorIdentifier   Option = 60 // string
	OptClientIdentifier   Option = 61 // string
	OptFQDN               Option = 81 // string

	// You shouldn't need to use the following directly. Instead,
	// refer to the fields in the Packet struct, and Marshal/Unmarshal
	// will handle encoding for you.

	OptTFTPServer      Option = 66 // string
	OptBootFile        Option = 67 // string
	OptDHCPMessageType Option = 53 // byte
)

// Options stores DHCP options.
type Options map[Option][]byte

var (
	errOptionNotPresent = errors.New("option not present in Options")
	errOptionWrongSize  = errors.New("option value is the wrong size")
)

// Unmarshal parses DHCP options into o.
func (o Options) Unmarshal(bs []byte) error {
	for len(bs) > 0 {
		opt := Option(bs[0])
		switch opt {
		case 0:
			// Padding byte
			bs = bs[1:]
		case 255:
			// End of options
			return nil
		default:
			// In theory, DHCP permits multiple instances of the same
			// option in a packet, as a way to have option values >255
			// bytes. Unfortunately, this is very loosely specified as
			// "up to individual options", and AFAICT, isn't at all
			// used in the wild.
			//
			// So, for now, seeing the same option twice in a packet
			// is going to be an error, until I get a bug report about
			// something that actually does it.
			if _, ok := o[opt]; ok {
				// Okay fine option 56 can be duped.
				if opt != 56 {
					return fmt.Errorf("packet has duplicate option %d (please file a bug with a pcap!)", opt)
				}
			}
			if len(bs) < 2 {
				return fmt.Errorf("option %d has no length byte", opt)
			}
			l := int(bs[1])
			if len(bs[2:]) < l {
				return fmt.Errorf("option %d claims to have %d bytes of payload, but only has %d bytes", opt, l, len(bs[2:]))
			}
			o[opt] = bs[2 : 2+l]
			bs = bs[2+l:]
		}
	}

	return errors.New("options are not terminated by a 255 byte")
}

// Marshal returns the wire encoding of o.
func (o Options) Marshal() ([]byte, error) {
	var ret bytes.Buffer
	opts, err := o.marshalLimited(&ret, 0, false)
	if err != nil {
		return nil, err
	}
	if len(opts) > 0 {
		return nil, errors.New("some options not written, but no limit was given (please file a bug)")
	}

	return ret.Bytes(), nil
}

// Copy returns a shallow copy of o.
func (o Options) Copy() Options {
	ret := make(Options, len(o))
	for k, v := range o {
		ret[k] = v
	}
	return ret
}

// marshalLimited serializes o into w. If nBytes > 0, as many options
// as possible are packed into that many bytes, inserting padding as
// needed, and the remaining unwritten options are returned.
func (o Options) marshalLimited(w io.Writer, nBytes int, skip52 bool) (Options, error) {
	ks := make([]int, 0, len(o))
	for n := range o {
		if n <= 0 || n >= 255 {
			return nil, fmt.Errorf("invalid DHCP option number %d", n)
		}
		ks = append(ks, int(n))
	}
	sort.Ints(ks)

	ret := make(Options)
	for _, n := range ks {
		opt := o[Option(n)]
		if len(opt) > 255 {
			return nil, fmt.Errorf("DHCP option %d has value >255 bytes", n)
		}

		// If space is limited, verify that we can fit the option plus
		// the final end-of-options marker.
		if nBytes > 0 && ((skip52 && n == 52) || len(opt)+3 > nBytes) {
			ret[Option(n)] = opt
			continue
		}

		w.Write([]byte{byte(n), byte(len(opt))})
		w.Write(opt)
		nBytes -= len(opt) + 2
	}

	w.Write([]byte{255})
	nBytes--
	if nBytes > 0 {
		w.Write(make([]byte, nBytes))
	}

	return ret, nil
}

// Bytes returns the value of option n as a byte slice.
func (o Options) Bytes(n Option) ([]byte, error) {
	bs := o[n]
	if bs == nil {
		return nil, errOptionNotPresent
	}
	return bs, nil
}

// String returns the value of option n as a string.
func (o Options) String(n Option) (string, error) {
	bs, err := o.Bytes(n)
	if err != nil {
		return "", err
	}
	return string(bs), err
}

// Byte returns the value of option n as a byte.
func (o Options) Byte(n Option) (byte, error) {
	bs, err := o.Bytes(n)
	if err != nil {
		return 0, err
	}
	if len(bs) != 1 {
		return 0, errOptionWrongSize
	}
	return bs[0], nil
}

// Uint16 returns the value of option n as a uint16.
func (o Options) Uint16(n Option) (uint16, error) {
	bs, err := o.Bytes(n)
	if err != nil {
		return 0, err
	}
	if len(bs) != 2 {
		return 0, errOptionWrongSize
	}
	return binary.BigEndian.Uint16(bs), nil
}

// Uint32 returns the value of option n as a uint32.
func (o Options) Uint32(n Option) (uint32, error) {
	bs, err := o.Bytes(n)
	if err != nil {
		return 0, err
	}
	if len(bs) != 4 {
		return 0, errOptionWrongSize
	}
	return binary.BigEndian.Uint32(bs), nil
}

// Int32 returns the value of option n as an int32.
func (o Options) Int32(n Option) (int32, error) {
	bs, err := o.Bytes(n)
	if err != nil {
		return 0, err
	}
	if len(bs) != 4 {
		return 0, errOptionWrongSize
	}
	return int32(binary.BigEndian.Uint32(bs)), nil
}

// IPs returns the value of option n as a list of IPv4 addresses.
func (o Options) IPs(n Option) ([]net.IP, error) {
	bs, err := o.Bytes(n)
	if err != nil {
		return nil, err
	}
	if len(bs) < 4 || len(bs)%4 != 0 {
		return nil, errOptionWrongSize
	}
	ret := make([]net.IP, 0, len(bs)/4)
	for i := 0; i < len(bs); i += 4 {
		ret = append(ret, net.IP(bs[i:i+4]))
	}
	return ret, nil
}

// IP returns the value of option n as an IPv4 address.
func (o Options) IP(n Option) (net.IP, error) {
	ips, err := o.IPs(n)
	if err != nil {
		return nil, err
	}
	if len(ips) != 1 {
		return nil, errOptionWrongSize
	}
	return ips[0], nil
}

// IPMask returns the value of option n as a net.IPMask.
func (o Options) IPMask(n Option) (net.IPMask, error) {
	bs := o[n]
	if bs == nil {
		return nil, fmt.Errorf("option %d not found", n)
	}
	if len(bs) != 4 {
		return nil, fmt.Errorf("option %d is the wrong size for an IPMask", n)
	}
	return net.IPMask(bs), nil
}
<|endoftext|>


<commit_before>package director

import (
	"errors"
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

func Test_TimedLooper(t *testing.T) {
	Convey("TimedLooper", t, func() {
		looper := NewTimedLooper(1, 1*time.Nanosecond, make(chan error))

		Convey("Sends a nil on the DoneChan when everything was kosher", func() {
			go looper.Done(nil)

			result := looper.Wait()
			So(result, ShouldBeNil)
		})

		Convey("Sends the error on the DoneChan when everything exploded", func() {
			err := errors.New("Borked!")
			go looper.Done(err)

			result := looper.Wait()
			So(result, ShouldEqual, err)
		})

		Convey("The loop executes the function", func() {
			run := false
			go looper.Loop(func() error { run = true; return nil })
			<-looper.DoneChan

			So(run, ShouldBeTrue)
		})

		Convey("The loop executes the correct number of times", func() {
			count := 0
			looper.Count = 5
			go looper.Loop(func() error { count++; return nil })
			looper.Wait()

			So(count, ShouldEqual, 5)
		})

		Convey("The loop returns an error on the DoneChan", func() {
			err := errors.New("Borked!")
			go looper.Loop(func() error { return err })
			So(looper.Wait(), ShouldEqual, err)
		})

		Convey("The loop exits when told to quit", func() {
			looper.Count = FOREVER
			count := 0

			go looper.Loop(func() error { count++; time.Sleep(2 * time.Nanosecond); return nil })
			looper.Quit()

			So(looper.Wait(), ShouldBeNil)
			So(count, ShouldBeLessThan, 2)
		})
	})
}

func Test_NewImmediateTimedLooper(t *testing.T) {
	Convey("ImmediateTimedLooper", t, func() {
		looper := NewImmediateTimedLooper(10, 1*time.Nanosecond, make(chan error))

		Convey("Immediate looper must have immediate set to true", func() {
			So(looper.Immediate, ShouldBeTrue)
		})
	})
}

func Test_FreeLooper(t *testing.T) {
	Convey("FreeLooper", t, func() {
		looper := NewFreeLooper(1, make(chan error))

		Convey("Sends a nil on the DoneChan when everything was kosher", func() {
			go looper.Done(nil)

			result := looper.Wait()
			So(result, ShouldBeNil)
		})

		Convey("Sends the error on the DoneChan when everything exploded", func() {
			err := errors.New("Borked!")
			go looper.Done(err)

			result := looper.Wait()
			So(result, ShouldEqual, err)
		})

		Convey("The loop executes the function", func() {
			run := false
			go looper.Loop(func() error { run = true; return nil })
			looper.Wait()

			So(run, ShouldBeTrue)
		})

		Convey("The loop executes the correct number of times", func() {
			count := 0
			looper.Count = 5
			go looper.Loop(func() error { count++; return nil })
			looper.Wait()

			So(count, ShouldEqual, 5)
		})

		Convey("The loop returns an error on the DoneChan", func() {
			err := errors.New("Borked!")
			go looper.Loop(func() error { return err })
			So(looper.Wait(), ShouldEqual, err)
		})

		Convey("The loop exits when told to quit", func() {
			looper.Count = FOREVER
			count := 0

			go looper.Loop(func() error { count++; time.Sleep(2 * time.Nanosecond); return nil })
			looper.Quit()

			So(looper.Wait(), ShouldBeNil)
			So(count, ShouldBeLessThan, 2)
		})
	})
}

func Example_TimedLooperWithoutQuit() {
	looper := NewTimedLooper(5, 1*time.Nanosecond, make(chan error))

	runner := func(looper Looper) {
		x := 0
		looper.Loop(func() error {
			println(x)
			x++
			return nil
		})
	}

	go runner(looper)
	<-looper.DoneChan

	// Output:
}
<commit_msg>Better examples in Godocs.<commit_after>package director

import (
	"errors"
	"fmt"
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

func Test_TimedLooper(t *testing.T) {
	Convey("TimedLooper", t, func() {
		looper := NewTimedLooper(1, 1*time.Nanosecond, make(chan error))

		Convey("Sends a nil on the DoneChan when everything was kosher", func() {
			go looper.Done(nil)

			result := looper.Wait()
			So(result, ShouldBeNil)
		})

		Convey("Sends the error on the DoneChan when everything exploded", func() {
			err := errors.New("Borked!")
			go looper.Done(err)

			result := looper.Wait()
			So(result, ShouldEqual, err)
		})

		Convey("The loop executes the function", func() {
			run := false
			go looper.Loop(func() error { run = true; return nil })
			<-looper.DoneChan

			So(run, ShouldBeTrue)
		})

		Convey("The loop executes the correct number of times", func() {
			count := 0
			looper.Count = 5
			go looper.Loop(func() error { count++; return nil })
			looper.Wait()

			So(count, ShouldEqual, 5)
		})

		Convey("The loop returns an error on the DoneChan", func() {
			err := errors.New("Borked!")
			go looper.Loop(func() error { return err })
			So(looper.Wait(), ShouldEqual, err)
		})

		Convey("The loop exits when told to quit", func() {
			looper.Count = FOREVER
			count := 0

			go looper.Loop(func() error { count++; time.Sleep(2 * time.Nanosecond); return nil })
			looper.Quit()

			So(looper.Wait(), ShouldBeNil)
			So(count, ShouldBeLessThan, 2)
		})
	})
}

func Test_NewImmediateTimedLooper(t *testing.T) {
	Convey("ImmediateTimedLooper", t, func() {
		looper := NewImmediateTimedLooper(10, 1*time.Nanosecond, make(chan error))

		Convey("Immediate looper must have immediate set to true", func() {
			So(looper.Immediate, ShouldBeTrue)
		})
	})
}

func Test_FreeLooper(t *testing.T) {
	Convey("FreeLooper", t, func() {
		looper := NewFreeLooper(1, make(chan error))

		Convey("Sends a nil on the DoneChan when everything was kosher", func() {
			go looper.Done(nil)

			result := looper.Wait()
			So(result, ShouldBeNil)
		})

		Convey("Sends the error on the DoneChan when everything exploded", func() {
			err := errors.New("Borked!")
			go looper.Done(err)

			result := looper.Wait()
			So(result, ShouldEqual, err)
		})

		Convey("The loop executes the function", func() {
			run := false
			go looper.Loop(func() error { run = true; return nil })
			looper.Wait()

			So(run, ShouldBeTrue)
		})

		Convey("The loop executes the correct number of times", func() {
			count := 0
			looper.Count = 5
			go looper.Loop(func() error { count++; return nil })
			looper.Wait()

			So(count, ShouldEqual, 5)
		})

		Convey("The loop returns an error on the DoneChan", func() {
			err := errors.New("Borked!")
			go looper.Loop(func() error { return err })
			So(looper.Wait(), ShouldEqual, err)
		})

		Convey("The loop exits when told to quit", func() {
			looper.Count = FOREVER
			count := 0

			go looper.Loop(func() error { count++; time.Sleep(2 * time.Nanosecond); return nil })
			looper.Quit()

			So(looper.Wait(), ShouldBeNil)
			So(count, ShouldBeLessThan, 2)
		})
	})
}

// In this example, we run a really fast TimedLooper for a
// fixed number of runs.
func ExampleTimedLooper() {
	looper := NewTimedLooper(5, 1*time.Nanosecond, make(chan error))

	runner := func(looper Looper) {
		x := 0
		looper.Loop(func() error {
			fmt.Println(x)
			x++
			return nil
		})
	}

	go runner(looper)
	<-looper.DoneChan

	// Output:
	// 0
	// 1
	// 2
	// 3
	// 4
}

// In this example we run a really fast TimedLooper for a fixed
// number of runs, but we interrupt it with a Quit() call so
// it only completes one run.
func ExampleTimedLooper_Quit() {
	looper := NewTimedLooper(5, 50*time.Millisecond, make(chan error))

	runner := func(looper Looper) {
		x := 0
		looper.Loop(func() error {
			fmt.Println(x)
			x++
			return nil
		})
	}

	go runner(looper)
	looper.Quit()
	<-looper.DoneChan

	// Output:
	// 0
}

// In this example, we are going to run a FreeLooper with 5 iterations.
// In the course of running, an error is generated, which the parent
// function captures and outputs. As a result of the error only 3
// of the 5 iterations are completed and the output reflects this.
func Example() {
	looper := NewFreeLooper(5, make(chan error))

	runner := func(looper Looper) {
		x := 0
		looper.Loop(func() error {
			fmt.Println(x)
			x++
			if x == 3 {
				return errors.New("Uh oh")
			}
			return nil
		})
	}

	go runner(looper)
	err := <-looper.DoneChan

	if err != nil {
		fmt.Printf("I got an error: %s\n", err.Error())
	}

	// Output:
	// 0
	// 1
	// 2
	// I got an error: Uh oh
}
<|endoftext|>


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

package distmv

import (
	"math"
	"math/rand"
	"sync"

	"github.com/gonum/floats"
	"github.com/gonum/matrix/mat64"
)

// Normal is a multivariate normal distribution (also known as the multivariate
// Gaussian distribution). Its pdf in k dimensions is given by
//  (2 π)^(-k/2) |Σ|^(-1/2) exp(-1/2 (x-μ)'Σ^-1(x-μ))
// where μ is the mean vector and Σ the covariance matrix. Σ must be symmetric
// and positive definite. Use NewNormal to construct.
type Normal struct {
	mu []float64

	once  sync.Once
	sigma *mat64.SymDense // only stored if needed

	chol       mat64.Cholesky
	lower      mat64.TriDense
	logSqrtDet float64
	dim        int

	src *rand.Rand
}

// NewNormal creates a new Normal with the given mean and covariance matrix.
// NewNormal panics if len(mu) == 0, or if len(mu) != sigma.N. If the covariance
// matrix is not positive-definite, the returned boolean is false.
func NewNormal(mu []float64, sigma mat64.Symmetric, src *rand.Rand) (*Normal, bool) {
	if len(mu) == 0 {
		panic(badZeroDimension)
	}
	dim := sigma.Symmetric()
	if dim != len(mu) {
		panic(badSizeMismatch)
	}
	n := &Normal{
		src: src,
		dim: dim,
		mu:  make([]float64, dim),
	}
	copy(n.mu, mu)
	ok := n.chol.Factorize(sigma)
	if !ok {
		return nil, false
	}
	n.lower.LFromCholesky(&n.chol)
	n.logSqrtDet = 0.5 * n.chol.LogDet()
	return n, true
}

// NewNormalChol creates a new Normal distribution with the given mean and
// covariance matrix represented by its Cholesky decomposition. NewNormalChol
// panics if len(mu) is not equal to chol.Size().
func NewNormalChol(mu []float64, chol *mat64.Cholesky, src *rand.Rand) *Normal {
	dim := len(mu)
	if dim != chol.Size() {
		panic(badSizeMismatch)
	}
	n := &Normal{
		src: src,
		dim: dim,
		mu:  make([]float64, dim),
	}
	copy(n.mu, mu)
	n.lower.LFromCholesky(chol)
	n.logSqrtDet = 0.5 * n.chol.LogDet()
	return n
}

// ConditionNormal returns the Normal distribution that is the receiver conditioned
// on the input evidence. The returned multivariate normal has dimension
// n - len(observed), where n is the dimension of the original receiver. The updated
// mean and covariance are
//  mu = mu_un + sigma_{ob,un}^T * sigma_{ob,ob}^-1 (v - mu_ob)
//  sigma = sigma_{un,un} - sigma_{ob,un}^T * sigma_{ob,ob}^-1 * sigma_{ob,un}
// where mu_un and mu_ob are the original means of the unobserved and observed
// variables respectively, sigma_{un,un} is the unobserved subset of the covariance
// matrix, sigma_{ob,ob} is the observed subset of the covariance matrix, and
// sigma_{un,ob} are the cross terms. The elements of x_2 have been observed with
// values v. The dimension order is preserved during conditioning, so if the value
// of dimension 1 is observed, the returned normal represents dimensions {0, 2, ...}
// of the original Normal distribution.
//
// ConditionNormal returns {nil, false} if there is a failure during the update.
// Mathematically this is impossible, but can occur with finite precision arithmetic.
func (n *Normal) ConditionNormal(observed []int, values []float64, src *rand.Rand) (*Normal, bool) {
	if len(observed) == 0 {
		panic("normal: no observed value")
	}
	if len(observed) != len(values) {
		panic("normal: input slice length mismatch")
	}
	for _, v := range observed {
		if v < 0 || v >= n.Dim() {
			panic("normal: observed value out of bounds")
		}
	}

	ob := len(observed)
	unob := n.Dim() - ob
	obMap := make(map[int]struct{})
	for _, v := range observed {
		if _, ok := obMap[v]; ok {
			panic("normal: observed dimension occurs twice")
		}
		obMap[v] = struct{}{}
	}
	if len(observed) == n.Dim() {
		panic("normal: all dimensions observed")
	}
	unobserved := make([]int, 0, unob)
	for i := 0; i < n.Dim(); i++ {
		if _, ok := obMap[i]; !ok {
			unobserved = append(unobserved, i)
		}
	}
	mu1 := make([]float64, unob)
	for i, v := range unobserved {
		mu1[i] = n.mu[v]
	}
	mu2 := make([]float64, ob) // really v - mu2
	for i, v := range observed {
		mu2[i] = values[i] - n.mu[v]
	}

	n.setSigma()

	var sigma11, sigma22 mat64.SymDense
	sigma11.SubsetSym(n.sigma, unobserved)
	sigma22.SubsetSym(n.sigma, observed)

	sigma21 := mat64.NewDense(ob, unob, nil)
	for i, r := range observed {
		for j, c := range unobserved {
			v := n.sigma.At(r, c)
			sigma21.Set(i, j, v)
		}
	}

	var chol mat64.Cholesky
	ok := chol.Factorize(&sigma22)
	if !ok {
		return nil, ok
	}

	// Compute sigma_{2,1}^T * sigma_{2,2}^-1 (v - mu_2).
	v := mat64.NewVector(ob, mu2)
	var tmp, tmp2 mat64.Vector
	err := tmp.SolveCholeskyVec(&chol, v)
	if err != nil {
		return nil, false
	}
	tmp2.MulVec(sigma21.T(), &tmp)

	// Compute sigma_{2,1}^T * sigma_{2,2}^-1 * sigma_{2,1}.
	// TODO(btracey): Should this be a method of SymDense?
	var tmp3, tmp4 mat64.Dense
	err = tmp3.SolveCholesky(&chol, sigma21)
	if err != nil {
		return nil, false
	}
	tmp4.Mul(sigma21.T(), &tmp3)

	for i := range mu1 {
		mu1[i] += tmp2.At(i, 0)
	}

	// TODO(btracey): If tmp2 can constructed with a method, then this can be
	// replaced with SubSym.
	for i := 0; i < len(unobserved); i++ {
		for j := i; j < len(unobserved); j++ {
			v := sigma11.At(i, j)
			sigma11.SetSym(i, j, v-tmp4.At(i, j))
		}
	}
	return NewNormal(mu1, &sigma11, src)
}

// CovarianceMatrix returns the covariance matrix of the distribution. Upon
// return, the value at element {i, j} of the covariance matrix is equal to
// the covariance of the i^th and j^th variables.
//  covariance(i, j) = E[(x_i - E[x_i])(x_j - E[x_j])]
// If the input matrix is nil a new matrix is allocated, otherwise the result
// is stored in-place into the input.
func (n *Normal) CovarianceMatrix(s *mat64.SymDense) *mat64.SymDense {
	if s == nil {
		s = mat64.NewSymDense(n.Dim(), nil)
	}
	sn := s.Symmetric()
	if sn != n.Dim() {
		panic("normal: input matrix size mismatch")
	}
	n.setSigma()
	s.CopySym(n.sigma)
	return s
}

// Dim returns the dimension of the distribution.
func (n *Normal) Dim() int {
	return n.dim
}

// Entropy returns the differential entropy of the distribution.
func (n *Normal) Entropy() float64 {
	return float64(n.dim)/2*(1+logTwoPi) + n.logSqrtDet
}

// LogProb computes the log of the pdf of the point x.
func (n *Normal) LogProb(x []float64) float64 {
	dim := n.dim
	if len(x) != dim {
		panic(badSizeMismatch)
	}
	// Compute the normalization constant
	c := -0.5*float64(dim)*logTwoPi - n.logSqrtDet

	// Compute (x-mu)'Sigma^-1 (x-mu)
	xMinusMu := make([]float64, dim)
	floats.SubTo(xMinusMu, x, n.mu)
	d := mat64.NewVector(dim, xMinusMu)
	tmp := make([]float64, dim)
	tmpVec := mat64.NewVector(dim, tmp)
	tmpVec.SolveCholeskyVec(&n.chol, d)
	return c - 0.5*floats.Dot(tmp, xMinusMu)
}

// MarginalNormal returns the marginal distribution of the given input variables.
// That is, MarginalNormal returns
//  p(x_i) = \int_{x_o} p(x_i | x_o) p(x_o) dx_o
// where x_i are the dimensions in the input, and x_o are the remaining dimensions.
// The input src is passed to the call to NewNormal.
func (n *Normal) MarginalNormal(vars []int, src *rand.Rand) (*Normal, bool) {
	newMean := make([]float64, len(vars))
	for i, v := range vars {
		newMean[i] = n.mu[v]
	}
	n.setSigma()
	var s mat64.SymDense
	s.SubsetSym(n.sigma, vars)
	return NewNormal(newMean, &s, src)
}

// Mean returns the mean of the probability distribution at x. If the
// input argument is nil, a new slice will be allocated, otherwise the result
// will be put in-place into the receiver.
func (n *Normal) Mean(x []float64) []float64 {
	x = reuseAs(x, n.dim)
	copy(x, n.mu)
	return x
}

// Prob computes the value of the probability density function at x.
func (n *Normal) Prob(x []float64) float64 {
	return math.Exp(n.LogProb(x))
}

// Rand generates a random number according to the distributon.
// If the input slice is nil, new memory is allocated, otherwise the result is stored
// in place.
func (n *Normal) Rand(x []float64) []float64 {
	x = reuseAs(x, n.dim)
	tmp := make([]float64, n.dim)
	if n.src == nil {
		for i := range x {
			tmp[i] = rand.NormFloat64()
		}
	} else {
		for i := range x {
			tmp[i] = n.src.NormFloat64()
		}
	}
	tmpVec := mat64.NewVector(n.dim, tmp)
	xVec := mat64.NewVector(n.dim, x)
	xVec.MulVec(&n.lower, tmpVec)
	floats.Add(x, n.mu)
	return x
}

// SetMean changes the mean of the normal distribution. SetMean panics if len(mu)
// does not equal the dimension of the normal distribution.
func (n *Normal) SetMean(mu []float64) {
	if len(mu) != n.Dim() {
		panic(badSizeMismatch)
	}
	copy(n.mu, mu)
}

// setSigma computes and stores the covariance matrix of the distribution.
func (n *Normal) setSigma() {
	n.once.Do(func() {
		n.sigma = mat64.NewSymDense(n.Dim(), nil)
		n.sigma.FromCholesky(&n.chol)
	})
}
<commit_msg>Add NewNormalPrecision for constructing a normal distribution based on its precision matrix<commit_after>// Copyright ©2015 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package distmv

import (
	"math"
	"math/rand"
	"sync"

	"github.com/gonum/floats"
	"github.com/gonum/matrix/mat64"
)

// Normal is a multivariate normal distribution (also known as the multivariate
// Gaussian distribution). Its pdf in k dimensions is given by
//  (2 π)^(-k/2) |Σ|^(-1/2) exp(-1/2 (x-μ)'Σ^-1(x-μ))
// where μ is the mean vector and Σ the covariance matrix. Σ must be symmetric
// and positive definite. Use NewNormal to construct.
type Normal struct {
	mu []float64

	once  sync.Once
	sigma *mat64.SymDense // only stored if needed

	chol       mat64.Cholesky
	lower      mat64.TriDense
	logSqrtDet float64
	dim        int

	src *rand.Rand
}

// NewNormal creates a new Normal with the given mean and covariance matrix.
// NewNormal panics if len(mu) == 0, or if len(mu) != sigma.N. If the covariance
// matrix is not positive-definite, the returned boolean is false.
func NewNormal(mu []float64, sigma mat64.Symmetric, src *rand.Rand) (*Normal, bool) {
	if len(mu) == 0 {
		panic(badZeroDimension)
	}
	dim := sigma.Symmetric()
	if dim != len(mu) {
		panic(badSizeMismatch)
	}
	n := &Normal{
		src: src,
		dim: dim,
		mu:  make([]float64, dim),
	}
	copy(n.mu, mu)
	ok := n.chol.Factorize(sigma)
	if !ok {
		return nil, false
	}
	n.lower.LFromCholesky(&n.chol)
	n.logSqrtDet = 0.5 * n.chol.LogDet()
	return n, true
}

// NewNormalChol creates a new Normal distribution with the given mean and
// covariance matrix represented by its Cholesky decomposition. NewNormalChol
// panics if len(mu) is not equal to chol.Size().
func NewNormalChol(mu []float64, chol *mat64.Cholesky, src *rand.Rand) *Normal {
	dim := len(mu)
	if dim != chol.Size() {
		panic(badSizeMismatch)
	}
	n := &Normal{
		src: src,
		dim: dim,
		mu:  make([]float64, dim),
	}
	copy(n.mu, mu)
	n.lower.LFromCholesky(chol)
	n.logSqrtDet = 0.5 * n.chol.LogDet()
	return n
}

// NewNormalPrecision creates a new Normal distribution with the given mean and
// precision matrix (inverse of the covariance matrix). NewNormalPrecision
// panics if len(mu) is not equal to prec.Symmetric(). If the precision matrix
// is not positive-definite, NewNormalPrecision returns nil for norm and ok
// for false.
func NewNormalPrecision(mu []float64, prec *mat64.SymDense, src *rand.Rand) (norm *Normal, ok bool) {
	if len(mu) == 0 {
		panic(badZeroDimension)
	}
	dim := prec.Symmetric()
	if dim != len(mu) {
		panic(badSizeMismatch)
	}
	// TODO(btracey): Computing a matrix inverse is generally numerically instable.
	// This only has to compute the inverse of a positive definite matrix, which
	// is much better, but this still loses precision. It is worth considering if
	// instead the precision matrix should be stored explicitly and used instead
	// of the Cholesky decomposition of the covariance matrix where appropriate.
	var chol mat64.Cholesky
	ok = chol.Factorize(prec)
	if !ok {
		return nil, false
	}
	var sigma mat64.SymDense
	sigma.InverseCholesky(&chol)
	return NewNormal(mu, &sigma, src)
}

// ConditionNormal returns the Normal distribution that is the receiver conditioned
// on the input evidence. The returned multivariate normal has dimension
// n - len(observed), where n is the dimension of the original receiver. The updated
// mean and covariance are
//  mu = mu_un + sigma_{ob,un}^T * sigma_{ob,ob}^-1 (v - mu_ob)
//  sigma = sigma_{un,un} - sigma_{ob,un}^T * sigma_{ob,ob}^-1 * sigma_{ob,un}
// where mu_un and mu_ob are the original means of the unobserved and observed
// variables respectively, sigma_{un,un} is the unobserved subset of the covariance
// matrix, sigma_{ob,ob} is the observed subset of the covariance matrix, and
// sigma_{un,ob} are the cross terms. The elements of x_2 have been observed with
// values v. The dimension order is preserved during conditioning, so if the value
// of dimension 1 is observed, the returned normal represents dimensions {0, 2, ...}
// of the original Normal distribution.
//
// ConditionNormal returns {nil, false} if there is a failure during the update.
// Mathematically this is impossible, but can occur with finite precision arithmetic.
func (n *Normal) ConditionNormal(observed []int, values []float64, src *rand.Rand) (*Normal, bool) {
	if len(observed) == 0 {
		panic("normal: no observed value")
	}
	if len(observed) != len(values) {
		panic("normal: input slice length mismatch")
	}
	for _, v := range observed {
		if v < 0 || v >= n.Dim() {
			panic("normal: observed value out of bounds")
		}
	}

	ob := len(observed)
	unob := n.Dim() - ob
	obMap := make(map[int]struct{})
	for _, v := range observed {
		if _, ok := obMap[v]; ok {
			panic("normal: observed dimension occurs twice")
		}
		obMap[v] = struct{}{}
	}
	if len(observed) == n.Dim() {
		panic("normal: all dimensions observed")
	}
	unobserved := make([]int, 0, unob)
	for i := 0; i < n.Dim(); i++ {
		if _, ok := obMap[i]; !ok {
			unobserved = append(unobserved, i)
		}
	}
	mu1 := make([]float64, unob)
	for i, v := range unobserved {
		mu1[i] = n.mu[v]
	}
	mu2 := make([]float64, ob) // really v - mu2
	for i, v := range observed {
		mu2[i] = values[i] - n.mu[v]
	}

	n.setSigma()

	var sigma11, sigma22 mat64.SymDense
	sigma11.SubsetSym(n.sigma, unobserved)
	sigma22.SubsetSym(n.sigma, observed)

	sigma21 := mat64.NewDense(ob, unob, nil)
	for i, r := range observed {
		for j, c := range unobserved {
			v := n.sigma.At(r, c)
			sigma21.Set(i, j, v)
		}
	}

	var chol mat64.Cholesky
	ok := chol.Factorize(&sigma22)
	if !ok {
		return nil, ok
	}

	// Compute sigma_{2,1}^T * sigma_{2,2}^-1 (v - mu_2).
	v := mat64.NewVector(ob, mu2)
	var tmp, tmp2 mat64.Vector
	err := tmp.SolveCholeskyVec(&chol, v)
	if err != nil {
		return nil, false
	}
	tmp2.MulVec(sigma21.T(), &tmp)

	// Compute sigma_{2,1}^T * sigma_{2,2}^-1 * sigma_{2,1}.
	// TODO(btracey): Should this be a method of SymDense?
	var tmp3, tmp4 mat64.Dense
	err = tmp3.SolveCholesky(&chol, sigma21)
	if err != nil {
		return nil, false
	}
	tmp4.Mul(sigma21.T(), &tmp3)

	for i := range mu1 {
		mu1[i] += tmp2.At(i, 0)
	}

	// TODO(btracey): If tmp2 can constructed with a method, then this can be
	// replaced with SubSym.
	for i := 0; i < len(unobserved); i++ {
		for j := i; j < len(unobserved); j++ {
			v := sigma11.At(i, j)
			sigma11.SetSym(i, j, v-tmp4.At(i, j))
		}
	}
	return NewNormal(mu1, &sigma11, src)
}

// CovarianceMatrix returns the covariance matrix of the distribution. Upon
// return, the value at element {i, j} of the covariance matrix is equal to
// the covariance of the i^th and j^th variables.
//  covariance(i, j) = E[(x_i - E[x_i])(x_j - E[x_j])]
// If the input matrix is nil a new matrix is allocated, otherwise the result
// is stored in-place into the input.
func (n *Normal) CovarianceMatrix(s *mat64.SymDense) *mat64.SymDense {
	if s == nil {
		s = mat64.NewSymDense(n.Dim(), nil)
	}
	sn := s.Symmetric()
	if sn != n.Dim() {
		panic("normal: input matrix size mismatch")
	}
	n.setSigma()
	s.CopySym(n.sigma)
	return s
}

// Dim returns the dimension of the distribution.
func (n *Normal) Dim() int {
	return n.dim
}

// Entropy returns the differential entropy of the distribution.
func (n *Normal) Entropy() float64 {
	return float64(n.dim)/2*(1+logTwoPi) + n.logSqrtDet
}

// LogProb computes the log of the pdf of the point x.
func (n *Normal) LogProb(x []float64) float64 {
	dim := n.dim
	if len(x) != dim {
		panic(badSizeMismatch)
	}
	// Compute the normalization constant
	c := -0.5*float64(dim)*logTwoPi - n.logSqrtDet

	// Compute (x-mu)'Sigma^-1 (x-mu)
	xMinusMu := make([]float64, dim)
	floats.SubTo(xMinusMu, x, n.mu)
	d := mat64.NewVector(dim, xMinusMu)
	tmp := make([]float64, dim)
	tmpVec := mat64.NewVector(dim, tmp)
	tmpVec.SolveCholeskyVec(&n.chol, d)
	return c - 0.5*floats.Dot(tmp, xMinusMu)
}

// MarginalNormal returns the marginal distribution of the given input variables.
// That is, MarginalNormal returns
//  p(x_i) = \int_{x_o} p(x_i | x_o) p(x_o) dx_o
// where x_i are the dimensions in the input, and x_o are the remaining dimensions.
// The input src is passed to the call to NewNormal.
func (n *Normal) MarginalNormal(vars []int, src *rand.Rand) (*Normal, bool) {
	newMean := make([]float64, len(vars))
	for i, v := range vars {
		newMean[i] = n.mu[v]
	}
	n.setSigma()
	var s mat64.SymDense
	s.SubsetSym(n.sigma, vars)
	return NewNormal(newMean, &s, src)
}

// Mean returns the mean of the probability distribution at x. If the
// input argument is nil, a new slice will be allocated, otherwise the result
// will be put in-place into the receiver.
func (n *Normal) Mean(x []float64) []float64 {
	x = reuseAs(x, n.dim)
	copy(x, n.mu)
	return x
}

// Prob computes the value of the probability density function at x.
func (n *Normal) Prob(x []float64) float64 {
	return math.Exp(n.LogProb(x))
}

// Rand generates a random number according to the distributon.
// If the input slice is nil, new memory is allocated, otherwise the result is stored
// in place.
func (n *Normal) Rand(x []float64) []float64 {
	x = reuseAs(x, n.dim)
	tmp := make([]float64, n.dim)
	if n.src == nil {
		for i := range x {
			tmp[i] = rand.NormFloat64()
		}
	} else {
		for i := range x {
			tmp[i] = n.src.NormFloat64()
		}
	}
	tmpVec := mat64.NewVector(n.dim, tmp)
	xVec := mat64.NewVector(n.dim, x)
	xVec.MulVec(&n.lower, tmpVec)
	floats.Add(x, n.mu)
	return x
}

// SetMean changes the mean of the normal distribution. SetMean panics if len(mu)
// does not equal the dimension of the normal distribution.
func (n *Normal) SetMean(mu []float64) {
	if len(mu) != n.Dim() {
		panic(badSizeMismatch)
	}
	copy(n.mu, mu)
}

// setSigma computes and stores the covariance matrix of the distribution.
func (n *Normal) setSigma() {
	n.once.Do(func() {
		n.sigma = mat64.NewSymDense(n.Dim(), nil)
		n.sigma.FromCholesky(&n.chol)
	})
}
<|endoftext|>