<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package tests_test

import (
	"flag"
	"time"

	"github.com/google/goexpect"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"kubevirt.io/kubevirt/pkg/api/v1"
	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/tests"
)

var _ = Describe("Configurations", func() {

	flag.Parse()

	virtClient, err := kubecli.GetKubevirtClient()
	tests.PanicOnError(err)

	BeforeEach(func() {
		tests.BeforeTestCleanup()
	})

	Context("New VM with different cpu topologies give", func() {

		var vm *v1.VirtualMachine

		BeforeEach(func() {
			vm = tests.NewRandomVMWithEphemeralDisk("kubevirt/alpine-registry-disk-demo:devel")
		})
		It("should report 3 cpu cores", func() {
			vm.Spec.Domain.CPU = &v1.CPU{
				Cores: 3,
			}

			vm, err = virtClient.VM(tests.NamespaceTestDefault).Create(vm)
			Expect(err).ToNot(HaveOccurred())
			tests.WaitForSuccessfulVMStart(vm)

			expecter, _, err := tests.NewConsoleExpecter(virtClient, vm, "serial0", 10*time.Second)
			Expect(err).ToNot(HaveOccurred())
			defer expecter.Close()
			_, err = expecter.ExpectBatch([]expect.Batcher{
				&expect.BExp{R: "Welcome to Alpine"},
				&expect.BSnd{S: "\n"},
				&expect.BExp{R: "login"},
				&expect.BSnd{S: "root\n"},
				&expect.BExp{R: "#"},
				&expect.BSnd{S: "grep -c ^processor /proc/cpuinfo\n"},
				&expect.BExp{R: "3"},
			}, 250*time.Second)

			Expect(err).ToNot(HaveOccurred())
		}, 300)
	})

	Context("New VM with explicitly set VirtIO drives", func() {

		var vm *v1.VirtualMachine
		var diskDev v1.DiskDevice

		BeforeEach(func() {
			diskDev = v1.DiskDevice{
				Disk: &v1.DiskTarget{
					Bus: "virtio",
				},
			}
			vm = tests.NewRandomVMWithDirectLunAndDevice(2, false, diskDev)
		})
		It("should have /dev/vda node", func() {
			vm, err = virtClient.VM(tests.NamespaceTestDefault).Create(vm)
			Expect(err).ToNot(HaveOccurred())
			tests.WaitForSuccessfulVMStart(vm)

			expecter, _, err := tests.NewConsoleExpecter(virtClient, vm, "serial0", 10*time.Second)
			Expect(err).ToNot(HaveOccurred())
			defer expecter.Close()
			_, err = expecter.ExpectBatch([]expect.Batcher{
				&expect.BExp{R: "Welcome to Alpine"},
				&expect.BSnd{S: "\n"},
				&expect.BExp{R: "login"},
				&expect.BSnd{S: "root\n"},
				&expect.BExp{R: "#"},
				&expect.BSnd{S: "ls /dev/vda\n"},
				&expect.BExp{R: "/dev/vda"},
			}, 150*time.Second)

			Expect(err).ToNot(HaveOccurred())
		})
	})

	Context("New VM with explicitly set SATA drives", func() {

		var vm *v1.VirtualMachine
		var diskDev v1.DiskDevice

		BeforeEach(func() {
			diskDev = v1.DiskDevice{
				Disk: &v1.DiskTarget{
					Bus: "sata",
				},
			}
			vm = tests.NewRandomVMWithDirectLunAndDevice(2, false, diskDev)
		})
		It("should have /dev/vda node", func() {
			vm, err = virtClient.VM(tests.NamespaceTestDefault).Create(vm)
			Expect(err).ToNot(HaveOccurred())
			tests.WaitForSuccessfulVMStart(vm)

			expecter, _, err := tests.NewConsoleExpecter(virtClient, vm, "serial0", 10*time.Second)
			Expect(err).ToNot(HaveOccurred())
			defer expecter.Close()
			_, err = expecter.ExpectBatch([]expect.Batcher{
				&expect.BExp{R: "Welcome to Alpine"},
				&expect.BSnd{S: "\n"},
				&expect.BExp{R: "login"},
				&expect.BSnd{S: "root\n"},
				&expect.BExp{R: "#"},
				&expect.BSnd{S: "ls /dev/sda\n"},
				&expect.BExp{R: "/dev/sda"},
			}, 150*time.Second)

			Expect(err).ToNot(HaveOccurred())
		})
	})

	Context("New VM with all supported drives", func() {

		var vm *v1.VirtualMachine

		BeforeEach(func() {
			// ordering:
			// virtio - added by NewRandomVMWithEphemeralDisk
			containerImage := "kubevirt/cirros-registry-disk-demo:devel"
			vm = tests.NewRandomVMWithEphemeralDisk(containerImage)
			// sata
			tests.AddEphemeralDisk(vm, "disk1", "sata", containerImage)
			// ide
			tests.AddEphemeralDisk(vm, "disk2", "ide", containerImage)
			// floppy
			tests.AddEphemeralDisk(vm, "disk3", "floppy", containerImage)
			// NOTE: we have one disk per bus, so we expect vda, sda, hda, fda
		})
		It("should have all the device nodes", func() {
			vm, err = virtClient.VM(tests.NamespaceTestDefault).Create(vm)
			Expect(err).ToNot(HaveOccurred())
			tests.WaitForSuccessfulVMStart(vm)

			expecter, _, err := tests.NewConsoleExpecter(virtClient, vm, "serial0", 10*time.Second)
			Expect(err).ToNot(HaveOccurred())
			defer expecter.Close()
			_, err = expecter.ExpectBatch([]expect.Batcher{
				&expect.BExp{R: "Welcome to Alpine"},
				&expect.BSnd{S: "\n"},
				&expect.BExp{R: "login"},
				&expect.BSnd{S: "root\n"},
				&expect.BExp{R: "#"},
				// keep the ordering!
				&expect.BSnd{S: "ls /dev/fda /dev/hda /dev/sda /dev/vda\n"},
				&expect.BExp{R: "/dev/fda /dev/hda /dev/sda /dev/vda"},
			}, 150*time.Second)

			Expect(err).ToNot(HaveOccurred())
		})
	})

})
<commit_msg>fix typo in the test expectation message<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package tests_test

import (
	"flag"
	"time"

	"github.com/google/goexpect"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"kubevirt.io/kubevirt/pkg/api/v1"
	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/tests"
)

var _ = Describe("Configurations", func() {

	flag.Parse()

	virtClient, err := kubecli.GetKubevirtClient()
	tests.PanicOnError(err)

	BeforeEach(func() {
		tests.BeforeTestCleanup()
	})

	Context("New VM with different cpu topologies give", func() {

		var vm *v1.VirtualMachine

		BeforeEach(func() {
			vm = tests.NewRandomVMWithEphemeralDisk("kubevirt/alpine-registry-disk-demo:devel")
		})
		It("should report 3 cpu cores", func() {
			vm.Spec.Domain.CPU = &v1.CPU{
				Cores: 3,
			}

			vm, err = virtClient.VM(tests.NamespaceTestDefault).Create(vm)
			Expect(err).ToNot(HaveOccurred())
			tests.WaitForSuccessfulVMStart(vm)

			expecter, _, err := tests.NewConsoleExpecter(virtClient, vm, "serial0", 10*time.Second)
			Expect(err).ToNot(HaveOccurred())
			defer expecter.Close()
			_, err = expecter.ExpectBatch([]expect.Batcher{
				&expect.BExp{R: "Welcome to Alpine"},
				&expect.BSnd{S: "\n"},
				&expect.BExp{R: "login"},
				&expect.BSnd{S: "root\n"},
				&expect.BExp{R: "#"},
				&expect.BSnd{S: "grep -c ^processor /proc/cpuinfo\n"},
				&expect.BExp{R: "3"},
			}, 250*time.Second)

			Expect(err).ToNot(HaveOccurred())
		}, 300)
	})

	Context("New VM with explicitly set VirtIO drives", func() {

		var vm *v1.VirtualMachine
		var diskDev v1.DiskDevice

		BeforeEach(func() {
			diskDev = v1.DiskDevice{
				Disk: &v1.DiskTarget{
					Bus: "virtio",
				},
			}
			vm = tests.NewRandomVMWithDirectLunAndDevice(2, false, diskDev)
		})
		It("should have /dev/vda node", func() {
			vm, err = virtClient.VM(tests.NamespaceTestDefault).Create(vm)
			Expect(err).ToNot(HaveOccurred())
			tests.WaitForSuccessfulVMStart(vm)

			expecter, _, err := tests.NewConsoleExpecter(virtClient, vm, "serial0", 10*time.Second)
			Expect(err).ToNot(HaveOccurred())
			defer expecter.Close()
			_, err = expecter.ExpectBatch([]expect.Batcher{
				&expect.BExp{R: "Welcome to Alpine"},
				&expect.BSnd{S: "\n"},
				&expect.BExp{R: "login"},
				&expect.BSnd{S: "root\n"},
				&expect.BExp{R: "#"},
				&expect.BSnd{S: "ls /dev/vda\n"},
				&expect.BExp{R: "/dev/vda"},
			}, 150*time.Second)

			Expect(err).ToNot(HaveOccurred())
		})
	})

	Context("New VM with explicitly set SATA drives", func() {

		var vm *v1.VirtualMachine
		var diskDev v1.DiskDevice

		BeforeEach(func() {
			diskDev = v1.DiskDevice{
				Disk: &v1.DiskTarget{
					Bus: "sata",
				},
			}
			vm = tests.NewRandomVMWithDirectLunAndDevice(2, false, diskDev)
		})
		It("should have /dev/sda node", func() {
			vm, err = virtClient.VM(tests.NamespaceTestDefault).Create(vm)
			Expect(err).ToNot(HaveOccurred())
			tests.WaitForSuccessfulVMStart(vm)

			expecter, _, err := tests.NewConsoleExpecter(virtClient, vm, "serial0", 10*time.Second)
			Expect(err).ToNot(HaveOccurred())
			defer expecter.Close()
			_, err = expecter.ExpectBatch([]expect.Batcher{
				&expect.BExp{R: "Welcome to Alpine"},
				&expect.BSnd{S: "\n"},
				&expect.BExp{R: "login"},
				&expect.BSnd{S: "root\n"},
				&expect.BExp{R: "#"},
				&expect.BSnd{S: "ls /dev/sda\n"},
				&expect.BExp{R: "/dev/sda"},
			}, 150*time.Second)

			Expect(err).ToNot(HaveOccurred())
		})
	})

	Context("New VM with all supported drives", func() {

		var vm *v1.VirtualMachine

		BeforeEach(func() {
			// ordering:
			// virtio - added by NewRandomVMWithEphemeralDisk
			containerImage := "kubevirt/cirros-registry-disk-demo:devel"
			vm = tests.NewRandomVMWithEphemeralDisk(containerImage)
			// sata
			tests.AddEphemeralDisk(vm, "disk1", "sata", containerImage)
			// ide
			tests.AddEphemeralDisk(vm, "disk2", "ide", containerImage)
			// floppy
			tests.AddEphemeralDisk(vm, "disk3", "floppy", containerImage)
			// NOTE: we have one disk per bus, so we expect vda, sda, hda, fda
		})
		It("should have all the device nodes", func() {
			vm, err = virtClient.VM(tests.NamespaceTestDefault).Create(vm)
			Expect(err).ToNot(HaveOccurred())
			tests.WaitForSuccessfulVMStart(vm)

			expecter, _, err := tests.NewConsoleExpecter(virtClient, vm, "serial0", 10*time.Second)
			Expect(err).ToNot(HaveOccurred())
			defer expecter.Close()
			_, err = expecter.ExpectBatch([]expect.Batcher{
				&expect.BExp{R: "Welcome to Alpine"},
				&expect.BSnd{S: "\n"},
				&expect.BExp{R: "login"},
				&expect.BSnd{S: "root\n"},
				&expect.BExp{R: "#"},
				// keep the ordering!
				&expect.BSnd{S: "ls /dev/fda /dev/hda /dev/sda /dev/vda\n"},
				&expect.BExp{R: "/dev/fda /dev/hda /dev/sda /dev/vda"},
			}, 150*time.Second)

			Expect(err).ToNot(HaveOccurred())
		})
	})

})
<|endoftext|>


<commit_before>package worker

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"
	"sync"
	"time"

	gocontext "context"

	"github.com/sirupsen/logrus"
	"github.com/streadway/amqp"
	"github.com/travis-ci/worker/context"
)

type amqpLogPart struct {
	JobID      uint64     `json:"id"`
	Content    string     `json:"log"`
	Number     int        `json:"number"`
	UUID       string     `json:"uuid"`
	Final      bool       `json:"final"`
	ReceivedAt *time.Time `json:"received_at,omitempty"`
}

type amqpLogWriter struct {
	ctx     gocontext.Context
	jobID   uint64
	sharded bool

	closeChan chan struct{}

	bufferMutex   sync.Mutex
	buffer        *bytes.Buffer
	logPartNumber int

	bytesWritten int
	maxLength    int

	amqpChanMutex sync.RWMutex
	amqpChan      *amqp.Channel

	timer   *time.Timer
	timeout time.Duration
}

func newAMQPLogWriter(ctx gocontext.Context, logWriterChan *amqp.Channel, jobID uint64, timeout time.Duration, sharded bool) (*amqpLogWriter, error) {

	writer := &amqpLogWriter{
		ctx:       context.FromComponent(ctx, "log_writer"),
		amqpChan:  logWriterChan,
		jobID:     jobID,
		closeChan: make(chan struct{}),
		buffer:    new(bytes.Buffer),
		timer:     time.NewTimer(time.Hour),
		timeout:   timeout,
		sharded:   sharded,
	}

	context.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"writer": writer,
		"job_id": jobID,
	}).Debug("created new log writer")

	go writer.flushRegularly(ctx)

	return writer, nil
}

func (w *amqpLogWriter) Write(p []byte) (int, error) {
	if w.closed() {
		return 0, fmt.Errorf("attempted write to closed log")
	}

	logger := context.LoggerFromContext(w.ctx).WithFields(logrus.Fields{
		"self": "amqp_log_writer",
		"inst": fmt.Sprintf("%p", w),
	})

	logger.WithFields(logrus.Fields{
		"length": len(p),
		"bytes":  string(p),
	}).Debug("writing bytes")

	w.timer.Reset(w.timeout)

	w.bytesWritten += len(p)
	if w.bytesWritten > w.maxLength {
		_, err := w.WriteAndClose([]byte(fmt.Sprintf("\n\nThe log length has exceeded the limit of %d MB (this usually means that the test suite is raising the same exception over and over).\n\nThe job has been terminated\n", w.maxLength/1000/1000)))
		if err != nil {
			logger.WithField("err", err).Error("couldn't write 'log length exceeded' error message to log")
		}
		return 0, ErrWrotePastMaxLogLength
	}

	w.bufferMutex.Lock()
	defer w.bufferMutex.Unlock()
	return w.buffer.Write(p)
}

func (w *amqpLogWriter) Close() error {
	if w.closed() {
		return nil
	}

	w.timer.Stop()

	close(w.closeChan)
	w.flush()

	part := amqpLogPart{
		JobID:  w.jobID,
		Number: w.logPartNumber,
		Final:  true,
	}
	w.logPartNumber++

	err := w.publishLogPart(part)
	return err
}

func (w *amqpLogWriter) Timeout() <-chan time.Time {
	return w.timer.C
}

func (w *amqpLogWriter) SetMaxLogLength(bytes int) {
	w.maxLength = bytes
}

// WriteAndClose works like a Write followed by a Close, but ensures that no
// other Writes are allowed in between.
func (w *amqpLogWriter) WriteAndClose(p []byte) (int, error) {
	if w.closed() {
		return 0, fmt.Errorf("log already closed")
	}

	w.timer.Stop()

	close(w.closeChan)

	w.bufferMutex.Lock()
	n, err := w.buffer.Write(p)
	w.bufferMutex.Unlock()
	if err != nil {
		return n, err
	}

	w.flush()

	part := amqpLogPart{
		JobID:  w.jobID,
		Number: w.logPartNumber,
		Final:  true,
	}
	w.logPartNumber++

	err = w.publishLogPart(part)
	return n, err
}

func (w *amqpLogWriter) closed() bool {
	select {
	case <-w.closeChan:
		return true
	default:
		return false
	}
}

func (w *amqpLogWriter) flushRegularly(ctx gocontext.Context) {
	ticker := time.NewTicker(LogWriterTick)
	defer ticker.Stop()
	for {
		select {
		case <-w.closeChan:
			return
		case <-ticker.C:
			w.flush()
		case <-ctx.Done():
			return
		}
	}
}

func (w *amqpLogWriter) flush() {
	if w.buffer.Len() <= 0 {
		return
	}

	buf := make([]byte, LogChunkSize)
	logger := context.LoggerFromContext(w.ctx).WithFields(logrus.Fields{
		"self": "amqp_log_writer",
		"inst": fmt.Sprintf("%p", w),
	})

	for w.buffer.Len() > 0 {
		w.bufferMutex.Lock()
		n, err := w.buffer.Read(buf)
		w.bufferMutex.Unlock()
		if err != nil {
			// According to documentation, err should only be non-nil if
			// there's no data in the buffer. We've checked for this, so
			// this means that err should never be nil. Something is very
			// wrong if this happens, so let's abort!
			panic("non-empty buffer shouldn't return an error on Read")
		}

		part := amqpLogPart{
			JobID:   w.jobID,
			Content: string(buf[0:n]),
			Number:  w.logPartNumber,
		}
		w.logPartNumber++

		err = w.publishLogPart(part)
		if err != nil {
			logger.WithField("err", err).Error("couldn't publish log part")
		}
	}
}

func (w *amqpLogWriter) publishLogPart(part amqpLogPart) error {
	part.UUID, _ = context.UUIDFromContext(w.ctx)

	if w.logPartNumber == 1 && w.ctx.Value("processedAt") != nil {
		processedAt := w.ctx.Value("processedAt").(time.Time)
		part.ReceivedAt = &processedAt
	}

	partBody, err := json.Marshal(part)
	if err != nil {
		return err
	}

	w.amqpChanMutex.RLock()
	var exchange string
	var routingKey string
	if w.sharded {
		exchange = "reporting.jobs.logs_sharded"
		routingKey = strconv.FormatUint(w.jobID, 10)
	} else {
		exchange = "reporting"
		routingKey = "reporting.jobs.logs"
	}
	err = w.amqpChan.Publish(exchange, routingKey, false, false, amqp.Publishing{
		ContentType:  "application/json",
		DeliveryMode: amqp.Persistent,
		Timestamp:    time.Now(),
		Type:         "job:test:log",
		Body:         partBody,
	})
	w.amqpChanMutex.RUnlock()

	return err
}
<commit_msg>check part number as assigned to part<commit_after>package worker

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"
	"sync"
	"time"

	gocontext "context"

	"github.com/sirupsen/logrus"
	"github.com/streadway/amqp"
	"github.com/travis-ci/worker/context"
)

type amqpLogPart struct {
	JobID      uint64     `json:"id"`
	Content    string     `json:"log"`
	Number     int        `json:"number"`
	UUID       string     `json:"uuid"`
	Final      bool       `json:"final"`
	ReceivedAt *time.Time `json:"received_at,omitempty"`
}

type amqpLogWriter struct {
	ctx     gocontext.Context
	jobID   uint64
	sharded bool

	closeChan chan struct{}

	bufferMutex   sync.Mutex
	buffer        *bytes.Buffer
	logPartNumber int

	bytesWritten int
	maxLength    int

	amqpChanMutex sync.RWMutex
	amqpChan      *amqp.Channel

	timer   *time.Timer
	timeout time.Duration
}

func newAMQPLogWriter(ctx gocontext.Context, logWriterChan *amqp.Channel, jobID uint64, timeout time.Duration, sharded bool) (*amqpLogWriter, error) {

	writer := &amqpLogWriter{
		ctx:       context.FromComponent(ctx, "log_writer"),
		amqpChan:  logWriterChan,
		jobID:     jobID,
		closeChan: make(chan struct{}),
		buffer:    new(bytes.Buffer),
		timer:     time.NewTimer(time.Hour),
		timeout:   timeout,
		sharded:   sharded,
	}

	context.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"writer": writer,
		"job_id": jobID,
	}).Debug("created new log writer")

	go writer.flushRegularly(ctx)

	return writer, nil
}

func (w *amqpLogWriter) Write(p []byte) (int, error) {
	if w.closed() {
		return 0, fmt.Errorf("attempted write to closed log")
	}

	logger := context.LoggerFromContext(w.ctx).WithFields(logrus.Fields{
		"self": "amqp_log_writer",
		"inst": fmt.Sprintf("%p", w),
	})

	logger.WithFields(logrus.Fields{
		"length": len(p),
		"bytes":  string(p),
	}).Debug("writing bytes")

	w.timer.Reset(w.timeout)

	w.bytesWritten += len(p)
	if w.bytesWritten > w.maxLength {
		_, err := w.WriteAndClose([]byte(fmt.Sprintf("\n\nThe log length has exceeded the limit of %d MB (this usually means that the test suite is raising the same exception over and over).\n\nThe job has been terminated\n", w.maxLength/1000/1000)))
		if err != nil {
			logger.WithField("err", err).Error("couldn't write 'log length exceeded' error message to log")
		}
		return 0, ErrWrotePastMaxLogLength
	}

	w.bufferMutex.Lock()
	defer w.bufferMutex.Unlock()
	return w.buffer.Write(p)
}

func (w *amqpLogWriter) Close() error {
	if w.closed() {
		return nil
	}

	w.timer.Stop()

	close(w.closeChan)
	w.flush()

	part := amqpLogPart{
		JobID:  w.jobID,
		Number: w.logPartNumber,
		Final:  true,
	}
	w.logPartNumber++

	err := w.publishLogPart(part)
	return err
}

func (w *amqpLogWriter) Timeout() <-chan time.Time {
	return w.timer.C
}

func (w *amqpLogWriter) SetMaxLogLength(bytes int) {
	w.maxLength = bytes
}

// WriteAndClose works like a Write followed by a Close, but ensures that no
// other Writes are allowed in between.
func (w *amqpLogWriter) WriteAndClose(p []byte) (int, error) {
	if w.closed() {
		return 0, fmt.Errorf("log already closed")
	}

	w.timer.Stop()

	close(w.closeChan)

	w.bufferMutex.Lock()
	n, err := w.buffer.Write(p)
	w.bufferMutex.Unlock()
	if err != nil {
		return n, err
	}

	w.flush()

	part := amqpLogPart{
		JobID:  w.jobID,
		Number: w.logPartNumber,
		Final:  true,
	}
	w.logPartNumber++

	err = w.publishLogPart(part)
	return n, err
}

func (w *amqpLogWriter) closed() bool {
	select {
	case <-w.closeChan:
		return true
	default:
		return false
	}
}

func (w *amqpLogWriter) flushRegularly(ctx gocontext.Context) {
	ticker := time.NewTicker(LogWriterTick)
	defer ticker.Stop()
	for {
		select {
		case <-w.closeChan:
			return
		case <-ticker.C:
			w.flush()
		case <-ctx.Done():
			return
		}
	}
}

func (w *amqpLogWriter) flush() {
	if w.buffer.Len() <= 0 {
		return
	}

	buf := make([]byte, LogChunkSize)
	logger := context.LoggerFromContext(w.ctx).WithFields(logrus.Fields{
		"self": "amqp_log_writer",
		"inst": fmt.Sprintf("%p", w),
	})

	for w.buffer.Len() > 0 {
		w.bufferMutex.Lock()
		n, err := w.buffer.Read(buf)
		w.bufferMutex.Unlock()
		if err != nil {
			// According to documentation, err should only be non-nil if
			// there's no data in the buffer. We've checked for this, so
			// this means that err should never be nil. Something is very
			// wrong if this happens, so let's abort!
			panic("non-empty buffer shouldn't return an error on Read")
		}

		part := amqpLogPart{
			JobID:   w.jobID,
			Content: string(buf[0:n]),
			Number:  w.logPartNumber,
		}
		w.logPartNumber++

		err = w.publishLogPart(part)
		if err != nil {
			logger.WithField("err", err).Error("couldn't publish log part")
		}
	}
}

func (w *amqpLogWriter) publishLogPart(part amqpLogPart) error {
	part.UUID, _ = context.UUIDFromContext(w.ctx)

	if part.Number == 0 && w.ctx.Value("processedAt") != nil {
		processedAt := w.ctx.Value("processedAt").(time.Time)
		part.ReceivedAt = &processedAt
	}

	partBody, err := json.Marshal(part)
	if err != nil {
		return err
	}

	w.amqpChanMutex.RLock()
	var exchange string
	var routingKey string
	if w.sharded {
		exchange = "reporting.jobs.logs_sharded"
		routingKey = strconv.FormatUint(w.jobID, 10)
	} else {
		exchange = "reporting"
		routingKey = "reporting.jobs.logs"
	}
	err = w.amqpChan.Publish(exchange, routingKey, false, false, amqp.Publishing{
		ContentType:  "application/json",
		DeliveryMode: amqp.Persistent,
		Timestamp:    time.Now(),
		Type:         "job:test:log",
		Body:         partBody,
	})
	w.amqpChanMutex.RUnlock()

	return err
}
<|endoftext|>


<commit_before>// Package circonusgometrics provides instrumentation for your applications in the form
// of counters, gauges and histograms and allows you to publish them to
// Circonus
//
// Counters
//
// A counter is a monotonically-increasing, unsigned, 64-bit integer used to
// represent the number of times an event has occurred. By tracking the deltas
// between measurements of a counter over intervals of time, an aggregation
// layer can derive rates, acceleration, etc.
//
// Gauges
//
// A gauge returns instantaneous measurements of something using signed, 64-bit
// integers. This value does not need to be monotonic.
//
// Histograms
//
// A histogram tracks the distribution of a stream of values (e.g. the number of
// seconds it takes to handle requests).  Circonus can calculate complex
// analytics on these.
//
// Reporting
//
// A period push to a Circonus httptrap is confgurable.

package circonusgometrics

import (
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"sync"
	"time"
)

const (
	// a few sensible defaults
	defaultApiHost  = "api.circonus.com"
	defaultApiApp   = "circonus-gometrics"
	defaultInterval = 10 * time.Second
)

// a few words about: "BrokerGroupId"
//
// calling it this because the instructions for how to get into the UI and FIND this value are more straight-forward:
//
// log into ui
// navigate to brokers page
// identify which broker you need to use
// click the little down arrow in the circle on the right-hand side of the line for the broker you'd like to use
// use the value from the "GROUP ID:" field under "Broker Details" in the drop-down afetr clicking the down arrow
//
// ... or ...
//
// log into ui
// navigate to brokers page
// identify which broker you need to use
// click the hamburger menu icon (three lines to the left of the broker name)
// click "view API object" from the drop-down menu
// look for "_cid" field, use integer value after "/broker/" e.g. "/broker/35" would be 35
//

type CirconusMetrics struct {
	ApiToken      string
	SubmissionUrl string
	CheckId       int
	ApiApp        string
	ApiHost       string
	InstanceId    string
	SearchTag     string
	BrokerGroupId int
	Tags          []string
	CheckSecret   string

	Interval time.Duration
	Log      *log.Logger
	Debug    bool

	// internals
	ready   bool
	trapUrl string
	trapCN  string
	trapSSL bool
	trapmu  sync.Mutex

	certPool      *x509.CertPool
	cert          []byte
	checkBundle   *CheckBundle
	activeMetrics map[string]bool
	checkType     string

	counters map[string]uint64
	cm       sync.Mutex

	counterFuncs map[string]func() uint64
	cfm          sync.Mutex

	gauges map[string]int64
	gm     sync.Mutex

	gaugeFuncs map[string]func() int64
	gfm        sync.Mutex

	histograms map[string]*Histogram
	hm         sync.Mutex
}

// return new CirconusMetrics instance
func NewCirconusMetrics() *CirconusMetrics {
	_, an := path.Split(os.Args[0])
	hn, err := os.Hostname()
	if err != nil {
		hn = "unknown"
	}

	return &CirconusMetrics{
		InstanceId:    fmt.Sprintf("%s:%s", hn, an),
		SearchTag:     fmt.Sprintf("service:%s", an),
		ApiHost:       defaultApiHost,
		ApiApp:        defaultApiApp,
		Interval:      defaultInterval,
		Log:           log.New(ioutil.Discard, "", log.LstdFlags),
		Debug:         false,
		ready:         false,
		trapUrl:       "",
		activeMetrics: make(map[string]bool),
		counters:      make(map[string]uint64),
		counterFuncs:  make(map[string]func() uint64),
		gauges:        make(map[string]int64),
		gaugeFuncs:    make(map[string]func() int64),
		histograms:    make(map[string]*Histogram),
		certPool:      x509.NewCertPool(),
		checkType:     "httptrap",
	}

}

// Start initializes the CirconusMetrics instance based on
// configuration settings and sets the httptrap check url to
// which metrics should be sent. It then starts a perdiodic
// submission process of all metrics collected.
func (m *CirconusMetrics) Start() error {
	if m.Debug {
		m.Log = log.New(os.Stderr, "", log.LstdFlags)
	}
	if !m.ready {
		if err := m.initializeTrap(); err != nil {
			return err
		}
	}

	go func() {
		for _ = range time.NewTicker(m.Interval).C {
			m.Flush()
		}
	}()

	return nil
}

// Flush metrics kicks off the process of sending metrics to Circonus
func (m *CirconusMetrics) Flush() {
	m.Log.Println("Flushing")
	if !m.ready {
		if err := m.initializeTrap(); err != nil {
			m.Log.Printf("Unable to initialize check, NOT flushing metrics. %s\n", err)
		}
		return
	}

	// check for new metrics and enable them automatically
	newMetrics := make(map[string]*CheckBundleMetric)

	counters, gauges, histograms := m.snapshot()
	output := make(map[string]interface{})
	for name, value := range counters {
		output[name] = map[string]interface{}{
			"_type":  "n",
			"_value": value,
		}
		if _, ok := m.activeMetrics[name]; !ok {
			newMetrics[name] = &CheckBundleMetric{
				Name:   name,
				Type:   "numeric",
				Status: "active",
			}
		}
	}

	for name, value := range gauges {
		output[name] = map[string]interface{}{
			"_type":  "n",
			"_value": value,
		}
		if _, ok := m.activeMetrics[name]; !ok {
			newMetrics[name] = &CheckBundleMetric{
				Name:   name,
				Type:   "numeric",
				Status: "active",
			}
		}
	}

	for name, value := range histograms {
		output[name] = map[string]interface{}{
			"_type":  "n",
			"_value": value.DecStrings(),
		}
		if _, ok := m.activeMetrics[name]; !ok {
			newMetrics[name] = &CheckBundleMetric{
				Name:   name,
				Type:   "histogram",
				Status: "active",
			}
		}
	}

	m.submit(output, newMetrics)
}
<commit_msg>Add MaxSubmissionUrlAge (default 60 seconds) * errors submititng metrics will refresh trapUrl (from api) only after this amount of time. * add flushing flag so that only one flush is running at a time.<commit_after>// Package circonusgometrics provides instrumentation for your applications in the form
// of counters, gauges and histograms and allows you to publish them to
// Circonus
//
// Counters
//
// A counter is a monotonically-increasing, unsigned, 64-bit integer used to
// represent the number of times an event has occurred. By tracking the deltas
// between measurements of a counter over intervals of time, an aggregation
// layer can derive rates, acceleration, etc.
//
// Gauges
//
// A gauge returns instantaneous measurements of something using signed, 64-bit
// integers. This value does not need to be monotonic.
//
// Histograms
//
// A histogram tracks the distribution of a stream of values (e.g. the number of
// seconds it takes to handle requests).  Circonus can calculate complex
// analytics on these.
//
// Reporting
//
// A period push to a Circonus httptrap is confgurable.

package circonusgometrics

import (
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"sync"
	"time"
)

const (
	// a few sensible defaults
	defaultApiHost             = "api.circonus.com"
	defaultApiApp              = "circonus-gometrics"
	defaultInterval            = 10 * time.Second
	defaultMaxSubmissionUrlAge = 60 * time.Second
)

// a few words about: "BrokerGroupId"
//
// calling it this because the instructions for how to get into the UI and FIND this value are more straight-forward:
//
// log into ui
// navigate to brokers page
// identify which broker you need to use
// click the little down arrow in the circle on the right-hand side of the line for the broker you'd like to use
// use the value from the "GROUP ID:" field under "Broker Details" in the drop-down afetr clicking the down arrow
//
// ... or ...
//
// log into ui
// navigate to brokers page
// identify which broker you need to use
// click the hamburger menu icon (three lines to the left of the broker name)
// click "view API object" from the drop-down menu
// look for "_cid" field, use integer value after "/broker/" e.g. "/broker/35" would be 35
//

type CirconusMetrics struct {
	ApiToken      string
	SubmissionUrl string
	CheckId       int
	ApiApp        string
	ApiHost       string
	InstanceId    string
	SearchTag     string
	BrokerGroupId int
	Tags          []string
	CheckSecret   string

	Interval time.Duration
	// if the submission url returns errors
	// this gates the amount of time to keep the current
	// submission url before attempting to retrieve it
	// again from the api
	MaxSubmissionUrlAge time.Duration

	Log   *log.Logger
	Debug bool

	// internals
	flushing bool
	flushmu  sync.Mutex

	ready          bool
	trapUrl        string
	trapCN         string
	trapSSL        bool
	trapLastUpdate time.Time
	trapmu         sync.Mutex

	certPool      *x509.CertPool
	cert          []byte
	checkBundle   *CheckBundle
	activeMetrics map[string]bool
	checkType     string

	counters map[string]uint64
	cm       sync.Mutex

	counterFuncs map[string]func() uint64
	cfm          sync.Mutex

	gauges map[string]int64
	gm     sync.Mutex

	gaugeFuncs map[string]func() int64
	gfm        sync.Mutex

	histograms map[string]*Histogram
	hm         sync.Mutex
}

// return new CirconusMetrics instance
func NewCirconusMetrics() *CirconusMetrics {
	_, an := path.Split(os.Args[0])
	hn, err := os.Hostname()
	if err != nil {
		hn = "unknown"
	}

	return &CirconusMetrics{
		InstanceId:          fmt.Sprintf("%s:%s", hn, an),
		SearchTag:           fmt.Sprintf("service:%s", an),
		ApiHost:             defaultApiHost,
		ApiApp:              defaultApiApp,
		Interval:            defaultInterval,
		MaxSubmissionUrlAge: defaultMaxSubmissionUrlAge,
		Log:                 log.New(ioutil.Discard, "", log.LstdFlags),
		Debug:               false,
		ready:               false,
		trapUrl:             "",
		activeMetrics:       make(map[string]bool),
		counters:            make(map[string]uint64),
		counterFuncs:        make(map[string]func() uint64),
		gauges:              make(map[string]int64),
		gaugeFuncs:          make(map[string]func() int64),
		histograms:          make(map[string]*Histogram),
		certPool:            x509.NewCertPool(),
		checkType:           "httptrap",
	}

}

// Start initializes the CirconusMetrics instance based on
// configuration settings and sets the httptrap check url to
// which metrics should be sent. It then starts a perdiodic
// submission process of all metrics collected.
func (m *CirconusMetrics) Start() error {
	if m.Debug {
		m.Log = log.New(os.Stderr, "", log.LstdFlags)
	}
	if !m.ready {
		if err := m.initializeTrap(); err != nil {
			return err
		}
	}

	go func() {
		for _ = range time.NewTicker(m.Interval).C {
			m.Flush()
		}
	}()

	return nil
}

// Flush metrics kicks off the process of sending metrics to Circonus
func (m *CirconusMetrics) Flush() {
	if m.flushing {
		m.Log.Println("Flush already active.")
		return
	}
	m.flushmu.Lock()
	m.flushing = true
	m.flushmu.Unlock()

	if !m.ready {
		m.Log.Println("Initializing trap")
		if err := m.initializeTrap(); err != nil {
			m.Log.Printf("Unable to initialize check, NOT flushing metrics. %s\n", err)
			return
		}
	}

	m.Log.Println("Flushing")

	// check for new metrics and enable them automatically
	newMetrics := make(map[string]*CheckBundleMetric)

	counters, gauges, histograms := m.snapshot()
	output := make(map[string]interface{})
	for name, value := range counters {
		output[name] = map[string]interface{}{
			"_type":  "n",
			"_value": value,
		}
		if _, ok := m.activeMetrics[name]; !ok {
			newMetrics[name] = &CheckBundleMetric{
				Name:   name,
				Type:   "numeric",
				Status: "active",
			}
		}
	}

	for name, value := range gauges {
		output[name] = map[string]interface{}{
			"_type":  "n",
			"_value": value,
		}
		if _, ok := m.activeMetrics[name]; !ok {
			newMetrics[name] = &CheckBundleMetric{
				Name:   name,
				Type:   "numeric",
				Status: "active",
			}
		}
	}

	for name, value := range histograms {
		output[name] = map[string]interface{}{
			"_type":  "n",
			"_value": value.DecStrings(),
		}
		if _, ok := m.activeMetrics[name]; !ok {
			newMetrics[name] = &CheckBundleMetric{
				Name:   name,
				Type:   "histogram",
				Status: "active",
			}
		}
	}

	m.submit(output, newMetrics)

	m.flushmu.Lock()
	m.flushing = false
	m.flushmu.Unlock()
}
<|endoftext|>


<commit_before>package auth

import (
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/atc/db"

	"net/url"

	"golang.org/x/net/context"
	"golang.org/x/oauth2"
)

type OAuthCallbackHandler struct {
	logger             lager.Logger
	providerFactory    ProviderFactory
	privateKey         *rsa.PrivateKey
	authTokenGenerator AuthTokenGenerator
	csrfTokenGenerator CSRFTokenGenerator
	teamFactory        db.TeamFactory
	expire             time.Duration
	isTLSEnabled       bool
	versionHandler     oauthCallbackHandler
}

func NewOAuthCallbackHandler(
	logger lager.Logger,
	providerFactory ProviderFactory,
	privateKey *rsa.PrivateKey,
	teamFactory db.TeamFactory,
	expire time.Duration,
	isTLSEnabled bool,
	versionHandler oauthCallbackHandler,
) http.Handler {
	return &OAuthCallbackHandler{
		logger:             logger,
		providerFactory:    providerFactory,
		privateKey:         privateKey,
		authTokenGenerator: NewAuthTokenGenerator(privateKey),
		csrfTokenGenerator: NewCSRFTokenGenerator(),
		teamFactory:        teamFactory,
		expire:             expire,
		isTLSEnabled:       isTLSEnabled,
		versionHandler:     versionHandler,
	}
}

func (handler *OAuthCallbackHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	hLog := handler.logger.Session("callback")
	providerName := r.FormValue(":provider")
	paramState := r.FormValue("state")

	cookieState, err := r.Cookie(OAuthStateCookie)
	if err != nil {
		hLog.Info("no-state-cookie", lager.Data{
			"error": err.Error(),
		})
		http.Error(w, "state cookie not set", http.StatusUnauthorized)
		return
	}

	if !handler.versionHandler.ValidState(cookieState.Value, paramState) {
		hLog.Info("state-cookie-mismatch", lager.Data{
			"param-state":  paramState,
			"cookie-state": cookieState.Value,
		})

		http.Error(w, "state cookie does not match param", http.StatusUnauthorized)
		return
	}

	stateJSON, err := base64.RawURLEncoding.DecodeString(cookieState.Value)
	if err != nil {
		hLog.Info("failed-to-decode-state", lager.Data{
			"error": err.Error(),
		})
		http.Error(w, "state value invalid base64", http.StatusUnauthorized)
		return
	}

	var oauthState OAuthState
	err = json.Unmarshal(stateJSON, &oauthState)
	if err != nil {
		hLog.Info("failed-to-unmarshal-state", lager.Data{
			"error": err.Error(),
		})
		http.Error(w, "state value invalid JSON", http.StatusUnauthorized)
		return
	}

	teamName := oauthState.TeamName
	team, found, err := handler.teamFactory.FindTeam(teamName)

	if err != nil {
		hLog.Error("failed-to-get-team", err)
		http.Error(w, "failed to get team", http.StatusInternalServerError)
		return
	}
	if !found {
		hLog.Info("failed-to-find-team", lager.Data{
			"teamName": teamName,
		})
		http.Error(w, "failed to find team", http.StatusNotFound)
		return
	}

	provider, found, err := handler.providerFactory.GetProvider(team, providerName)
	if err != nil {
		handler.logger.Error("failed-to-get-provider", err, lager.Data{
			"provider": providerName,
			"teamName": teamName,
		})

		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	if !found {
		handler.logger.Info("provider-not-found-for-team", lager.Data{
			"provider": providerName,
			"teamName": teamName,
		})

		w.WriteHeader(http.StatusNotFound)
		return
	}

	preTokenClient, err := provider.PreTokenClient()
	if err != nil {
		handler.logger.Error("failed-to-construct-pre-token-client", err, lager.Data{
			"provider": providerName,
			"teamName": teamName,
		})

		http.Error(w, "unable to connect to provider: "+err.Error(), http.StatusInternalServerError)
		return
	}

	ctx := context.WithValue(oauth2.NoContext, oauth2.HTTPClient, preTokenClient)
	ctx = context.WithValue(ctx, "request", r)

	token, err := provider.Exchange(ctx, r.FormValue("code"))
	if err != nil {
		hLog.Error("failed-to-exchange-token", err)
		http.Error(w, "failed to exchange token", http.StatusInternalServerError)
		return
	}

	httpClient := provider.Client(ctx, token)

	verified, err := provider.Verify(hLog.Session("verify"), httpClient)
	if err != nil {
		hLog.Error("failed-to-verify-token", err)
		http.Error(w, "failed to verify token", http.StatusInternalServerError)
		return
	}

	if !verified {
		hLog.Info("verification-failed")
		http.Error(w, "verification failed", http.StatusUnauthorized)
		return
	}

	exp := time.Now().Add(handler.expire)

	csrfToken, err := handler.csrfTokenGenerator.GenerateToken()
	if err != nil {
		hLog.Error("generate-csrf-token", err)
		http.Error(w, "failed to generate csrf token", http.StatusInternalServerError)
		return
	}

	tokenType, signedToken, err := handler.authTokenGenerator.GenerateToken(exp, team.Name(), team.Admin(), csrfToken)
	if err != nil {
		hLog.Error("failed-to-sign-token", err)
		http.Error(w, "failed to generate auth token", http.StatusInternalServerError)
		return
	}

	tokenStr := string(tokenType) + " " + string(signedToken)

	authCookie := &http.Cookie{
		Name:     AuthCookieName,
		Value:    tokenStr,
		Path:     "/",
		Expires:  exp,
		HttpOnly: true,
	}
	if handler.isTLSEnabled {
		authCookie.Secure = true
	}
	// TODO: Add SameSite once Golang supports it
	// https://github.com/golang/go/issues/15867
	http.SetCookie(w, authCookie)

	// Deletes the oauth state cookie to avoid CSRF attacks
	http.SetCookie(w, &http.Cookie{
		Name:   cookieState.Name,
		Path:   "/",
		MaxAge: -1,
	})

	w.Header().Set(CSRFHeaderName, csrfToken)

	if oauthState.Redirect != "" && !strings.HasPrefix(oauthState.Redirect, "/") {
		hLog.Info("invalid-redirect")
		http.Error(w, "invalid redirect", http.StatusBadRequest)
		return
	}

	if oauthState.Redirect != "" {
		redirectURL, err := url.Parse(oauthState.Redirect)
		if err != nil {
			hLog.Info("invalid-redirect")
			http.Error(w, "invalid redirect", http.StatusBadRequest)
			return
		}
		queryParams := redirectURL.Query()
		queryParams.Set("csrf_token", csrfToken)
		redirectURL.RawQuery = queryParams.Encode()
		http.Redirect(w, r, redirectURL.String(), http.StatusTemporaryRedirect)
		return
	}

	if oauthState.FlyLocalPort == "" {
		// Old login flow
		fmt.Fprintln(w, tokenStr)
	} else {
		encodedToken := url.QueryEscape(tokenStr)
		http.Redirect(w, r, fmt.Sprintf("http://127.0.0.1:%s/oauth/callback?token=%s", oauthState.FlyLocalPort, encodedToken), http.StatusTemporaryRedirect)
	}
}

type oauthCallbackHandler interface {
	ValidState(cookieState string, paramState string) bool
}

type oauthCallbackHandlerV1 struct{}
type oauthCallbackHandlerV2 struct{}

func (oauthCallbackHandlerV1) ValidState(cookieState string, paramState string) bool {
	return true
}

func (oauthCallbackHandlerV2) ValidState(cookieState string, paramState string) bool {
	return cookieState == paramState
}
<commit_msg>Add comment explaining the use of the state cookie<commit_after>package auth

import (
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/atc/db"

	"net/url"

	"golang.org/x/net/context"
	"golang.org/x/oauth2"
)

type OAuthCallbackHandler struct {
	logger             lager.Logger
	providerFactory    ProviderFactory
	privateKey         *rsa.PrivateKey
	authTokenGenerator AuthTokenGenerator
	csrfTokenGenerator CSRFTokenGenerator
	teamFactory        db.TeamFactory
	expire             time.Duration
	isTLSEnabled       bool
	versionHandler     oauthCallbackHandler
}

func NewOAuthCallbackHandler(
	logger lager.Logger,
	providerFactory ProviderFactory,
	privateKey *rsa.PrivateKey,
	teamFactory db.TeamFactory,
	expire time.Duration,
	isTLSEnabled bool,
	versionHandler oauthCallbackHandler,
) http.Handler {
	return &OAuthCallbackHandler{
		logger:             logger,
		providerFactory:    providerFactory,
		privateKey:         privateKey,
		authTokenGenerator: NewAuthTokenGenerator(privateKey),
		csrfTokenGenerator: NewCSRFTokenGenerator(),
		teamFactory:        teamFactory,
		expire:             expire,
		isTLSEnabled:       isTLSEnabled,
		versionHandler:     versionHandler,
	}
}

func (handler *OAuthCallbackHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	hLog := handler.logger.Session("callback")
	providerName := r.FormValue(":provider")
	paramState := r.FormValue("state")

	cookieState, err := r.Cookie(OAuthStateCookie)
	if err != nil {
		hLog.Info("no-state-cookie", lager.Data{
			"error": err.Error(),
		})
		http.Error(w, "state cookie not set", http.StatusUnauthorized)
		return
	}

	if !handler.versionHandler.ValidState(cookieState.Value, paramState) {
		hLog.Info("state-cookie-mismatch", lager.Data{
			"param-state":  paramState,
			"cookie-state": cookieState.Value,
		})

		http.Error(w, "state cookie does not match param", http.StatusUnauthorized)
		return
	}

	// Read the state from the cookie instead of the param, as the param
	// will be empty if this is an OAuth 1 request. For OAuth 2, we already
	// made sure that the cookie and the param contain the same state.
	stateJSON, err := base64.RawURLEncoding.DecodeString(cookieState.Value)
	if err != nil {
		hLog.Info("failed-to-decode-state", lager.Data{
			"error": err.Error(),
		})
		http.Error(w, "state value invalid base64", http.StatusUnauthorized)
		return
	}

	var oauthState OAuthState
	err = json.Unmarshal(stateJSON, &oauthState)
	if err != nil {
		hLog.Info("failed-to-unmarshal-state", lager.Data{
			"error": err.Error(),
		})
		http.Error(w, "state value invalid JSON", http.StatusUnauthorized)
		return
	}

	teamName := oauthState.TeamName
	team, found, err := handler.teamFactory.FindTeam(teamName)

	if err != nil {
		hLog.Error("failed-to-get-team", err)
		http.Error(w, "failed to get team", http.StatusInternalServerError)
		return
	}
	if !found {
		hLog.Info("failed-to-find-team", lager.Data{
			"teamName": teamName,
		})
		http.Error(w, "failed to find team", http.StatusNotFound)
		return
	}

	provider, found, err := handler.providerFactory.GetProvider(team, providerName)
	if err != nil {
		handler.logger.Error("failed-to-get-provider", err, lager.Data{
			"provider": providerName,
			"teamName": teamName,
		})

		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	if !found {
		handler.logger.Info("provider-not-found-for-team", lager.Data{
			"provider": providerName,
			"teamName": teamName,
		})

		w.WriteHeader(http.StatusNotFound)
		return
	}

	preTokenClient, err := provider.PreTokenClient()
	if err != nil {
		handler.logger.Error("failed-to-construct-pre-token-client", err, lager.Data{
			"provider": providerName,
			"teamName": teamName,
		})

		http.Error(w, "unable to connect to provider: "+err.Error(), http.StatusInternalServerError)
		return
	}

	ctx := context.WithValue(oauth2.NoContext, oauth2.HTTPClient, preTokenClient)
	ctx = context.WithValue(ctx, "request", r)

	token, err := provider.Exchange(ctx, r.FormValue("code"))
	if err != nil {
		hLog.Error("failed-to-exchange-token", err)
		http.Error(w, "failed to exchange token", http.StatusInternalServerError)
		return
	}

	httpClient := provider.Client(ctx, token)

	verified, err := provider.Verify(hLog.Session("verify"), httpClient)
	if err != nil {
		hLog.Error("failed-to-verify-token", err)
		http.Error(w, "failed to verify token", http.StatusInternalServerError)
		return
	}

	if !verified {
		hLog.Info("verification-failed")
		http.Error(w, "verification failed", http.StatusUnauthorized)
		return
	}

	exp := time.Now().Add(handler.expire)

	csrfToken, err := handler.csrfTokenGenerator.GenerateToken()
	if err != nil {
		hLog.Error("generate-csrf-token", err)
		http.Error(w, "failed to generate csrf token", http.StatusInternalServerError)
		return
	}

	tokenType, signedToken, err := handler.authTokenGenerator.GenerateToken(exp, team.Name(), team.Admin(), csrfToken)
	if err != nil {
		hLog.Error("failed-to-sign-token", err)
		http.Error(w, "failed to generate auth token", http.StatusInternalServerError)
		return
	}

	tokenStr := string(tokenType) + " " + string(signedToken)

	authCookie := &http.Cookie{
		Name:     AuthCookieName,
		Value:    tokenStr,
		Path:     "/",
		Expires:  exp,
		HttpOnly: true,
	}
	if handler.isTLSEnabled {
		authCookie.Secure = true
	}
	// TODO: Add SameSite once Golang supports it
	// https://github.com/golang/go/issues/15867
	http.SetCookie(w, authCookie)

	// Deletes the oauth state cookie to avoid CSRF attacks
	http.SetCookie(w, &http.Cookie{
		Name:   cookieState.Name,
		Path:   "/",
		MaxAge: -1,
	})

	w.Header().Set(CSRFHeaderName, csrfToken)

	if oauthState.Redirect != "" && !strings.HasPrefix(oauthState.Redirect, "/") {
		hLog.Info("invalid-redirect")
		http.Error(w, "invalid redirect", http.StatusBadRequest)
		return
	}

	if oauthState.Redirect != "" {
		redirectURL, err := url.Parse(oauthState.Redirect)
		if err != nil {
			hLog.Info("invalid-redirect")
			http.Error(w, "invalid redirect", http.StatusBadRequest)
			return
		}
		queryParams := redirectURL.Query()
		queryParams.Set("csrf_token", csrfToken)
		redirectURL.RawQuery = queryParams.Encode()
		http.Redirect(w, r, redirectURL.String(), http.StatusTemporaryRedirect)
		return
	}

	if oauthState.FlyLocalPort == "" {
		// Old login flow
		fmt.Fprintln(w, tokenStr)
	} else {
		encodedToken := url.QueryEscape(tokenStr)
		http.Redirect(w, r, fmt.Sprintf("http://127.0.0.1:%s/oauth/callback?token=%s", oauthState.FlyLocalPort, encodedToken), http.StatusTemporaryRedirect)
	}
}

type oauthCallbackHandler interface {
	ValidState(cookieState string, paramState string) bool
}

type oauthCallbackHandlerV1 struct{}
type oauthCallbackHandlerV2 struct{}

func (oauthCallbackHandlerV1) ValidState(cookieState string, paramState string) bool {
	return true
}

func (oauthCallbackHandlerV2) ValidState(cookieState string, paramState string) bool {
	return cookieState == paramState
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package kubelet

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"k8s.io/api/core/v1"
	rbac "k8s.io/api/rbac/v1"
	apierrs "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/strategicpatch"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	kubeadmapi "k8s.io/kubernetes/cmd/kubeadm/app/apis/kubeadm"
	kubeadmconstants "k8s.io/kubernetes/cmd/kubeadm/app/constants"
	kubeadmutil "k8s.io/kubernetes/cmd/kubeadm/app/util"
	"k8s.io/kubernetes/cmd/kubeadm/app/util/apiclient"
	kubeconfigutil "k8s.io/kubernetes/cmd/kubeadm/app/util/kubeconfig"
	rbachelper "k8s.io/kubernetes/pkg/apis/rbac/v1"
	kubeletconfigscheme "k8s.io/kubernetes/pkg/kubelet/apis/kubeletconfig/scheme"
	kubeletconfigv1alpha1 "k8s.io/kubernetes/pkg/kubelet/apis/kubeletconfig/v1alpha1"
)

// CreateBaseKubeletConfiguration creates base kubelet configuration for dynamic kubelet configuration feature.
func CreateBaseKubeletConfiguration(cfg *kubeadmapi.MasterConfiguration, client clientset.Interface) error {
	fmt.Printf("[kubelet] Uploading a ConfigMap %q in namespace %s with base configuration for the kubelets in the cluster",
		kubeadmconstants.KubeletBaseConfigurationConfigMap, metav1.NamespaceSystem)

	_, kubeletCodecs, err := kubeletconfigscheme.NewSchemeAndCodecs()
	if err != nil {
		return err
	}
	kubeletBytes, err := kubeadmutil.MarshalToYamlForCodecs(cfg.KubeletConfiguration.BaseConfig, kubeletconfigv1alpha1.SchemeGroupVersion, *kubeletCodecs)
	if err != nil {
		return err
	}

	if err = apiclient.CreateOrUpdateConfigMap(client, &v1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      kubeadmconstants.KubeletBaseConfigurationConfigMap,
			Namespace: metav1.NamespaceSystem,
		},
		Data: map[string]string{
			kubeadmconstants.KubeletBaseConfigurationConfigMapKey: string(kubeletBytes),
		},
	}); err != nil {
		return err
	}

	if err := createKubeletBaseConfigMapRBACRules(client); err != nil {
		return fmt.Errorf("error creating base kubelet configmap RBAC rules: %v", err)
	}

	return updateNodeWithConfigMap(client, cfg.NodeName)
}

// ConsumeBaseKubeletConfiguration consumes base kubelet configuration for dynamic kubelet configuration feature.
func ConsumeBaseKubeletConfiguration(nodeName string) error {
	client, err := getLocalNodeTLSBootstrappedClient()
	if err != nil {
		return err
	}

	kubeletCfg, err := client.CoreV1().ConfigMaps(metav1.NamespaceSystem).Get(kubeadmconstants.KubeletBaseConfigurationConfigMap, metav1.GetOptions{})
	if err != nil {
		return err
	}

	if err := writeInitKubeletConfigToDisk([]byte(kubeletCfg.Data[kubeadmconstants.KubeletBaseConfigurationConfigMapKey])); err != nil {
		return fmt.Errorf("failed to write initial remote configuration of kubelet to disk for node %s: %v", nodeName, err)
	}

	return updateNodeWithConfigMap(client, nodeName)
}

// updateNodeWithConfigMap updates node ConfigSource with KubeletBaseConfigurationConfigMap
func updateNodeWithConfigMap(client clientset.Interface, nodeName string) error {
	fmt.Printf("[kubelet] Using Dynamic Kubelet Config for node %q; config sourced from ConfigMap %q in namespace %s",
		nodeName, kubeadmconstants.KubeletBaseConfigurationConfigMap, metav1.NamespaceSystem)

	// Loop on every falsy return. Return with an error if raised. Exit successfully if true is returned.
	return wait.Poll(kubeadmconstants.APICallRetryInterval, kubeadmconstants.UpdateNodeTimeout, func() (bool, error) {
		node, err := client.CoreV1().Nodes().Get(nodeName, metav1.GetOptions{})
		if err != nil {
			return false, nil
		}

		oldData, err := json.Marshal(node)
		if err != nil {
			return false, err
		}

		kubeletCfg, err := client.CoreV1().ConfigMaps(metav1.NamespaceSystem).Get(kubeadmconstants.KubeletBaseConfigurationConfigMap, metav1.GetOptions{})
		if err != nil {
			return false, nil
		}

		node.Spec.ConfigSource = &v1.NodeConfigSource{
			ConfigMapRef: &v1.ObjectReference{
				Name:      kubeadmconstants.KubeletBaseConfigurationConfigMap,
				Namespace: metav1.NamespaceSystem,
				UID:       kubeletCfg.UID,
			},
		}

		newData, err := json.Marshal(node)
		if err != nil {
			return false, err
		}

		patchBytes, err := strategicpatch.CreateTwoWayMergePatch(oldData, newData, v1.Node{})
		if err != nil {
			return false, err
		}

		if _, err := client.CoreV1().Nodes().Patch(node.Name, types.StrategicMergePatchType, patchBytes); err != nil {
			if apierrs.IsConflict(err) {
				fmt.Println("Temporarily unable to update node metadata due to conflict (will retry)")
				return false, nil
			}
			return false, err
		}

		return true, nil
	})
}

// createKubeletBaseConfigMapRBACRules creates the RBAC rules for exposing the base kubelet ConfigMap in the kube-system namespace to unauthenticated users
func createKubeletBaseConfigMapRBACRules(client clientset.Interface) error {
	if err := apiclient.CreateOrUpdateRole(client, &rbac.Role{
		ObjectMeta: metav1.ObjectMeta{
			Name:      kubeadmconstants.KubeletBaseConfigMapRoleName,
			Namespace: metav1.NamespaceSystem,
		},
		Rules: []rbac.PolicyRule{
			rbachelper.NewRule("get").Groups("").Resources("configmaps").Names(kubeadmconstants.KubeletBaseConfigurationConfigMap).RuleOrDie(),
		},
	}); err != nil {
		return err
	}

	return apiclient.CreateOrUpdateRoleBinding(client, &rbac.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Name:      kubeadmconstants.KubeletBaseConfigMapRoleName,
			Namespace: metav1.NamespaceSystem,
		},
		RoleRef: rbac.RoleRef{
			APIGroup: rbac.GroupName,
			Kind:     "Role",
			Name:     kubeadmconstants.KubeletBaseConfigMapRoleName,
		},
		Subjects: []rbac.Subject{
			{
				Kind: rbac.GroupKind,
				Name: kubeadmconstants.NodesGroup,
			},
			{
				Kind: rbac.GroupKind,
				Name: kubeadmconstants.NodeBootstrapTokenAuthGroup,
			},
		},
	})
}

// getLocalNodeTLSBootstrappedClient waits for the kubelet to perform the TLS bootstrap
// and then creates a client from config file /etc/kubernetes/kubelet.conf
func getLocalNodeTLSBootstrappedClient() (clientset.Interface, error) {
	fmt.Println("[tlsbootstrap] Waiting for the kubelet to perform the TLS Bootstrap...")

	kubeletKubeConfig := filepath.Join(kubeadmconstants.KubernetesDir, kubeadmconstants.KubeletKubeConfigFileName)

	// Loop on every falsy return. Return with an error if raised. Exit successfully if true is returned.
	err := wait.PollImmediateInfinite(kubeadmconstants.APICallRetryInterval, func() (bool, error) {
		_, err := os.Stat(kubeletKubeConfig)
		return (err == nil), nil
	})
	if err != nil {
		return nil, err
	}

	return kubeconfigutil.ClientSetFromFile(kubeletKubeConfig)
}

// WriteInitKubeletConfigToDiskOnMaster writes base kubelet configuration to disk on master.
func WriteInitKubeletConfigToDiskOnMaster(cfg *kubeadmapi.MasterConfiguration) error {
	fmt.Printf("[kubelet] Writing base configuration of kubelets to disk on master node %s", cfg.NodeName)

	_, kubeletCodecs, err := kubeletconfigscheme.NewSchemeAndCodecs()
	if err != nil {
		return err
	}

	kubeletBytes, err := kubeadmutil.MarshalToYamlForCodecs(cfg.KubeletConfiguration.BaseConfig, kubeletconfigv1alpha1.SchemeGroupVersion, *kubeletCodecs)
	if err != nil {
		return err
	}

	if err := writeInitKubeletConfigToDisk(kubeletBytes); err != nil {
		return fmt.Errorf("failed to write base configuration of kubelet to disk on master node %s: %v", cfg.NodeName, err)
	}

	return nil
}

func writeInitKubeletConfigToDisk(kubeletConfig []byte) error {
	baseCongfigFile := filepath.Join(kubeadmconstants.KubeletBaseConfigurationDir, kubeadmconstants.KubeletBaseConfigurationFile)
	if err := ioutil.WriteFile(baseCongfigFile, kubeletConfig, 0644); err != nil {
		return fmt.Errorf("failed to write initial remote configuration of kubelet into file %q: %v", baseCongfigFile, err)
	}
	return nil
}
<commit_msg>Should make dir before writing file.<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package kubelet

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"k8s.io/api/core/v1"
	rbac "k8s.io/api/rbac/v1"
	apierrs "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/strategicpatch"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	kubeadmapi "k8s.io/kubernetes/cmd/kubeadm/app/apis/kubeadm"
	kubeadmconstants "k8s.io/kubernetes/cmd/kubeadm/app/constants"
	kubeadmutil "k8s.io/kubernetes/cmd/kubeadm/app/util"
	"k8s.io/kubernetes/cmd/kubeadm/app/util/apiclient"
	kubeconfigutil "k8s.io/kubernetes/cmd/kubeadm/app/util/kubeconfig"
	rbachelper "k8s.io/kubernetes/pkg/apis/rbac/v1"
	kubeletconfigscheme "k8s.io/kubernetes/pkg/kubelet/apis/kubeletconfig/scheme"
	kubeletconfigv1alpha1 "k8s.io/kubernetes/pkg/kubelet/apis/kubeletconfig/v1alpha1"
)

// CreateBaseKubeletConfiguration creates base kubelet configuration for dynamic kubelet configuration feature.
func CreateBaseKubeletConfiguration(cfg *kubeadmapi.MasterConfiguration, client clientset.Interface) error {
	fmt.Printf("[kubelet] Uploading a ConfigMap %q in namespace %s with base configuration for the kubelets in the cluster",
		kubeadmconstants.KubeletBaseConfigurationConfigMap, metav1.NamespaceSystem)

	_, kubeletCodecs, err := kubeletconfigscheme.NewSchemeAndCodecs()
	if err != nil {
		return err
	}
	kubeletBytes, err := kubeadmutil.MarshalToYamlForCodecs(cfg.KubeletConfiguration.BaseConfig, kubeletconfigv1alpha1.SchemeGroupVersion, *kubeletCodecs)
	if err != nil {
		return err
	}

	if err = apiclient.CreateOrUpdateConfigMap(client, &v1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      kubeadmconstants.KubeletBaseConfigurationConfigMap,
			Namespace: metav1.NamespaceSystem,
		},
		Data: map[string]string{
			kubeadmconstants.KubeletBaseConfigurationConfigMapKey: string(kubeletBytes),
		},
	}); err != nil {
		return err
	}

	if err := createKubeletBaseConfigMapRBACRules(client); err != nil {
		return fmt.Errorf("error creating base kubelet configmap RBAC rules: %v", err)
	}

	return updateNodeWithConfigMap(client, cfg.NodeName)
}

// ConsumeBaseKubeletConfiguration consumes base kubelet configuration for dynamic kubelet configuration feature.
func ConsumeBaseKubeletConfiguration(nodeName string) error {
	client, err := getLocalNodeTLSBootstrappedClient()
	if err != nil {
		return err
	}

	kubeletCfg, err := client.CoreV1().ConfigMaps(metav1.NamespaceSystem).Get(kubeadmconstants.KubeletBaseConfigurationConfigMap, metav1.GetOptions{})
	if err != nil {
		return err
	}

	if err := writeInitKubeletConfigToDisk([]byte(kubeletCfg.Data[kubeadmconstants.KubeletBaseConfigurationConfigMapKey])); err != nil {
		return fmt.Errorf("failed to write initial remote configuration of kubelet to disk for node %s: %v", nodeName, err)
	}

	return updateNodeWithConfigMap(client, nodeName)
}

// updateNodeWithConfigMap updates node ConfigSource with KubeletBaseConfigurationConfigMap
func updateNodeWithConfigMap(client clientset.Interface, nodeName string) error {
	fmt.Printf("[kubelet] Using Dynamic Kubelet Config for node %q; config sourced from ConfigMap %q in namespace %s",
		nodeName, kubeadmconstants.KubeletBaseConfigurationConfigMap, metav1.NamespaceSystem)

	// Loop on every falsy return. Return with an error if raised. Exit successfully if true is returned.
	return wait.Poll(kubeadmconstants.APICallRetryInterval, kubeadmconstants.UpdateNodeTimeout, func() (bool, error) {
		node, err := client.CoreV1().Nodes().Get(nodeName, metav1.GetOptions{})
		if err != nil {
			return false, nil
		}

		oldData, err := json.Marshal(node)
		if err != nil {
			return false, err
		}

		kubeletCfg, err := client.CoreV1().ConfigMaps(metav1.NamespaceSystem).Get(kubeadmconstants.KubeletBaseConfigurationConfigMap, metav1.GetOptions{})
		if err != nil {
			return false, nil
		}

		node.Spec.ConfigSource = &v1.NodeConfigSource{
			ConfigMapRef: &v1.ObjectReference{
				Name:      kubeadmconstants.KubeletBaseConfigurationConfigMap,
				Namespace: metav1.NamespaceSystem,
				UID:       kubeletCfg.UID,
			},
		}

		newData, err := json.Marshal(node)
		if err != nil {
			return false, err
		}

		patchBytes, err := strategicpatch.CreateTwoWayMergePatch(oldData, newData, v1.Node{})
		if err != nil {
			return false, err
		}

		if _, err := client.CoreV1().Nodes().Patch(node.Name, types.StrategicMergePatchType, patchBytes); err != nil {
			if apierrs.IsConflict(err) {
				fmt.Println("Temporarily unable to update node metadata due to conflict (will retry)")
				return false, nil
			}
			return false, err
		}

		return true, nil
	})
}

// createKubeletBaseConfigMapRBACRules creates the RBAC rules for exposing the base kubelet ConfigMap in the kube-system namespace to unauthenticated users
func createKubeletBaseConfigMapRBACRules(client clientset.Interface) error {
	if err := apiclient.CreateOrUpdateRole(client, &rbac.Role{
		ObjectMeta: metav1.ObjectMeta{
			Name:      kubeadmconstants.KubeletBaseConfigMapRoleName,
			Namespace: metav1.NamespaceSystem,
		},
		Rules: []rbac.PolicyRule{
			rbachelper.NewRule("get").Groups("").Resources("configmaps").Names(kubeadmconstants.KubeletBaseConfigurationConfigMap).RuleOrDie(),
		},
	}); err != nil {
		return err
	}

	return apiclient.CreateOrUpdateRoleBinding(client, &rbac.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Name:      kubeadmconstants.KubeletBaseConfigMapRoleName,
			Namespace: metav1.NamespaceSystem,
		},
		RoleRef: rbac.RoleRef{
			APIGroup: rbac.GroupName,
			Kind:     "Role",
			Name:     kubeadmconstants.KubeletBaseConfigMapRoleName,
		},
		Subjects: []rbac.Subject{
			{
				Kind: rbac.GroupKind,
				Name: kubeadmconstants.NodesGroup,
			},
			{
				Kind: rbac.GroupKind,
				Name: kubeadmconstants.NodeBootstrapTokenAuthGroup,
			},
		},
	})
}

// getLocalNodeTLSBootstrappedClient waits for the kubelet to perform the TLS bootstrap
// and then creates a client from config file /etc/kubernetes/kubelet.conf
func getLocalNodeTLSBootstrappedClient() (clientset.Interface, error) {
	fmt.Println("[tlsbootstrap] Waiting for the kubelet to perform the TLS Bootstrap...")

	kubeletKubeConfig := filepath.Join(kubeadmconstants.KubernetesDir, kubeadmconstants.KubeletKubeConfigFileName)

	// Loop on every falsy return. Return with an error if raised. Exit successfully if true is returned.
	err := wait.PollImmediateInfinite(kubeadmconstants.APICallRetryInterval, func() (bool, error) {
		_, err := os.Stat(kubeletKubeConfig)
		return (err == nil), nil
	})
	if err != nil {
		return nil, err
	}

	return kubeconfigutil.ClientSetFromFile(kubeletKubeConfig)
}

// WriteInitKubeletConfigToDiskOnMaster writes base kubelet configuration to disk on master.
func WriteInitKubeletConfigToDiskOnMaster(cfg *kubeadmapi.MasterConfiguration) error {
	fmt.Printf("[kubelet] Writing base configuration of kubelets to disk on master node %s", cfg.NodeName)

	_, kubeletCodecs, err := kubeletconfigscheme.NewSchemeAndCodecs()
	if err != nil {
		return err
	}

	kubeletBytes, err := kubeadmutil.MarshalToYamlForCodecs(cfg.KubeletConfiguration.BaseConfig, kubeletconfigv1alpha1.SchemeGroupVersion, *kubeletCodecs)
	if err != nil {
		return err
	}

	if err := writeInitKubeletConfigToDisk(kubeletBytes); err != nil {
		return fmt.Errorf("failed to write base configuration of kubelet to disk on master node %s: %v", cfg.NodeName, err)
	}

	return nil
}

func writeInitKubeletConfigToDisk(kubeletConfig []byte) error {
	if err := os.MkdirAll(kubeadmconstants.KubeletBaseConfigurationDir, 0644); err != nil {
		return fmt.Errorf("failed to create directory %q: %v", kubeadmconstants.KubeletBaseConfigurationDir, err)
	}
	baseConfigFile := filepath.Join(kubeadmconstants.KubeletBaseConfigurationDir, kubeadmconstants.KubeletBaseConfigurationFile)
	if err := ioutil.WriteFile(baseConfigFile, kubeletConfig, 0644); err != nil {
		return fmt.Errorf("failed to write initial remote configuration of kubelet into file %q: %v", baseConfigFile, err)
	}
	return nil
}
<|endoftext|>


<commit_before>package anonlib

import (
	"encoding/binary"
	"fmt"
	"log"
	"net"
	"sync"

	"../common"
)

type anonConnection struct {
	sync.Mutex
}

const (
	TYPE_ACK byte = iota
	TYPE_DATA
	TYPE_BEAT
)

//Some Global variables
var (
	BindAddr string
	BindPort string
)

// GetTheClientData builds the frame-id:bool value byte string
func BuildClientFrameData(framid string, enable bool) []byte {

	enableStr := "0"
	if enable == true {
		enableStr = "1"
	}
	tempStr := framid + ":" + enableStr
	return []byte(tempStr)

}

// BuildTheClientDataFromMap locks and read the map and build calls the BuildClientFrameData
func BuildTheClientDataFromMap(M map[string]bool) []byte {

	common.ToAnon.Lck.Lock()
	defer common.ToAnon.Lck.Unlock()

	result := []byte{}
	for key, value := range M {
		result = append(result, BuildClientFrameData(key, value)...)
	}

	return result
}

func GenerateDataMsg(msgType byte, data []byte, totalMsgCount int) []byte {

	//calculate the total byte lenght
	totalLength := int(len(data) + 4 + 4 + 1)
	//create a dummy byte array
	resultByteArray := make([]byte, totalLength)

	//encode msg type in the first byte
	resultByteArray[0] = msgType
	//encode the total msg length
	binary.BigEndian.PutUint32(resultByteArray[1:5], uint32(len(data)+(totalMsgCount*2)))
	//encode the total frames id sent
	binary.BigEndian.PutUint32(resultByteArray[5:9], uint32(totalMsgCount))
	//next start from 9th index
	copy(resultByteArray[9:], data)
	return resultByteArray
}

func NewaAnonConnection() *anonConnection {
	return &anonConnection{}
}

func BindAndStartListener() {
	//Bind to all the network inteface int the system on the same port
	tcpListener, err := net.Listen("tcp", ":"+BindPort)
	if err != nil {
		fmt.Println(" BindAndStartListener", err)
	}
	HandleClientConnection(tcpListener)
}

//listen for incommng connection
func HandleClientConnection(tcpListener net.Listener) {

	for {
		conn, err := tcpListener.Accept()
		if err != nil {
			fmt.Println("Unable to accept connection")
		}
		newAnonClientConn := NewaAnonConnection()
		go newAnonClientConn.SendMsg(conn)
		go newAnonClientConn.RecvMsg(conn)
	}
}

// SendMsg will wait on the Common.ToAnon
func (annonClient *anonConnection) SendMsg(conn net.Conn) {
	for {
		<-common.ToAnon.Ch
		result := BuildTheClientDataFromMap(common.ToAnon.M)
		result = GenerateDataMsg(TYPE_DATA, result, len(common.ToAnon.M))
		fmt.Println(result, "sending to cleint")
		annonClient.Lock()
		n, err := conn.Write(result)
		if err != nil {
			fmt.Println("Unable to send data to client", err)
			annonClient.Unlock()
			return
		}
		log.Printf("Sent %d bytes", n)
		annonClient.Unlock()
	}

}

func (annonClient *anonConnection) RecvMsg(conn net.Conn) {

	for {
		localClientBuf := make([]byte, 4096)
		//defer annonClient.Unlock()
		_, err := conn.Read(localClientBuf)
		if err != nil {
			fmt.Println("Unable to read from the client", err)
			return
		}
		msgType := localClientBuf[0:1]
		if msgType[0] == TYPE_BEAT {
			annonClient.Lock()
			annonClient.SendAck(conn)
			annonClient.Unlock()
		} else {
			fmt.Println("Unknown msg from client", msgType)
		}
	}

}

func (annonClient *anonConnection) SendAck(conn net.Conn) {
	_, err := conn.Write([]byte{TYPE_ACK})
	if err != nil {
		fmt.Println("Failed to send ack", err)
	}

}

func Run(inBindAddress string, inBindPort string) {
	BindAddr = inBindAddress
	BindPort = inBindPort

	go BindAndStartListener()
}
<commit_msg>Changed information logging to use log package<commit_after>package anonlib

import (
	"encoding/binary"
	"log"
	"net"
	"sync"

	"../common"
)

type anonConnection struct {
	sync.Mutex
}

const (
	TYPE_ACK byte = iota
	TYPE_DATA
	TYPE_BEAT
)

//Some Global variables
var (
	BindAddr string
	BindPort string
)

// GetTheClientData builds the frame-id:bool value byte string
func BuildClientFrameData(framid string, enable bool) []byte {

	enableStr := "0"
	if enable == true {
		enableStr = "1"
	}
	tempStr := framid + ":" + enableStr
	return []byte(tempStr)

}

// BuildTheClientDataFromMap locks and read the map and build calls the BuildClientFrameData
func BuildTheClientDataFromMap(M map[string]bool) []byte {

	common.ToAnon.Lck.Lock()
	defer common.ToAnon.Lck.Unlock()

	result := []byte{}
	for key, value := range M {
		result = append(result, BuildClientFrameData(key, value)...)
	}

	return result
}

func GenerateDataMsg(msgType byte, data []byte, totalMsgCount int) []byte {

	//calculate the total byte lenght
	totalLength := int(len(data) + 4 + 4 + 1)
	//create a dummy byte array
	resultByteArray := make([]byte, totalLength)

	//encode msg type in the first byte
	resultByteArray[0] = msgType
	//encode the total msg length
	binary.BigEndian.PutUint32(resultByteArray[1:5], uint32(len(data)+(totalMsgCount*2)))
	//encode the total frames id sent
	binary.BigEndian.PutUint32(resultByteArray[5:9], uint32(totalMsgCount))
	//next start from 9th index
	copy(resultByteArray[9:], data)
	return resultByteArray
}

func NewaAnonConnection() *anonConnection {
	return &anonConnection{}
}

func BindAndStartListener() {
	//Bind to all the network inteface int the system on the same port
	tcpListener, err := net.Listen("tcp", ":"+BindPort)
	if err != nil {
		log.Println(" BindAndStartListener", err)
	}
	HandleClientConnection(tcpListener)
}

//listen for incommng connection
func HandleClientConnection(tcpListener net.Listener) {

	for {
		conn, err := tcpListener.Accept()
		if err != nil {
			log.Println("Unable to accept connection")
		}
		newAnonClientConn := NewaAnonConnection()
		go newAnonClientConn.SendMsg(conn)
		go newAnonClientConn.RecvMsg(conn)
	}
}

// SendMsg will wait on the Common.ToAnon
func (annonClient *anonConnection) SendMsg(conn net.Conn) {
	for {
		<-common.ToAnon.Ch
		result := BuildTheClientDataFromMap(common.ToAnon.M)
		result = GenerateDataMsg(TYPE_DATA, result, len(common.ToAnon.M))
		log.Println(result, "sending to client")
		annonClient.Lock()
		n, err := conn.Write(result)
		if err != nil {
			log.Println("Unable to send data to client", err)
			annonClient.Unlock()
			return
		}
		log.Printf("Sent %d bytes", n)
		annonClient.Unlock()
	}

}

func (annonClient *anonConnection) RecvMsg(conn net.Conn) {

	for {
		localClientBuf := make([]byte, 4096)
		//defer annonClient.Unlock()
		_, err := conn.Read(localClientBuf)
		if err != nil {
			log.Println("Unable to read from the client", err)
			return
		}
		msgType := localClientBuf[0:1]
		if msgType[0] == TYPE_BEAT {
			annonClient.Lock()
			annonClient.SendAck(conn)
			annonClient.Unlock()
		} else {
			log.Println("Unknown msg from client", msgType)
		}
	}

}

func (annonClient *anonConnection) SendAck(conn net.Conn) {
	_, err := conn.Write([]byte{TYPE_ACK})
	if err != nil {
		log.Println("Failed to send ack", err)
	}

}

func Run(inBindAddress string, inBindPort string) {
	BindAddr = inBindAddress
	BindPort = inBindPort

	go BindAndStartListener()
}
<|endoftext|>


<commit_before>// Daedalean-specific configuration file, defining a bunch of constants that are company-specific. Create your own and
// mark it with your build tag, then remove the !ddln tag below.
// +build ddln !ddln

package config

// Project name
const ProjectName = "Reqtraq"

type RequirementLevel int

// Requirement levels according to DO-178C (do not change!)
const (
	SYSTEM RequirementLevel = iota
	HIGH
	LOW
	CODE
)

// Document types:
// ORD - Overall (aka System) Requirement Document
// SRD - Software Requirements Data
// SDD - Software Design Description
// HRD - Hardware Requirements Data
// HDD - Hardware Design Description

// Requirement types:
// SYS - System/overall requirements (defined in ORD documents)
// SWH - Sofware high-level requirements (defined in SRD documents)
// SWL - Software low-level requirements (defined in SDD documents)
// HWH - Hardware high-level requirements (defined in HRD documents)
// HWL - Hardware low-level requirements (defined in HDD documents)

// Map from requirement type to requirement level.
var ReqTypeToReqLevel = map[string]RequirementLevel{
	"SYS": SYSTEM,
	"SWH": HIGH,
	"HWH": HIGH,
	"SWL": LOW,
	"HWL": LOW,
}

// Map from document type to requirement type.
var DocTypeToReqType = map[string]string{
	"ORD": "SYS",
	"SRD": "SWH",
	"HRD": "HWH",
	"SDD": "SWL",
	"HDD": "HWL",
}

// Map from requirement type to document type.
var ReqTypeToDocType = map[string]string{
	"SYS": "ORD",
	"SWH": "SRD",
	"SWL": "SDD",
	"HWH": "HRD",
	"HWL": "HDD",
}

// Map from document type to document ID.
// https://a.daedalean.ai/organisation-of-documentation
var DocTypeToDocId = map[string]string{
	"H":       "0",
	"DS":      "1",
	"CLSRS":   "5",
	"RS":      "6",
	"SDS":     "7",
	"CS":      "8",
	"HRS":     "9",
	"HCS":     "10",
	"HDS":     "11",
	"HVVS":    "12",
	"HAS":     "13",
	"HCMS":    "14",
	"PDAS":    "15",
	"CMP":     "20",
	"CLCMP":   "21",
	"PAP":     "22",
	"CLPAP":   "23",
	"CLTQP":   "25",
	"SCMP":    "26",
	"CLSCMP":  "27",
	"SQAP":    "28",
	"CLSQAP":  "29",
	"SDP":     "30",
	"HPAP":    "32",
	"CLHPAP":  "33",
	"TPPR":    "50",
	"TPSQAR":  "51",
	"ORD":     "100",
	"ICD":     "101",
	"CP":      "102",
	"DP":      "103",
	"DD":      "104",
	"VAP":     "105",
	"VEP":     "106",
	"CI":      "107",
	"FHA":     "108",
	"SFHA":    "109",
	"PSSA":    "110",
	"SSA":     "111",
	"CCA":     "112",
	"SPP":     "113",
	"VAD":     "114",
	"VED":     "115",
	"ECM":     "116",
	"EPA":     "117",
	"CSCR":    "118",
	"PSAC":    "134",
	"TQP":     "135",
	"SVP":     "136",
	"SRD":     "137",
	"SDD":     "138",
	"SVCP":    "139",
	"SVR":     "140",
	"SLECI":   "141",
	"SCI":     "142",
	"SAS":     "143",
	"STD":     "144",
	"SQARI":   "145",
	"PHAC":    "167",
	"HDP":     "168",
	"HVEP":    "169",
	"ECMP":    "170",
	"HVAP":    "171",
	"HCMP":    "172",
	"HECI":    "173",
	"HCI":     "174",
	"HRD":     "175",
	"HDD":     "176",
	"HTD":     "177",
	"HRAP":    "178",
	"HRAR":    "179",
	"HTP":     "180",
	"HTR":     "181",
	"HATC":    "182",
	"HACS":    "183",
	"FFPA":    "184",
	"TPORD":   "200",
	"TPICD":   "201",
	"TPCP":    "202",
	"TPDP":    "203",
	"TPDD":    "204",
	"TPVAP":   "205",
	"TPVEP":   "206",
	"TPCI":    "207",
	"TPFHA":   "208",
	"TPSFHA":  "209",
	"TPPSSA":  "210",
	"TPSSA":   "211",
	"TPCCA":   "212",
	"TPSPP":   "213",
	"TPVAD":   "214",
	"TPVED":   "215",
	"TPECM":   "216",
	"TPEPA":   "217",
	"TPCSCR":  "218",
	"TPPSAC":  "234",
	"TPSDP":   "235",
	"TPSVP":   "236",
	"TPSRD":   "237",
	"TPSDD":   "238",
	"TPSVCP":  "239",
	"TPSVR":   "240",
	"TPSLECI": "241",
	"TPSCI":   "242",
	"TPSAS":   "243",
	"TPSTD":   "244",
	"TPPHAC":  "267",
	"TPHDP":   "268",
	"TPHVEP":  "269",
	"TPECMP":  "270",
	"TPHVAP":  "271",
	"TPHCMP":  "272",
	"TPHECI":  "273",
	"TPHCI":   "274",
	"TPHRD":   "275",
	"TPHDD":   "276",
	"TPHTD":   "277",
	"TPHRAP":  "278",
	"TPHRAR":  "279",
	"TPHTP":   "280",
	"TPHTR":   "281",
	"TPHATC":  "282",
	"TPHACS":  "283",
	"TPFFPA":  "284",
	"CLORD":   "300",
	"CLICD":   "301",
	"CLCP":    "302",
	"CLDP":    "303",
	"CLDD":    "304",
	"CLVAP":   "305",
	"CLVEP":   "306",
	"CLCI":    "307",
	"CLFHA":   "308",
	"CLSFHA":  "309",
	"CLPSSA":  "310",
	"CLSSA":   "311",
	"CLCCA":   "312",
	"CLSPP":   "313",
	"CLVAD":   "314",
	"CLVED":   "315",
	"CLECM":   "316",
	"CLEPA":   "317",
	"CLCSCR":  "318",
	"CLPSAC":  "334",
	"CLSDP":   "335",
	"CLSVP":   "336",
	"CLSRD":   "337",
	"CLSDD":   "338",
	"CLSVCP":  "339",
	"CLSVR":   "340",
	"CLSLECI": "341",
	"CLSCI":   "342",
	"CLSAS":   "343",
	"CLSTD":   "344",
	"CLSOI1":  "345",
	"CLSOI2":  "346",
	"CLSOI3":  "347",
	"CLSOI4":  "348",
	"CLSPR":   "349",
	"CLRA":    "350",
	"CLSCR":   "351",
	"CLPHAC":  "367",
	"CLHDP":   "368",
	"CLHVEP":  "369",
	"CLECMP":  "370",
	"CLHVAP":  "371",
	"CLHCMP":  "372",
	"CLHECI":  "373",
	"CLHCI":   "374",
	"CLHRD":   "375",
	"CLHDD":   "376",
	"CLHTD":   "377",
	"CLHRAP":  "378",
	"CLHRAR":  "379",
	"CLHTP":   "380",
	"CLHTR":   "381",
	"CLHATC":  "382",
	"CLHACS":  "383",
	"CLFFPA":  "384",
}
<commit_msg>Updating allowable document names to reflect verification needs<commit_after>// Daedalean-specific configuration file, defining a bunch of constants that are company-specific. Create your own and
// mark it with your build tag, then remove the !ddln tag below.
// +build ddln !ddln

package config

// Project name
const ProjectName = "Reqtraq"

type RequirementLevel int

// Requirement levels according to DO-178C (do not change!)
const (
	SYSTEM RequirementLevel = iota
	HIGH
	LOW
	CODE
)

// Document types:
// ORD - Overall (aka System) Requirement Document
// SRD - Software Requirements Data
// SDD - Software Design Description
// HRD - Hardware Requirements Data
// HDD - Hardware Design Description

// Requirement types:
// SYS - System/overall requirements (defined in ORD documents)
// SWH - Sofware high-level requirements (defined in SRD documents)
// SWL - Software low-level requirements (defined in SDD documents)
// HWH - Hardware high-level requirements (defined in HRD documents)
// HWL - Hardware low-level requirements (defined in HDD documents)

// Map from requirement type to requirement level.
var ReqTypeToReqLevel = map[string]RequirementLevel{
	"SYS": SYSTEM,
	"SWH": HIGH,
	"HWH": HIGH,
	"SWL": LOW,
	"HWL": LOW,
}

// Map from document type to requirement type.
var DocTypeToReqType = map[string]string{
	"ORD": "SYS",
	"SRD": "SWH",
	"HRD": "HWH",
	"SDD": "SWL",
	"HDD": "HWL",
}

// Map from requirement type to document type.
var ReqTypeToDocType = map[string]string{
	"SYS": "ORD",
	"SWH": "SRD",
	"SWL": "SDD",
	"HWH": "HRD",
	"HWL": "HDD",
}

// Map from document type to document ID.
// https://a.daedalean.ai/organisation-of-documentation
var DocTypeToDocId = map[string]string{
	"H":       "0",
	"DS":      "1",
	"CLSRS":   "5",
	"RS":      "6",
	"SDS":     "7",
	"CS":      "8",
	"HRS":     "9",
	"HCS":     "10",
	"HDS":     "11",
	"HVVS":    "12",
	"HAS":     "13",
	"HCMS":    "14",
	"PDAS":    "15",
	"CMP":     "20",
	"CLCMP":   "21",
	"PAP":     "22",
	"CLPAP":   "23",
	"CLTQP":   "25",
	"SCMP":    "26",
	"CLSCMP":  "27",
	"SQAP":    "28",
	"CLSQAP":  "29",
	"SDP":     "30",
	"HPAP":    "32",
	"CLHPAP":  "33",
	"TPPR":    "50",
	"TPSQAR":  "51",
	"ORD":     "100",
	"ICD":     "101",
	"CP":      "102",
	"DP":      "103",
	"DD":      "104",
	"VAP":     "105",
	"VEP":     "106",
	"CI":      "107",
	"FHA":     "108",
	"SFHA":    "109",
	"PSSA":    "110",
	"SSA":     "111",
	"CCA":     "112",
	"SPP":     "113",
	"VAD":     "114",
	"VED":     "115",
	"ECM":     "116",
	"EPA":     "117",
	"CSCR":    "118",
	"PSAC":    "134",
	"TQP":     "135",
	"SVP":     "136",
	"SRD":     "137",
	"SDD":     "138",
	"SLECI":   "141",
	"SCI":     "142",
	"SAS":     "143",
	"STD":     "144",
	"SQARI":   "145",
	"SPVP":    "146",
	"SVVDVP":  "147",
	"STCP":    "148",
	"SPVR":    "149",
	"SVVDVR":  "150",
	"STR":     "151",
	"PHAC":    "167",
	"HDP":     "168",
	"HVEP":    "169",
	"ECMP":    "170",
	"HVAP":    "171",
	"HCMP":    "172",
	"HECI":    "173",
	"HCI":     "174",
	"HRD":     "175",
	"HDD":     "176",
	"HTD":     "177",
	"HRAP":    "178",
	"HRAR":    "179",
	"HTP":     "180",
	"HTR":     "181",
	"HATC":    "182",
	"HACS":    "183",
	"FFPA":    "184",
	"TPORD":   "200",
	"TPICD":   "201",
	"TPCP":    "202",
	"TPDP":    "203",
	"TPDD":    "204",
	"TPVAP":   "205",
	"TPVEP":   "206",
	"TPCI":    "207",
	"TPFHA":   "208",
	"TPSFHA":  "209",
	"TPPSSA":  "210",
	"TPSSA":   "211",
	"TPCCA":   "212",
	"TPSPP":   "213",
	"TPVAD":   "214",
	"TPVED":   "215",
	"TPECM":   "216",
	"TPEPA":   "217",
	"TPCSCR":  "218",
	"TPPSAC":  "234",
	"TPSDP":   "235",
	"TPSVP":   "236",
	"TPSRD":   "237",
	"TPSDD":   "238",
	"TPSVCP":  "239",
	"TPSVR":   "240",
	"TPSLECI": "241",
	"TPSCI":   "242",
	"TPSAS":   "243",
	"TPSTD":   "244",
	"TPPHAC":  "267",
	"TPHDP":   "268",
	"TPHVEP":  "269",
	"TPECMP":  "270",
	"TPHVAP":  "271",
	"TPHCMP":  "272",
	"TPHECI":  "273",
	"TPHCI":   "274",
	"TPHRD":   "275",
	"TPHDD":   "276",
	"TPHTD":   "277",
	"TPHRAP":  "278",
	"TPHRAR":  "279",
	"TPHTP":   "280",
	"TPHTR":   "281",
	"TPHATC":  "282",
	"TPHACS":  "283",
	"TPFFPA":  "284",
	"CLORD":   "300",
	"CLICD":   "301",
	"CLCP":    "302",
	"CLDP":    "303",
	"CLDD":    "304",
	"CLVAP":   "305",
	"CLVEP":   "306",
	"CLCI":    "307",
	"CLFHA":   "308",
	"CLSFHA":  "309",
	"CLPSSA":  "310",
	"CLSSA":   "311",
	"CLCCA":   "312",
	"CLSPP":   "313",
	"CLVAD":   "314",
	"CLVED":   "315",
	"CLECM":   "316",
	"CLEPA":   "317",
	"CLCSCR":  "318",
	"CLPSAC":  "334",
	"CLSDP":   "335",
	"CLSVP":   "336",
	"CLSRD":   "337",
	"CLSDD":   "338",
	"CLSVCP":  "339",
	"CLSVR":   "340",
	"CLSLECI": "341",
	"CLSCI":   "342",
	"CLSAS":   "343",
	"CLSTD":   "344",
	"CLSOI1":  "345",
	"CLSOI2":  "346",
	"CLSOI3":  "347",
	"CLSOI4":  "348",
	"CLSPR":   "349",
	"CLRA":    "350",
	"CLSCR":   "351",
	"CLPHAC":  "367",
	"CLHDP":   "368",
	"CLHVEP":  "369",
	"CLECMP":  "370",
	"CLHVAP":  "371",
	"CLHCMP":  "372",
	"CLHECI":  "373",
	"CLHCI":   "374",
	"CLHRD":   "375",
	"CLHDD":   "376",
	"CLHTD":   "377",
	"CLHRAP":  "378",
	"CLHRAR":  "379",
	"CLHTP":   "380",
	"CLHTR":   "381",
	"CLHATC":  "382",
	"CLHACS":  "383",
	"CLFFPA":  "384",
}
<|endoftext|>


<commit_before>package agent

import (
	"context"
	"github.com/fnproject/fn/api/common"
	"sync"
)

// TODO this should expose:
// * hot containers active
// * memory used / available

// global statistics
type stats struct {
	mu sync.Mutex
	// statistics for all functions combined
	queue    uint64
	running  uint64
	complete uint64
	failed   uint64
	// statistics for individual functions, keyed by function path
	functionStatsMap map[string]*functionStats
}

// statistics for an individual function
type functionStats struct {
	queue    uint64
	running  uint64
	complete uint64
	failed   uint64
}

// Stats hold the statistics for all functions combined
// and the statistics for each individual function
type Stats struct {
	Queue    uint64
	Running  uint64
	Complete uint64
	Failed   uint64
	// statistics for individual functions, keyed by function path
	FunctionStatsMap map[string]*FunctionStats
}

// FunctionStats holds the statistics for an individual function
type FunctionStats struct {
	Queue    uint64
	Running  uint64
	Complete uint64
	Failed   uint64
}

func (s *stats) getStatsForFunction(path string) *functionStats {
	if s.functionStatsMap == nil {
		s.functionStatsMap = make(map[string]*functionStats)
	}
	thisFunctionStats, found := s.functionStatsMap[path]
	if !found {
		thisFunctionStats = &functionStats{}
		s.functionStatsMap[path] = thisFunctionStats
	}

	return thisFunctionStats
}

func (s *stats) Enqueue(ctx context.Context, app string, path string) {
	s.mu.Lock()

	s.queue++
	s.getStatsForFunction(path).queue++
	common.IncrementGauge(ctx, queuedMetricName)

	common.IncrementCounter(ctx, callsMetricName)

	s.mu.Unlock()
}

// Call when a function has been queued but cannot be started because of an error
func (s *stats) Dequeue(ctx context.Context, app string, path string) {
	s.mu.Lock()

	s.queue--
	s.getStatsForFunction(path).queue--
	common.DecrementGauge(ctx, queuedMetricName)

	s.mu.Unlock()
}

func (s *stats) DequeueAndStart(ctx context.Context, app string, path string) {
	s.mu.Lock()

	s.queue--
	s.getStatsForFunction(path).queue--
	common.DecrementGauge(ctx, queuedMetricName)

	s.running++
	s.getStatsForFunction(path).running++
	common.IncrementGauge(ctx, runningMetricName)

	s.mu.Unlock()
}

func (s *stats) Complete(ctx context.Context, app string, path string) {
	s.mu.Lock()

	s.running--
	s.getStatsForFunction(path).running--
	common.DecrementGauge(ctx, runningMetricName)

	s.complete++
	s.getStatsForFunction(path).complete++
	common.IncrementCounter(ctx, completedMetricName)

	s.mu.Unlock()
}

func (s *stats) Failed(ctx context.Context, app string, path string) {
	s.mu.Lock()

	s.running--
	s.getStatsForFunction(path).running--
	common.DecrementGauge(ctx, runningMetricName)

	s.failed++
	s.getStatsForFunction(path).failed++
	common.IncrementCounter(ctx, failedMetricName)

	s.mu.Unlock()
}

func (s *stats) DequeueAndFail(ctx context.Context, app string, path string) {
	s.mu.Lock()

	s.queue--
	s.getStatsForFunction(path).queue--
	common.DecrementGauge(ctx, queuedMetricName)

	s.failed++
	s.getStatsForFunction(path).failed++
	common.IncrementCounter(ctx, failedMetricName)

	s.mu.Unlock()
}

func (s *stats) IncrementTimedout(ctx context.Context) {
	common.IncrementCounter(ctx, timedoutMetricName)
}

func (s *stats) IncrementErrors(ctx context.Context) {
	common.IncrementCounter(ctx, errorsMetricName)
}

func (s *stats) Stats() Stats {
	var stats Stats
	s.mu.Lock()
	stats.Running = s.running
	stats.Complete = s.complete
	stats.Queue = s.queue
	stats.Failed = s.failed
	stats.FunctionStatsMap = make(map[string]*FunctionStats)
	for key, value := range s.functionStatsMap {
		thisFunctionStats := &FunctionStats{Queue: value.queue, Running: value.running, Complete: value.complete, Failed: value.failed}
		stats.FunctionStatsMap[key] = thisFunctionStats
	}
	s.mu.Unlock()
	return stats
}

const (
	queuedMetricName    = "queued"
	callsMetricName     = "calls"
	runningMetricName   = "running"
	completedMetricName = "completed"
	failedMetricName    = "failed"
	timedoutMetricName  = "timedout"
	errorsMetricName    = "errors"
)
<commit_msg>Change timedout to timeouts (#709)<commit_after>package agent

import (
	"context"
	"github.com/fnproject/fn/api/common"
	"sync"
)

// TODO this should expose:
// * hot containers active
// * memory used / available

// global statistics
type stats struct {
	mu sync.Mutex
	// statistics for all functions combined
	queue    uint64
	running  uint64
	complete uint64
	failed   uint64
	// statistics for individual functions, keyed by function path
	functionStatsMap map[string]*functionStats
}

// statistics for an individual function
type functionStats struct {
	queue    uint64
	running  uint64
	complete uint64
	failed   uint64
}

// Stats hold the statistics for all functions combined
// and the statistics for each individual function
type Stats struct {
	Queue    uint64
	Running  uint64
	Complete uint64
	Failed   uint64
	// statistics for individual functions, keyed by function path
	FunctionStatsMap map[string]*FunctionStats
}

// FunctionStats holds the statistics for an individual function
type FunctionStats struct {
	Queue    uint64
	Running  uint64
	Complete uint64
	Failed   uint64
}

func (s *stats) getStatsForFunction(path string) *functionStats {
	if s.functionStatsMap == nil {
		s.functionStatsMap = make(map[string]*functionStats)
	}
	thisFunctionStats, found := s.functionStatsMap[path]
	if !found {
		thisFunctionStats = &functionStats{}
		s.functionStatsMap[path] = thisFunctionStats
	}

	return thisFunctionStats
}

func (s *stats) Enqueue(ctx context.Context, app string, path string) {
	s.mu.Lock()

	s.queue++
	s.getStatsForFunction(path).queue++
	common.IncrementGauge(ctx, queuedMetricName)

	common.IncrementCounter(ctx, callsMetricName)

	s.mu.Unlock()
}

// Call when a function has been queued but cannot be started because of an error
func (s *stats) Dequeue(ctx context.Context, app string, path string) {
	s.mu.Lock()

	s.queue--
	s.getStatsForFunction(path).queue--
	common.DecrementGauge(ctx, queuedMetricName)

	s.mu.Unlock()
}

func (s *stats) DequeueAndStart(ctx context.Context, app string, path string) {
	s.mu.Lock()

	s.queue--
	s.getStatsForFunction(path).queue--
	common.DecrementGauge(ctx, queuedMetricName)

	s.running++
	s.getStatsForFunction(path).running++
	common.IncrementGauge(ctx, runningMetricName)

	s.mu.Unlock()
}

func (s *stats) Complete(ctx context.Context, app string, path string) {
	s.mu.Lock()

	s.running--
	s.getStatsForFunction(path).running--
	common.DecrementGauge(ctx, runningMetricName)

	s.complete++
	s.getStatsForFunction(path).complete++
	common.IncrementCounter(ctx, completedMetricName)

	s.mu.Unlock()
}

func (s *stats) Failed(ctx context.Context, app string, path string) {
	s.mu.Lock()

	s.running--
	s.getStatsForFunction(path).running--
	common.DecrementGauge(ctx, runningMetricName)

	s.failed++
	s.getStatsForFunction(path).failed++
	common.IncrementCounter(ctx, failedMetricName)

	s.mu.Unlock()
}

func (s *stats) DequeueAndFail(ctx context.Context, app string, path string) {
	s.mu.Lock()

	s.queue--
	s.getStatsForFunction(path).queue--
	common.DecrementGauge(ctx, queuedMetricName)

	s.failed++
	s.getStatsForFunction(path).failed++
	common.IncrementCounter(ctx, failedMetricName)

	s.mu.Unlock()
}

func (s *stats) IncrementTimedout(ctx context.Context) {
	common.IncrementCounter(ctx, timedoutMetricName)
}

func (s *stats) IncrementErrors(ctx context.Context) {
	common.IncrementCounter(ctx, errorsMetricName)
}

func (s *stats) Stats() Stats {
	var stats Stats
	s.mu.Lock()
	stats.Running = s.running
	stats.Complete = s.complete
	stats.Queue = s.queue
	stats.Failed = s.failed
	stats.FunctionStatsMap = make(map[string]*FunctionStats)
	for key, value := range s.functionStatsMap {
		thisFunctionStats := &FunctionStats{Queue: value.queue, Running: value.running, Complete: value.complete, Failed: value.failed}
		stats.FunctionStatsMap[key] = thisFunctionStats
	}
	s.mu.Unlock()
	return stats
}

const (
	queuedMetricName    = "queued"
	callsMetricName     = "calls"
	runningMetricName   = "running"
	completedMetricName = "completed"
	failedMetricName    = "failed"
	timedoutMetricName  = "timeouts"
	errorsMetricName    = "errors"
)
<|endoftext|>


<commit_before>package builder

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/gorilla/websocket"

	"github.com/velocity-ci/velocity/backend/pkg/architect"
	"github.com/velocity-ci/velocity/backend/pkg/domain/builder"
	"github.com/velocity-ci/velocity/backend/pkg/velocity"
)

type Builder struct {
}

func (b *Builder) Start() {
	address := getArchitectAddress()
	secret := getBuilderSecret()
	client := &http.Client{
		Timeout: time.Second * 10,
	}
	if !waitForService(client, address) {
		logrus.Fatalf("Could not connect to: %s", address)
	}

	ws := connectToArchitect(address, secret)

	logrus.Infof("connected to %s", address)

	go monitorCommands(ws)
}

func (b *Builder) Stop() error {
	return nil
}

func New() architect.App {
	velocity.SetLogLevel()
	return &Builder{}
}

func getArchitectAddress() string {
	address := os.Getenv("ARCHITECT_ADDRESS") // http://architect || https://architect
	if address == "" {
		logrus.Fatal("Missing ARCHITECT_ADDRESS environment variable")
	}

	if address[:5] != "https" {
		logrus.Info("WARNING: Builds are not protected by TLS.")
	}

	return address
}

func getBuilderSecret() string {
	secret := os.Getenv("BUILDER_SECRET")
	if secret == "" {
		logrus.Fatal("Missing BUILDER_SECRET environment variable")
	}

	return secret
}

func waitForService(client *http.Client, address string) bool {

	for i := 0; i < 6; i++ {
		logrus.Infof("attempting connection to %s", address)
		_, err := client.Get(address)
		if err != nil {
			logrus.Debugf("connection error: %v", err)
		} else {
			logrus.Infof("%s is alive!", address)
			return true
		}
		time.Sleep(5 * time.Second)
	}

	return false
}

func connectToArchitect(address string, secret string) *websocket.Conn {
	wsAddress := strings.Replace(address, "http", "ws", 1)
	headers := http.Header{}
	headers.Set("Authorization", secret)
	var dialer *websocket.Dialer
	conn, _, err := dialer.Dial(
		fmt.Sprintf("%s/builder/ws", wsAddress),
		headers,
	)

	if err != nil {
		logrus.Fatal(err)
	}

	return conn
}

func monitorCommands(ws *websocket.Conn) {
	for {
		command := &builder.BuilderCtrlMessage{}
		err := ws.ReadJSON(command)
		if err != nil {
			log.Println(err)
			log.Println("Closing WebSocket")
			ws.Close()
			return
		}

		if command.Command == builder.CommandBuild {
			logrus.Infof("Got Build: %v", command.Payload)
			runBuild(command.Payload.(*builder.BuildCtrl), ws)
		} else if command.Command == builder.CommandKnownHosts {
			logrus.Infof("Got known hosts: %v", command.Payload)
			updateKnownHosts(command.Payload.(*builder.KnownHostCtrl))
		}
	}
}
<commit_msg>[backend] restarting builder when websocket connection fails<commit_after>package builder

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/gorilla/websocket"

	"github.com/velocity-ci/velocity/backend/pkg/architect"
	"github.com/velocity-ci/velocity/backend/pkg/domain/builder"
	"github.com/velocity-ci/velocity/backend/pkg/velocity"
)

type Builder struct {
	run bool
}

func (b *Builder) Start() {
	address := getArchitectAddress()
	secret := getBuilderSecret()
	client := &http.Client{
		Timeout: time.Second * 10,
	}

	for b.run {
		if !waitForService(client, address) {
			logrus.Fatalf("Could not connect to: %s", address)
		}

		ws := connectToArchitect(address, secret)

		logrus.Infof("connected to %s", address)

		monitorCommands(ws)
	}
}

func (b *Builder) Stop() error {
	b.run = false
	return nil
}

func New() architect.App {
	velocity.SetLogLevel()
	return &Builder{run: true}
}

func getArchitectAddress() string {
	address := os.Getenv("ARCHITECT_ADDRESS") // http://architect || https://architect
	if address == "" {
		logrus.Fatal("Missing ARCHITECT_ADDRESS environment variable")
	}

	if address[:5] != "https" {
		logrus.Info("WARNING: Builds are not protected by TLS.")
	}

	return address
}

func getBuilderSecret() string {
	secret := os.Getenv("BUILDER_SECRET")
	if secret == "" {
		logrus.Fatal("Missing BUILDER_SECRET environment variable")
	}

	return secret
}

func waitForService(client *http.Client, address string) bool {

	for i := 0; i < 6; i++ {
		logrus.Infof("attempting connection to %s", address)
		_, err := client.Get(address)
		if err != nil {
			logrus.Debugf("connection error: %v", err)
		} else {
			logrus.Infof("%s is alive!", address)
			return true
		}
		time.Sleep(5 * time.Second)
	}

	return false
}

func connectToArchitect(address string, secret string) *websocket.Conn {
	wsAddress := strings.Replace(address, "http", "ws", 1)
	headers := http.Header{}
	headers.Set("Authorization", secret)
	var dialer *websocket.Dialer
	conn, _, err := dialer.Dial(
		fmt.Sprintf("%s/builder/ws", wsAddress),
		headers,
	)

	if err != nil {
		logrus.Fatal(err)
	}

	return conn
}

func monitorCommands(ws *websocket.Conn) {
	for {
		command := &builder.BuilderCtrlMessage{}
		err := ws.ReadJSON(command)
		if err != nil {
			log.Println(err)
			log.Println("Closing WebSocket")
			ws.Close()
			return
		}

		if command.Command == builder.CommandBuild {
			logrus.Infof("Got Build: %v", command.Payload)
			runBuild(command.Payload.(*builder.BuildCtrl), ws)
		} else if command.Command == builder.CommandKnownHosts {
			logrus.Infof("Got known hosts: %v", command.Payload)
			updateKnownHosts(command.Payload.(*builder.KnownHostCtrl))
		}
	}
}
<|endoftext|>