<commit_before>package server

import (
	"bytes"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

	"go-rest/server/context"

	"github.com/gorilla/mux"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
)

type Resource struct {
	Foo string `json:"foo"`
}

type MockResourceHandler struct {
	mock.Mock
}

func (m *MockResourceHandler) ResourceName() string {
	args := m.Mock.Called()
	return args.String(0)
}

func (m *MockResourceHandler) CreateResource(r context.RequestContext, data map[string]interface{}) (interface{}, error) {
	args := m.Mock.Called()
	resource := args.Get(0)
	if resource != nil {
		resource = resource.(*Resource)
	}
	return resource, args.Error(1)
}

func (m *MockResourceHandler) ReadResource(r context.RequestContext, id string) (interface{}, error) {
	args := m.Mock.Called()
	resource := args.Get(0)
	if resource != nil {
		resource = resource.(*Resource)
	}
	return resource, args.Error(1)
}

func (m *MockResourceHandler) UpdateResource(r context.RequestContext, id string, data map[string]interface{}) (interface{}, error) {
	args := m.Mock.Called()
	return args.Get(0).(*Resource), args.Error(1)
}

func (m *MockResourceHandler) DeleteResource(r context.RequestContext, id string) (interface{}, error) {
	args := m.Mock.Called()
	return args.Get(0).(*Resource), args.Error(1)
}

// Ensures that the create handler returns a Not Implemented code if an invalid response format is provided.
func TestHandleCreateBadFormat(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")

	RegisterResourceHandler(router, handler)
	createHandler := router.Get("create").GetHandler()

	payload := []byte(`{"foo": "bar"}`)
	r := bytes.NewReader(payload)
	req, _ := http.NewRequest("POST", "http://foo.com/api/v0.1/foo?format=blah", r)
	resp := httptest.NewRecorder()

	createHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusNotImplemented, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"error":"Format not implemented: blah","success":false}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}

// Ensures that the create handler returns an Internal Server Error code when the createFunc returns an error.
func TestHandleCreateBadCreate(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")
	handler.On("CreateResource").Return(nil, fmt.Errorf("couldn't create"))

	RegisterResourceHandler(router, handler)
	createHandler := router.Get("create").GetHandler()

	payload := []byte(`{"foo": "bar"}`)
	r := bytes.NewReader(payload)
	req, _ := http.NewRequest("POST", "http://foo.com/api/v0.1/foo", r)
	resp := httptest.NewRecorder()

	createHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusInternalServerError, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"error":"couldn't create","success":false}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}

// Ensures that the create handler returns the serialized resource and Created code when createFunc succeeds.
func TestHandleCreateHappyPath(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")
	handler.On("CreateResource").Return(&Resource{Foo: "bar"}, nil)

	RegisterResourceHandler(router, handler)
	createHandler := router.Get("create").GetHandler()

	payload := []byte(`{"foo": "bar"}`)
	r := bytes.NewReader(payload)
	req, _ := http.NewRequest("POST", "http://foo.com/api/v0.1/foo", r)
	resp := httptest.NewRecorder()

	createHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusCreated, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"result":{"foo":"bar"},"success":true}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}

// Ensures that the read handler returns a Not Implemented code if an invalid response format is provided.
func TestHandleReadBadFormat(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")

	RegisterResourceHandler(router, handler)
	readHandler := router.Get("read").GetHandler()

	req, _ := http.NewRequest("GET", "http://foo.com/api/v0.1/foo/1?format=blah", nil)
	resp := httptest.NewRecorder()

	readHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusNotImplemented, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"error":"Format not implemented: blah","success":false}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}

// Ensures that the read handler returns an Internal Server Error code when the readFunc returns an error.
func TestHandleReadBadRead(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")
	handler.On("ReadResource").Return(nil, fmt.Errorf("no resource"))

	RegisterResourceHandler(router, handler)
	readHandler := router.Get("read").GetHandler()

	req, _ := http.NewRequest("GET", "http://foo.com/api/v0.1/foo/1", nil)
	resp := httptest.NewRecorder()

	readHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusInternalServerError, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"error":"no resource","success":false}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}

// Ensures that the read handler returns the serialized resource and OK code when readFunc succeeds.
func TestHandleReadHappyPath(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")
	handler.On("ReadResource").Return(&Resource{Foo: "hello"}, nil)

	RegisterResourceHandler(router, handler)
	readHandler := router.Get("read").GetHandler()

	req, _ := http.NewRequest("GET", "http://foo.com/api/v0.1/foo/1", nil)
	resp := httptest.NewRecorder()

	readHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusOK, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"result":{"foo":"hello"},"success":true}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}
<commit_msg>Add unit tests for update handler<commit_after>package server

import (
	"bytes"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

	"go-rest/server/context"

	"github.com/gorilla/mux"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
)

type Resource struct {
	Foo string `json:"foo"`
}

type MockResourceHandler struct {
	mock.Mock
}

func (m *MockResourceHandler) ResourceName() string {
	args := m.Mock.Called()
	return args.String(0)
}

func (m *MockResourceHandler) CreateResource(r context.RequestContext, data map[string]interface{}) (interface{}, error) {
	args := m.Mock.Called()
	resource := args.Get(0)
	if resource != nil {
		resource = resource.(*Resource)
	}
	return resource, args.Error(1)
}

func (m *MockResourceHandler) ReadResource(r context.RequestContext, id string) (interface{}, error) {
	args := m.Mock.Called()
	resource := args.Get(0)
	if resource != nil {
		resource = resource.(*Resource)
	}
	return resource, args.Error(1)
}

func (m *MockResourceHandler) UpdateResource(r context.RequestContext, id string, data map[string]interface{}) (interface{}, error) {
	args := m.Mock.Called()
	resource := args.Get(0)
	if resource != nil {
		resource = resource.(*Resource)
	}
	return resource, args.Error(1)
}

func (m *MockResourceHandler) DeleteResource(r context.RequestContext, id string) (interface{}, error) {
	args := m.Mock.Called()
	return args.Get(0).(*Resource), args.Error(1)
}

// Ensures that the create handler returns a Not Implemented code if an invalid response format is provided.
func TestHandleCreateBadFormat(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")

	RegisterResourceHandler(router, handler)
	createHandler := router.Get("create").GetHandler()

	payload := []byte(`{"foo": "bar"}`)
	r := bytes.NewReader(payload)
	req, _ := http.NewRequest("POST", "http://foo.com/api/v0.1/foo?format=blah", r)
	resp := httptest.NewRecorder()

	createHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusNotImplemented, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"error":"Format not implemented: blah","success":false}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}

// Ensures that the create handler returns an Internal Server Error code when the createFunc returns an error.
func TestHandleCreateBadCreate(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")
	handler.On("CreateResource").Return(nil, fmt.Errorf("couldn't create"))

	RegisterResourceHandler(router, handler)
	createHandler := router.Get("create").GetHandler()

	payload := []byte(`{"foo": "bar"}`)
	r := bytes.NewReader(payload)
	req, _ := http.NewRequest("POST", "http://foo.com/api/v0.1/foo", r)
	resp := httptest.NewRecorder()

	createHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusInternalServerError, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"error":"couldn't create","success":false}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}

// Ensures that the create handler returns the serialized resource and Created code when createFunc succeeds.
func TestHandleCreateHappyPath(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")
	handler.On("CreateResource").Return(&Resource{Foo: "bar"}, nil)

	RegisterResourceHandler(router, handler)
	createHandler := router.Get("create").GetHandler()

	payload := []byte(`{"foo": "bar"}`)
	r := bytes.NewReader(payload)
	req, _ := http.NewRequest("POST", "http://foo.com/api/v0.1/foo", r)
	resp := httptest.NewRecorder()

	createHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusCreated, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"result":{"foo":"bar"},"success":true}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}

// Ensures that the read handler returns a Not Implemented code if an invalid response format is provided.
func TestHandleReadBadFormat(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")

	RegisterResourceHandler(router, handler)
	readHandler := router.Get("read").GetHandler()

	req, _ := http.NewRequest("GET", "http://foo.com/api/v0.1/foo/1?format=blah", nil)
	resp := httptest.NewRecorder()

	readHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusNotImplemented, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"error":"Format not implemented: blah","success":false}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}

// Ensures that the read handler returns an Internal Server Error code when the readFunc returns an error.
func TestHandleReadBadRead(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")
	handler.On("ReadResource").Return(nil, fmt.Errorf("no resource"))

	RegisterResourceHandler(router, handler)
	readHandler := router.Get("read").GetHandler()

	req, _ := http.NewRequest("GET", "http://foo.com/api/v0.1/foo/1", nil)
	resp := httptest.NewRecorder()

	readHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusInternalServerError, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"error":"no resource","success":false}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}

// Ensures that the read handler returns the serialized resource and OK code when readFunc succeeds.
func TestHandleReadHappyPath(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")
	handler.On("ReadResource").Return(&Resource{Foo: "hello"}, nil)

	RegisterResourceHandler(router, handler)
	readHandler := router.Get("read").GetHandler()

	req, _ := http.NewRequest("GET", "http://foo.com/api/v0.1/foo/1", nil)
	resp := httptest.NewRecorder()

	readHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusOK, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"result":{"foo":"hello"},"success":true}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}

// Ensures that the update handler returns a Not Implemented code if an invalid response format is provided.
func TestHandleUpdateBadFormat(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")

	RegisterResourceHandler(router, handler)
	updateHandler := router.Get("update").GetHandler()

	payload := []byte(`{"foo": "bar"}`)
	r := bytes.NewReader(payload)
	req, _ := http.NewRequest("PUT", "http://foo.com/api/v0.1/foo/1?format=blah", r)
	resp := httptest.NewRecorder()

	updateHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusNotImplemented, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"error":"Format not implemented: blah","success":false}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}

// Ensures that the update handler returns an Internal Server Error code when the updateFunc returns an error.
func TestHandleUpdateBadUpdate(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")
	handler.On("UpdateResource").Return(nil, fmt.Errorf("couldn't update"))

	RegisterResourceHandler(router, handler)
	updateHandler := router.Get("update").GetHandler()

	payload := []byte(`{"foo": "bar"}`)
	r := bytes.NewReader(payload)
	req, _ := http.NewRequest("PUT", "http://foo.com/api/v0.1/foo/1", r)
	resp := httptest.NewRecorder()

	updateHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusInternalServerError, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"error":"couldn't update","success":false}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}

// Ensures that the update handler returns the serialized resource and OK code when updateFunc succeeds.
func TestHandleUpdateHappyPath(t *testing.T) {
	assert := assert.New(t)
	handler := new(MockResourceHandler)
	router := mux.NewRouter()

	handler.On("ResourceName").Return("foo")
	handler.On("UpdateResource").Return(&Resource{Foo: "bar"}, nil)

	RegisterResourceHandler(router, handler)
	updateHandler := router.Get("update").GetHandler()

	payload := []byte(`{"foo": "bar"}`)
	r := bytes.NewReader(payload)
	req, _ := http.NewRequest("PUT", "http://foo.com/api/v0.1/foo/1", r)
	resp := httptest.NewRecorder()

	updateHandler.ServeHTTP(resp, req)

	handler.Mock.AssertExpectations(t)
	assert.Equal(http.StatusOK, resp.Code, "Incorrect response code")
	assert.Equal(
		`{"result":{"foo":"bar"},"success":true}`,
		resp.Body.String(),
		"Incorrect response string",
	)
}
<|endoftext|>


<commit_before><commit_msg>Добавил хелперы для возврата ошибок из методов API на Gin<commit_after><|endoftext|>


<commit_before><commit_msg>修改配置文件<commit_after><|endoftext|>


<commit_before><commit_msg>Added some more info to a TODO<commit_after><|endoftext|>


<commit_before><commit_msg>Adds passthrough for "--compilation_mode" and "-c" (#348)<commit_after><|endoftext|>


<commit_before><commit_msg>Honor nat 1to1 config when creating agent<commit_after><|endoftext|>


<commit_before><commit_msg>Fix 8bit ADD/ADC (thanks z80)<commit_after><|endoftext|>


<commit_before>// +build linux

package socket

import (
	"fmt"
	"io"
	"sync"
	"unsafe"

	"github.com/pkg/errors"
	"golang.org/x/sys/unix"
)

func ioR(t, nr, size uintptr) uintptr {
	return (2 << 30) | (t << 8) | nr | (size << 16)
}

func ioW(t, nr, size uintptr) uintptr {
	return (1 << 30) | (t << 8) | nr | (size << 16)
}

func ioctl(fd, op, arg uintptr) error {
	if _, _, ep := unix.Syscall(unix.SYS_IOCTL, fd, op, arg); ep != 0 {
		return ep
	}
	return nil
}

const (
	ioctlSize     = 4
	hciMaxDevices = 16
	typHCI        = 72 // 'H'
)

var (
	hciUpDevice      = ioW(typHCI, 201, ioctlSize) // HCIDEVUP
	hciDownDevice    = ioW(typHCI, 202, ioctlSize) // HCIDEVDOWN
	hciResetDevice   = ioW(typHCI, 203, ioctlSize) // HCIDEVRESET
	hciGetDeviceList = ioR(typHCI, 210, ioctlSize) // HCIGETDEVLIST
	hciGetDeviceInfo = ioR(typHCI, 211, ioctlSize) // HCIGETDEVINFO
)

type devListRequest struct {
	devNum     uint16
	devRequest [hciMaxDevices]struct {
		id  uint16
		opt uint32
	}
}

// Socket implements a HCI User Channel as ReadWriteCloser.
type Socket struct {
	fd     int
	closed chan struct{}
	rmu    sync.Mutex
	wmu    sync.Mutex
}

// NewSocket returns a HCI User Channel of specified device id.
// If id is -1, the first available HCI device is returned.
func NewSocket(id int) (*Socket, error) {
	var err error
	// Create RAW HCI Socket.
	fd, err := unix.Socket(unix.AF_BLUETOOTH, unix.SOCK_RAW, unix.BTPROTO_HCI)
	if err != nil {
		return nil, errors.Wrap(err, "can't create socket")
	}

	if id != -1 {
		return open(fd, id)
	}

	req := devListRequest{devNum: hciMaxDevices}
	if err = ioctl(uintptr(fd), hciGetDeviceList, uintptr(unsafe.Pointer(&req))); err != nil {
		return nil, errors.Wrap(err, "can't get device list")
	}
	var msg string
	for id := 0; id < int(req.devNum); id++ {
		s, err := open(fd, id)
		if err == nil {
			return s, nil
		}
		msg = msg + fmt.Sprintf("(hci%d: %s)", id, err)
	}
	return nil, errors.Errorf("no devices available: %s", msg)
}

func open(fd, id int) (*Socket, error) {
	// Reset the device in case previous session didn't cleanup properly.
	if err := ioctl(uintptr(fd), hciDownDevice, uintptr(id)); err != nil {
		return nil, errors.Wrap(err, "can't down device")
	}
	if err := ioctl(uintptr(fd), hciUpDevice, uintptr(id)); err != nil {
		return nil, errors.Wrap(err, "can't up device")
	}

	// HCI User Channel requires exclusive access to the device.
	// The device has to be down at the time of binding.
	if err := ioctl(uintptr(fd), hciDownDevice, uintptr(id)); err != nil {
		return nil, errors.Wrap(err, "can't down device")
	}

	// Bind the RAW socket to HCI User Channel
	sa := unix.SockaddrHCI{Dev: uint16(id), Channel: unix.HCI_CHANNEL_USER}
	if err := unix.Bind(fd, &sa); err != nil {
		return nil, errors.Wrap(err, "can't bind socket to hci user channel")
	}

	// poll for 20ms to see if any data becomes available, then clear it
	pfds := []unix.PollFd{unix.PollFd{Fd: int32(fd), Events: unix.POLLIN}}
	unix.Poll(pfds, 20)
	if pfds[0].Revents&unix.POLLIN > 0 {
		b := make([]byte, 100)
		unix.Read(fd, b)
	}

	return &Socket{fd: fd, closed: make(chan struct{})}, nil
}

func (s *Socket) Read(p []byte) (int, error) {
	select {
	case <-s.closed:
		return 0, io.EOF
	default:
	}
	s.rmu.Lock()
	defer s.rmu.Unlock()
	n, err := unix.Read(s.fd, p)
	return n, errors.Wrap(err, "can't read hci socket")
}

func (s *Socket) Write(p []byte) (int, error) {
	s.wmu.Lock()
	defer s.wmu.Unlock()
	n, err := unix.Write(s.fd, p)
	return n, errors.Wrap(err, "can't write hci socket")
}

func (s *Socket) Close() error {
	close(s.closed)
	s.Write([]byte{0x01, 0x09, 0x10, 0x00})
	s.rmu.Lock()
	defer s.rmu.Unlock()
	return errors.Wrap(unix.Close(s.fd), "can't close hci socket")
}
<commit_msg>gofmt socket.go<commit_after>// +build linux

package socket

import (
	"fmt"
	"io"
	"sync"
	"unsafe"

	"github.com/pkg/errors"
	"golang.org/x/sys/unix"
)

func ioR(t, nr, size uintptr) uintptr {
	return (2 << 30) | (t << 8) | nr | (size << 16)
}

func ioW(t, nr, size uintptr) uintptr {
	return (1 << 30) | (t << 8) | nr | (size << 16)
}

func ioctl(fd, op, arg uintptr) error {
	if _, _, ep := unix.Syscall(unix.SYS_IOCTL, fd, op, arg); ep != 0 {
		return ep
	}
	return nil
}

const (
	ioctlSize     = 4
	hciMaxDevices = 16
	typHCI        = 72 // 'H'
)

var (
	hciUpDevice      = ioW(typHCI, 201, ioctlSize) // HCIDEVUP
	hciDownDevice    = ioW(typHCI, 202, ioctlSize) // HCIDEVDOWN
	hciResetDevice   = ioW(typHCI, 203, ioctlSize) // HCIDEVRESET
	hciGetDeviceList = ioR(typHCI, 210, ioctlSize) // HCIGETDEVLIST
	hciGetDeviceInfo = ioR(typHCI, 211, ioctlSize) // HCIGETDEVINFO
)

type devListRequest struct {
	devNum     uint16
	devRequest [hciMaxDevices]struct {
		id  uint16
		opt uint32
	}
}

// Socket implements a HCI User Channel as ReadWriteCloser.
type Socket struct {
	fd     int
	closed chan struct{}
	rmu    sync.Mutex
	wmu    sync.Mutex
}

// NewSocket returns a HCI User Channel of specified device id.
// If id is -1, the first available HCI device is returned.
func NewSocket(id int) (*Socket, error) {
	var err error
	// Create RAW HCI Socket.
	fd, err := unix.Socket(unix.AF_BLUETOOTH, unix.SOCK_RAW, unix.BTPROTO_HCI)
	if err != nil {
		return nil, errors.Wrap(err, "can't create socket")
	}

	if id != -1 {
		return open(fd, id)
	}

	req := devListRequest{devNum: hciMaxDevices}
	if err = ioctl(uintptr(fd), hciGetDeviceList, uintptr(unsafe.Pointer(&req))); err != nil {
		return nil, errors.Wrap(err, "can't get device list")
	}
	var msg string
	for id := 0; id < int(req.devNum); id++ {
		s, err := open(fd, id)
		if err == nil {
			return s, nil
		}
		msg = msg + fmt.Sprintf("(hci%d: %s)", id, err)
	}
	return nil, errors.Errorf("no devices available: %s", msg)
}

func open(fd, id int) (*Socket, error) {
	// Reset the device in case previous session didn't cleanup properly.
	if err := ioctl(uintptr(fd), hciDownDevice, uintptr(id)); err != nil {
		return nil, errors.Wrap(err, "can't down device")
	}
	if err := ioctl(uintptr(fd), hciUpDevice, uintptr(id)); err != nil {
		return nil, errors.Wrap(err, "can't up device")
	}

	// HCI User Channel requires exclusive access to the device.
	// The device has to be down at the time of binding.
	if err := ioctl(uintptr(fd), hciDownDevice, uintptr(id)); err != nil {
		return nil, errors.Wrap(err, "can't down device")
	}

	// Bind the RAW socket to HCI User Channel
	sa := unix.SockaddrHCI{Dev: uint16(id), Channel: unix.HCI_CHANNEL_USER}
	if err := unix.Bind(fd, &sa); err != nil {
		return nil, errors.Wrap(err, "can't bind socket to hci user channel")
	}

	// poll for 20ms to see if any data becomes available, then clear it
	pfds := []unix.PollFd{{Fd: int32(fd), Events: unix.POLLIN}}
	unix.Poll(pfds, 20)
	if pfds[0].Revents&unix.POLLIN > 0 {
		b := make([]byte, 100)
		unix.Read(fd, b)
	}

	return &Socket{fd: fd, closed: make(chan struct{})}, nil
}

func (s *Socket) Read(p []byte) (int, error) {
	select {
	case <-s.closed:
		return 0, io.EOF
	default:
	}
	s.rmu.Lock()
	defer s.rmu.Unlock()
	n, err := unix.Read(s.fd, p)
	return n, errors.Wrap(err, "can't read hci socket")
}

func (s *Socket) Write(p []byte) (int, error) {
	s.wmu.Lock()
	defer s.wmu.Unlock()
	n, err := unix.Write(s.fd, p)
	return n, errors.Wrap(err, "can't write hci socket")
}

func (s *Socket) Close() error {
	close(s.closed)
	s.Write([]byte{0x01, 0x09, 0x10, 0x00})
	s.rmu.Lock()
	defer s.rmu.Unlock()
	return errors.Wrap(unix.Close(s.fd), "can't close hci socket")
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/x509"
	"flag"
        "fmt"
        "log"
	"io/ioutil"
        "math/rand"
	"net/http"
	"os/exec"
	"time"
        "encoding/pem"
)
var db = map[string]int{}
var csrLocation = "/var/csr/"
var crtLocation = "/var/crt/"
var confLocation = "/opt/pollendina/openssl-ca.conf"

func main() {

	flag.Parse()

	http.HandleFunc("/v1/authorize", Authorize)
	http.HandleFunc("/v1/sign", Sign)

	// Placeholder for authentication / authorization middleware on authorize call.

	err := http.ListenAndServe(":33004", nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}

}

func Authorize(w http.ResponseWriter, req *http.Request) {
	log.Println("Received authorize call.")
	// Parse input
	sn := req.FormValue("cn")
	ttl := req.FormValue("ttl")

	// Construct ttl
	d, err := time.ParseDuration(ttl)
        if err != nil {
            log.Println(err)
            w.WriteHeader(http.StatusBadRequest)
            return
        }
	expires := time.Now().Add(d)

	// queue for write to map
	// ...

        log.Println("Service: " + sn)
        log.Println("Expires: " + expires.String())

	req.Body.Close()
}

func Sign(w http.ResponseWriter, req *http.Request) {
	log.Println("Received sign call.")

	// Upload the CSR and copy it to some known location
        body, err := ioutil.ReadAll(req.Body)
        if err != nil {
                log.Println(err)
                w.WriteHeader(http.StatusBadRequest)
                return
        }

        randoName := fmt.Sprintf("%d.csr", rand.Int63())
        csrFilename := csrLocation + randoName
        err = ioutil.WriteFile(csrFilename, body, 0777)
        if err != nil {
                log.Println(err)
                w.WriteHeader(http.StatusInternalServerError)
                return
        }
        log.Println("File uploaded.")

	// Parse the CSR
	rawCSR, _ := ioutil.ReadFile(csrFilename)
        decodedCSR, _ := pem.Decode(rawCSR)
	csr, err := x509.ParseCertificateRequest(decodedCSR.Bytes)
	if err != nil {
		log.Println(err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	log.Println("Received CSR for: " + csr.Subject.CommonName)
	// check authorization for the provided commonname
	// TODO ...

	// Build the command for exec
	// openssl ca -config openssl-ca.cnf -policy signing_policy -extensions signing_req -out servercert.pem -infiles servercert.csr
	app := "openssl"
	command := "ca"
	c_flag := "-config"
	p_flag := "-policy"
	p_value := "signing_policy"
	e_flag := "-extensions"
	e_value := "signing_req"
	o_flag := "-out"
	outputFile := crtLocation + randoName + ".crt"
	i_flag := "-infiles"
        b_flag := "-batch"

	// Sign the CSR with OpenSSL
	cmd := exec.Command(app, command, b_flag, c_flag, confLocation, p_flag, p_value, e_flag, e_value, o_flag, outputFile, i_flag, csrFilename)
        args := fmt.Sprintf("%s %s %s %s %s %s %s %s %s %s %s %s %s", app, command, b_flag, c_flag, confLocation, p_flag, p_value, e_flag, e_value, o_flag, outputFile, i_flag, csrFilename)
	fmt.Println(args)
        stdOut, err := cmd.Output()
	if err != nil {
                log.Println("STDOUT: " + string(stdOut))
		log.Println("STDERR: " + err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	// Open the output file for reading and stream it back on the response
        outputData, err := ioutil.ReadFile(outputFile)
	w.Write(outputData)
}
<commit_msg>Added cn and token authorization. Records to an inmemory DB. Ready to be validated by the Sign method.<commit_after>package main

import (
	"crypto/x509"
	"flag"
        "fmt"
        "log"
	"io/ioutil"
        "math/rand"
	"net/http"
	"os/exec"
        "encoding/pem"
)
var db = map[string]string{}
var csrLocation = "/var/csr/"
var crtLocation = "/var/crt/"
var confLocation = "/opt/pollendina/openssl-ca.conf"

type Tuple struct { CN, Token string }
var updates = make(chan Tuple)

func main() {

	flag.Parse()

	http.HandleFunc("/v1/authorize", Authorize)
	http.HandleFunc("/v1/sign", Sign)

        go MapWriter()

	// Placeholder for authentication / authorization middleware on authorize call.

	err := http.ListenAndServe(":33004", nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}

func MapWriter() {
      for {
              select {
                      case t,ok := <-updates:
                          if (!ok) {
                               log.Printf("Publisher channel closed. Stopping.")
                               return;
                          }
                          log.Printf("Setting key %s to value %s", t.Token, t.CN)
                          db[t.Token] = t.CN
              }
      }
}

func Authorize(w http.ResponseWriter, req *http.Request) {
	log.Println("Received authorize call.")
	// Parse input
	cn := req.FormValue("cn")
	token := req.FormValue("token")

	// queue for write to map
	// ...
        t := Tuple{cn, token}
        updates <- t

        log.Println("Service: " + cn)
        log.Println("Token: " + token)

	req.Body.Close()
}

func Sign(w http.ResponseWriter, req *http.Request) {
	log.Println("Received sign call.")

	// Upload the CSR and copy it to some known location
        body, err := ioutil.ReadAll(req.Body)
        if err != nil {
                log.Println(err)
                w.WriteHeader(http.StatusBadRequest)
                return
        }

        randoName := fmt.Sprintf("%d.csr", rand.Int63())
        csrFilename := csrLocation + randoName
        err = ioutil.WriteFile(csrFilename, body, 0777)
        if err != nil {
                log.Println(err)
                w.WriteHeader(http.StatusInternalServerError)
                return
        }
        log.Println("File uploaded.")

	// Parse the CSR
	rawCSR, _ := ioutil.ReadFile(csrFilename)
        decodedCSR, _ := pem.Decode(rawCSR)
	csr, err := x509.ParseCertificateRequest(decodedCSR.Bytes)
	if err != nil {
		log.Println(err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	log.Println("Received CSR for: " + csr.Subject.CommonName)
	// check authorization for the provided commonname
	// TODO ...

	// Build the command for exec
	// openssl ca -config openssl-ca.cnf -policy signing_policy -extensions signing_req -out servercert.pem -infiles servercert.csr
	app := "openssl"
	command := "ca"
	c_flag := "-config"
	p_flag := "-policy"
	p_value := "signing_policy"
	e_flag := "-extensions"
	e_value := "signing_req"
	o_flag := "-out"
	outputFile := crtLocation + randoName + ".crt"
	i_flag := "-infiles"
        b_flag := "-batch"

	// Sign the CSR with OpenSSL
	cmd := exec.Command(app, command, b_flag, c_flag, confLocation, p_flag, p_value, e_flag, e_value, o_flag, outputFile, i_flag, csrFilename)
        args := fmt.Sprintf("%s %s %s %s %s %s %s %s %s %s %s %s %s", app, command, b_flag, c_flag, confLocation, p_flag, p_value, e_flag, e_value, o_flag, outputFile, i_flag, csrFilename)
	fmt.Println(args)
        stdOut, err := cmd.Output()
	if err != nil {
                log.Println("STDOUT: " + string(stdOut))
		log.Println("STDERR: " + err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	// Open the output file for reading and stream it back on the response
        outputData, err := ioutil.ReadFile(outputFile)
	w.Write(outputData)
}
<|endoftext|>


<commit_before>package simra

import (
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"testing"
)

type mock struct {
	m map[string]interface{}
}

func (db *mock) Open(dirpath string) error {
	if db.m == nil {
		db.m = map[string]interface{}{}
	}
	return nil
}

func (db *mock) Close() {
	db.m = nil
}

func (db *mock) Put(key string, value interface{}) {
	if db.m == nil {
		log.Fatal("database is not opened yet")
		return
	}
	db.m[key] = value
}

func (db *mock) Get(key string) interface{} {
	if db.m == nil {
		log.Fatal("database is not opened yet")
		return nil
	}
	return db.m[key]
}

func TestMock(t *testing.T) {
	db := OpenDB(&mock{}, filepath.Join("."))
	if db == nil {
		t.Error("failed to open database")
	}
	defer db.Close()

	db.Put("key1", "value1")
	value1 := db.Get("key1")
	if value1 != "value1" {
		t.Error("failed to open database")
	}
}

func TestBolt(t *testing.T) {
	tmpdir, err := ioutil.TempDir(filepath.Join("."), "simradb_")
	if err != nil {
		t.Fatalf("failed to create temporary working directory. err: %s", err.Error())
	}
	defer func() {
		_ = os.RemoveAll(tmpdir)
	}()

	db := OpenDB(&Boltdb{}, filepath.Join(tmpdir))
	if db == nil {
		t.Error("failed to open database")
	}
	defer db.Close()

	db.Put("key1", "value1")
	fetched := db.Get("key1")
	if bytes, ok := fetched.([]uint8); ok {
		v := string(bytes)
		if v != "value1" {
			t.Errorf("unexpected fetched value. expected: %s, fetched: %s", "value1", v)
		}
	}
}
<commit_msg>[#113] fix build failure<commit_after>package simra

import (
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"testing"

	"github.com/pankona/gomo-simra/simra/database"
)

type mock struct {
	m map[string]interface{}
}

func (db *mock) Open(dirpath string) error {
	if db.m == nil {
		db.m = map[string]interface{}{}
	}
	return nil
}

func (db *mock) Close() {
	db.m = nil
}

func (db *mock) Put(key string, value interface{}) {
	if db.m == nil {
		log.Fatal("database is not opened yet")
		return
	}
	db.m[key] = value
}

func (db *mock) Get(key string) interface{} {
	if db.m == nil {
		log.Fatal("database is not opened yet")
		return nil
	}
	return db.m[key]
}

func TestMock(t *testing.T) {
	db := OpenDB(&mock{}, filepath.Join("."))
	if db == nil {
		t.Error("failed to open database")
	}
	defer db.Close()

	db.Put("key1", "value1")
	value1 := db.Get("key1")
	if value1 != "value1" {
		t.Error("failed to open database")
	}
}

func TestBolt(t *testing.T) {
	tmpdir, err := ioutil.TempDir(filepath.Join("."), "simradb_")
	if err != nil {
		t.Fatalf("failed to create temporary working directory. err: %s", err.Error())
	}
	defer func() {
		_ = os.RemoveAll(tmpdir)
	}()

	db := OpenDB(&database.Boltdb{}, filepath.Join(tmpdir))
	if db == nil {
		t.Error("failed to open database")
	}
	defer db.Close()

	db.Put("key1", "value1")
	fetched := db.Get("key1")
	if bytes, ok := fetched.([]uint8); ok {
		v := string(bytes)
		if v != "value1" {
			t.Errorf("unexpected fetched value. expected: %s, fetched: %s", "value1", v)
		}
	}
}
<|endoftext|>


<commit_before>package saltboot

import (
    "log"
    "github.com/sequenceiq/salt-bootstrap/saltboot/model"
    "net/http"
    "encoding/json"
    "fmt"
    "gopkg.in/yaml.v2"
    "io/ioutil"
    "os"
    "strings"
)

type SaltServerSetupRequest struct {
    Password string     `json:"password,omitempty"`
}

type SaltActionRequest struct {
    Minions []SaltMinion   `json:"minions,omitempty"`
    Server  string         `json:"server,omitempty"`
    Action  string         `json:"action"`
}

type SaltMinion struct {
    Address string     `json:"address"`
    Roles   []string   `json:"roles,omitempty"`
    Server  string     `json:"server,omitempty"`
}

type SaltPillar struct {
    Path string                 `json:"path"`
    Json map[string]interface{} `json:"json"`
}

func (saltMinion SaltMinion) AsByteArray() []byte {
    b, _ := json.Marshal(saltMinion)
    return b
}

//consul:
//    advertise_addr: 10.0.0.24
//    recursors:
//        - 10.0.0.2
//        - 8.8.8.8
//hostgroup: hostgroup_3
//roles:
//    - ambari_server
//    - ambari_agent

type ConsulGrainConfig struct {
    AdvertiseAddr string     `json:"advertise_addr" yaml:"advertise_addr"`
    DNSRecursors  []string   `json:"recursors" yaml:"recursors"`
}

type GrainConfig struct {
    Consul ConsulGrainConfig       `json:"consul" yaml:"consul"`
    //    HostGroup           string                  `json:"hostgroup yaml:"hostgroup"`
    Roles  []string                `json:"roles" yaml:"roles"`
}

func (r SaltActionRequest) String() string {
    b, _ := json.Marshal(r)
    return fmt.Sprintf(string(b))
}

func (r SaltActionRequest) distributeAction(user string, pass string) (result []model.Response) {
    log.Printf("[distributeAction] distribute salt state command to targets: %s", r.String())
    var targets []string
    var payloads []Payload
    for _, minion := range r.Minions {
        targets = append(targets, minion.Address)
        if minion.Server == "" {
            minion.Server = r.Server
        }
        payloads = append(payloads, minion)
    }

    for res := range DistributePayload(targets, payloads, SaltMinionEp + "/" + r.Action, user, pass) {
        result = append(result, res)
    }
    if len(r.Server) > 0 {
        result = append(result, <-Distribute([]string{r.Server}, nil, SaltServerEp + "/" + r.Action, user, pass))
    }
    return result
}

func SaltMinionRunRequestHandler(w http.ResponseWriter, req *http.Request) {
    log.Printf("[SaltMinionRunRequestHandler] execute salt run request")

    decoder := json.NewDecoder(req.Body)
    var saltMinion SaltMinion
    err := decoder.Decode(&saltMinion)
    if err != nil {
        log.Printf("[SaltMinionRunRequestHandler] [ERROR] couldn't decode json: %s", err)
        model.Response{Status: err.Error()}.WriteBadRequestHttp(w)
        return
    }

    recursors := DetermineDNSRecursors([]string{})

    grainConfig := GrainConfig{
        Consul:            ConsulGrainConfig{
            AdvertiseAddr:      saltMinion.Address,
            DNSRecursors:       recursors,
        },
        Roles:             saltMinion.Roles,
    }

    grainYaml, err := yaml.Marshal(grainConfig)
    var resp model.Response;
    if err != nil {
        resp = model.Response{ErrorText: err.Error(), StatusCode:http.StatusInternalServerError}
        resp.WriteHttp(w)
        return
    }

    err = os.MkdirAll("/etc/salt/minion.d", 0755)
    if err != nil {
        resp = model.Response{ErrorText: err.Error(), StatusCode:http.StatusInternalServerError}
        resp.WriteHttp(w)
        return
    }

    masterConf := []byte("master: " + saltMinion.Server)
    err = ioutil.WriteFile("/etc/salt/minion.d/master.conf", masterConf, 0644)
    if err != nil {
        resp = model.Response{ErrorText: err.Error(), StatusCode:http.StatusInternalServerError}
        resp.WriteHttp(w)
        return
    }

    err = ioutil.WriteFile("/etc/salt/grains", grainYaml, 0644)
    if err != nil {
        resp = model.Response{ErrorText: err.Error(), StatusCode:http.StatusInternalServerError}
        resp.WriteHttp(w)
        return
    }
    resp, _ = LaunchService("salt-minion")
    resp.WriteHttp(w)
}

func SaltMinionStopRequestHandler(w http.ResponseWriter, req *http.Request) {
    log.Printf("[SaltMinionStopRequestHandler] execute salt minion stop request")

    decoder := json.NewDecoder(req.Body)
    var saltMinion SaltMinion
    err := decoder.Decode(&saltMinion)
    if err != nil {
        log.Printf("[SaltMinionRunRequestHandler] [ERROR] couldn't decode json: %s", err)
        model.Response{Status: err.Error()}.WriteBadRequestHttp(w)
        return
    }

    resp, _ := StopService("salt-minion")
    resp.WriteHttp(w)
}

func SaltServerRunRequestHandler(w http.ResponseWriter, req *http.Request) {
    log.Printf("[SaltServerRunRequestHandler] execute salt run request")
    resp, err := LaunchService("salt-master")
    resp.WriteHttp(w)
    if err != nil {
        return
    }
    resp, _ = LaunchService("salt-api")
    resp.WriteHttp(w)
}

func SaltServerStopRequestHandler(w http.ResponseWriter, req *http.Request) {
    log.Printf("[SaltServerStopRequestHandler] execute salt master stop request")
    resp, err := StopService("salt-master")
    resp.WriteHttp(w)
    if err != nil {
        return
    }
    resp, _ = StopService("salt-api")
    resp.WriteHttp(w)
}

func SaltServerSetupRequestHandler(w http.ResponseWriter, req *http.Request) {

}

func (pillar SaltPillar) WritePillar() (outStr string, err error) {
    file := "/srv/pillar" + pillar.Path
    dir := file[0:strings.LastIndex(file, "/")]

    log.Printf("[SaltPillar.WritePillar] mkdir %s", dir)
    err = os.MkdirAll(dir, 0644)
    if err != nil {
        return "Failed to create dir " + dir, err
    }

    yml, _ := yaml.Marshal(pillar.Json)
    log.Printf("[SaltPillarRequestHandler] generated yaml from json %s", string(yml))
    err = ioutil.WriteFile(file, yml, 0644)
    if err != nil {
        return "Failed to write to " + file, err
    }
    return "Salt pillar successfully written to " + file, err
}

func SaltPillarRequestHandler(w http.ResponseWriter, req *http.Request) {
    log.Printf("[SaltPillarRequestHandler] execute salt pillar save request")
    decoder := json.NewDecoder(req.Body)
    var saltPillar SaltPillar
    err := decoder.Decode(&saltPillar)
    if err != nil {
        log.Printf("[SaltPillarRequestHandler] [ERROR] couldn't decode json: %s", err)
        model.Response{Status: err.Error()}.WriteBadRequestHttp(w)
        return
    }

    if !strings.HasSuffix(saltPillar.Path, ".sls") {
        log.Printf("[SaltPillarRequestHandler] [ERROR] path is not ending with '.sls' suffix %s", saltPillar.Path)
        model.Response{Status: "path is not ending with '.sls' suffix"}.WriteBadRequestHttp(w)
        return
    }
    if !strings.HasPrefix(saltPillar.Path, "/") {
        log.Printf("[SaltPillarRequestHandler] [ERROR] path is not starting with '/' %s", saltPillar.Path)
        model.Response{Status: "path is not starting with '/'"}.WriteBadRequestHttp(w)
        return
    }
    if strings.Contains(saltPillar.Path, "..") {
        log.Printf("[SaltPillarRequestHandler] [ERROR] path cannot contain '..' charachters %s", saltPillar.Path)
        model.Response{Status: "path cannot contain '..' charachters"}.WriteBadRequestHttp(w)
        return
    }

    jsonString, _ := json.Marshal(saltPillar.Json)
    log.Printf("[SaltPillarRequestHandler] Recieved arbitrary json: %s", jsonString)

    outStr, err := saltPillar.WritePillar()
    if err != nil {
        log.Printf("[SaltPillarRequestHandler] failed to execute salt pillar save config: %s", err.Error())
        model.Response{ErrorText: err.Error(), StatusCode:http.StatusInternalServerError}.WriteHttp(w)
    } else {
        cResp := model.Response{Status: outStr}.WriteHttp(w)
        log.Printf("[SaltPillarRequestHandler] save salt pillar request executed: %s", cResp.String())
    }
}

func SaltActionDistributeRequestHandler(w http.ResponseWriter, req *http.Request) {
    log.Printf("[SaltActionDistributeRequestHandler] execute Salt state distribute request")

    decoder := json.NewDecoder(req.Body)
    var saltActionRequest SaltActionRequest
    err := decoder.Decode(&saltActionRequest)
    if err != nil {
        log.Printf("[SaltActionDistributeRequestHandler] [ERROR] couldn't decode json: %s", err)
        model.Response{Status: err.Error()}.WriteBadRequestHttp(w)
        return
    }

    user, pass := GetAuthUserPass(req)
    result := saltActionRequest.distributeAction(user, pass)
    cResp := model.Responses{Responses:result}
    log.Printf("[SaltActionDistributeRequestHandler] distribute salt state command request executed: %s", cResp.String())
    json.NewEncoder(w).Encode(cResp)
}<commit_msg>CLOUD-56738 include host group in grains<commit_after>package saltboot

import (
    "log"
    "github.com/sequenceiq/salt-bootstrap/saltboot/model"
    "net/http"
    "encoding/json"
    "fmt"
    "gopkg.in/yaml.v2"
    "io/ioutil"
    "os"
    "strings"
)

type SaltServerSetupRequest struct {
    Password string     `json:"password,omitempty"`
}

type SaltActionRequest struct {
    Minions []SaltMinion   `json:"minions,omitempty"`
    Server  string         `json:"server,omitempty"`
    Action  string         `json:"action"`
}

type SaltMinion struct {
    Address   string     `json:"address"`
    Roles     []string   `json:"roles,omitempty"`
    Server    string     `json:"server,omitempty"`
    HostGroup string     `json:"hostGroup,omitempty"`
}

type SaltPillar struct {
    Path string                 `json:"path"`
    Json map[string]interface{} `json:"json"`
}

func (saltMinion SaltMinion) AsByteArray() []byte {
    b, _ := json.Marshal(saltMinion)
    return b
}

type GrainConfig struct {
    HostGroup string     `json:"hostgroup" yaml:"hostgroup"`
    Roles     []string   `json:"roles" yaml:"roles"`
}

func (r SaltActionRequest) String() string {
    b, _ := json.Marshal(r)
    return fmt.Sprintf(string(b))
}

func (r SaltActionRequest) distributeAction(user string, pass string) (result []model.Response) {
    log.Printf("[distributeAction] distribute salt state command to targets: %s", r.String())
    var targets []string
    var payloads []Payload
    for _, minion := range r.Minions {
        targets = append(targets, minion.Address)
        if minion.Server == "" {
            minion.Server = r.Server
        }
        payloads = append(payloads, minion)
    }

    for res := range DistributePayload(targets, payloads, SaltMinionEp + "/" + r.Action, user, pass) {
        result = append(result, res)
    }
    if len(r.Server) > 0 {
        result = append(result, <-Distribute([]string{r.Server}, nil, SaltServerEp + "/" + r.Action, user, pass))
    }
    return result
}

func SaltMinionRunRequestHandler(w http.ResponseWriter, req *http.Request) {
    log.Printf("[SaltMinionRunRequestHandler] execute salt run request")

    decoder := json.NewDecoder(req.Body)
    var saltMinion SaltMinion
    err := decoder.Decode(&saltMinion)
    if err != nil {
        log.Printf("[SaltMinionRunRequestHandler] [ERROR] couldn't decode json: %s", err)
        model.Response{Status: err.Error()}.WriteBadRequestHttp(w)
        return
    }

    grainConfig := GrainConfig{Roles:saltMinion.Roles, HostGroup:saltMinion.HostGroup}
    grainYaml, err := yaml.Marshal(grainConfig)
    var resp model.Response;
    if err != nil {
        resp = model.Response{ErrorText: err.Error(), StatusCode:http.StatusInternalServerError}
        resp.WriteHttp(w)
        return
    }

    err = os.MkdirAll("/etc/salt/minion.d", 0755)
    if err != nil {
        resp = model.Response{ErrorText: err.Error(), StatusCode:http.StatusInternalServerError}
        resp.WriteHttp(w)
        return
    }

    masterConf := []byte("master: " + saltMinion.Server)
    err = ioutil.WriteFile("/etc/salt/minion.d/master.conf", masterConf, 0644)
    if err != nil {
        resp = model.Response{ErrorText: err.Error(), StatusCode:http.StatusInternalServerError}
        resp.WriteHttp(w)
        return
    }

    err = ioutil.WriteFile("/etc/salt/grains", grainYaml, 0644)
    if err != nil {
        resp = model.Response{ErrorText: err.Error(), StatusCode:http.StatusInternalServerError}
        resp.WriteHttp(w)
        return
    }
    resp, _ = LaunchService("salt-minion")
    resp.WriteHttp(w)
}

func SaltMinionStopRequestHandler(w http.ResponseWriter, req *http.Request) {
    log.Printf("[SaltMinionStopRequestHandler] execute salt minion stop request")

    decoder := json.NewDecoder(req.Body)
    var saltMinion SaltMinion
    err := decoder.Decode(&saltMinion)
    if err != nil {
        log.Printf("[SaltMinionRunRequestHandler] [ERROR] couldn't decode json: %s", err)
        model.Response{Status: err.Error()}.WriteBadRequestHttp(w)
        return
    }

    resp, _ := StopService("salt-minion")
    resp.WriteHttp(w)
}

func SaltServerRunRequestHandler(w http.ResponseWriter, req *http.Request) {
    log.Printf("[SaltServerRunRequestHandler] execute salt run request")
    resp, err := LaunchService("salt-master")
    resp.WriteHttp(w)
    if err != nil {
        return
    }
    resp, _ = LaunchService("salt-api")
    resp.WriteHttp(w)
}

func SaltServerStopRequestHandler(w http.ResponseWriter, req *http.Request) {
    log.Printf("[SaltServerStopRequestHandler] execute salt master stop request")
    resp, err := StopService("salt-master")
    resp.WriteHttp(w)
    if err != nil {
        return
    }
    resp, _ = StopService("salt-api")
    resp.WriteHttp(w)
}

func SaltServerSetupRequestHandler(w http.ResponseWriter, req *http.Request) {

}

func (pillar SaltPillar) WritePillar() (outStr string, err error) {
    file := "/srv/pillar" + pillar.Path
    dir := file[0:strings.LastIndex(file, "/")]

    log.Printf("[SaltPillar.WritePillar] mkdir %s", dir)
    err = os.MkdirAll(dir, 0644)
    if err != nil {
        return "Failed to create dir " + dir, err
    }

    yml, _ := yaml.Marshal(pillar.Json)
    log.Printf("[SaltPillarRequestHandler] generated yaml from json %s", string(yml))
    err = ioutil.WriteFile(file, yml, 0644)
    if err != nil {
        return "Failed to write to " + file, err
    }
    return "Salt pillar successfully written to " + file, err
}

func SaltPillarRequestHandler(w http.ResponseWriter, req *http.Request) {
    log.Printf("[SaltPillarRequestHandler] execute salt pillar save request")
    decoder := json.NewDecoder(req.Body)
    var saltPillar SaltPillar
    err := decoder.Decode(&saltPillar)
    if err != nil {
        log.Printf("[SaltPillarRequestHandler] [ERROR] couldn't decode json: %s", err)
        model.Response{Status: err.Error()}.WriteBadRequestHttp(w)
        return
    }

    if !strings.HasSuffix(saltPillar.Path, ".sls") {
        log.Printf("[SaltPillarRequestHandler] [ERROR] path is not ending with '.sls' suffix %s", saltPillar.Path)
        model.Response{Status: "path is not ending with '.sls' suffix"}.WriteBadRequestHttp(w)
        return
    }
    if !strings.HasPrefix(saltPillar.Path, "/") {
        log.Printf("[SaltPillarRequestHandler] [ERROR] path is not starting with '/' %s", saltPillar.Path)
        model.Response{Status: "path is not starting with '/'"}.WriteBadRequestHttp(w)
        return
    }
    if strings.Contains(saltPillar.Path, "..") {
        log.Printf("[SaltPillarRequestHandler] [ERROR] path cannot contain '..' charachters %s", saltPillar.Path)
        model.Response{Status: "path cannot contain '..' charachters"}.WriteBadRequestHttp(w)
        return
    }

    jsonString, _ := json.Marshal(saltPillar.Json)
    log.Printf("[SaltPillarRequestHandler] Recieved arbitrary json: %s", jsonString)

    outStr, err := saltPillar.WritePillar()
    if err != nil {
        log.Printf("[SaltPillarRequestHandler] failed to execute salt pillar save config: %s", err.Error())
        model.Response{ErrorText: err.Error(), StatusCode:http.StatusInternalServerError}.WriteHttp(w)
    } else {
        cResp := model.Response{Status: outStr}.WriteHttp(w)
        log.Printf("[SaltPillarRequestHandler] save salt pillar request executed: %s", cResp.String())
    }
}

func SaltActionDistributeRequestHandler(w http.ResponseWriter, req *http.Request) {
    log.Printf("[SaltActionDistributeRequestHandler] execute Salt state distribute request")

    decoder := json.NewDecoder(req.Body)
    var saltActionRequest SaltActionRequest
    err := decoder.Decode(&saltActionRequest)
    if err != nil {
        log.Printf("[SaltActionDistributeRequestHandler] [ERROR] couldn't decode json: %s", err)
        model.Response{Status: err.Error()}.WriteBadRequestHttp(w)
        return
    }

    user, pass := GetAuthUserPass(req)
    result := saltActionRequest.distributeAction(user, pass)
    cResp := model.Responses{Responses:result}
    log.Printf("[SaltActionDistributeRequestHandler] distribute salt state command request executed: %s", cResp.String())
    json.NewEncoder(w).Encode(cResp)
}<|endoftext|>


<commit_before>package sentinel

import (
	"github.com/mdevilliers/redishappy/services/logger"
	"github.com/mdevilliers/redishappy/services/redis"
	"github.com/mdevilliers/redishappy/types"
	"testing"
	"time"
)

type TestRedisConnection struct {
	// RedisClient *TestRedisClient
}

type TestRedisClient struct{}

type TestRedisReply struct{}

func (c TestRedisConnection) Dial(protocol, uri string) (redis.RedisClient, error) {
	return &TestRedisClient{}, nil
}

func (c *TestRedisClient) Cmd(cmd string, args ...interface{}) redis.RedisReply {
	return &TestRedisReply{}
}

func (c *TestRedisReply) String() string {
	return "PIG"
}

func (c *TestRedisReply) Err() error {
	return nil
}

type TestManager struct {
	NotifyCalled                   int
	ScheduleNewHealthCheckerCalled int
}

func (tm *TestManager) Notify(event SentinelEvent) {
	tm.NotifyCalled++
}
func (*TestManager) GetState(request TopologyRequest) {

}
func (*TestManager) NewSentinelMonitor(types.Sentinel) (*SentinelHealthCheckerClient, error) {
	return nil, nil
}
func (tm *TestManager) ScheduleNewHealthChecker(sentinel types.Sentinel) {
	tm.ScheduleNewHealthCheckerCalled++
}

func TestNewClientWillFailWhenPingUnsucessful(t *testing.T) {
	logger.InitLogging("../log")

	sentinel := types.Sentinel{}
	sentinelManager := &TestManager{}
	redisConnection := &TestRedisConnection{}

	client, _ := NewHealthCheckerClient(sentinel, sentinelManager, redisConnection)
	client.Start()

	time.Sleep(time.Second)

	if sentinelManager.NotifyCalled != 1 {
		t.Error("Notify should have been called!")
	}
}
<commit_msg>healthchecker tests updated for all of the conditions a healthcher 'should' encounter<commit_after>package sentinel

import (
	"errors"
	"github.com/mdevilliers/redishappy/services/logger"
	"github.com/mdevilliers/redishappy/services/redis"
	"github.com/mdevilliers/redishappy/types"
	"reflect"
	"testing"
	"time"
)

// MOCKS
type TestRedisConnection struct {
	RedisClient *TestRedisClient
}

type TestRedisClient struct {
	RedisReply *TestRedisReply
}

type TestRedisReply struct {
	Reply string
	Error error
}

func (c TestRedisConnection) Dial(protocol, uri string) (redis.RedisClient, error) {

	//fail to connect
	if uri == "DOESNOTEXIST:1234" {
		return nil, errors.New("CannotConnect")
	}
	return c.RedisClient, nil
}

func (c *TestRedisClient) Cmd(cmd string, args ...interface{}) redis.RedisReply {
	return c.RedisReply
}

func (c *TestRedisReply) String() string {
	return c.Reply
}

func (c *TestRedisReply) Err() error {
	return c.Error
}

type TestManager struct {
	NotifyCalledWithSentinelPing   int
	NotifyCalledWithSentinelLost   int
	NotifyCalledWithSentinelAdded  int
	ScheduleNewHealthCheckerCalled int
}

func (tm *TestManager) Notify(event SentinelEvent) {
	t := reflect.TypeOf(event).String()

	if t == "*sentinel.SentinelLost" {
		tm.NotifyCalledWithSentinelLost++
	}
	if t == "*sentinel.SentinelPing" {
		tm.NotifyCalledWithSentinelPing++
	}
	if t == "*sentinel.SentinelAdded" {
		tm.NotifyCalledWithSentinelAdded++
	}
}
func (*TestManager) GetState(request TopologyRequest) {

}
func (*TestManager) NewSentinelMonitor(types.Sentinel) (*SentinelHealthCheckerClient, error) {
	return nil, nil
}
func (tm *TestManager) ScheduleNewHealthChecker(sentinel types.Sentinel) {
	tm.ScheduleNewHealthCheckerCalled++
}

func TestNewClientWillGetASuccessfulPing(t *testing.T) {
	logger.InitLogging("../log")

	sentinel := types.Sentinel{}
	sentinelManager := &TestManager{}
	redisConnection := &TestRedisConnection{RedisClient: &TestRedisClient{RedisReply: &TestRedisReply{Reply: "PONG"}}}

	client, _ := NewHealthCheckerClient(sentinel, sentinelManager, redisConnection)
	client.Start()

	time.Sleep(time.Second)

	if sentinelManager.NotifyCalledWithSentinelPing != 1 {
		t.Error("Notify should have been called with a SentinelPing event!")
	}
}

func TestNewClientWillFailWhenPingUnsucessful(t *testing.T) {
	logger.InitLogging("../log")

	sentinel := types.Sentinel{}
	sentinelManager := &TestManager{}
	redisConnection := &TestRedisConnection{RedisClient: &TestRedisClient{RedisReply: &TestRedisReply{Reply: "ERROR"}}}

	client, _ := NewHealthCheckerClient(sentinel, sentinelManager, redisConnection)
	client.Start()

	time.Sleep(time.Second)

	if sentinelManager.NotifyCalledWithSentinelLost != 1 {
		t.Error("Notify should have been called with a SentinelPing event!")
	}
}

func TestNewClientWillFailWhenErrorOnPing(t *testing.T) {
	logger.InitLogging("../log")

	sentinel := types.Sentinel{}
	sentinelManager := &TestManager{}
	redisConnection := &TestRedisConnection{RedisClient: &TestRedisClient{RedisReply: &TestRedisReply{Error: errors.New("BOOYAH!")}}}

	client, _ := NewHealthCheckerClient(sentinel, sentinelManager, redisConnection)
	client.Start()

	time.Sleep(time.Second)

	if sentinelManager.NotifyCalledWithSentinelLost != 1 {
		t.Error("Notify should have been called with a SentinelPing event!")
	}
}

func TestNewClientWillWillSignalSentinelLostIfCanNotConnect(t *testing.T) {
	logger.InitLogging("../log")

	sentinel := types.Sentinel{Host: "DOESNOTEXIST", Port: 1234} // mock coded to not connect
	sentinelManager := &TestManager{}
	redisConnection := &TestRedisConnection{}

	_, _ = NewHealthCheckerClient(sentinel, sentinelManager, redisConnection)

	if sentinelManager.NotifyCalledWithSentinelLost != 1 {
		t.Error("Notify should have been called!")
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2017, Boise State University All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"context"
	"crypto/rand"
	"errors"
	"fmt"
	"log"
	"path"
	"regexp"
	"sync"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/nat"
)

const (
	// jupyterNotebookImageMatch matches images from jupyter/docker-stacks, for
	// example: jupyter/minimal-notebook
	jupyterNotebookImageMatch = `[a-zA-Z0-9]+/[a-zA-Z0-9]+-notebook[:]{0,1}[a-zA-Z0-9]*`

	// allImageMatch applies no filter
	allImageMatch = `.*`

	// defaultContainerLifetime is used if a lifetime is not provided
	defaultContainerLifetime = time.Minute * 10

	// defaultMaxContainers governs the port set size and triggers reclamation
	defaultMaxContainers = 100

	// collectionFraction is the fraction of lifetime to collect containers.  For
	// example 4 collects every 1/4 of the container lifetime.
	collectionFraction = 4
)

// tempNotebook holds context for a single container
type tempNotebook struct {
	// guard struct (only used for lastAccessed right now)
	sync.Mutex
	// id is the docker container id.
	id string
	// hash is  a random generated hash that is used in the path of the server.
	hash string
	// imageName is the name of the image used to start the container
	imageName string
	// lastAccessed is when the container was used last.
	lastAccessed time.Time
	// port is the passthrough port for the reverse proxy.
	port int
}

// notebookPool holds data regarding running notebooks.
type notebookPool struct {
	// guards the entire struct
	sync.Mutex

	// availableImages is a list of docker images that installed on the machine,
	// and match the imageMatch expression
	availableImages map[string]struct{}

	// imageMatch filters available images by name
	imageMatch *regexp.Regexp

	// containerMap is stores the contexts for the containers.
	containerMap map[string]*tempNotebook

	// portSet holds free ports
	portSet *portRange

	// maxContainers governs the port set size and resource reclamation.
	maxContainers int

	// containerLifetime governs when the container resources are reclaimed.
	containerLifetime time.Duration

	// token is the security token for auto-auth
	token string

	// killCollection stops the automated resource reclamation
	killCollection chan struct{}

	// lastCollMu guards the time for reclamation.  It is used infrequently, and
	// we don't need to lock the whole object.
	lastCollMu sync.Mutex

	// lastCollection is the timestamp the last time the containers were
	// reclaimed.
	lastCollection time.Time

	// deregisterMux is a channel for sending a path that needs to be
	// de-registered from the server mux.
	deregisterMux chan string
}

// errNotebookPoolFull indicates the pool is at maxContainers
var errNotebookPoolFull = errors.New("container pool hit max size limit")

// newNotebookPool creates a notebookPool and sets defaults, overriding some
// with passed arguments.
func newNotebookPool(imageRegexp string, maxContainers int, lifetime time.Duration) (*notebookPool, error) {
	if imageRegexp == "" {
		imageRegexp = jupyterNotebookImageMatch
	}
	if int64(lifetime) <= 0 {
		lifetime = defaultContainerLifetime
	}
	if maxContainers < 1 {
		maxContainers = defaultMaxContainers
	}
	imageMatch, err := regexp.Compile(imageRegexp)
	if err != nil {
		return nil, err
	}
	imageMap := map[string]struct{}{}
	cli, err := client.NewEnvClient()
	if err != nil {
		return nil, err
	}
	images, err := cli.ImageList(context.Background(), types.ImageListOptions{})
	if err != nil {
		return nil, err
	}
	for _, image := range images {
		if len(image.RepoTags) < 1 || !imageMatch.MatchString(image.RepoTags[0]) {
			continue
		}
		log.Printf("found image %s", image.RepoTags[0])
		imageMap[image.RepoTags[0]] = struct{}{}
	}
	pool := &notebookPool{
		availableImages:   imageMap,
		imageMatch:        imageMatch,
		containerMap:      make(map[string]*tempNotebook),
		portSet:           newPortRange(8000, maxContainers),
		maxContainers:     maxContainers,
		containerLifetime: lifetime,
		killCollection:    make(chan struct{}),
		deregisterMux:     make(chan string),
	}
	pool.startCollector(time.Duration(int64(lifetime) / collectionFraction))
	pool.lastCollMu.Lock()
	pool.lastCollection = time.Now()
	pool.lastCollMu.Unlock()
	return pool, nil
}

// defaultHashSize is used for the unique hash generation
const defaultHashSize = 32

// newHash makes a n byte hash and returns the hex encoding.
func newHash(n int) string {
	b := make([]byte, n)
	_, err := rand.Read(b[:])
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", b)
}

// newNotebook initializes and sets values for a new notebook.
func (p *notebookPool) newNotebook(image string, pull bool) (*tempNotebook, error) {
	ctx := context.Background()
	cli, err := client.NewEnvClient()
	if err != nil {
		return nil, err
	}

	// TODO(kyle): possibly provide tag support
	if pull {
		log.Printf("pulling container %s", image)
		ctx, cancel := context.WithTimeout(ctx, 45*time.Second)
		defer cancel()
		_, err = cli.ImagePull(ctx, image, types.ImagePullOptions{})
		if err != nil {
			return nil, err
		}
		log.Printf("successfully pulled")
	}

	hash := newHash(defaultHashSize)

	port, err := p.portSet.Acquire()
	if err != nil {
		return nil, err
	}
	portString := fmt.Sprintf("%d", port)

	var pSet = nat.PortSet{}
	pt, err := nat.NewPort("tcp", portString)
	pSet[pt] = struct{}{}
	containerConfig := container.Config{
		Hostname: "0.0.0.0",
		User:     "jovyan",
		Cmd: []string{`jupyter`,
			`notebook`,
			`--no-browser`,
			`--port`,
			portString,
			`--ip=0.0.0.0`,
			fmt.Sprintf("--NotebookApp.base_url=%s", path.Join("/book", hash)),
			`--NotebookApp.port_retries=0`,
			fmt.Sprintf(`--NotebookApp.token="%s"`, p.token),
			`--NotebookApp.disable_check_xsrf=True`,
		},
		Env:          []string{fmt.Sprintf("CONFIGPROXY_AUTH_TOKEN=%s", p.token)},
		Image:        image,
		ExposedPorts: pSet,
	}

	hostConfig := container.HostConfig{
		NetworkMode: "host",
	}

	resp, err := cli.ContainerCreate(ctx, &containerConfig, &hostConfig, nil, "")
	if err != nil {
		p.portSet.Drop(port)
		return nil, err
	}

	if err := cli.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{}); err != nil {
		p.portSet.Drop(port)
		return nil, err
	}
	log.Printf("created container: %s", resp.ID)
	t := &tempNotebook{
		id:           resp.ID,
		hash:         hash,
		imageName:    image,
		lastAccessed: time.Now(),
		port:         port,
	}
	err = p.addNotebook(t)
	if err != nil {
		log.Print(err)
		p.portSet.Drop(port)
		return nil, err
	}
	return t, nil
}

// size returns the size of the containerMap with appropriate locks
func (p *notebookPool) size() int {
	p.Lock()
	n := len(p.containerMap)
	p.Unlock()
	return n
}

// addNotebook adds a tempNotebook to the containerMap, if there is room.
func (p *notebookPool) addNotebook(t *tempNotebook) error {
	n := p.size()
	log.Printf("pool size: %d of %d", n+1, p.maxContainers)
	if p.size()+1 > p.maxContainers {
		p.releaseContainers(false)
	}
	if p.size()+1 > p.maxContainers {
		return errNotebookPoolFull
	}
	p.Lock()
	p.containerMap[t.hash] = t
	p.Unlock()
	return nil
}

// stopAndKillContainer requests the stopping (docker stop) and the removal of
// the container (docker rm).  Errors are logged, but not returned and rm is
// always called.
func (p *notebookPool) stopAndKillContainer(id string) {
	d := time.Minute
	cli, err := client.NewEnvClient()
	if err != nil {
		log.Print(err)
	}
	ctx := context.Background()
	if err := cli.ContainerStop(ctx, id, &d); err != nil {
		log.Print(err)
	}
	if err := cli.ContainerRemove(ctx, id, types.ContainerRemoveOptions{Force: true}); err != nil {
		log.Print(err)
	}
}

// activeNotebooks fetchs copies of the tempNotebooks and returns them as a
// slice.
func (p *notebookPool) activeNotebooks() []tempNotebook {
	nbs := make([]tempNotebook, p.size())
	p.Lock()
	i := 0
	for _, nb := range p.containerMap {
		nbs[i] = *nb
		i++
	}
	p.Unlock()
	return nbs
}

// zombieNotebooks queries docker for containers that aren't under our
// supervision.  These can block ports assigned to our containers.
func (p *notebookPool) zombieContainers() ([]types.Container, error) {
	var cs []types.Container
	ids := map[string]struct{}{}
	p.Lock()
	for _, c := range p.containerMap {
		ids[c.id] = struct{}{}
	}
	p.Unlock()
	cli, err := client.NewEnvClient()
	opts := types.ContainerListOptions{}
	containers, err := cli.ContainerList(context.Background(), opts)
	if err != nil {
		return nil, err
	}
	for _, c := range containers {
		// If we manage it, leave it be
		if _, ok := ids[c.ID]; ok {
			continue
		}
		cs = append(cs, c)
	}
	return cs, nil
}

// nextCollection returns when the collector is run again
func (p *notebookPool) NextCollection() time.Time {
	p.lastCollMu.Lock()
	t := p.lastCollection.Add(p.containerLifetime / collectionFraction)
	p.lastCollMu.Unlock()
	return t
}

// startCollector launches a goroutine that checks for expired containers at
// interval d.  d is typically set to containerLifetime / collectionFraction.  Call
// stopCollector to stop the reclamation.
func (p *notebookPool) startCollector(d time.Duration) {
	go func() {
		ticker := time.NewTicker(d)
		for {
			select {
			case <-ticker.C:
				p.releaseContainers(false)
				p.lastCollMu.Lock()
				p.lastCollection = time.Now()
				p.lastCollMu.Unlock()
			case <-p.killCollection:
				ticker.Stop()
				return
			}
		}
	}()
}

// stopCollector sends a message on a channel to kill the auto reclamation.
func (p *notebookPool) stopCollector() {
	p.killCollection <- struct{}{}
}

// releaseContainers checks for expired containers and frees them from the
// containerMap.  It also frees the port in the portSet.  If force is true, age
// is ignored.
func (p *notebookPool) releaseContainers(force bool) error {
	p.Lock()
	trash := []tempNotebook{}
	for _, c := range p.containerMap {
		c.Lock()
		age := time.Now().Sub(c.lastAccessed)
		if age.Seconds() > p.containerLifetime.Seconds() || force {
			log.Printf("age: %v\n", age)
			trash = append(trash, *c)
		}
		c.Unlock()
	}
	p.Unlock()
	for _, c := range trash {
		c := c
		go func() {
			log.Printf("attempting to release container %s last accessed at %v", c.id, c.lastAccessed)
			p.stopAndKillContainer(c.id)
			p.portSet.Drop(c.port)
			p.Lock()
			delete(p.containerMap, c.hash)
			p.Unlock()
			// This isn't very elegant, but we couldn't delete the pattern from the mux
			// before, but now we can with the vendored/updated copy in mux.go.  We add
			// a trailing slash when we register the path, so we must add it here too.
			p.deregisterMux <- path.Join("/book", c.hash) + "/"
		}()
	}
	return nil
}

// killZombieContainers stops and kills any docker containers that aren't under
// out supervision.
//
// FIXME(kyle): not currently called at any time, when, why, etc...
func (p *notebookPool) killZombieContainers() error {
	zombies, err := p.zombieContainers()
	if err != nil {
		return err
	}
	for _, c := range zombies {
		p.stopAndKillContainer(c.ID)
	}
	return nil
}
<commit_msg>add logging to pull command<commit_after>// Copyright (c) 2017, Boise State University All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"context"
	"crypto/rand"
	"errors"
	"fmt"
	"io"
	"log"
	"path"
	"regexp"
	"sync"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/nat"
)

const (
	// jupyterNotebookImageMatch matches images from jupyter/docker-stacks, for
	// example: jupyter/minimal-notebook
	jupyterNotebookImageMatch = `[a-zA-Z0-9]+/[a-zA-Z0-9]+-notebook[:]{0,1}[a-zA-Z0-9]*`

	// allImageMatch applies no filter
	allImageMatch = `.*`

	// defaultContainerLifetime is used if a lifetime is not provided
	defaultContainerLifetime = time.Minute * 10

	// defaultMaxContainers governs the port set size and triggers reclamation
	defaultMaxContainers = 100

	// collectionFraction is the fraction of lifetime to collect containers.  For
	// example 4 collects every 1/4 of the container lifetime.
	collectionFraction = 4
)

// tempNotebook holds context for a single container
type tempNotebook struct {
	// guard struct (only used for lastAccessed right now)
	sync.Mutex
	// id is the docker container id.
	id string
	// hash is  a random generated hash that is used in the path of the server.
	hash string
	// imageName is the name of the image used to start the container
	imageName string
	// lastAccessed is when the container was used last.
	lastAccessed time.Time
	// port is the passthrough port for the reverse proxy.
	port int
}

// notebookPool holds data regarding running notebooks.
type notebookPool struct {
	// guards the entire struct
	sync.Mutex

	// availableImages is a list of docker images that installed on the machine,
	// and match the imageMatch expression
	availableImages map[string]struct{}

	// imageMatch filters available images by name
	imageMatch *regexp.Regexp

	// containerMap is stores the contexts for the containers.
	containerMap map[string]*tempNotebook

	// portSet holds free ports
	portSet *portRange

	// maxContainers governs the port set size and resource reclamation.
	maxContainers int

	// containerLifetime governs when the container resources are reclaimed.
	containerLifetime time.Duration

	// token is the security token for auto-auth
	token string

	// killCollection stops the automated resource reclamation
	killCollection chan struct{}

	// lastCollMu guards the time for reclamation.  It is used infrequently, and
	// we don't need to lock the whole object.
	lastCollMu sync.Mutex

	// lastCollection is the timestamp the last time the containers were
	// reclaimed.
	lastCollection time.Time

	// deregisterMux is a channel for sending a path that needs to be
	// de-registered from the server mux.
	deregisterMux chan string
}

// errNotebookPoolFull indicates the pool is at maxContainers
var errNotebookPoolFull = errors.New("container pool hit max size limit")

// newNotebookPool creates a notebookPool and sets defaults, overriding some
// with passed arguments.
func newNotebookPool(imageRegexp string, maxContainers int, lifetime time.Duration) (*notebookPool, error) {
	if imageRegexp == "" {
		imageRegexp = jupyterNotebookImageMatch
	}
	if int64(lifetime) <= 0 {
		lifetime = defaultContainerLifetime
	}
	if maxContainers < 1 {
		maxContainers = defaultMaxContainers
	}
	imageMatch, err := regexp.Compile(imageRegexp)
	if err != nil {
		return nil, err
	}
	imageMap := map[string]struct{}{}
	cli, err := client.NewEnvClient()
	if err != nil {
		return nil, err
	}
	images, err := cli.ImageList(context.Background(), types.ImageListOptions{})
	if err != nil {
		return nil, err
	}
	for _, image := range images {
		if len(image.RepoTags) < 1 || !imageMatch.MatchString(image.RepoTags[0]) {
			continue
		}
		log.Printf("found image %s", image.RepoTags[0])
		imageMap[image.RepoTags[0]] = struct{}{}
	}
	pool := &notebookPool{
		availableImages:   imageMap,
		imageMatch:        imageMatch,
		containerMap:      make(map[string]*tempNotebook),
		portSet:           newPortRange(8000, maxContainers),
		maxContainers:     maxContainers,
		containerLifetime: lifetime,
		killCollection:    make(chan struct{}),
		deregisterMux:     make(chan string),
	}
	pool.startCollector(time.Duration(int64(lifetime) / collectionFraction))
	pool.lastCollMu.Lock()
	pool.lastCollection = time.Now()
	pool.lastCollMu.Unlock()
	return pool, nil
}

// defaultHashSize is used for the unique hash generation
const defaultHashSize = 32

// newHash makes a n byte hash and returns the hex encoding.
func newHash(n int) string {
	b := make([]byte, n)
	_, err := rand.Read(b[:])
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", b)
}

// newNotebook initializes and sets values for a new notebook.
func (p *notebookPool) newNotebook(image string, pull bool) (*tempNotebook, error) {
	ctx := context.Background()
	cli, err := client.NewEnvClient()
	if err != nil {
		return nil, err
	}

	// TODO(kyle): possibly provide tag support
	if pull {
		log.Printf("pulling container %s", image)
		ctx, cancel := context.WithTimeout(ctx, 45*time.Second)
		defer cancel()
		r, err := cli.ImagePull(ctx, image, types.ImagePullOptions{})
		if err != nil {
			return nil, err
		}
		defer r.Close()
		var b bytes.Buffer
		io.Copy(&b, r)
		log.Print(string(b.Bytes()))
	}

	hash := newHash(defaultHashSize)

	port, err := p.portSet.Acquire()
	if err != nil {
		return nil, err
	}
	portString := fmt.Sprintf("%d", port)

	var pSet = nat.PortSet{}
	pt, err := nat.NewPort("tcp", portString)
	pSet[pt] = struct{}{}
	containerConfig := container.Config{
		Hostname: "0.0.0.0",
		User:     "jovyan",
		Cmd: []string{`jupyter`,
			`notebook`,
			`--no-browser`,
			`--port`,
			portString,
			`--ip=0.0.0.0`,
			fmt.Sprintf("--NotebookApp.base_url=%s", path.Join("/book", hash)),
			`--NotebookApp.port_retries=0`,
			fmt.Sprintf(`--NotebookApp.token="%s"`, p.token),
			`--NotebookApp.disable_check_xsrf=True`,
		},
		Env:          []string{fmt.Sprintf("CONFIGPROXY_AUTH_TOKEN=%s", p.token)},
		Image:        image,
		ExposedPorts: pSet,
	}

	hostConfig := container.HostConfig{
		NetworkMode: "host",
	}

	resp, err := cli.ContainerCreate(ctx, &containerConfig, &hostConfig, nil, "")
	if err != nil {
		p.portSet.Drop(port)
		return nil, err
	}

	if err := cli.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{}); err != nil {
		p.portSet.Drop(port)
		return nil, err
	}
	log.Printf("created container: %s", resp.ID)
	t := &tempNotebook{
		id:           resp.ID,
		hash:         hash,
		imageName:    image,
		lastAccessed: time.Now(),
		port:         port,
	}
	err = p.addNotebook(t)
	if err != nil {
		log.Print(err)
		p.portSet.Drop(port)
		return nil, err
	}
	return t, nil
}

// size returns the size of the containerMap with appropriate locks
func (p *notebookPool) size() int {
	p.Lock()
	n := len(p.containerMap)
	p.Unlock()
	return n
}

// addNotebook adds a tempNotebook to the containerMap, if there is room.
func (p *notebookPool) addNotebook(t *tempNotebook) error {
	n := p.size()
	log.Printf("pool size: %d of %d", n+1, p.maxContainers)
	if p.size()+1 > p.maxContainers {
		p.releaseContainers(false)
	}
	if p.size()+1 > p.maxContainers {
		return errNotebookPoolFull
	}
	p.Lock()
	p.containerMap[t.hash] = t
	p.Unlock()
	return nil
}

// stopAndKillContainer requests the stopping (docker stop) and the removal of
// the container (docker rm).  Errors are logged, but not returned and rm is
// always called.
func (p *notebookPool) stopAndKillContainer(id string) {
	d := time.Minute
	cli, err := client.NewEnvClient()
	if err != nil {
		log.Print(err)
	}
	ctx := context.Background()
	if err := cli.ContainerStop(ctx, id, &d); err != nil {
		log.Print(err)
	}
	if err := cli.ContainerRemove(ctx, id, types.ContainerRemoveOptions{Force: true}); err != nil {
		log.Print(err)
	}
}

// activeNotebooks fetchs copies of the tempNotebooks and returns them as a
// slice.
func (p *notebookPool) activeNotebooks() []tempNotebook {
	nbs := make([]tempNotebook, p.size())
	p.Lock()
	i := 0
	for _, nb := range p.containerMap {
		nbs[i] = *nb
		i++
	}
	p.Unlock()
	return nbs
}

// zombieNotebooks queries docker for containers that aren't under our
// supervision.  These can block ports assigned to our containers.
func (p *notebookPool) zombieContainers() ([]types.Container, error) {
	var cs []types.Container
	ids := map[string]struct{}{}
	p.Lock()
	for _, c := range p.containerMap {
		ids[c.id] = struct{}{}
	}
	p.Unlock()
	cli, err := client.NewEnvClient()
	opts := types.ContainerListOptions{}
	containers, err := cli.ContainerList(context.Background(), opts)
	if err != nil {
		return nil, err
	}
	for _, c := range containers {
		// If we manage it, leave it be
		if _, ok := ids[c.ID]; ok {
			continue
		}
		cs = append(cs, c)
	}
	return cs, nil
}

// nextCollection returns when the collector is run again
func (p *notebookPool) NextCollection() time.Time {
	p.lastCollMu.Lock()
	t := p.lastCollection.Add(p.containerLifetime / collectionFraction)
	p.lastCollMu.Unlock()
	return t
}

// startCollector launches a goroutine that checks for expired containers at
// interval d.  d is typically set to containerLifetime / collectionFraction.  Call
// stopCollector to stop the reclamation.
func (p *notebookPool) startCollector(d time.Duration) {
	go func() {
		ticker := time.NewTicker(d)
		for {
			select {
			case <-ticker.C:
				p.releaseContainers(false)
				p.lastCollMu.Lock()
				p.lastCollection = time.Now()
				p.lastCollMu.Unlock()
			case <-p.killCollection:
				ticker.Stop()
				return
			}
		}
	}()
}

// stopCollector sends a message on a channel to kill the auto reclamation.
func (p *notebookPool) stopCollector() {
	p.killCollection <- struct{}{}
}

// releaseContainers checks for expired containers and frees them from the
// containerMap.  It also frees the port in the portSet.  If force is true, age
// is ignored.
func (p *notebookPool) releaseContainers(force bool) error {
	p.Lock()
	trash := []tempNotebook{}
	for _, c := range p.containerMap {
		c.Lock()
		age := time.Now().Sub(c.lastAccessed)
		if age.Seconds() > p.containerLifetime.Seconds() || force {
			log.Printf("age: %v\n", age)
			trash = append(trash, *c)
		}
		c.Unlock()
	}
	p.Unlock()
	for _, c := range trash {
		c := c
		go func() {
			log.Printf("attempting to release container %s last accessed at %v", c.id, c.lastAccessed)
			p.stopAndKillContainer(c.id)
			p.portSet.Drop(c.port)
			p.Lock()
			delete(p.containerMap, c.hash)
			p.Unlock()
			// This isn't very elegant, but we couldn't delete the pattern from the mux
			// before, but now we can with the vendored/updated copy in mux.go.  We add
			// a trailing slash when we register the path, so we must add it here too.
			p.deregisterMux <- path.Join("/book", c.hash) + "/"
		}()
	}
	return nil
}

// killZombieContainers stops and kills any docker containers that aren't under
// out supervision.
//
// FIXME(kyle): not currently called at any time, when, why, etc...
func (p *notebookPool) killZombieContainers() error {
	zombies, err := p.zombieContainers()
	if err != nil {
		return err
	}
	for _, c := range zombies {
		p.stopAndKillContainer(c.ID)
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 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 nodeup

import (
	"k8s.io/klog/v2"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/nodeup/nodetasks"
)

type Loader struct {
	Builders []fi.ModelBuilder
}

// Build is responsible for running the build tasks for nodeup
func (l *Loader) Build() (map[string]fi.Task, error) {
	tasks := make(map[string]fi.Task)
	for _, builder := range l.Builders {
		context := &fi.ModelBuilderContext{
			Tasks: tasks,
		}
		err := builder.Build(context)
		if err != nil {
			return nil, err
		}
		tasks = context.Tasks
	}

	// If there is a package task, we need an update packages task
	for _, t := range tasks {
		if _, ok := t.(*nodetasks.Package); ok {
			klog.Infof("Package task found; adding UpdatePackages task")
			tasks["UpdatePackages"] = nodetasks.NewUpdatePackages()
			break
		}
	}
	if tasks["UpdatePackages"] == nil {
		klog.Infof("No package task found; won't update packages")
	}

	return tasks, nil
}
<commit_msg>Include type of builder in error message<commit_after>/*
Copyright 2019 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 nodeup

import (
	"fmt"
	"reflect"

	"k8s.io/klog/v2"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/nodeup/nodetasks"
)

type Loader struct {
	Builders []fi.ModelBuilder
}

// Build is responsible for running the build tasks for nodeup
func (l *Loader) Build() (map[string]fi.Task, error) {
	tasks := make(map[string]fi.Task)
	for _, builder := range l.Builders {
		context := &fi.ModelBuilderContext{
			Tasks: tasks,
		}
		err := builder.Build(context)
		if err != nil {
			return nil, fmt.Errorf("building %s: %v", reflect.TypeOf(builder), err)
		}
		tasks = context.Tasks
	}

	// If there is a package task, we need an update packages task
	for _, t := range tasks {
		if _, ok := t.(*nodetasks.Package); ok {
			klog.Infof("Package task found; adding UpdatePackages task")
			tasks["UpdatePackages"] = nodetasks.NewUpdatePackages()
			break
		}
	}
	if tasks["UpdatePackages"] == nil {
		klog.Infof("No package task found; won't update packages")
	}

	return tasks, nil
}
<|endoftext|>


<commit_before>package veneur

import (
	"testing"

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

func TestCounterEmpty(t *testing.T) {

	c := NewCounter("a.b.c", []string{"a:b"})
	c.Sample(1, 1.0)

	assert.Equal(t, "a.b.c", c.name, "Name")
	assert.Len(t, c.tags, 1, "Tag length")
	assert.Equal(t, c.tags[0], "a:b", "Tag contents")

	metrics := c.Flush()
	assert.Len(t, metrics, 1, "Flushes 1 metric")

	m1 := metrics[0]
	assert.Equal(t, int32(10), m1.Interval, "Interval")
	assert.Equal(t, "rate", m1.MetricType, "Type")
	assert.Len(t, c.tags, 1, "Tag length")
	assert.Equal(t, c.tags[0], "a:b", "Tag contents")
	// The counter returns an array with a single tuple of timestamp,value
	assert.Equal(t, 0.1, m1.Value[0][1], "Metric value")
}

func TestCounterRate(t *testing.T) {

	c := NewCounter("a.b.c", []string{"a:b"})

	c.Sample(5, 1.0)

	// The counter returns an array with a single tuple of timestamp,value
	metrics := c.Flush()
	assert.Equal(t, 0.5, metrics[0].Value[0][1], "Metric value")
}

func TestCounterSampleRate(t *testing.T) {

	c := NewCounter("a.b.c", []string{"a:b"})

	c.Sample(5, 0.5)

	// The counter returns an array with a single tuple of timestamp,value
	metrics := c.Flush()
	assert.Equal(t, float64(1), metrics[0].Value[0][1], "Metric value")
}

func TestGauge(t *testing.T) {

	g := NewGauge("a.b.c", []string{"a:b"})

	assert.Equal(t, "a.b.c", g.name, "Name")
	assert.Len(t, g.tags, 1, "Tag length")
	assert.Equal(t, g.tags[0], "a:b", "Tag contents")

	g.Sample(5, 1.0)

	metrics := g.Flush()
	assert.Len(t, metrics, 1, "Flushed metric count")

	m1 := metrics[0]
	// Interval is not meaningful for this
	assert.Equal(t, int32(0), m1.Interval, "Interval")
	assert.Equal(t, "gauge", m1.MetricType, "Type")
	tags := m1.Tags
	assert.Len(t, tags, 1, "Tag length")
	assert.Equal(t, tags[0], "a:b", "Tag contents")

	// The counter returns an array with a single tuple of timestamp,value
	assert.Equal(t, float64(5), m1.Value[0][1], "Value")
}

func TestSet(t *testing.T) {

	s := NewSet("a.b.c", []string{"a:b"}, 1000, 0.99)

	assert.Equal(t, "a.b.c", s.name, "Name")
	assert.Len(t, s.tags, 1, "Tag count")
	assert.Equal(t, "a:b", s.tags[0], "First tag")

	s.Sample(5, 1.0)

	s.Sample(5, 1.0)

	s.Sample(123, 1.0)

	metrics := s.Flush()
	assert.Len(t, metrics, 1, "Flush")

	m1 := metrics[0]
	// Interval is not meaningful for this
	assert.Equal(t, int32(0), m1.Interval, "Interval")
	assert.Equal(t, "set", m1.MetricType, "Type")
	assert.Len(t, m1.Tags, 1, "Tag count")
	assert.Equal(t, "a:b", m1.Tags[0], "First tag")
	assert.Equal(t, float64(2), m1.Value[0][1], "Value")
}

func TestHisto(t *testing.T) {

	h := NewHist("a.b.c", []string{"a:b"}, []float64{0.50})

	assert.Equal(t, "a.b.c", h.name, "Name")
	assert.Len(t, h.tags, 1, "Tag count")
	assert.Equal(t, "a:b", h.tags[0], "First tag")

	h.Sample(5, 1.0)
	h.Sample(10, 1.0)
	h.Sample(15, 1.0)
	h.Sample(20, 1.0)
	h.Sample(25, 1.0)

	metrics := h.Flush()
	// We get lots of metrics back for histograms!
	assert.Len(t, metrics, 4, "Flushed metrics length")

	// First the count
	m1 := metrics[0]
	assert.Equal(t, "a.b.c.count", m1.Name, "Name")
	assert.Equal(t, int32(10), m1.Interval, "Interval")
	assert.Equal(t, "rate", m1.MetricType, "Type")
	assert.Len(t, m1.Tags, 1, "Tag count")
	assert.Equal(t, "a:b", m1.Tags[0], "First tag")

	// The counter returns an array with a single tuple of timestamp,value
	assert.Equal(t, float64(0.5), m1.Value[0][1], "Value")

	// Now the max
	m2 := metrics[1]
	assert.Equal(t, "a.b.c.max", m2.Name, "Name")
	assert.Equal(t, int32(0), m2.Interval, "Interval")
	assert.Equal(t, "gauge", m2.MetricType, "Type")
	assert.Len(t, m2.Tags, 1, "Tag count")
	assert.Equal(t, "a:b", m2.Tags[0], "First tag")
	// The counter returns an array with a single tuple of timestamp,value
	assert.Equal(t, float64(25), m2.Value[0][1], "Value")

	// Now the min
	m3 := metrics[2]
	assert.Equal(t, "a.b.c.min", m3.Name, "Name")
	assert.Equal(t, int32(0), m3.Interval, "Interval")
	assert.Equal(t, "gauge", m3.MetricType, "Type")
	assert.Len(t, m3.Tags, 1, "Tag count")
	assert.Equal(t, "a:b", m3.Tags[0], "First tag")
	// The counter returns an array with a single tuple of timestamp,value
	assert.Equal(t, float64(5), m3.Value[0][1], "Value")

	// And the percentile
	m4 := metrics[3]
	assert.Equal(t, "a.b.c.50percentile", m4.Name, "Name")
	assert.Equal(t, int32(0), m4.Interval, "Interval")
	assert.Equal(t, "gauge", m4.MetricType, "Type")
	assert.Len(t, m4.Tags, 1, "Tag count")
	assert.Equal(t, "a:b", m4.Tags[0], "First tag")
	// The counter returns an array with a single tuple of timestamp,value
	assert.Equal(t, float64(15), m4.Value[0][1], "Value")
}

func TestHistoSampleRate(t *testing.T) {

	h := NewHist("a.b.c", []string{"a:b"}, []float64{0.50})

	assert.Equal(t, "a.b.c", h.name, "Name")
	assert.Len(t, h.tags, 1, "Tag length")
	assert.Equal(t, h.tags[0], "a:b", "Tag contents")

	h.Sample(5, 0.5)
	h.Sample(10, 0.5)
	h.Sample(15, 0.5)
	h.Sample(20, 0.5)
	h.Sample(25, 0.5)

	metrics := h.Flush()
	assert.Len(t, metrics, 4, "Metrics flush length")

	// First the count
	m1 := metrics[0]
	assert.Equal(t, "a.b.c.count", m1.Name, "Count name")
	assert.Equal(t, float64(1), m1.Value[0][1], "Sampled count as rate")
}
<commit_msg>Fix type test for set<commit_after>package veneur

import (
	"testing"

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

func TestCounterEmpty(t *testing.T) {

	c := NewCounter("a.b.c", []string{"a:b"})
	c.Sample(1, 1.0)

	assert.Equal(t, "a.b.c", c.name, "Name")
	assert.Len(t, c.tags, 1, "Tag length")
	assert.Equal(t, c.tags[0], "a:b", "Tag contents")

	metrics := c.Flush()
	assert.Len(t, metrics, 1, "Flushes 1 metric")

	m1 := metrics[0]
	assert.Equal(t, int32(10), m1.Interval, "Interval")
	assert.Equal(t, "rate", m1.MetricType, "Type")
	assert.Len(t, c.tags, 1, "Tag length")
	assert.Equal(t, c.tags[0], "a:b", "Tag contents")
	// The counter returns an array with a single tuple of timestamp,value
	assert.Equal(t, 0.1, m1.Value[0][1], "Metric value")
}

func TestCounterRate(t *testing.T) {

	c := NewCounter("a.b.c", []string{"a:b"})

	c.Sample(5, 1.0)

	// The counter returns an array with a single tuple of timestamp,value
	metrics := c.Flush()
	assert.Equal(t, 0.5, metrics[0].Value[0][1], "Metric value")
}

func TestCounterSampleRate(t *testing.T) {

	c := NewCounter("a.b.c", []string{"a:b"})

	c.Sample(5, 0.5)

	// The counter returns an array with a single tuple of timestamp,value
	metrics := c.Flush()
	assert.Equal(t, float64(1), metrics[0].Value[0][1], "Metric value")
}

func TestGauge(t *testing.T) {

	g := NewGauge("a.b.c", []string{"a:b"})

	assert.Equal(t, "a.b.c", g.name, "Name")
	assert.Len(t, g.tags, 1, "Tag length")
	assert.Equal(t, g.tags[0], "a:b", "Tag contents")

	g.Sample(5, 1.0)

	metrics := g.Flush()
	assert.Len(t, metrics, 1, "Flushed metric count")

	m1 := metrics[0]
	// Interval is not meaningful for this
	assert.Equal(t, int32(0), m1.Interval, "Interval")
	assert.Equal(t, "gauge", m1.MetricType, "Type")
	tags := m1.Tags
	assert.Len(t, tags, 1, "Tag length")
	assert.Equal(t, tags[0], "a:b", "Tag contents")

	// The counter returns an array with a single tuple of timestamp,value
	assert.Equal(t, float64(5), m1.Value[0][1], "Value")
}

func TestSet(t *testing.T) {

	s := NewSet("a.b.c", []string{"a:b"}, 1000, 0.99)

	assert.Equal(t, "a.b.c", s.name, "Name")
	assert.Len(t, s.tags, 1, "Tag count")
	assert.Equal(t, "a:b", s.tags[0], "First tag")

	s.Sample(5, 1.0)

	s.Sample(5, 1.0)

	s.Sample(123, 1.0)

	metrics := s.Flush()
	assert.Len(t, metrics, 1, "Flush")

	m1 := metrics[0]
	// Interval is not meaningful for this
	assert.Equal(t, int32(0), m1.Interval, "Interval")
	assert.Equal(t, "gauge", m1.MetricType, "Type")
	assert.Len(t, m1.Tags, 1, "Tag count")
	assert.Equal(t, "a:b", m1.Tags[0], "First tag")
	assert.Equal(t, float64(2), m1.Value[0][1], "Value")
}

func TestHisto(t *testing.T) {

	h := NewHist("a.b.c", []string{"a:b"}, []float64{0.50})

	assert.Equal(t, "a.b.c", h.name, "Name")
	assert.Len(t, h.tags, 1, "Tag count")
	assert.Equal(t, "a:b", h.tags[0], "First tag")

	h.Sample(5, 1.0)
	h.Sample(10, 1.0)
	h.Sample(15, 1.0)
	h.Sample(20, 1.0)
	h.Sample(25, 1.0)

	metrics := h.Flush()
	// We get lots of metrics back for histograms!
	assert.Len(t, metrics, 4, "Flushed metrics length")

	// First the count
	m1 := metrics[0]
	assert.Equal(t, "a.b.c.count", m1.Name, "Name")
	assert.Equal(t, int32(10), m1.Interval, "Interval")
	assert.Equal(t, "rate", m1.MetricType, "Type")
	assert.Len(t, m1.Tags, 1, "Tag count")
	assert.Equal(t, "a:b", m1.Tags[0], "First tag")

	// The counter returns an array with a single tuple of timestamp,value
	assert.Equal(t, float64(0.5), m1.Value[0][1], "Value")

	// Now the max
	m2 := metrics[1]
	assert.Equal(t, "a.b.c.max", m2.Name, "Name")
	assert.Equal(t, int32(0), m2.Interval, "Interval")
	assert.Equal(t, "gauge", m2.MetricType, "Type")
	assert.Len(t, m2.Tags, 1, "Tag count")
	assert.Equal(t, "a:b", m2.Tags[0], "First tag")
	// The counter returns an array with a single tuple of timestamp,value
	assert.Equal(t, float64(25), m2.Value[0][1], "Value")

	// Now the min
	m3 := metrics[2]
	assert.Equal(t, "a.b.c.min", m3.Name, "Name")
	assert.Equal(t, int32(0), m3.Interval, "Interval")
	assert.Equal(t, "gauge", m3.MetricType, "Type")
	assert.Len(t, m3.Tags, 1, "Tag count")
	assert.Equal(t, "a:b", m3.Tags[0], "First tag")
	// The counter returns an array with a single tuple of timestamp,value
	assert.Equal(t, float64(5), m3.Value[0][1], "Value")

	// And the percentile
	m4 := metrics[3]
	assert.Equal(t, "a.b.c.50percentile", m4.Name, "Name")
	assert.Equal(t, int32(0), m4.Interval, "Interval")
	assert.Equal(t, "gauge", m4.MetricType, "Type")
	assert.Len(t, m4.Tags, 1, "Tag count")
	assert.Equal(t, "a:b", m4.Tags[0], "First tag")
	// The counter returns an array with a single tuple of timestamp,value
	assert.Equal(t, float64(15), m4.Value[0][1], "Value")
}

func TestHistoSampleRate(t *testing.T) {

	h := NewHist("a.b.c", []string{"a:b"}, []float64{0.50})

	assert.Equal(t, "a.b.c", h.name, "Name")
	assert.Len(t, h.tags, 1, "Tag length")
	assert.Equal(t, h.tags[0], "a:b", "Tag contents")

	h.Sample(5, 0.5)
	h.Sample(10, 0.5)
	h.Sample(15, 0.5)
	h.Sample(20, 0.5)
	h.Sample(25, 0.5)

	metrics := h.Flush()
	assert.Len(t, metrics, 4, "Metrics flush length")

	// First the count
	m1 := metrics[0]
	assert.Equal(t, "a.b.c.count", m1.Name, "Count name")
	assert.Equal(t, float64(1), m1.Value[0][1], "Sampled count as rate")
}
<|endoftext|>