<commit_before>/*
Copyright 2017 The Nuclio 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 cmdrunner

import (
	"bytes"
	"fmt"
	"os/exec"
	"strings"
	"syscall"

	"github.com/nuclio/nuclio-sdk"
)

type CaptureOutputMode int

const (
	CaptureOutputModeCombined CaptureOutputMode = iota
	CaptureOutputModeStdout
)

// RunOptions specifies runOptions to CmdRunner.Run
type RunOptions struct {
	WorkingDir        *string
	Stdin             *string
	Env               map[string]string
	CaptureOutputMode CaptureOutputMode
}

type RunResult struct {
	Output   string
	Stderr   string
	ExitCode int
}

// CmdRunner specifies the interface to an underlying command runner
type CmdRunner interface {

	// Run runs a command, given runOptions
	Run(runOptions *RunOptions, format string, vars ...interface{}) (RunResult, error)
}

type ShellRunner struct {
	logger nuclio.Logger
	shell  string
}

func NewShellRunner(parentLogger nuclio.Logger) (*ShellRunner, error) {
	return &ShellRunner{
		logger: parentLogger.GetChild("runner"),
		shell:  "/bin/sh",
	}, nil
}

func (sr *ShellRunner) Run(runOptions *RunOptions, format string, vars ...interface{}) (RunResult, error) {

	// format the command
	formattedCommand := fmt.Sprintf(format, vars...)
	sr.logger.DebugWith("Executing", "command", formattedCommand, "runOptions", runOptions)

	// create a command
	cmd := exec.Command(sr.shell, "-c", formattedCommand)

	// if no run options passed, use default values
	if runOptions == nil {
		runOptions = &RunOptions{}
	}

	// if there are runOptions, set them
	if runOptions != nil {
		if runOptions.WorkingDir != nil {
			cmd.Dir = *runOptions.WorkingDir
		}

		// get environment variables if any
		if runOptions.Env != nil {
			cmd.Env = sr.getEnvFromOptions(runOptions)
		}

		if runOptions.Stdin != nil {
			cmd.Stdin = strings.NewReader(*runOptions.Stdin)
		}
	}

	runResult := RunResult{
		ExitCode: 0,
	}

	err := sr.runAndCaptureOutput(cmd, runOptions, &runResult)

	if err != nil {
		var exitCode int

		// Did the command fail because of an unsuccessful exit code
		if exitError, ok := err.(*exec.ExitError); ok {
			exitCode = exitError.Sys().(syscall.WaitStatus).ExitStatus()
		}

		runResult.ExitCode = exitCode

		sr.logger.DebugWith("Failed to execute command",
			"output", runResult.Output,
			"stderr", runResult.Stderr,
			"exitCode", runResult.ExitCode,
			"err", err)

		return runResult, err
	}

	sr.logger.DebugWith("Command executed successfully",
		"output", runResult.Output,
		"stderr", runResult.Stderr,
		"exitCode", runResult.ExitCode)

	return runResult, nil
}

func (sr *ShellRunner) SetShell(shell string) {
	sr.shell = shell
}

func (sr *ShellRunner) getEnvFromOptions(runOptions *RunOptions) []string {
	envs := []string{}

	for name, value := range runOptions.Env {
		envs = append(envs, fmt.Sprintf("%s=%s", name, value))
	}

	return envs
}

func (sr *ShellRunner) runAndCaptureOutput(cmd *exec.Cmd,
	runOptions *RunOptions,
	runResult *RunResult) error {

	switch runOptions.CaptureOutputMode {

	case CaptureOutputModeCombined:
		stdoutAndStderr, err := cmd.CombinedOutput()
		runResult.Output = string(stdoutAndStderr)
		return err

	case CaptureOutputModeStdout:
		var stdOut, stdErr bytes.Buffer
		cmd.Stdout = &stdOut
		cmd.Stderr = &stdErr

		err := cmd.Run()

		runResult.Output = stdOut.String()
		runResult.Stderr = stdErr.String()

		return err
	}

	return fmt.Errorf("Invalid output capture mode: %d", runOptions.CaptureOutputMode)
}
<commit_msg>ShellRunner.Run() Error Includes Stdout And Stderr (#436)<commit_after>/*
Copyright 2017 The Nuclio 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 cmdrunner

import (
	"bytes"
	"fmt"
	"os/exec"
	"strings"
	"syscall"

	"github.com/nuclio/nuclio/pkg/errors"

	"github.com/nuclio/nuclio-sdk"
)

type CaptureOutputMode int

const (
	CaptureOutputModeCombined CaptureOutputMode = iota
	CaptureOutputModeStdout
)

// RunOptions specifies runOptions to CmdRunner.Run
type RunOptions struct {
	WorkingDir        *string
	Stdin             *string
	Env               map[string]string
	CaptureOutputMode CaptureOutputMode
}

type RunResult struct {
	Output   string
	Stderr   string
	ExitCode int
}

// CmdRunner specifies the interface to an underlying command runner
type CmdRunner interface {
	// Run runs a command, given runOptions
	Run(runOptions *RunOptions, format string, vars ...interface{}) (RunResult, error)
}

type ShellRunner struct {
	logger nuclio.Logger
	shell  string
}

func NewShellRunner(parentLogger nuclio.Logger) (*ShellRunner, error) {
	return &ShellRunner{
		logger: parentLogger.GetChild("runner"),
		shell:  "/bin/sh",
	}, nil
}

func (sr *ShellRunner) Run(runOptions *RunOptions, format string, vars ...interface{}) (RunResult, error) {

	// format the command
	formattedCommand := fmt.Sprintf(format, vars...)
	sr.logger.DebugWith("Executing", "command", formattedCommand, "runOptions", runOptions)

	// create a command
	cmd := exec.Command(sr.shell, "-c", formattedCommand)

	// if no run options passed, use default values
	if runOptions == nil {
		runOptions = &RunOptions{}
	}

	// if there are runOptions, set them
	if runOptions != nil {
		if runOptions.WorkingDir != nil {
			cmd.Dir = *runOptions.WorkingDir
		}

		// get environment variables if any
		if runOptions.Env != nil {
			cmd.Env = sr.getEnvFromOptions(runOptions)
		}

		if runOptions.Stdin != nil {
			cmd.Stdin = strings.NewReader(*runOptions.Stdin)
		}
	}

	runResult := RunResult{
		ExitCode: 0,
	}

	err := sr.runAndCaptureOutput(cmd, runOptions, &runResult)

	if err != nil {
		var exitCode int

		// Did the command fail because of an unsuccessful exit code
		if exitError, ok := err.(*exec.ExitError); ok {
			exitCode = exitError.Sys().(syscall.WaitStatus).ExitStatus()
		}

		runResult.ExitCode = exitCode

		sr.logger.DebugWith("Failed to execute command",
			"output", runResult.Output,
			"stderr", runResult.Stderr,
			"exitCode", runResult.ExitCode,
			"err", err)

		err = errors.Wrapf(err, "stdout:\n%s\nstderr:\n%s", runResult.Output, runResult.Stderr)

		return runResult, err
	}

	sr.logger.DebugWith("Command executed successfully",
		"output", runResult.Output,
		"stderr", runResult.Stderr,
		"exitCode", runResult.ExitCode)

	return runResult, nil
}

func (sr *ShellRunner) SetShell(shell string) {
	sr.shell = shell
}

func (sr *ShellRunner) getEnvFromOptions(runOptions *RunOptions) []string {
	envs := []string{}

	for name, value := range runOptions.Env {
		envs = append(envs, fmt.Sprintf("%s=%s", name, value))
	}

	return envs
}

func (sr *ShellRunner) runAndCaptureOutput(cmd *exec.Cmd,
	runOptions *RunOptions,
	runResult *RunResult) error {

	switch runOptions.CaptureOutputMode {

	case CaptureOutputModeCombined:
		stdoutAndStderr, err := cmd.CombinedOutput()
		runResult.Output = string(stdoutAndStderr)
		return err

	case CaptureOutputModeStdout:
		var stdOut, stdErr bytes.Buffer
		cmd.Stdout = &stdOut
		cmd.Stderr = &stdErr

		err := cmd.Run()

		runResult.Output = stdOut.String()
		runResult.Stderr = stdErr.String()

		return err
	}

	return fmt.Errorf("Invalid output capture mode: %d", runOptions.CaptureOutputMode)
}
<|endoftext|>


<commit_before>package executor

import (
	log "github.com/Sirupsen/logrus"
	. "github.com/smartystreets/goconvey/convey"
	"testing"
)

const (
	fifoTestDirTemplate = "/tmp/swan_local_test.XXXXXXXXXXX"
	fifoTestName        = "swan_fifo"
)

// TestLocal
func TestLocal(t *testing.T) {
	log.SetLevel(log.ErrorLevel)

	Convey("While using Local Shell", t, func() {
		l := NewLocal()

		Convey("When blocking infinitive sleep command "+
			"is executed", func() {
			task, err := l.Execute("sleep inf")

			Convey("There should be no error", func() {
				So(err, ShouldBeNil)

				task.Stop()
			})

			Convey("Task should be still running and status should be nil", func() {
				taskState, taskStatus := task.Status()
				So(taskState, ShouldEqual, RUNNING)
				So(taskStatus, ShouldBeNil)

				task.Stop()
			})

			Convey("When we wait for task termination with the 1ms timeout", func() {
				isTaskTerminated := task.Wait(1)

				Convey("The timeout should exceed and the task not terminated ", func() {
					So(isTaskTerminated, ShouldBeFalse)
				})

				Convey("The task should be still running and status should be nil", func() {
					taskState, taskStatus := task.Status()
					So(taskState, ShouldEqual, RUNNING)
					So(taskStatus, ShouldBeNil)
				})

				task.Stop()
			})

			Convey("When we stop the task", func() {
				err := task.Stop()

				Convey("There should be no error", func() {
					So(err, ShouldBeNil)
				})

				Convey("The task should be terminated and the task status should be -1", func() {
					taskState, taskStatus := task.Status()
					So(taskState, ShouldEqual, TERMINATED)
					So(taskStatus.ExitCode, ShouldEqual, -1)
				})
			})
		})

		Convey("When command `echo output` is executed", func() {
			task, err := l.Execute("echo output")

			Convey("There should be no error", func() {
				So(err, ShouldBeNil)

				task.Stop()
			})

			Convey("When we wait for the task to terminate", func() {
				isTaskTerminated := task.Wait(500)

				Convey("Wait should states that task terminated", func() {
					So(isTaskTerminated, ShouldBeTrue)
				})

				taskState, taskStatus := task.Status()

				Convey("The task should be terminated", func() {
					So(taskState, ShouldEqual, TERMINATED)
				})

				Convey("And the exit status should be 0", func() {
					So(taskStatus.ExitCode, ShouldEqual, 0)
				})

				Convey("And command stdout needs to match 'output", func() {
					So(taskStatus.Stdout, ShouldEqual, "output\n")
				})
			})
		})

		Convey("When command which does not exists is executed", func() {
			task, err := l.Execute("commandThatDoesNotExists")

			Convey("There should be no error", func() {
				So(err, ShouldBeNil)

				task.Stop()
			})

			Convey("When we wait for the task to terminate", func() {
				isTaskTerminated := task.Wait(500)

				Convey("Wait should state that task terminated", func() {
					So(isTaskTerminated, ShouldBeTrue)
				})

				taskState, taskStatus := task.Status()

				Convey("The task should be terminated", func() {
					So(taskState, ShouldEqual, TERMINATED)
				})

				Convey("And the exit status should be 127", func() {
					So(taskStatus.ExitCode, ShouldEqual, 127)
				})
			})
		})

		Convey("When we execute two tasks in the same time", func() {
			task, err := l.Execute("echo output1")
			task2, err2 := l.Execute("echo output2")

			Convey("There should be no errors", func() {
				So(err, ShouldBeNil)
				So(err2, ShouldBeNil)
			})

			Convey("When we wait for the tasks to terminate", func() {
				isTaskTerminated := task.Wait(0)
				isTaskTerminated2 := task2.Wait(0)

				Convey("Wait should state that tasks are terminated", func() {
					So(isTaskTerminated, ShouldBeTrue)
					So(isTaskTerminated2, ShouldBeTrue)
				})

				taskState1, taskStatus1 := task.Status()
				taskState2, taskStatus2 := task2.Status()

				Convey("The tasks should be terminated", func() {
					So(taskState1, ShouldEqual, TERMINATED)
					So(taskState2, ShouldEqual, TERMINATED)
				})

				Convey("The commands stdouts needs to match 'output1' & 'output2'", func() {
					So(taskStatus1.Stdout, ShouldEqual, "output1\n")
					So(taskStatus2.Stdout, ShouldEqual, "output2\n")
				})

				Convey("Both exit statuses should be 0", func() {
					So(taskStatus1.ExitCode, ShouldEqual, 0)
					So(taskStatus2.ExitCode, ShouldEqual, 0)
				})
			})
		})
	})
}
<commit_msg>Issues addressed.<commit_after>package executor

import (
	log "github.com/Sirupsen/logrus"
	. "github.com/smartystreets/goconvey/convey"
	"testing"
)

// TestLocal tests the execution of process on local machine.
func TestLocal(t *testing.T) {
	log.SetLevel(log.ErrorLevel)

	Convey("While using Local Shell", t, func() {
		l := NewLocal()

		Convey("When blocking infinitively sleep command "+
			"is executed", func() {
			task, err := l.Execute("sleep inf")

			Convey("There should be no error", func() {
				So(err, ShouldBeNil)

				task.Stop()
			})

			Convey("Task should be still running and status should be nil", func() {
				taskState, taskStatus := task.Status()
				So(taskState, ShouldEqual, RUNNING)
				So(taskStatus, ShouldBeNil)

				task.Stop()
			})

			Convey("When we wait for task termination with the 1ms timeout", func() {
				isTaskTerminated := task.Wait(1)

				Convey("The timeout should exceed and the task not terminated ", func() {
					So(isTaskTerminated, ShouldBeFalse)
				})

				Convey("The task should be still running and status should be nil", func() {
					taskState, taskStatus := task.Status()
					So(taskState, ShouldEqual, RUNNING)
					So(taskStatus, ShouldBeNil)
				})

				task.Stop()
			})

			Convey("When we stop the task", func() {
				err := task.Stop()

				Convey("There should be no error", func() {
					So(err, ShouldBeNil)
				})

				Convey("The task should be terminated and the task status should be -1", func() {
					taskState, taskStatus := task.Status()
					So(taskState, ShouldEqual, TERMINATED)
					So(taskStatus.ExitCode, ShouldEqual, -1)
				})
			})
		})

		Convey("When command `echo output` is executed", func() {
			task, err := l.Execute("echo output")

			Convey("There should be no error", func() {
				So(err, ShouldBeNil)

				task.Stop()
			})

			Convey("When we wait for the task to terminate", func() {
				isTaskTerminated := task.Wait(500)

				Convey("Wait should states that task terminated", func() {
					So(isTaskTerminated, ShouldBeTrue)
				})

				taskState, taskStatus := task.Status()

				Convey("The task should be terminated", func() {
					So(taskState, ShouldEqual, TERMINATED)
				})

				Convey("And the exit status should be 0", func() {
					So(taskStatus.ExitCode, ShouldEqual, 0)
				})

				Convey("And command stdout needs to match 'output", func() {
					So(taskStatus.Stdout, ShouldEqual, "output\n")
				})
			})
		})

		Convey("When command which does not exists is executed", func() {
			task, err := l.Execute("commandThatDoesNotExists")

			Convey("There should be no error", func() {
				So(err, ShouldBeNil)

				task.Stop()
			})

			Convey("When we wait for the task to terminate", func() {
				isTaskTerminated := task.Wait(500)

				Convey("Wait should state that task terminated", func() {
					So(isTaskTerminated, ShouldBeTrue)
				})

				taskState, taskStatus := task.Status()

				Convey("The task should be terminated", func() {
					So(taskState, ShouldEqual, TERMINATED)
				})

				Convey("And the exit status should be 127", func() {
					So(taskStatus.ExitCode, ShouldEqual, 127)
				})
			})
		})

		Convey("When we execute two tasks in the same time", func() {
			task, err := l.Execute("echo output1")
			task2, err2 := l.Execute("echo output2")

			Convey("There should be no errors", func() {
				So(err, ShouldBeNil)
				So(err2, ShouldBeNil)
			})

			Convey("When we wait for the tasks to terminate", func() {
				isTaskTerminated := task.Wait(0)
				isTaskTerminated2 := task2.Wait(0)

				Convey("Wait should state that tasks are terminated", func() {
					So(isTaskTerminated, ShouldBeTrue)
					So(isTaskTerminated2, ShouldBeTrue)
				})

				taskState1, taskStatus1 := task.Status()
				taskState2, taskStatus2 := task2.Status()

				Convey("The tasks should be terminated", func() {
					So(taskState1, ShouldEqual, TERMINATED)
					So(taskState2, ShouldEqual, TERMINATED)
				})

				Convey("The commands stdouts needs to match 'output1' & 'output2'", func() {
					So(taskStatus1.Stdout, ShouldEqual, "output1\n")
					So(taskStatus2.Stdout, ShouldEqual, "output2\n")
				})

				Convey("Both exit statuses should be 0", func() {
					So(taskStatus1.ExitCode, ShouldEqual, 0)
					So(taskStatus2.ExitCode, ShouldEqual, 0)
				})
			})
		})
	})
}
<|endoftext|>


<commit_before>// Copyright (c) 2015-2020 The usbtmc developers. All rights reserved.
// Project site: https://github.com/gotmc/usbtmc
// Use of this source code is governed by a MIT-style license that
// can be found in the LICENSE.txt file for the project.

package usbtmc

import (
	"bytes"
	"fmt"
	"io"
	"strings"

	"github.com/gotmc/usbtmc/driver"
)

// Device models a USBTMC device, which includes a USB device and the required
// USBTMC attributes and methods.
type Device struct {
	usbDevice       driver.USBDevice
	bTag            byte
	termChar        byte
	termCharEnabled bool
}

// Write creates the appropriate USBMTC header, writes the header and data on
// the bulk out endpoint, and returns the number of bytes written and any
// errors.
func (d *Device) Write(p []byte) (n int, err error) {
	// FIXME(mdr): I need to change this so that I look at the size of the buf
	// being written to see if it can truly fit into one transfer, and if not
	// split it into multiple transfers.
	d.bTag = nextbTag(d.bTag)
	header := encodeBulkOutHeader(d.bTag, uint32(len(p)), true)
	data := append(header[:], p...)
	if moduloFour := len(data) % 4; moduloFour > 0 {
		numAlignment := 4 - moduloFour
		alignment := bytes.Repeat([]byte{0x00}, numAlignment)
		data = append(data, alignment...)
	}
	return d.usbDevice.Write(data)
}

// Read creates and sends the header on the bulk out endpoint and then reads
// from the bulk in endpoint per USBTMC standard.
func (d *Device) Read(p []byte) (n int, err error) {
	d.bTag = nextbTag(d.bTag)
	header := encodeMsgInBulkOutHeader(d.bTag, uint32(len(p)), d.termCharEnabled, d.termChar)
	_, err = d.usbDevice.Write(header[:])
	return d.readRemoveHeader(p)
}

func (d *Device) readRemoveHeader(p []byte) (n int, err error) {
	// FIXME(mdr): Seems like I shouldn't use 512 as a magic number or as a hard
	// size limit. I should grab the max size of the bulk in endpoint.
	usbtmcHeaderLen := 12
	temp := make([]byte, 512)
	n, err = d.usbDevice.Read(temp)
	// Remove the USBMTC Bulk-IN Header from the data and the number of bytes
	if n < usbtmcHeaderLen {
		return 0, err
	}
	reader := bytes.NewReader(temp)
	_, err = reader.ReadAt(p, int64(usbtmcHeaderLen))
	if err != nil && err != io.EOF {
		return n - usbtmcHeaderLen, err
	}
	return n - usbtmcHeaderLen, nil
}

func (d *Device) readKeepHeader(p []byte) (n int, err error) {
	return d.usbDevice.Read(p)
}

// Close closes the underlying USB device.
func (d *Device) Close() error {
	return d.usbDevice.Close()
}

// WriteString writes a string using the underlying USB device. A newline
// terminator is not automatically added.
func (d *Device) WriteString(s string) (n int, err error) {
	return d.Write([]byte(s))
}

// Command sends the SCPI/ASCII command to the underlying USB device. A newline
// character is automatically added to the end of the string.
func (d *Device) Command(format string, a ...interface{}) error {
	cmd := format
	if a != nil {
		cmd = fmt.Sprintf(format, a...)
	}
	_, err := d.Write([]byte(strings.TrimSpace(cmd) + "\n"))
	return err
}

// Query writes the given string to the USBTMC device and returns the returned
// value as a string. A newline character is automatically added to the query
// command sent to the instrument.
func (d *Device) Query(s string) (string, error) {
	err := d.Command(s)
	if err != nil {
		return "", err
	}
	p := make([]byte, 512)
	n, err := d.Read(p)
	if err != nil {
		return "", err
	}
	s = fmt.Sprintf("%s", p[:n])
	p = nil
	return s, nil
}
<commit_msg>Use WriteString instead of Write for Commands<commit_after>// Copyright (c) 2015-2020 The usbtmc developers. All rights reserved.
// Project site: https://github.com/gotmc/usbtmc
// Use of this source code is governed by a MIT-style license that
// can be found in the LICENSE.txt file for the project.

package usbtmc

import (
	"bytes"
	"fmt"
	"io"
	"strings"

	"github.com/gotmc/usbtmc/driver"
)

// Device models a USBTMC device, which includes a USB device and the required
// USBTMC attributes and methods.
type Device struct {
	usbDevice       driver.USBDevice
	bTag            byte
	termChar        byte
	termCharEnabled bool
}

// Write creates the appropriate USBMTC header, writes the header and data on
// the bulk out endpoint, and returns the number of bytes written and any
// errors.
func (d *Device) Write(p []byte) (n int, err error) {
	// FIXME(mdr): I need to change this so that I look at the size of the buf
	// being written to see if it can truly fit into one transfer, and if not
	// split it into multiple transfers.
	d.bTag = nextbTag(d.bTag)
	header := encodeBulkOutHeader(d.bTag, uint32(len(p)), true)
	data := append(header[:], p...)
	if moduloFour := len(data) % 4; moduloFour > 0 {
		numAlignment := 4 - moduloFour
		alignment := bytes.Repeat([]byte{0x00}, numAlignment)
		data = append(data, alignment...)
	}
	return d.usbDevice.Write(data)
}

// Read creates and sends the header on the bulk out endpoint and then reads
// from the bulk in endpoint per USBTMC standard.
func (d *Device) Read(p []byte) (n int, err error) {
	d.bTag = nextbTag(d.bTag)
	header := encodeMsgInBulkOutHeader(d.bTag, uint32(len(p)), d.termCharEnabled, d.termChar)
	_, err = d.usbDevice.Write(header[:])
	return d.readRemoveHeader(p)
}

func (d *Device) readRemoveHeader(p []byte) (n int, err error) {
	// FIXME(mdr): Seems like I shouldn't use 512 as a magic number or as a hard
	// size limit. I should grab the max size of the bulk in endpoint.
	usbtmcHeaderLen := 12
	temp := make([]byte, 512)
	n, err = d.usbDevice.Read(temp)
	// Remove the USBMTC Bulk-IN Header from the data and the number of bytes
	if n < usbtmcHeaderLen {
		return 0, err
	}
	reader := bytes.NewReader(temp)
	_, err = reader.ReadAt(p, int64(usbtmcHeaderLen))
	if err != nil && err != io.EOF {
		return n - usbtmcHeaderLen, err
	}
	return n - usbtmcHeaderLen, nil
}

func (d *Device) readKeepHeader(p []byte) (n int, err error) {
	return d.usbDevice.Read(p)
}

// Close closes the underlying USB device.
func (d *Device) Close() error {
	return d.usbDevice.Close()
}

// WriteString writes a string using the underlying USB device. A newline
// terminator is not automatically added.
func (d *Device) WriteString(s string) (n int, err error) {
	return d.Write([]byte(s))
}

// Command sends the SCPI/ASCII command to the underlying USB device. A newline
// character is automatically added to the end of the string.
func (d *Device) Command(format string, a ...interface{}) error {
	cmd := format
	if a != nil {
		cmd = fmt.Sprintf(format, a...)
	}
	_, err := d.WriteString(strings.TrimSpace(cmd) + "\n")
	return err
}

// Query writes the given string to the USBTMC device and returns the returned
// value as a string. A newline character is automatically added to the query
// command sent to the instrument.
func (d *Device) Query(s string) (string, error) {
	err := d.Command(s)
	if err != nil {
		return "", err
	}
	p := make([]byte, 512)
	n, err := d.Read(p)
	if err != nil {
		return "", err
	}
	s = fmt.Sprintf("%s", p[:n])
	p = nil
	return s, nil
}
<|endoftext|>


<commit_before>package acme

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"strings"

	"golang.org/x/crypto/acme"
	api "k8s.io/api/core/v1"
	k8sErrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/jetstack/cert-manager/pkg/apis/certmanager/v1alpha1"
	"github.com/jetstack/cert-manager/pkg/controller"
)

const (
	acmeAccountPrivateKeyKey = "key.pem"
)

type account struct {
	ctx    *controller.Context
	issuer *v1alpha1.Issuer
}

func (a *account) uri() string {
	return a.issuer.Spec.ACME.URI
}

func (a *account) email() string {
	return a.issuer.Spec.ACME.Email
}

func (a *account) server() string {
	return a.issuer.Spec.ACME.Server
}

// privateKey returns the private key for this account from the given context,
// or an error
// TODO (@munnerz): how can we support different types of private keys other
// than rsa?
func (a *account) privateKey() (*rsa.PrivateKey, error) {
	keyName := a.issuer.Spec.ACME.PrivateKey
	keySecret, err := a.ctx.InformerFactory.Core().V1().Secrets().Lister().Secrets(a.issuer.Namespace).Get(keyName)

	if err != nil {
		// we return the plain error here so k8sErrors.IsNotFound can be used
		return nil, err
	}

	keyBytes, okkey := keySecret.Data[acmeAccountPrivateKeyKey]

	// TODO: should we automatically recover from this situation by creating the key?
	if !okkey {
		return nil, fmt.Errorf("no '%s' key set in account secret", api.TLSPrivateKeyKey)
	}

	// decode the private key pem
	block, _ := pem.Decode(keyBytes)
	if block == nil {
		return nil, fmt.Errorf("error decoding private key PEM block in '%s'", keyName)
	}
	// parse the private key
	key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("error parsing private key in '%s': %s", keyName, err.Error())
	}
	// validate the private key
	if err = key.Validate(); err != nil {
		return nil, fmt.Errorf("private key failed validation in '%s': %s", keyName, err.Error())
	}

	return key, nil
}

// verify verifies an acme account is valid with the acme server
func (a *account) verify() error {
	if a.issuer.Spec.ACME.Server == "" {
		a.ctx.Logger.Printf("acme server url must be set")
		return nil
	}

	privateKey, err := a.privateKey()

	if err != nil {
		a.issuer.Status.Ready = false
		return err
	}

	a.ctx.Logger.Printf("using acme server '%s' for verification", a.issuer.Spec.ACME.Server)
	cl := acme.Client{
		Key:          privateKey,
		DirectoryURL: a.issuer.Spec.ACME.Server,
	}

	_, err = cl.GetReg(context.Background(), a.issuer.Spec.ACME.URI)

	if err != nil {
		return fmt.Errorf("error getting acme registration: %s", err.Error())
	}

	// TODO: come up with some way to verify the private key is valid for this
	// account

	return nil
}

// register will register an account with the acme server and store the account
// details in the context
func (a *account) register() error {
	if a.issuer.Spec.ACME.Server == "" {
		a.ctx.Logger.Printf("acme server url must be set")
		return nil
	}

	privateKey, err := a.privateKey()

	if err != nil {
		if !k8sErrors.IsNotFound(err) {
			return fmt.Errorf("error getting private key: %s", err.Error())
		}

		// TODO (@munnerz): allow changing the keysize
		privateKey, err = a.generatePrivateKey(2048)

		if err != nil {
			return fmt.Errorf("error generating private key: %s", err.Error())
		}

		_, err = a.ctx.Client.Secrets(a.issuer.Namespace).Create(&api.Secret{
			ObjectMeta: metav1.ObjectMeta{
				Name:      a.issuer.Spec.ACME.PrivateKey,
				Namespace: a.issuer.Namespace,
			},
			Data: map[string][]byte{
				acmeAccountPrivateKeyKey: pem.EncodeToMemory(
					&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)},
				),
			},
		})

		if err != nil {
			return fmt.Errorf("error saving private key: %s", err.Error())
		}
	}

	a.ctx.Logger.Printf("using acme server '%s' for registration", a.issuer.Spec.ACME.Server)
	cl := acme.Client{
		Key:          privateKey,
		DirectoryURL: a.issuer.Spec.ACME.Server,
	}

	acc := &acme.Account{
		Contact: []string{fmt.Sprintf("mailto:%s", strings.ToLower(a.issuer.Spec.ACME.Email))},
	}

	// todo (@munnerz): don't use ctx.Background() here
	account, err := cl.Register(context.Background(), acc, acme.AcceptTOS)

	if err != nil {
		var acmeErr *acme.Error
		var ok bool
		if acmeErr, ok = err.(*acme.Error); !ok || (acmeErr.StatusCode != 409) {
			return fmt.Errorf("error registering acme account: %s", err.Error())
		}

		if a.issuer.Spec.ACME.URI == "" {
			return fmt.Errorf("private key already registered but user URI not found. delete existing private key or set acme account URI")
		}

		if account, err = cl.UpdateReg(context.Background(), acc); err != nil {
			return fmt.Errorf("error updating acme account registration: %s", err.Error())
		}
	}

	a.issuer.Spec.ACME.URI = account.URI
	return nil
}

func (a *account) generatePrivateKey(bits int) (*rsa.PrivateKey, error) {
	return rsa.GenerateKey(rand.Reader, bits)
}
<commit_msg>Return privateKeyPem in generatePrivateKey<commit_after>package acme

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"strings"

	"golang.org/x/crypto/acme"
	api "k8s.io/api/core/v1"
	k8sErrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/jetstack/cert-manager/pkg/apis/certmanager/v1alpha1"
	"github.com/jetstack/cert-manager/pkg/controller"
)

const (
	acmeAccountPrivateKeyKey = "key.pem"
)

type account struct {
	ctx    *controller.Context
	issuer *v1alpha1.Issuer
}

func (a *account) uri() string {
	return a.issuer.Spec.ACME.URI
}

func (a *account) email() string {
	return a.issuer.Spec.ACME.Email
}

func (a *account) server() string {
	return a.issuer.Spec.ACME.Server
}

// privateKey returns the private key for this account from the given context,
// or an error
// TODO (@munnerz): how can we support different types of private keys other
// than rsa?
func (a *account) privateKey() (*rsa.PrivateKey, error) {
	keyName := a.issuer.Spec.ACME.PrivateKey
	keySecret, err := a.ctx.InformerFactory.Core().V1().Secrets().Lister().Secrets(a.issuer.Namespace).Get(keyName)

	if err != nil {
		// we return the plain error here so k8sErrors.IsNotFound can be used
		return nil, err
	}

	keyBytes, okkey := keySecret.Data[acmeAccountPrivateKeyKey]

	// TODO: should we automatically recover from this situation by creating the key?
	if !okkey {
		return nil, fmt.Errorf("no '%s' key set in account secret", api.TLSPrivateKeyKey)
	}

	// decode the private key pem
	block, _ := pem.Decode(keyBytes)
	if block == nil {
		return nil, fmt.Errorf("error decoding private key PEM block in '%s'", keyName)
	}
	// parse the private key
	key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("error parsing private key in '%s': %s", keyName, err.Error())
	}
	// validate the private key
	if err = key.Validate(); err != nil {
		return nil, fmt.Errorf("private key failed validation in '%s': %s", keyName, err.Error())
	}

	return key, nil
}

// verify verifies an acme account is valid with the acme server
func (a *account) verify() error {
	if a.issuer.Spec.ACME.Server == "" {
		a.ctx.Logger.Printf("acme server url must be set")
		return nil
	}

	privateKey, err := a.privateKey()

	if err != nil {
		a.issuer.Status.Ready = false
		return err
	}

	a.ctx.Logger.Printf("using acme server '%s' for verification", a.issuer.Spec.ACME.Server)
	cl := acme.Client{
		Key:          privateKey,
		DirectoryURL: a.issuer.Spec.ACME.Server,
	}

	_, err = cl.GetReg(context.Background(), a.issuer.Spec.ACME.URI)

	if err != nil {
		return fmt.Errorf("error getting acme registration: %s", err.Error())
	}

	// TODO: come up with some way to verify the private key is valid for this
	// account

	return nil
}

// register will register an account with the acme server and store the account
// details in the context
func (a *account) register() error {
	if a.issuer.Spec.ACME.Server == "" {
		a.ctx.Logger.Printf("acme server url must be set")
		return nil
	}

	privateKey, err := a.privateKey()
	var privateKeyPem []byte
	if err != nil {
		if !k8sErrors.IsNotFound(err) {
			return fmt.Errorf("error getting private key: %s", err.Error())
		}

		// TODO (@munnerz): allow changing the keysize
		privateKeyPem, privateKey, err = generatePrivateKey(2048)

		if err != nil {
			return fmt.Errorf("error generating private key: %s", err.Error())
		}

		_, err = a.ctx.Client.Secrets(a.issuer.Namespace).Create(&api.Secret{
			ObjectMeta: metav1.ObjectMeta{
				Name:      a.issuer.Spec.ACME.PrivateKey,
				Namespace: a.issuer.Namespace,
			},
			Data: map[string][]byte{
				acmeAccountPrivateKeyKey: privateKeyPem,
			},
		})

		if err != nil {
			return fmt.Errorf("error saving private key: %s", err.Error())
		}
	}

	a.ctx.Logger.Printf("using acme server '%s' for registration", a.issuer.Spec.ACME.Server)
	cl := acme.Client{
		Key:          privateKey,
		DirectoryURL: a.issuer.Spec.ACME.Server,
	}

	acc := &acme.Account{
		Contact: []string{fmt.Sprintf("mailto:%s", strings.ToLower(a.issuer.Spec.ACME.Email))},
	}

	// todo (@munnerz): don't use ctx.Background() here
	account, err := cl.Register(context.Background(), acc, acme.AcceptTOS)

	if err != nil {
		var acmeErr *acme.Error
		var ok bool
		if acmeErr, ok = err.(*acme.Error); !ok || (acmeErr.StatusCode != 409) {
			return fmt.Errorf("error registering acme account: %s", err.Error())
		}

		if a.issuer.Spec.ACME.URI == "" {
			return fmt.Errorf("private key already registered but user URI not found. delete existing private key or set acme account URI")
		}

		if account, err = cl.UpdateReg(context.Background(), acc); err != nil {
			return fmt.Errorf("error updating acme account registration: %s", err.Error())
		}
	}

	a.issuer.Spec.ACME.URI = account.URI
	return nil
}

func generatePrivateKey(keySize int) ([]byte, *rsa.PrivateKey, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, keySize)
	if err != nil {
		return []byte{}, nil, err
	}

	block := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)}

	return pem.EncodeToMemory(block), privateKey, nil
}
<|endoftext|>


<commit_before>package github

import (
	"net/http"
	"sync"

	"github.com/dgrijalva/jwt-go"
	"github.com/google/go-github/github"
	"github.com/hellofresh/janus/pkg/config"
	"github.com/hellofresh/janus/pkg/jwt/provider"
	"github.com/pkg/errors"
)

func init() {
	provider.Register("github", &Provider{})
}

// Provider abstracts the authentication for github
type Provider struct {
	provider.Verifier

	teams  []Team
	config config.Credentials
}

// Build acts like the constructor for a provider
func (gp *Provider) Build(config config.Credentials) provider.Provider {
	client := NewClient()
	teams := gp.teamConfigsToTeam(config.Github.Teams)

	return &Provider{
		Verifier: provider.NewVerifierBasket(
			NewTeamVerifier(teams, client),
			NewOrganizationVerifier(config.Github.Organizations, client),
		),
		teams:  teams,
		config: config,
	}
}

// GetClaims returns a JWT Map Claim
func (gp *Provider) GetClaims(httpClient *http.Client) (jwt.MapClaims, error) {
	client := NewClient()

	var (
		wg            sync.WaitGroup
		user          *github.User
		usersOrgTeams OrganizationTeams
		err           error
	)

	wg.Add(2)

	go func() {
		defer wg.Done()
		user, err = client.CurrentUser(httpClient)
		if err != nil {
			err = errors.Wrap(err, "failed to get github users")
		}
	}()

	go func() {
		defer wg.Done()
		usersOrgTeams, err = client.Teams(httpClient)
		if err != nil {
			err = errors.Wrap(err, "failed to get github teams")
		}
	}()

	wg.Wait()

	if err != nil {
		return nil, err
	}

	return jwt.MapClaims{
		"sub":      *user.Login,
		"is_admin": gp.isAdmin(usersOrgTeams),
	}, nil
}

func (gp *Provider) teamConfigsToTeam(configTeams map[string]string) []Team {
	teams := []Team{}
	for org, team := range configTeams {
		teams = append(teams, Team{
			Name:         team,
			Organization: org,
		})
	}
	return teams
}

func (gp *Provider) isAdmin(usersOrgTeams OrganizationTeams) bool {
	for _, team := range gp.teams {
		if teams, ok := usersOrgTeams[team.Organization]; ok {
			for _, teamUserBelongsTo := range teams {
				if teamUserBelongsTo == gp.config.JanusAdminTeam {
					return true
				}
			}
		}
	}
	return false
}
<commit_msg>Fix race condition<commit_after>package github

import (
	"net/http"
	"sync"

	"github.com/dgrijalva/jwt-go"
	"github.com/google/go-github/github"
	"github.com/hellofresh/janus/pkg/config"
	"github.com/hellofresh/janus/pkg/jwt/provider"
	"github.com/pkg/errors"
)

func init() {
	provider.Register("github", &Provider{})
}

// Provider abstracts the authentication for github
type Provider struct {
	provider.Verifier

	teams  []Team
	config config.Credentials
}

// Build acts like the constructor for a provider
func (gp *Provider) Build(config config.Credentials) provider.Provider {
	client := NewClient()
	teams := gp.teamConfigsToTeam(config.Github.Teams)

	return &Provider{
		Verifier: provider.NewVerifierBasket(
			NewTeamVerifier(teams, client),
			NewOrganizationVerifier(config.Github.Organizations, client),
		),
		teams:  teams,
		config: config,
	}
}

// GetClaims returns a JWT Map Claim
func (gp *Provider) GetClaims(httpClient *http.Client) (jwt.MapClaims, error) {
	client := NewClient()

	var (
		wg            sync.WaitGroup
		user          *github.User
		usersOrgTeams OrganizationTeams
		errs          []error
	)

	wg.Add(2)
	go func(user *github.User, client Client, httpClient *http.Client) {
		defer wg.Done()
		user, err := client.CurrentUser(httpClient)
		if err != nil {
			errs = append(errs, errors.Wrap(err, "failed to get github users"))
		}
	}(user, client, httpClient)

	go func(usersOrgTeams OrganizationTeams, client Client, httpClient *http.Client) {
		defer wg.Done()
		usersOrgTeams, err := client.Teams(httpClient)
		if err != nil {
			errs = append(errs, errors.Wrap(err, "failed to get github teams"))
		}
	}(usersOrgTeams, client, httpClient)

	wg.Wait()

	if len(errs) > 0 {
		return nil, errs[0]
	}

	return jwt.MapClaims{
		"sub":      *user.Login,
		"is_admin": gp.isAdmin(usersOrgTeams),
	}, nil
}

func (gp *Provider) teamConfigsToTeam(configTeams map[string]string) []Team {
	teams := []Team{}
	for org, team := range configTeams {
		teams = append(teams, Team{
			Name:         team,
			Organization: org,
		})
	}
	return teams
}

func (gp *Provider) isAdmin(usersOrgTeams OrganizationTeams) bool {
	for _, team := range gp.teams {
		if teams, ok := usersOrgTeams[team.Organization]; ok {
			for _, teamUserBelongsTo := range teams {
				if teamUserBelongsTo == gp.config.JanusAdminTeam {
					return true
				}
			}
		}
	}
	return false
}
<|endoftext|>


<commit_before>package events

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/pborman/uuid"

	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/cancel"
	"github.com/lxc/lxd/shared/logger"
)

// EventSource indicates the source of an event.
type EventSource uint8

// EventSourceLocal indicates the event was generated locally.
const EventSourceLocal = 0

// EventSourcePull indicates the event was received from an outbound event listener stream.
const EventSourcePull = 1

// EventSourcePush indicates the event was received from an event listener client connected to us.
const EventSourcePush = 2

// InjectFunc is used to inject an event received by a listener into the local events dispatcher.
type InjectFunc func(event api.Event, eventSource EventSource)

// NotifyFunc is called when an event is dispatched.
type NotifyFunc func(event api.Event)

// Server represents an instance of an event server.
type Server struct {
	serverCommon

	listeners map[string]*Listener
	notify    NotifyFunc
	location  string
}

// NewServer returns a new event server.
func NewServer(debug bool, verbose bool, notify NotifyFunc) *Server {
	server := &Server{
		serverCommon: serverCommon{
			debug:   debug,
			verbose: verbose,
		},
		listeners: map[string]*Listener{},
		notify:    notify,
	}

	return server
}

// SetLocalLocation sets the local location of this member.
// This value will be added to the Location event field if not populated from another member.
func (s *Server) SetLocalLocation(location string) {
	s.lock.Lock()
	defer s.lock.Unlock()

	s.location = location
}

// AddListener creates and returns a new event listener.
func (s *Server) AddListener(projectName string, allProjects bool, connection EventListenerConnection, messageTypes []string, excludeSources []EventSource, recvFunc EventHandler, excludeLocations []string) (*Listener, error) {
	if allProjects && projectName != "" {
		return nil, fmt.Errorf("Cannot specify project name when listening for events on all projects")
	}

	listener := &Listener{
		listenerCommon: listenerCommon{
			EventListenerConnection: connection,
			messageTypes:            messageTypes,
			done:                    cancel.New(context.Background()),
			id:                      uuid.New(),
			recvFunc:                recvFunc,
		},

		allProjects:      allProjects,
		projectName:      projectName,
		excludeSources:   excludeSources,
		excludeLocations: excludeLocations,
	}

	s.lock.Lock()
	defer s.lock.Unlock()

	if s.listeners[listener.id] != nil {
		return nil, fmt.Errorf("A listener with ID %q already exists", listener.id)
	}

	s.listeners[listener.id] = listener

	go listener.start()

	return listener, nil
}

// SendLifecycle broadcasts a lifecycle event.
func (s *Server) SendLifecycle(projectName string, event api.EventLifecycle) {
	_ = s.Send(projectName, api.EventTypeLifecycle, event)
}

// Send broadcasts a custom event.
func (s *Server) Send(projectName string, eventType string, eventMessage any) error {
	encodedMessage, err := json.Marshal(eventMessage)
	if err != nil {
		return err
	}

	event := api.Event{
		Type:      eventType,
		Timestamp: time.Now(),
		Metadata:  encodedMessage,
		Project:   projectName,
	}

	return s.broadcast(event, EventSourceLocal)
}

// Inject an event from another member into the local events dispatcher.
// eventSource is used to indicate where this event was received from.
func (s *Server) Inject(event api.Event, eventSource EventSource) {
	if event.Type == api.EventTypeLogging {
		// Parse the message
		logEntry := api.EventLogging{}
		err := json.Unmarshal(event.Metadata, &logEntry)
		if err != nil {
			return
		}

		if !s.debug && logEntry.Level == "debug" {
			return
		}

		if !s.debug && !s.verbose && logEntry.Level == "info" {
			return
		}
	}

	err := s.broadcast(event, eventSource)
	if err != nil {
		logger.Warn("Failed to forward event from member", logger.Ctx{"member": event.Location, "err": err})
	}
}

func (s *Server) broadcast(event api.Event, eventSource EventSource) error {
	sourceInSlice := func(source EventSource, sources []EventSource) bool {
		for _, i := range sources {
			if source == i {
				return true
			}
		}

		return false
	}

	s.lock.Lock()

	// Set the Location for local events to the local serverName if not already populated (do it here rather
	// than in Send as the lock to read s.location has been taken here already).
	if eventSource == EventSourceLocal && event.Location == "" {
		event.Location = s.location
	}

	// If a notifcation hook is present, then call it for locally produced events.
	// This can be used to send local events to another target (such as an event-hub member).
	if s.notify != nil && eventSource == EventSourceLocal {
		s.notify(event)
	}

	listeners := s.listeners
	for _, listener := range listeners {
		// If the event is project specific, check if the listener is requesting events from that project.
		if event.Project != "" && !listener.allProjects && event.Project != listener.projectName {
			continue
		}

		if sourceInSlice(eventSource, listener.excludeSources) {
			continue
		}

		if !shared.StringInSlice(event.Type, listener.messageTypes) {
			continue
		}

		// If the event doesn't come from this member and has been excluded by listener, don't deliver it.
		if eventSource != EventSourceLocal && shared.StringInSlice(event.Location, listener.excludeLocations) {
			continue
		}

		go func(listener *Listener, event api.Event) {
			// Check that the listener still exists
			if listener == nil {
				return
			}

			// Make sure we're not done already
			if listener.IsClosed() {
				return
			}

			err := listener.WriteJSON(event)
			if err != nil {
				// Remove the listener from the list
				s.lock.Lock()
				delete(s.listeners, listener.id)
				s.lock.Unlock()

				listener.Close()
			}
		}(listener, event)
	}

	s.lock.Unlock()

	return nil
}

// Listener describes an event listener.
type Listener struct {
	listenerCommon

	allProjects      bool
	projectName      string
	excludeSources   []EventSource
	excludeLocations []string
}
<commit_msg>lxd/events: Remove closed listeners<commit_after>package events

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/pborman/uuid"

	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/cancel"
	"github.com/lxc/lxd/shared/logger"
)

// EventSource indicates the source of an event.
type EventSource uint8

// EventSourceLocal indicates the event was generated locally.
const EventSourceLocal = 0

// EventSourcePull indicates the event was received from an outbound event listener stream.
const EventSourcePull = 1

// EventSourcePush indicates the event was received from an event listener client connected to us.
const EventSourcePush = 2

// InjectFunc is used to inject an event received by a listener into the local events dispatcher.
type InjectFunc func(event api.Event, eventSource EventSource)

// NotifyFunc is called when an event is dispatched.
type NotifyFunc func(event api.Event)

// Server represents an instance of an event server.
type Server struct {
	serverCommon

	listeners map[string]*Listener
	notify    NotifyFunc
	location  string
}

// NewServer returns a new event server.
func NewServer(debug bool, verbose bool, notify NotifyFunc) *Server {
	server := &Server{
		serverCommon: serverCommon{
			debug:   debug,
			verbose: verbose,
		},
		listeners: map[string]*Listener{},
		notify:    notify,
	}

	return server
}

// SetLocalLocation sets the local location of this member.
// This value will be added to the Location event field if not populated from another member.
func (s *Server) SetLocalLocation(location string) {
	s.lock.Lock()
	defer s.lock.Unlock()

	s.location = location
}

// AddListener creates and returns a new event listener.
func (s *Server) AddListener(projectName string, allProjects bool, connection EventListenerConnection, messageTypes []string, excludeSources []EventSource, recvFunc EventHandler, excludeLocations []string) (*Listener, error) {
	if allProjects && projectName != "" {
		return nil, fmt.Errorf("Cannot specify project name when listening for events on all projects")
	}

	listener := &Listener{
		listenerCommon: listenerCommon{
			EventListenerConnection: connection,
			messageTypes:            messageTypes,
			done:                    cancel.New(context.Background()),
			id:                      uuid.New(),
			recvFunc:                recvFunc,
		},

		allProjects:      allProjects,
		projectName:      projectName,
		excludeSources:   excludeSources,
		excludeLocations: excludeLocations,
	}

	s.lock.Lock()
	defer s.lock.Unlock()

	if s.listeners[listener.id] != nil {
		return nil, fmt.Errorf("A listener with ID %q already exists", listener.id)
	}

	s.listeners[listener.id] = listener

	go listener.start()

	return listener, nil
}

// SendLifecycle broadcasts a lifecycle event.
func (s *Server) SendLifecycle(projectName string, event api.EventLifecycle) {
	_ = s.Send(projectName, api.EventTypeLifecycle, event)
}

// Send broadcasts a custom event.
func (s *Server) Send(projectName string, eventType string, eventMessage any) error {
	encodedMessage, err := json.Marshal(eventMessage)
	if err != nil {
		return err
	}

	event := api.Event{
		Type:      eventType,
		Timestamp: time.Now(),
		Metadata:  encodedMessage,
		Project:   projectName,
	}

	return s.broadcast(event, EventSourceLocal)
}

// Inject an event from another member into the local events dispatcher.
// eventSource is used to indicate where this event was received from.
func (s *Server) Inject(event api.Event, eventSource EventSource) {
	if event.Type == api.EventTypeLogging {
		// Parse the message
		logEntry := api.EventLogging{}
		err := json.Unmarshal(event.Metadata, &logEntry)
		if err != nil {
			return
		}

		if !s.debug && logEntry.Level == "debug" {
			return
		}

		if !s.debug && !s.verbose && logEntry.Level == "info" {
			return
		}
	}

	err := s.broadcast(event, eventSource)
	if err != nil {
		logger.Warn("Failed to forward event from member", logger.Ctx{"member": event.Location, "err": err})
	}
}

func (s *Server) broadcast(event api.Event, eventSource EventSource) error {
	sourceInSlice := func(source EventSource, sources []EventSource) bool {
		for _, i := range sources {
			if source == i {
				return true
			}
		}

		return false
	}

	s.lock.Lock()

	// Set the Location for local events to the local serverName if not already populated (do it here rather
	// than in Send as the lock to read s.location has been taken here already).
	if eventSource == EventSourceLocal && event.Location == "" {
		event.Location = s.location
	}

	// If a notifcation hook is present, then call it for locally produced events.
	// This can be used to send local events to another target (such as an event-hub member).
	if s.notify != nil && eventSource == EventSourceLocal {
		s.notify(event)
	}

	listeners := s.listeners
	for _, listener := range listeners {
		// If the event is project specific, check if the listener is requesting events from that project.
		if event.Project != "" && !listener.allProjects && event.Project != listener.projectName {
			continue
		}

		if sourceInSlice(eventSource, listener.excludeSources) {
			continue
		}

		if !shared.StringInSlice(event.Type, listener.messageTypes) {
			continue
		}

		// If the event doesn't come from this member and has been excluded by listener, don't deliver it.
		if eventSource != EventSourceLocal && shared.StringInSlice(event.Location, listener.excludeLocations) {
			continue
		}

		go func(listener *Listener, event api.Event) {
			// Check that the listener still exists
			if listener == nil {
				return
			}

			// Make sure we're not done already
			if listener.IsClosed() {
				// Remove the listener from the list
				s.lock.Lock()
				delete(s.listeners, listener.id)
				s.lock.Unlock()
				return
			}

			err := listener.WriteJSON(event)
			if err != nil {
				// Remove the listener from the list
				s.lock.Lock()
				delete(s.listeners, listener.id)
				s.lock.Unlock()

				listener.Close()
			}
		}(listener, event)
	}

	s.lock.Unlock()

	return nil
}

// Listener describes an event listener.
type Listener struct {
	listenerCommon

	allProjects      bool
	projectName      string
	excludeSources   []EventSource
	excludeLocations []string
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 provision

import (
	"bytes"
	"fmt"
	"path"
	"text/template"
	"time"

	"github.com/docker/machine/libmachine/auth"
	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/engine"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/provision"
	"github.com/docker/machine/libmachine/provision/pkgaction"
	"github.com/docker/machine/libmachine/swarm"
	"k8s.io/minikube/pkg/util"
)

type BuildrootProvisioner struct {
	provision.SystemdProvisioner
}

func init() {
	provision.Register("Buildroot", &provision.RegisteredProvisioner{
		New: NewBuildrootProvisioner,
	})
}

func NewBuildrootProvisioner(d drivers.Driver) provision.Provisioner {
	return &BuildrootProvisioner{
		provision.NewSystemdProvisioner("buildroot", d),
	}
}

func (p *BuildrootProvisioner) String() string {
	return "buildroot"
}

func (p *BuildrootProvisioner) GenerateDockerOptions(dockerPort int) (*provision.DockerOptions, error) {
	var engineCfg bytes.Buffer

	driverNameLabel := fmt.Sprintf("provider=%s", p.Driver.DriverName())
	p.EngineOptions.Labels = append(p.EngineOptions.Labels, driverNameLabel)

	engineConfigTmpl := `[Unit]
Description=Docker Application Container Engine
Documentation=https://docs.docker.com
After=network.target docker.socket
Requires=docker.socket

[Service]
Type=notify

# DOCKER_RAMDISK disables pivot_root in Docker, using MS_MOVE instead.
Environment=DOCKER_RAMDISK=yes
{{range .EngineOptions.Env}}Environment={{.}}
{{end}}

# This file is a systemd drop-in unit that inherits from the base dockerd configuration.
# The base configuration already specifies an 'ExecStart=...' command. The first directive
# here is to clear out that command inherited from the base configuration. Without this,
# the command from the base configuration and the command specified here are treated as
# a sequence of commands, which is not the desired behavior, nor is it valid -- systemd
# will catch this invalid input and refuse to start the service with an error like:
#  Service has more than one ExecStart= setting, which is only allowed for Type=oneshot services.
ExecStart=
ExecStart=/usr/bin/docker daemon -H tcp://0.0.0.0:{{.DockerPort}} -H unix:///var/run/docker.sock --tlsverify --tlscacert {{.AuthOptions.CaCertRemotePath}} --tlscert {{.AuthOptions.ServerCertRemotePath}} --tlskey {{.AuthOptions.ServerKeyRemotePath}} {{ range .EngineOptions.Labels }}--label {{.}} {{ end }}{{ range .EngineOptions.InsecureRegistry }}--insecure-registry {{.}} {{ end }}{{ range .EngineOptions.RegistryMirror }}--registry-mirror {{.}} {{ end }}{{ range .EngineOptions.ArbitraryFlags }}--{{.}} {{ end }}
ExecReload=/bin/kill -s HUP $MAINPID

# Having non-zero Limit*s causes performance problems due to accounting overhead
# in the kernel. We recommend using cgroups to do container-local accounting.
LimitNOFILE=infinity
LimitNPROC=infinity
LimitCORE=infinity

# Uncomment TasksMax if your systemd version supports it.
# Only systemd 226 and above support this version.
TasksMax=infinity
TimeoutStartSec=0

# set delegate yes so that systemd does not reset the cgroups of docker containers
Delegate=yes

# kill only the docker process, not all processes in the cgroup
KillMode=process

[Install]
WantedBy=multi-user.target
`
	t, err := template.New("engineConfig").Parse(engineConfigTmpl)
	if err != nil {
		return nil, err
	}

	engineConfigContext := provision.EngineConfigContext{
		DockerPort:    dockerPort,
		AuthOptions:   p.AuthOptions,
		EngineOptions: p.EngineOptions,
	}

	t.Execute(&engineCfg, engineConfigContext)

	return &provision.DockerOptions{
		EngineOptions:     engineCfg.String(),
		EngineOptionsPath: p.DaemonOptionsFile,
	}, nil
}

func (p *BuildrootProvisioner) Package(name string, action pkgaction.PackageAction) error {
	return nil
}

func (p *BuildrootProvisioner) Provision(swarmOptions swarm.Options, authOptions auth.Options, engineOptions engine.Options) error {
	p.SwarmOptions = swarmOptions
	p.AuthOptions = authOptions
	p.EngineOptions = engineOptions

	log.Debugf("setting hostname %q", p.Driver.GetMachineName())
	if err := p.SetHostname(p.Driver.GetMachineName()); err != nil {
		return err
	}

	p.AuthOptions = setRemoteAuthOptions(p)
	log.Debugf("set auth options %+v", p.AuthOptions)

	log.Debugf("setting up certificates")

	configureAuth := func() error {
		if err := provision.ConfigureAuth(p); err != nil {
			return &util.RetriableError{Err: err}
		}
		return nil
	}

	err := util.RetryAfter(5, configureAuth, time.Second*10)
	if err != nil {
		log.Debugf("Error configuring auth during provisioning %v", err)
		return err
	}

	return nil
}

func setRemoteAuthOptions(p provision.Provisioner) auth.Options {
	dockerDir := p.GetDockerOptionsDir()
	authOptions := p.GetAuthOptions()

	// due to windows clients, we cannot use filepath.Join as the paths
	// will be mucked on the linux hosts
	authOptions.CaCertRemotePath = path.Join(dockerDir, "ca.pem")
	authOptions.ServerCertRemotePath = path.Join(dockerDir, "server.pem")
	authOptions.ServerKeyRemotePath = path.Join(dockerDir, "server-key.pem")

	return authOptions
}
<commit_msg>docker daemon must be started as 'dockerd' in newer versions of docker<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 provision

import (
	"bytes"
	"fmt"
	"path"
	"text/template"
	"time"

	"github.com/docker/machine/libmachine/auth"
	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/engine"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/provision"
	"github.com/docker/machine/libmachine/provision/pkgaction"
	"github.com/docker/machine/libmachine/swarm"
	"k8s.io/minikube/pkg/util"
)

type BuildrootProvisioner struct {
	provision.SystemdProvisioner
}

func init() {
	provision.Register("Buildroot", &provision.RegisteredProvisioner{
		New: NewBuildrootProvisioner,
	})
}

func NewBuildrootProvisioner(d drivers.Driver) provision.Provisioner {
	return &BuildrootProvisioner{
		provision.NewSystemdProvisioner("buildroot", d),
	}
}

func (p *BuildrootProvisioner) String() string {
	return "buildroot"
}

func (p *BuildrootProvisioner) GenerateDockerOptions(dockerPort int) (*provision.DockerOptions, error) {
	var engineCfg bytes.Buffer

	driverNameLabel := fmt.Sprintf("provider=%s", p.Driver.DriverName())
	p.EngineOptions.Labels = append(p.EngineOptions.Labels, driverNameLabel)

	engineConfigTmpl := `[Unit]
Description=Docker Application Container Engine
Documentation=https://docs.docker.com
After=network.target docker.socket
Requires=docker.socket

[Service]
Type=notify

# DOCKER_RAMDISK disables pivot_root in Docker, using MS_MOVE instead.
Environment=DOCKER_RAMDISK=yes
{{range .EngineOptions.Env}}Environment={{.}}
{{end}}

# This file is a systemd drop-in unit that inherits from the base dockerd configuration.
# The base configuration already specifies an 'ExecStart=...' command. The first directive
# here is to clear out that command inherited from the base configuration. Without this,
# the command from the base configuration and the command specified here are treated as
# a sequence of commands, which is not the desired behavior, nor is it valid -- systemd
# will catch this invalid input and refuse to start the service with an error like:
#  Service has more than one ExecStart= setting, which is only allowed for Type=oneshot services.
ExecStart=
ExecStart=/usr/bin/dockerd -H tcp://0.0.0.0:{{.DockerPort}} -H unix:///var/run/docker.sock --tlsverify --tlscacert {{.AuthOptions.CaCertRemotePath}} --tlscert {{.AuthOptions.ServerCertRemotePath}} --tlskey {{.AuthOptions.ServerKeyRemotePath}} {{ range .EngineOptions.Labels }}--label {{.}} {{ end }}{{ range .EngineOptions.InsecureRegistry }}--insecure-registry {{.}} {{ end }}{{ range .EngineOptions.RegistryMirror }}--registry-mirror {{.}} {{ end }}{{ range .EngineOptions.ArbitraryFlags }}--{{.}} {{ end }}
ExecReload=/bin/kill -s HUP $MAINPID

# Having non-zero Limit*s causes performance problems due to accounting overhead
# in the kernel. We recommend using cgroups to do container-local accounting.
LimitNOFILE=infinity
LimitNPROC=infinity
LimitCORE=infinity

# Uncomment TasksMax if your systemd version supports it.
# Only systemd 226 and above support this version.
TasksMax=infinity
TimeoutStartSec=0

# set delegate yes so that systemd does not reset the cgroups of docker containers
Delegate=yes

# kill only the docker process, not all processes in the cgroup
KillMode=process

[Install]
WantedBy=multi-user.target
`
	t, err := template.New("engineConfig").Parse(engineConfigTmpl)
	if err != nil {
		return nil, err
	}

	engineConfigContext := provision.EngineConfigContext{
		DockerPort:    dockerPort,
		AuthOptions:   p.AuthOptions,
		EngineOptions: p.EngineOptions,
	}

	t.Execute(&engineCfg, engineConfigContext)

	return &provision.DockerOptions{
		EngineOptions:     engineCfg.String(),
		EngineOptionsPath: p.DaemonOptionsFile,
	}, nil
}

func (p *BuildrootProvisioner) Package(name string, action pkgaction.PackageAction) error {
	return nil
}

func (p *BuildrootProvisioner) Provision(swarmOptions swarm.Options, authOptions auth.Options, engineOptions engine.Options) error {
	p.SwarmOptions = swarmOptions
	p.AuthOptions = authOptions
	p.EngineOptions = engineOptions

	log.Debugf("setting hostname %q", p.Driver.GetMachineName())
	if err := p.SetHostname(p.Driver.GetMachineName()); err != nil {
		return err
	}

	p.AuthOptions = setRemoteAuthOptions(p)
	log.Debugf("set auth options %+v", p.AuthOptions)

	log.Debugf("setting up certificates")

	configureAuth := func() error {
		if err := provision.ConfigureAuth(p); err != nil {
			return &util.RetriableError{Err: err}
		}
		return nil
	}

	err := util.RetryAfter(5, configureAuth, time.Second*10)
	if err != nil {
		log.Debugf("Error configuring auth during provisioning %v", err)
		return err
	}

	return nil
}

func setRemoteAuthOptions(p provision.Provisioner) auth.Options {
	dockerDir := p.GetDockerOptionsDir()
	authOptions := p.GetAuthOptions()

	// due to windows clients, we cannot use filepath.Join as the paths
	// will be mucked on the linux hosts
	authOptions.CaCertRemotePath = path.Join(dockerDir, "ca.pem")
	authOptions.ServerCertRemotePath = path.Join(dockerDir, "server.pem")
	authOptions.ServerKeyRemotePath = path.Join(dockerDir, "server-key.pem")

	return authOptions
}
<|endoftext|>


<commit_before>// Copyright 2014 Wandoujia Inc. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package router

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/reborndb/reborn/pkg/utils"

	"github.com/reborndb/reborn/pkg/models"
	"github.com/reborndb/reborn/pkg/proxy/group"
	"github.com/reborndb/reborn/pkg/proxy/parser"
	"github.com/reborndb/reborn/pkg/proxy/router/topology"

	log "github.com/ngaut/logging"

	"github.com/juju/errors"
	topo "github.com/ngaut/go-zookeeper/zk"
	stats "github.com/ngaut/gostats"

	respcoding "github.com/ngaut/resp"
)

var blackList = []string{
	"KEYS", "MOVE", "OBJECT", "RENAME", "RENAMENX", "SORT", "SCAN", "BITOP" /*"MGET",*/ /* "MSET",*/, "MSETNX", "SCAN",
	"BLPOP", "BRPOP", "BRPOPLPUSH", "PSUBSCRIBE，PUBLISH", "PUNSUBSCRIBE", "SUBSCRIBE", "RANDOMKEY",
	"UNSUBSCRIBE", "DISCARD", "EXEC", "MULTI", "UNWATCH", "WATCH", "SCRIPT EXISTS", "SCRIPT FLUSH", "SCRIPT KILL",
	"SCRIPT LOAD" /*, "AUTH" , "ECHO"*/ /*"QUIT",*/ /*"SELECT",*/, "BGREWRITEAOF", "BGSAVE", "CLIENT KILL", "CLIENT LIST",
	"CONFIG GET", "CONFIG SET", "CONFIG RESETSTAT", "DBSIZE", "DEBUG OBJECT", "DEBUG SEGFAULT", "FLUSHALL", "FLUSHDB",
	"LASTSAVE", "MONITOR", "SAVE", "SHUTDOWN", "SLAVEOF", "SLOWLOG", "SYNC", "TIME", "SLOTSMGRTONE", "SLOTSMGRT",
	"SLOTSDEL",
}

var (
	blackListCommand = make(map[string]struct{})
	OK_BYTES         = []byte("+OK\r\n")
)

func init() {
	for _, k := range blackList {
		blackListCommand[k] = struct{}{}
	}
}

func allowOp(op string) bool {
	_, black := blackListCommand[op]
	return !black
}

func isMulOp(op string) bool {
	if op == "MGET" || op == "DEL" || op == "MSET" {
		return true
	}

	return false
}

func validSlot(i int) bool {
	if i < 0 || i >= models.DEFAULT_SLOT_NUM {
		return false
	}

	return true
}

func WriteMigrateKeyCmd(w io.Writer, addr string, timeoutMs int, key []byte) error {
	hostPort := strings.Split(addr, ":")
	if len(hostPort) != 2 {
		return errors.Errorf("invalid address " + addr)
	}
	respW := respcoding.NewRESPWriter(w)
	err := respW.WriteCommand("slotsmgrttagone", hostPort[0], hostPort[1],
		strconv.Itoa(int(timeoutMs)), string(key))
	return errors.Trace(err)
}

type DeadlineReadWriter interface {
	io.ReadWriter
	SetWriteDeadline(t time.Time) error
	SetReadDeadline(t time.Time) error
}

func handleSpecCommand(cmd string, keys [][]byte, timeout int) ([]byte, bool, bool, error) {
	var b []byte
	shouldClose := false
	switch cmd {
	case "PING":
		b = []byte("+PONG\r\n")
	case "QUIT":
		b = OK_BYTES
		shouldClose = true
	case "SELECT":
		b = OK_BYTES
	case "AUTH":
		b = OK_BYTES
	case "ECHO":
		if len(keys) > 0 {
			var err error
			b, err = respcoding.Marshal(string(keys[0]))
			if err != nil {
				return nil, true, false, errors.Trace(err)
			}
		} else {
			return nil, true, false, nil
		}
	}

	if len(b) > 0 {
		return b, shouldClose, true, nil
	}

	return b, shouldClose, false, nil
}

func write2Client(redisReader *bufio.Reader, clientWriter io.Writer) (redisErr error, clientErr error) {
	resp, err := parser.Parse(redisReader)
	if err != nil {
		return errors.Trace(err), errors.Trace(err)
	}

	b, err := resp.Bytes()
	if err != nil {
		return errors.Trace(err), errors.Trace(err)
	}

	_, err = clientWriter.Write(b)
	return nil, errors.Trace(err)
}

func write2Redis(resp *parser.Resp, redisWriter io.Writer) error {
	// get resp in bytes
	b, err := resp.Bytes()
	if err != nil {
		return errors.Trace(err)
	}

	return writeBytes2Redis(b, redisWriter)
}

func writeBytes2Redis(b []byte, redisWriter io.Writer) error {
	// write to redis
	_, err := redisWriter.Write(b)
	return errors.Trace(err)
}

type BufioDeadlineReadWriter interface {
	DeadlineReadWriter
	BufioReader() *bufio.Reader
}

func forward(c DeadlineReadWriter, redisConn BufioDeadlineReadWriter, resp *parser.Resp, timeout int) (redisErr error, clientErr error) {
	redisReader := redisConn.BufioReader()
	if err := redisConn.SetWriteDeadline(time.Now().Add(time.Duration(timeout) * time.Second)); err != nil {
		return errors.Trace(err), errors.Trace(err)
	}

	if err := write2Redis(resp, redisConn); err != nil {
		return errors.Trace(err), errors.Trace(err)
	}

	if err := redisConn.SetReadDeadline(time.Now().Add(time.Duration(timeout) * time.Second)); err != nil {
		return errors.Trace(err), errors.Trace(err)
	}

	if err := c.SetWriteDeadline(time.Now().Add(time.Duration(timeout) * time.Second)); err != nil {
		return nil, errors.Trace(err)
	}

	// read and parse redis response
	return write2Client(redisReader, c)
}

func StringsContain(s []string, key string) bool {
	for _, val := range s {
		if val == key { //need our resopnse
			return true
		}
	}

	return false
}

func getRespOpKeys(c *session) (*parser.Resp, []byte, [][]byte, error) {
	resp, err := parser.Parse(c.r) // read client request
	if err != nil {
		return nil, nil, nil, errors.Trace(err)
	}

	op, keys, err := resp.GetOpKeys()
	if err != nil {
		return nil, nil, nil, errors.Trace(err)
	}

	if len(keys) == 0 {
		keys = [][]byte{[]byte("fakeKey")}
	}

	return resp, op, keys, nil
}

func filter(opstr string, keys [][]byte, c *session, timeoutSec int) (rawresp []byte, next bool, err error) {
	if !allowOp(opstr) {
		errmsg, err := respcoding.Marshal(fmt.Errorf("%s not allowed", opstr))
		if err != nil {
			log.Fatal("should never happend", opstr)
		}
		return errmsg, false, errors.Trace(fmt.Errorf("%s not allowed", opstr))
	}

	buf, shouldClose, handled, err := handleSpecCommand(opstr, keys, timeoutSec)
	if shouldClose { //quit command
		return buf, false, errors.Trace(io.EOF)
	}
	if err != nil {
		return nil, false, errors.Trace(err)
	}

	if handled {
		return buf, false, nil
	}

	return nil, true, nil
}

func GetEventPath(evt interface{}) string {
	return evt.(topo.Event).Path
}

func CheckUlimit(min int) {
	ulimitN, err := exec.Command("/bin/sh", "-c", "ulimit -n").Output()
	if err != nil {
		log.Warning("get ulimit failed", err)
	}

	n, err := strconv.Atoi(strings.TrimSpace(string(ulimitN)))
	if err != nil || n < min {
		log.Fatalf("ulimit too small: %d, should be at least %d", n, min)
	}
}

func GetOriginError(err *errors.Err) error {
	if err != nil {
		if err.Cause() == nil && err.Underlying() == nil {
			return err
		} else {
			return err.Underlying()
		}
	}

	return err
}

func recordResponseTime(c *stats.Counters, d time.Duration) {
	switch {
	case d < 5:
		c.Add("0-5ms", 1)
	case d >= 5 && d < 10:
		c.Add("5-10ms", 1)
	case d >= 10 && d < 50:
		c.Add("10-50ms", 1)
	case d >= 50 && d < 200:
		c.Add("50-200ms", 1)
	case d >= 200 && d < 1000:
		c.Add("200-1000ms", 1)
	case d >= 1000 && d < 5000:
		c.Add("1000-5000ms", 1)
	case d >= 5000 && d < 10000:
		c.Add("5000-10000ms", 1)
	default:
		c.Add("10000ms+", 1)
	}
}

type killEvent struct {
	done chan error
}

type Conf struct {
	proxyId     string
	productName string
	zkAddr      string
	f           topology.ZkFactory
	netTimeout  int    //seconds
	proto       string //tcp or tcp4
	provider    string
}

func LoadConf(configFile string) (*Conf, error) {
	srvConf := &Conf{}
	conf, err := utils.InitConfigFromFile(configFile)
	if err != nil {
		log.Fatal(err)
	}

	srvConf.productName, _ = conf.ReadString("product", "test")
	if len(srvConf.productName) == 0 {
		log.Fatalf("invalid config: product entry is missing in %s", configFile)
	}
	srvConf.zkAddr, _ = conf.ReadString("zk", "")
	if len(srvConf.zkAddr) == 0 {
		log.Fatalf("invalid config: need zk entry is missing in %s", configFile)
	}
	srvConf.zkAddr = strings.TrimSpace(srvConf.zkAddr)

	srvConf.proxyId, _ = conf.ReadString("proxy_id", "")
	if len(srvConf.proxyId) == 0 {
		log.Fatalf("invalid config: need proxy_id entry is missing in %s", configFile)
	}

	srvConf.netTimeout, _ = conf.ReadInt("net_timeout", 5)
	srvConf.proto, _ = conf.ReadString("proto", "tcp")
	srvConf.provider, _ = conf.ReadString("coordinator", "zookeeper")
	log.Infof("%+v", srvConf)

	return srvConf, nil
}

type Slot struct {
	slotInfo    *models.Slot
	groupInfo   *models.ServerGroup
	dst         *group.Group
	migrateFrom *group.Group
}

type OnSuicideFun func() error

func needResponse(receivers []string, self models.ProxyInfo) bool {
	var pi models.ProxyInfo
	for _, v := range receivers {
		err := json.Unmarshal([]byte(v), &pi)
		if err != nil {
			//is it old version of dashboard
			if v == self.Id {
				return true
			}
			return false
		}

		if pi.Id == self.Id && pi.Pid == self.Pid && pi.StartAt == self.StartAt {
			return true
		}
	}

	return false
}
<commit_msg>Tiny clean up<commit_after>// Copyright 2014 Wandoujia Inc. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package router

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/reborndb/reborn/pkg/utils"

	"github.com/reborndb/reborn/pkg/models"
	"github.com/reborndb/reborn/pkg/proxy/group"
	"github.com/reborndb/reborn/pkg/proxy/parser"
	"github.com/reborndb/reborn/pkg/proxy/router/topology"

	log "github.com/ngaut/logging"

	"github.com/juju/errors"
	topo "github.com/ngaut/go-zookeeper/zk"
	stats "github.com/ngaut/gostats"

	respcoding "github.com/ngaut/resp"
)

var blackList = []string{
	"KEYS", "MOVE", "OBJECT", "RENAME", "RENAMENX", "SORT", "SCAN", "BITOP" /*"MGET",*/ /* "MSET",*/, "MSETNX", "SCAN",
	"BLPOP", "BRPOP", "BRPOPLPUSH", "PSUBSCRIBE，PUBLISH", "PUNSUBSCRIBE", "SUBSCRIBE", "RANDOMKEY",
	"UNSUBSCRIBE", "DISCARD", "EXEC", "MULTI", "UNWATCH", "WATCH", "SCRIPT EXISTS", "SCRIPT FLUSH", "SCRIPT KILL",
	"SCRIPT LOAD" /*, "AUTH" , "ECHO"*/ /*"QUIT",*/ /*"SELECT",*/, "BGREWRITEAOF", "BGSAVE", "CLIENT KILL", "CLIENT LIST",
	"CONFIG GET", "CONFIG SET", "CONFIG RESETSTAT", "DBSIZE", "DEBUG OBJECT", "DEBUG SEGFAULT", "FLUSHALL", "FLUSHDB",
	"LASTSAVE", "MONITOR", "SAVE", "SHUTDOWN", "SLAVEOF", "SLOWLOG", "SYNC", "TIME", "SLOTSMGRTONE", "SLOTSMGRT",
	"SLOTSDEL",
}

var (
	blackListCommand = make(map[string]struct{})
	OK_BYTES         = []byte("+OK\r\n")
)

func init() {
	for _, k := range blackList {
		blackListCommand[k] = struct{}{}
	}
}

func allowOp(op string) bool {
	_, black := blackListCommand[op]
	return !black
}

func isMulOp(op string) bool {
	if op == "MGET" || op == "DEL" || op == "MSET" {
		return true
	}

	return false
}

func validSlot(i int) bool {
	if i < 0 || i >= models.DEFAULT_SLOT_NUM {
		return false
	}

	return true
}

func WriteMigrateKeyCmd(w io.Writer, addr string, timeoutMs int, key []byte) error {
	hostPort := strings.Split(addr, ":")
	if len(hostPort) != 2 {
		return errors.Errorf("invalid address " + addr)
	}
	respW := respcoding.NewRESPWriter(w)
	err := respW.WriteCommand("slotsmgrttagone", hostPort[0], hostPort[1],
		strconv.Itoa(int(timeoutMs)), string(key))
	return errors.Trace(err)
}

type DeadlineReadWriter interface {
	io.ReadWriter
	SetWriteDeadline(t time.Time) error
	SetReadDeadline(t time.Time) error
}

func handleSpecCommand(cmd string, keys [][]byte, timeout int) ([]byte, bool, bool, error) {
	var b []byte
	shouldClose := false
	switch cmd {
	case "PING":
		b = []byte("+PONG\r\n")
	case "QUIT":
		b = OK_BYTES
		shouldClose = true
	case "SELECT":
		b = OK_BYTES
	case "AUTH":
		b = OK_BYTES
	case "ECHO":
		if len(keys) > 0 {
			var err error
			b, err = respcoding.Marshal(string(keys[0]))
			if err != nil {
				return nil, true, false, errors.Trace(err)
			}
		} else {
			return nil, true, false, nil
		}
	}

	if len(b) > 0 {
		return b, shouldClose, true, nil
	}

	return b, shouldClose, false, nil
}

func write2Client(redisReader *bufio.Reader, clientWriter io.Writer) (redisErr error, clientErr error) {
	resp, err := parser.Parse(redisReader)
	if err != nil {
		return errors.Trace(err), errors.Trace(err)
	}

	b, err := resp.Bytes()
	if err != nil {
		return errors.Trace(err), errors.Trace(err)
	}

	_, err = clientWriter.Write(b)
	return nil, errors.Trace(err)
}

func write2Redis(resp *parser.Resp, redisWriter io.Writer) error {
	// get resp in bytes
	b, err := resp.Bytes()
	if err != nil {
		return errors.Trace(err)
	}

	return writeBytes2Redis(b, redisWriter)
}

func writeBytes2Redis(b []byte, redisWriter io.Writer) error {
	// write to redis
	_, err := redisWriter.Write(b)
	return errors.Trace(err)
}

type BufioDeadlineReadWriter interface {
	DeadlineReadWriter
	BufioReader() *bufio.Reader
}

func forward(c DeadlineReadWriter, redisConn BufioDeadlineReadWriter, resp *parser.Resp, timeout int) (redisErr error, clientErr error) {
	redisReader := redisConn.BufioReader()
	if err := redisConn.SetWriteDeadline(time.Now().Add(time.Duration(timeout) * time.Second)); err != nil {
		return errors.Trace(err), errors.Trace(err)
	}

	if err := write2Redis(resp, redisConn); err != nil {
		return errors.Trace(err), errors.Trace(err)
	}

	if err := redisConn.SetReadDeadline(time.Now().Add(time.Duration(timeout) * time.Second)); err != nil {
		return errors.Trace(err), errors.Trace(err)
	}

	if err := c.SetWriteDeadline(time.Now().Add(time.Duration(timeout) * time.Second)); err != nil {
		return nil, errors.Trace(err)
	}

	// read and parse redis response
	return write2Client(redisReader, c)
}

func StringsContain(s []string, key string) bool {
	for _, val := range s {
		if val == key { //need our resopnse
			return true
		}
	}

	return false
}

func getRespOpKeys(c *session) (*parser.Resp, []byte, [][]byte, error) {
	resp, err := parser.Parse(c.r) // read client request
	if err != nil {
		return nil, nil, nil, errors.Trace(err)
	}

	op, keys, err := resp.GetOpKeys()
	if err != nil {
		return nil, nil, nil, errors.Trace(err)
	}

	if len(keys) == 0 {
		keys = [][]byte{[]byte("fakeKey")}
	}

	return resp, op, keys, nil
}

func filter(opstr string, keys [][]byte, c *session, timeoutSec int) (rawresp []byte, next bool, err error) {
	if !allowOp(opstr) {
		errmsg, err := respcoding.Marshal(fmt.Errorf("%s not allowed", opstr))
		if err != nil {
			log.Fatal("should never happend", opstr)
		}
		return errmsg, false, errors.New(string(errmsg))
	}

	buf, shouldClose, handled, err := handleSpecCommand(opstr, keys, timeoutSec)
	if shouldClose { //quit command
		return buf, false, errors.Trace(io.EOF)
	}
	if err != nil {
		return nil, false, errors.Trace(err)
	}

	if handled {
		return buf, false, nil
	}

	return nil, true, nil
}

func GetEventPath(evt interface{}) string {
	return evt.(topo.Event).Path
}

func CheckUlimit(min int) {
	ulimitN, err := exec.Command("/bin/sh", "-c", "ulimit -n").Output()
	if err != nil {
		log.Warning("get ulimit failed", err)
	}

	n, err := strconv.Atoi(strings.TrimSpace(string(ulimitN)))
	if err != nil || n < min {
		log.Fatalf("ulimit too small: %d, should be at least %d", n, min)
	}
}

func GetOriginError(err *errors.Err) error {
	if err != nil {
		if err.Cause() == nil && err.Underlying() == nil {
			return err
		} else {
			return err.Underlying()
		}
	}

	return err
}

func recordResponseTime(c *stats.Counters, d time.Duration) {
	switch {
	case d < 5:
		c.Add("0-5ms", 1)
	case d >= 5 && d < 10:
		c.Add("5-10ms", 1)
	case d >= 10 && d < 50:
		c.Add("10-50ms", 1)
	case d >= 50 && d < 200:
		c.Add("50-200ms", 1)
	case d >= 200 && d < 1000:
		c.Add("200-1000ms", 1)
	case d >= 1000 && d < 5000:
		c.Add("1000-5000ms", 1)
	case d >= 5000 && d < 10000:
		c.Add("5000-10000ms", 1)
	default:
		c.Add("10000ms+", 1)
	}
}

type killEvent struct {
	done chan error
}

type Conf struct {
	proxyId     string
	productName string
	zkAddr      string
	f           topology.ZkFactory
	netTimeout  int    //seconds
	proto       string //tcp or tcp4
	provider    string
}

func LoadConf(configFile string) (*Conf, error) {
	srvConf := &Conf{}
	conf, err := utils.InitConfigFromFile(configFile)
	if err != nil {
		log.Fatal(err)
	}

	srvConf.productName, _ = conf.ReadString("product", "test")
	if len(srvConf.productName) == 0 {
		log.Fatalf("invalid config: product entry is missing in %s", configFile)
	}
	srvConf.zkAddr, _ = conf.ReadString("zk", "")
	if len(srvConf.zkAddr) == 0 {
		log.Fatalf("invalid config: need zk entry is missing in %s", configFile)
	}
	srvConf.zkAddr = strings.TrimSpace(srvConf.zkAddr)

	srvConf.proxyId, _ = conf.ReadString("proxy_id", "")
	if len(srvConf.proxyId) == 0 {
		log.Fatalf("invalid config: need proxy_id entry is missing in %s", configFile)
	}

	srvConf.netTimeout, _ = conf.ReadInt("net_timeout", 5)
	srvConf.proto, _ = conf.ReadString("proto", "tcp")
	srvConf.provider, _ = conf.ReadString("coordinator", "zookeeper")
	log.Infof("%+v", srvConf)

	return srvConf, nil
}

type Slot struct {
	slotInfo    *models.Slot
	groupInfo   *models.ServerGroup
	dst         *group.Group
	migrateFrom *group.Group
}

type OnSuicideFun func() error

func needResponse(receivers []string, self models.ProxyInfo) bool {
	var pi models.ProxyInfo
	for _, v := range receivers {
		err := json.Unmarshal([]byte(v), &pi)
		if err != nil {
			//is it old version of dashboard
			if v == self.Id {
				return true
			}
			return false
		}

		if pi.Id == self.Id && pi.Pid == self.Pid && pi.StartAt == self.StartAt {
			return true
		}
	}

	return false
}
<|endoftext|>