<commit_before>// Copyright 2017 Mesosphere, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"errors"
	"fmt"
	"math"
	"regexp"
	"strconv"
)

// measurement represents each point in a time series
type measurement struct {
	Name  string            `json:"name"`
	Value float64           `json:"value"`
	Tags  map[string]string `json:"tags"`
	Time  int64             `json:"time"`
}

func newMeasurement() *measurement {
	return &measurement{
		Tags: make(map[string]string),
	}
}

func (m *measurement) floorTime(interval int64) {
	m.Time = m.Time / interval * interval
}

func (m *measurement) String() string {
	return m.Name
}

func (m *measurement) setValue(value interface{}) error {
	val, err := m.toFloat64(value)
	if err != nil {
		return fmt.Errorf("Could not set value: %v", err)
	}
	m.Value = val
	return nil
}

func (m *measurement) addTag(name string, value string) error {
	matched, err := regexp.MatchString(`^[-.:_\w]+\z{1,64}$`, name)
	if err != nil {
		return fmt.Errorf("Error occurred matching tag name: %v", err)
	}
	if !matched {
		return fmt.Errorf("Tag name '%s' is not valid", name)
	}
	matched, err = regexp.MatchString(`^[-.:_\\/\w ]{1,255}$`, value)
	if err != nil {
		return fmt.Errorf("Error occurred matching tag value: %v", err)
	}
	if !matched {
		return fmt.Errorf("Tag value '%s' is not valid", value)
	}
	m.Tags[name] = value
	return nil
}

func (m *measurement) toFloat64(value interface{}) (float64, error) {
	switch value := value.(type) {
	case float64:
		return value, nil
	case float32:
		return float64(value), nil
	case int64:
		return float64(value), nil
	case int32:
		return float64(value), nil
	case int:
		return float64(value), nil
	case string:
		return strconv.ParseFloat(value, 64)
	}
	return math.NaN(), fmt.Errorf("Could not convert %+v (%T) to float64", value, value)
}

func (m *measurement) validate() error {
	if len(m.Tags) == 0 {
		return errors.New("Must have at least one tag")
	}
	if len(m.Name) == 0 {
		return errors.New("Must have a non-empty name")
	}
	if m.Time <= 0 {
		return errors.New("Time must be >= 0")
	}
	return nil
}
<commit_msg>Refactored to use plugin.DatapointValueToFloat64.<commit_after>// Copyright 2017 Mesosphere, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"errors"
	"fmt"
	"regexp"

	"github.com/dcos/dcos-metrics/plugins"
)

// measurement represents each point in a time series
type measurement struct {
	Name  string            `json:"name"`
	Value float64           `json:"value"`
	Tags  map[string]string `json:"tags"`
	Time  int64             `json:"time"`
}

func newMeasurement() *measurement {
	return &measurement{
		Tags: make(map[string]string),
	}
}

func (m *measurement) floorTime(interval int64) {
	m.Time = m.Time / interval * interval
}

func (m *measurement) String() string {
	return m.Name
}

func (m *measurement) setValue(value interface{}) error {
	val, err := plugin.DatapointValueToFloat64(value)
	if err != nil {
		return fmt.Errorf("Could not set value: %v", err)
	}
	m.Value = val
	return nil
}

func (m *measurement) addTag(name string, value string) error {
	matched, err := regexp.MatchString(`^[-.:_\w]+\z{1,64}$`, name)
	if err != nil {
		return fmt.Errorf("Error occurred matching tag name: %v", err)
	}
	if !matched {
		return fmt.Errorf("Tag name '%s' is not valid", name)
	}
	matched, err = regexp.MatchString(`^[-.:_\\/\w ]{1,255}$`, value)
	if err != nil {
		return fmt.Errorf("Error occurred matching tag value: %v", err)
	}
	if !matched {
		return fmt.Errorf("Tag value '%s' is not valid", value)
	}
	m.Tags[name] = value
	return nil
}

func (m *measurement) validate() error {
	if len(m.Tags) == 0 {
		return errors.New("Must have at least one tag")
	}
	if len(m.Name) == 0 {
		return errors.New("Must have a non-empty name")
	}
	if m.Time <= 0 {
		return errors.New("Time must be >= 0")
	}
	return nil
}
<|endoftext|>


<commit_before>package portallocator

import (
	"bufio"
	"errors"
	"fmt"
	"net"
	"os"
	"sync"

	"github.com/Sirupsen/logrus"
)

const (
	// DefaultPortRangeStart indicates the first port in port range
	DefaultPortRangeStart = 49153
	// DefaultPortRangeEnd indicates the last port in port range
	DefaultPortRangeEnd = 65535
)

type ipMapping map[string]protoMap

var (
	// ErrAllPortsAllocated is returned when no more ports are available
	ErrAllPortsAllocated = errors.New("all ports are allocated")
	// ErrUnknownProtocol is returned when an unknown protocol was specified
	ErrUnknownProtocol = errors.New("unknown protocol")
	defaultIP          = net.ParseIP("0.0.0.0")
	once               sync.Once
	instance           *PortAllocator
	createInstance     = func() { instance = newInstance() }
)

// ErrPortAlreadyAllocated is the returned error information when a requested port is already being used
type ErrPortAlreadyAllocated struct {
	ip   string
	port int
}

func newErrPortAlreadyAllocated(ip string, port int) ErrPortAlreadyAllocated {
	return ErrPortAlreadyAllocated{
		ip:   ip,
		port: port,
	}
}

// IP returns the address to which the used port is associated
func (e ErrPortAlreadyAllocated) IP() string {
	return e.ip
}

// Port returns the value of the already used port
func (e ErrPortAlreadyAllocated) Port() int {
	return e.port
}

// IPPort returns the address and the port in the form ip:port
func (e ErrPortAlreadyAllocated) IPPort() string {
	return fmt.Sprintf("%s:%d", e.ip, e.port)
}

// Error is the implementation of error.Error interface
func (e ErrPortAlreadyAllocated) Error() string {
	return fmt.Sprintf("Bind for %s:%d failed: port is already allocated", e.ip, e.port)
}

type (
	// PortAllocator manages the transport ports database
	PortAllocator struct {
		mutex sync.Mutex
		ipMap ipMapping
		Begin int
		End   int
	}
	portMap struct {
		p          map[int]struct{}
		begin, end int
		last       int
	}
	protoMap map[string]*portMap
)

// Get returns the default instance of PortAllocator
func Get() *PortAllocator {
	// Port Allocator is a singleton
	// Note: Long term solution will be each PortAllocator will have access to
	// the OS so that it can have up to date view of the OS port allocation.
	// When this happens singleton behavior will be removed. Clients do not
	// need to worry about this, they will not see a change in behavior.
	once.Do(createInstance)
	return instance
}

func newInstance() *PortAllocator {
	start, end, err := getDynamicPortRange()
	if err != nil {
		logrus.Warn(err)
		start, end = DefaultPortRangeStart, DefaultPortRangeEnd
	}
	return &PortAllocator{
		ipMap: ipMapping{},
		Begin: start,
		End:   end,
	}
}

func getDynamicPortRange() (start int, end int, err error) {
	const portRangeKernelParam = "/proc/sys/net/ipv4/ip_local_port_range"
	portRangeFallback := fmt.Sprintf("using fallback port range %d-%d", DefaultPortRangeStart, DefaultPortRangeEnd)
	file, err := os.Open(portRangeKernelParam)
	if err != nil {
		return 0, 0, fmt.Errorf("port allocator - %s due to error: %v", portRangeFallback, err)
	}
	n, err := fmt.Fscanf(bufio.NewReader(file), "%d\t%d", &start, &end)
	if n != 2 || err != nil {
		if err == nil {
			err = fmt.Errorf("unexpected count of parsed numbers (%d)", n)
		}
		return 0, 0, fmt.Errorf("port allocator - failed to parse system ephemeral port range from %s - %s: %v", portRangeKernelParam, portRangeFallback, err)
	}
	return start, end, nil
}

// RequestPort requests new port from global ports pool for specified ip and proto.
// If port is 0 it returns first free port. Otherwise it checks port availability
// in pool and return that port or error if port is already busy.
func (p *PortAllocator) RequestPort(ip net.IP, proto string, port int) (int, error) {
	p.mutex.Lock()
	defer p.mutex.Unlock()

	if proto != "tcp" && proto != "udp" {
		return 0, ErrUnknownProtocol
	}

	if ip == nil {
		ip = defaultIP
	}
	ipstr := ip.String()
	protomap, ok := p.ipMap[ipstr]
	if !ok {
		protomap = protoMap{
			"tcp": p.newPortMap(),
			"udp": p.newPortMap(),
		}

		p.ipMap[ipstr] = protomap
	}
	mapping := protomap[proto]
	if port > 0 {
		if _, ok := mapping.p[port]; !ok {
			mapping.p[port] = struct{}{}
			return port, nil
		}
		return 0, newErrPortAlreadyAllocated(ipstr, port)
	}

	port, err := mapping.findPort()
	if err != nil {
		return 0, err
	}
	return port, nil
}

// ReleasePort releases port from global ports pool for specified ip and proto.
func (p *PortAllocator) ReleasePort(ip net.IP, proto string, port int) error {
	p.mutex.Lock()
	defer p.mutex.Unlock()

	if ip == nil {
		ip = defaultIP
	}
	protomap, ok := p.ipMap[ip.String()]
	if !ok {
		return nil
	}
	delete(protomap[proto].p, port)
	return nil
}

func (p *PortAllocator) newPortMap() *portMap {
	return &portMap{
		p:     map[int]struct{}{},
		begin: p.Begin,
		end:   p.End,
		last:  p.End,
	}
}

// ReleaseAll releases all ports for all ips.
func (p *PortAllocator) ReleaseAll() error {
	p.mutex.Lock()
	p.ipMap = ipMapping{}
	p.mutex.Unlock()
	return nil
}

func (pm *portMap) findPort() (int, error) {
	port := pm.last
	for i := 0; i <= pm.end-pm.begin; i++ {
		port++
		if port > pm.end {
			port = pm.begin
		}

		if _, ok := pm.p[port]; !ok {
			pm.p[port] = struct{}{}
			pm.last = port
			return port, nil
		}
	}
	return 0, ErrAllPortsAllocated
}
<commit_msg>Do not warn in packages<commit_after>package portallocator

import (
	"bufio"
	"errors"
	"fmt"
	"net"
	"os"
	"sync"
)

const (
	// DefaultPortRangeStart indicates the first port in port range
	DefaultPortRangeStart = 49153
	// DefaultPortRangeEnd indicates the last port in port range
	DefaultPortRangeEnd = 65535
)

type ipMapping map[string]protoMap

var (
	// ErrAllPortsAllocated is returned when no more ports are available
	ErrAllPortsAllocated = errors.New("all ports are allocated")
	// ErrUnknownProtocol is returned when an unknown protocol was specified
	ErrUnknownProtocol = errors.New("unknown protocol")
	defaultIP          = net.ParseIP("0.0.0.0")
	once               sync.Once
	instance           *PortAllocator
	createInstance     = func() { instance = newInstance() }
)

// ErrPortAlreadyAllocated is the returned error information when a requested port is already being used
type ErrPortAlreadyAllocated struct {
	ip   string
	port int
}

func newErrPortAlreadyAllocated(ip string, port int) ErrPortAlreadyAllocated {
	return ErrPortAlreadyAllocated{
		ip:   ip,
		port: port,
	}
}

// IP returns the address to which the used port is associated
func (e ErrPortAlreadyAllocated) IP() string {
	return e.ip
}

// Port returns the value of the already used port
func (e ErrPortAlreadyAllocated) Port() int {
	return e.port
}

// IPPort returns the address and the port in the form ip:port
func (e ErrPortAlreadyAllocated) IPPort() string {
	return fmt.Sprintf("%s:%d", e.ip, e.port)
}

// Error is the implementation of error.Error interface
func (e ErrPortAlreadyAllocated) Error() string {
	return fmt.Sprintf("Bind for %s:%d failed: port is already allocated", e.ip, e.port)
}

type (
	// PortAllocator manages the transport ports database
	PortAllocator struct {
		mutex sync.Mutex
		ipMap ipMapping
		Begin int
		End   int
	}
	portMap struct {
		p          map[int]struct{}
		begin, end int
		last       int
	}
	protoMap map[string]*portMap
)

// Get returns the default instance of PortAllocator
func Get() *PortAllocator {
	// Port Allocator is a singleton
	// Note: Long term solution will be each PortAllocator will have access to
	// the OS so that it can have up to date view of the OS port allocation.
	// When this happens singleton behavior will be removed. Clients do not
	// need to worry about this, they will not see a change in behavior.
	once.Do(createInstance)
	return instance
}

func newInstance() *PortAllocator {
	start, end, err := getDynamicPortRange()
	if err != nil {
		start, end = DefaultPortRangeStart, DefaultPortRangeEnd
	}
	return &PortAllocator{
		ipMap: ipMapping{},
		Begin: start,
		End:   end,
	}
}

func getDynamicPortRange() (start int, end int, err error) {
	const portRangeKernelParam = "/proc/sys/net/ipv4/ip_local_port_range"
	portRangeFallback := fmt.Sprintf("using fallback port range %d-%d", DefaultPortRangeStart, DefaultPortRangeEnd)
	file, err := os.Open(portRangeKernelParam)
	if err != nil {
		return 0, 0, fmt.Errorf("port allocator - %s due to error: %v", portRangeFallback, err)
	}
	n, err := fmt.Fscanf(bufio.NewReader(file), "%d\t%d", &start, &end)
	if n != 2 || err != nil {
		if err == nil {
			err = fmt.Errorf("unexpected count of parsed numbers (%d)", n)
		}
		return 0, 0, fmt.Errorf("port allocator - failed to parse system ephemeral port range from %s - %s: %v", portRangeKernelParam, portRangeFallback, err)
	}
	return start, end, nil
}

// RequestPort requests new port from global ports pool for specified ip and proto.
// If port is 0 it returns first free port. Otherwise it checks port availability
// in pool and return that port or error if port is already busy.
func (p *PortAllocator) RequestPort(ip net.IP, proto string, port int) (int, error) {
	p.mutex.Lock()
	defer p.mutex.Unlock()

	if proto != "tcp" && proto != "udp" {
		return 0, ErrUnknownProtocol
	}

	if ip == nil {
		ip = defaultIP
	}
	ipstr := ip.String()
	protomap, ok := p.ipMap[ipstr]
	if !ok {
		protomap = protoMap{
			"tcp": p.newPortMap(),
			"udp": p.newPortMap(),
		}

		p.ipMap[ipstr] = protomap
	}
	mapping := protomap[proto]
	if port > 0 {
		if _, ok := mapping.p[port]; !ok {
			mapping.p[port] = struct{}{}
			return port, nil
		}
		return 0, newErrPortAlreadyAllocated(ipstr, port)
	}

	port, err := mapping.findPort()
	if err != nil {
		return 0, err
	}
	return port, nil
}

// ReleasePort releases port from global ports pool for specified ip and proto.
func (p *PortAllocator) ReleasePort(ip net.IP, proto string, port int) error {
	p.mutex.Lock()
	defer p.mutex.Unlock()

	if ip == nil {
		ip = defaultIP
	}
	protomap, ok := p.ipMap[ip.String()]
	if !ok {
		return nil
	}
	delete(protomap[proto].p, port)
	return nil
}

func (p *PortAllocator) newPortMap() *portMap {
	return &portMap{
		p:     map[int]struct{}{},
		begin: p.Begin,
		end:   p.End,
		last:  p.End,
	}
}

// ReleaseAll releases all ports for all ips.
func (p *PortAllocator) ReleaseAll() error {
	p.mutex.Lock()
	p.ipMap = ipMapping{}
	p.mutex.Unlock()
	return nil
}

func (pm *portMap) findPort() (int, error) {
	port := pm.last
	for i := 0; i <= pm.end-pm.begin; i++ {
		port++
		if port > pm.end {
			port = pm.begin
		}

		if _, ok := pm.p[port]; !ok {
			pm.p[port] = struct{}{}
			pm.last = port
			return port, nil
		}
	}
	return 0, ErrAllPortsAllocated
}
<|endoftext|>


<commit_before>package ebpf

import (
	"fmt"
	"runtime"
	"syscall"
	"unsafe"

	"github.com/pkg/errors"
)

// MapSpec defines a Map.
type MapSpec struct {
	Type       MapType
	KeySize    uint32
	ValueSize  uint32
	MaxEntries uint32
	Flags      uint32
	// InnerMap is used as a template for ArrayOfMaps and HashOfMaps
	InnerMap *MapSpec
}

func (ms *MapSpec) String() string {
	return fmt.Sprintf("%s(keySize=%d, valueSize=%d, maxEntries=%d, flags=%d)", ms.Type, ms.KeySize, ms.ValueSize, ms.MaxEntries, ms.Flags)
}

// Map represents a Map file descriptor.
//
// Methods which take interface{} arguments by default encode
// them using binary.Read/Write in the machine's native endianness.
//
// Implement Marshaler on the arguments if you need custom encoding.
type Map struct {
	fd   uint32
	meta MapSpec
	// Per CPU maps return values larger than the size in the spec
	fullValueSize int
}

// NewMap creates a new Map
func NewMap(spec *MapSpec) (*Map, error) {
	if spec.Type != ArrayOfMaps && spec.Type != HashOfMaps {
		return newMap(spec, 0)
	}

	if spec.ValueSize != 0 {
		return nil, fmt.Errorf("ebpf: ValueSize must be zero for map of map")
	}
	if spec.InnerMap == nil {
		return nil, fmt.Errorf("ebpf: map of map requires InnerMap")
	}

	inner, err := newMap(spec.InnerMap, 0)
	if err != nil {
		return nil, err
	}
	defer inner.Close()

	outerSpec := *spec
	outerSpec.InnerMap = nil
	outerSpec.ValueSize = 4
	return newMap(&outerSpec, inner.fd)
}

func newMap(spec *MapSpec, inner uint32) (*Map, error) {
	if spec.InnerMap != nil {
		return nil, fmt.Errorf("ebpf: inner map not allowed for this type")
	}
	attr := mapCreateAttr{
		spec.Type,
		spec.KeySize,
		spec.ValueSize,
		spec.MaxEntries,
		spec.Flags,
		inner,
	}
	fd, err := bpfCall(_MapCreate, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
	if err != nil {
		return nil, errors.Wrap(err, "map create")
	}

	return alignMap(uint32(fd), spec)
}

func alignMap(fd uint32, spec *MapSpec) (*Map, error) {
	if !spec.Type.hasPerCPUValue() {
		return &Map{
			uint32(fd),
			*spec,
			int(spec.ValueSize),
		}, nil
	}

	possibleCPUs, err := possibleCPUs()
	if err != nil {
		return nil, err
	}

	fullValueSize := align(int(spec.ValueSize), 8) * possibleCPUs
	m := &Map{
		uint32(fd),
		*spec,
		fullValueSize,
	}
	runtime.SetFinalizer(m, (*Map).Close)
	return m, nil
}

func (m *Map) String() string {
	return fmt.Sprintf("%s#%d", m.meta.Type, m.fd)
}

// Get gets a value from a Map
func (m *Map) Get(key, valueOut interface{}) (bool, error) {
	valueBytes, err := m.GetBytes(key)
	if err != nil {
		return false, err
	}
	if valueBytes == nil {
		return false, nil
	}

	if m.meta.Type.hasPerCPUValue() {
		return true, unmarshalPerCPUValue(valueOut, int(m.meta.ValueSize), valueBytes)
	}

	return true, unmarshalBytes(valueOut, valueBytes)
}

// GetBytes gets a value from Map
func (m *Map) GetBytes(key interface{}) ([]byte, error) {
	keyBytes, err := marshalBytes(key, int(m.meta.KeySize))
	if err != nil {
		return nil, err
	}
	valueBytes := make([]byte, m.fullValueSize)
	attr := mapOpAttr{
		mapFd: m.fd,
		key:   newPtr(unsafe.Pointer(&keyBytes[0])),
		value: newPtr(unsafe.Pointer(&valueBytes[0])),
	}
	_, err = bpfCall(_MapLookupElem, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
	if errors.Cause(err) == syscall.ENOENT {
		return nil, nil
	}
	return valueBytes, err
}

// Create creates a new value in a map, failing if the key exists already
func (m *Map) Create(key, value interface{}) error {
	return m.put(key, value, _NoExist)
}

// Put replaces or creates a value in map
func (m *Map) Put(key, value interface{}) error {
	return m.put(key, value, _Any)
}

// Replace replaces a value in a map, failing if the value did not exist
func (m *Map) Replace(key, value interface{}) error {
	return m.put(key, value, _Exist)
}

// Delete removes a value.
//
// Use DeleteStrict if you desire an error if key does not exist.
func (m *Map) Delete(key interface{}) error {
	err := m.DeleteStrict(key)
	if err == syscall.ENOENT {
		return nil
	}
	return err
}

// DeleteStrict removes a key and returns an error if the
// key doesn't exist.
func (m *Map) DeleteStrict(key interface{}) error {
	keyBytes, err := marshalBytes(key, int(m.meta.KeySize))
	if err != nil {
		return err
	}
	attr := mapOpAttr{
		mapFd: m.fd,
		key:   newPtr(unsafe.Pointer(&keyBytes[0])),
	}
	_, err = bpfCall(_MapDeleteElem, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
	return err
}

// NextKey finds the key following an initial key.
//
// See NextKeyBytes for details.
func (m *Map) NextKey(key, nextKeyOut interface{}) (bool, error) {
	nextKeyBytes, err := m.NextKeyBytes(key)
	if err != nil {
		return false, err
	}
	if nextKeyBytes == nil {
		return false, nil
	}
	err = unmarshalBytes(nextKeyOut, nextKeyBytes)
	if err != nil {
		return false, err
	}
	return true, nil
}

// NextKeyBytes returns the key following an initial key as a byte slice.
//
// Passing nil will return the first key.
//
// Use Iterate if you want to traverse all entries in the map.
func (m *Map) NextKeyBytes(key interface{}) ([]byte, error) {
	var keyPtr syscallPtr
	if key != nil {
		keyBytes, err := marshalBytes(key, int(m.meta.KeySize))
		if err != nil {
			return nil, err
		}
		keyPtr = newPtr(unsafe.Pointer(&keyBytes[0]))
	}

	nextKey := make([]byte, m.meta.KeySize)
	attr := mapOpAttr{
		mapFd: m.fd,
		key:   keyPtr,
		value: newPtr(unsafe.Pointer(&nextKey[0])),
	}
	_, err := bpfCall(_MapGetNextKey, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
	if errors.Cause(err) == syscall.ENOENT {
		return nil, nil
	}
	return nextKey, err
}

// Iterate traverses a map.
//
// It's not possible to guarantee that all keys in a map will be
// returned if there are concurrent modifications to the map. If a
// map is modified too heavily iteration may abort.
func (m *Map) Iterate() *MapIterator {
	return &MapIterator{
		target: m,
	}
}

// Close removes a Map
func (m *Map) Close() error {
	runtime.SetFinalizer(m, nil)
	return syscall.Close(int(m.fd))
}

// FD gets the raw fd value of Map
func (m *Map) FD() int {
	return int(m.fd)
}

// Pin persists the map past the lifetime of the process that created it.
//
// This requires bpffs to be mounted above fileName. See http://cilium.readthedocs.io/en/doc-1.0/kubernetes/install/#mounting-the-bpf-fs-optional
func (m *Map) Pin(fileName string) error {
	return pinObject(fileName, m.fd)
}

// LoadPinnedMap load a Map from a BPF file.
//
// Requires at least Linux 4.13, use LoadPinnedMapExplicit on
// earlier versions.
func LoadPinnedMap(fileName string) (*Map, error) {
	fd, err := getObject(fileName)
	if err != nil {
		return nil, err
	}
	spec, err := getMapSpecByFD(fd)
	if err != nil {
		return nil, err
	}
	return alignMap(fd, spec)
}

// LoadPinnedMapExplicit loads a map with explicit parameters.
func LoadPinnedMapExplicit(fileName string, spec *MapSpec) (*Map, error) {
	fd, err := getObject(fileName)
	if err != nil {
		return nil, err
	}
	return alignMap(fd, spec)
}

func (m *Map) put(key, value interface{}, putType uint64) error {
	keyBytes, err := marshalBytes(key, int(m.meta.KeySize))
	if err != nil {
		return err
	}

	var valueBytes []byte
	if m.meta.Type.hasPerCPUValue() {
		valueBytes, err = marshalPerCPUValue(value, int(m.meta.ValueSize))
	} else {
		valueBytes, err = marshalBytes(value, int(m.meta.ValueSize))
	}
	if err != nil {
		return err
	}

	attr := mapOpAttr{
		mapFd: m.fd,
		key:   newPtr(unsafe.Pointer(&keyBytes[0])),
		value: newPtr(unsafe.Pointer(&valueBytes[0])),
		flags: putType,
	}
	_, err = bpfCall(_MapUpdateElem, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
	return err
}

// UnmarshalBinary implements BinaryUnmarshaler.
func (m *Map) UnmarshalBinary(buf []byte) error {
	if m.fd != 0 {
		return errors.New("ebpf: can't unmarshal into existing map")
	}
	if len(buf) != 4 {
		return errors.New("ebpf: map id requires uint32")
	}
	// Looking up an entry in a nested map or prog array returns an id,
	// not an fd.
	id := nativeEndian.Uint32(buf)
	fd, err := getMapFDByID(id)
	if err != nil {
		return err
	}
	meta, err := getMapSpecByFD(fd)
	if err != nil {
		return err
	}
	m.fd = fd
	m.meta = *meta
	return nil
}

// MarshalBinary implements BinaryMarshaler.
func (m *Map) MarshalBinary() ([]byte, error) {
	buf := make([]byte, 4)
	nativeEndian.PutUint32(buf, m.fd)
	return buf, nil
}

// MapIterator iterates a Map.
//
// See Map.Iterate.
type MapIterator struct {
	target *Map
	prev   interface{}
	done   bool
	err    error
}

// Next decodes the next key and value.
//
// Returns false if there are no more entries.
func (mi *MapIterator) Next(keyOut, valueOut interface{}) bool {
	if mi.err != nil || mi.done {
		return false
	}

	var nextBytes []byte
	for i := 0; i < 3; i++ {
		nextBytes, mi.err = mi.target.NextKeyBytes(mi.prev)
		if mi.err != nil {
			return false
		}

		if nextBytes == nil {
			mi.done = true
			return false
		}

		var ok bool
		ok, mi.err = mi.target.Get(nextBytes, valueOut)
		if mi.err != nil {
			return false
		}

		if ok {
			break
		}

		// The next key was deleted before we could retrieve
		// it's value. As of Linux 4.16 there is no safe API which
		// prevents this race.
		nextBytes = nil
	}

	if nextBytes == nil {
		// We still hit the race condition even though we retried.
		mi.err = errors.New("ebpf: can't retrieve next entry, map mutated too quickly")
		return false
	}

	mi.err = unmarshalBytes(keyOut, nextBytes)
	if mi.err != nil {
		return false
	}

	// The user can get access to nextBytes since marshalBytes
	// does not copy when unmarshaling into a []byte.
	// Make a copy to prevent accidental corruption of
	// iterator state.
	prevBytes := make([]byte, len(nextBytes))
	copy(prevBytes, nextBytes)
	mi.prev = prevBytes
	return true
}

// Err returns any encountered error.
//
// The method must be called after Next returns nil.
func (mi *MapIterator) Err() error {
	return mi.err
}
<commit_msg>fix Map.UnmarshalBinary not setting fullValueSize<commit_after>package ebpf

import (
	"fmt"
	"runtime"
	"syscall"
	"unsafe"

	"github.com/pkg/errors"
)

// MapSpec defines a Map.
type MapSpec struct {
	Type       MapType
	KeySize    uint32
	ValueSize  uint32
	MaxEntries uint32
	Flags      uint32
	// InnerMap is used as a template for ArrayOfMaps and HashOfMaps
	InnerMap *MapSpec
}

func (ms *MapSpec) String() string {
	return fmt.Sprintf("%s(keySize=%d, valueSize=%d, maxEntries=%d, flags=%d)", ms.Type, ms.KeySize, ms.ValueSize, ms.MaxEntries, ms.Flags)
}

// Map represents a Map file descriptor.
//
// Methods which take interface{} arguments by default encode
// them using binary.Read/Write in the machine's native endianness.
//
// Implement Marshaler on the arguments if you need custom encoding.
type Map struct {
	fd   uint32
	meta MapSpec
	// Per CPU maps return values larger than the size in the spec
	fullValueSize int
}

// NewMap creates a new Map
func NewMap(spec *MapSpec) (*Map, error) {
	if spec.Type != ArrayOfMaps && spec.Type != HashOfMaps {
		return createMap(spec, 0)
	}

	if spec.ValueSize != 0 {
		return nil, fmt.Errorf("ebpf: ValueSize must be zero for map of map")
	}
	if spec.InnerMap == nil {
		return nil, fmt.Errorf("ebpf: map of map requires InnerMap")
	}

	inner, err := createMap(spec.InnerMap, 0)
	if err != nil {
		return nil, err
	}
	defer inner.Close()

	outerSpec := *spec
	outerSpec.InnerMap = nil
	outerSpec.ValueSize = 4
	return createMap(&outerSpec, inner.fd)
}

func createMap(spec *MapSpec, inner uint32) (*Map, error) {
	if spec.InnerMap != nil {
		return nil, fmt.Errorf("ebpf: inner map not allowed for this type")
	}
	attr := mapCreateAttr{
		spec.Type,
		spec.KeySize,
		spec.ValueSize,
		spec.MaxEntries,
		spec.Flags,
		inner,
	}
	fd, err := bpfCall(_MapCreate, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
	if err != nil {
		return nil, errors.Wrap(err, "map create")
	}

	return newMap(uint32(fd), spec)
}

func newMap(fd uint32, spec *MapSpec) (*Map, error) {
	m := &Map{
		uint32(fd),
		*spec,
		int(spec.ValueSize),
	}
	runtime.SetFinalizer(m, (*Map).Close)

	if !spec.Type.hasPerCPUValue() {
		return m, nil
	}

	possibleCPUs, err := possibleCPUs()
	if err != nil {
		return nil, err
	}

	m.fullValueSize = align(int(spec.ValueSize), 8) * possibleCPUs
	return m, nil
}

func (m *Map) String() string {
	return fmt.Sprintf("%s#%d", m.meta.Type, m.fd)
}

// Get gets a value from a Map
func (m *Map) Get(key, valueOut interface{}) (bool, error) {
	valueBytes, err := m.GetBytes(key)
	if err != nil {
		return false, err
	}
	if valueBytes == nil {
		return false, nil
	}

	if m.meta.Type.hasPerCPUValue() {
		return true, unmarshalPerCPUValue(valueOut, int(m.meta.ValueSize), valueBytes)
	}

	return true, unmarshalBytes(valueOut, valueBytes)
}

// GetBytes gets a value from Map
func (m *Map) GetBytes(key interface{}) ([]byte, error) {
	keyBytes, err := marshalBytes(key, int(m.meta.KeySize))
	if err != nil {
		return nil, err
	}
	valueBytes := make([]byte, m.fullValueSize)
	attr := mapOpAttr{
		mapFd: m.fd,
		key:   newPtr(unsafe.Pointer(&keyBytes[0])),
		value: newPtr(unsafe.Pointer(&valueBytes[0])),
	}
	_, err = bpfCall(_MapLookupElem, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
	if errors.Cause(err) == syscall.ENOENT {
		return nil, nil
	}
	return valueBytes, err
}

// Create creates a new value in a map, failing if the key exists already
func (m *Map) Create(key, value interface{}) error {
	return m.put(key, value, _NoExist)
}

// Put replaces or creates a value in map
func (m *Map) Put(key, value interface{}) error {
	return m.put(key, value, _Any)
}

// Replace replaces a value in a map, failing if the value did not exist
func (m *Map) Replace(key, value interface{}) error {
	return m.put(key, value, _Exist)
}

// Delete removes a value.
//
// Use DeleteStrict if you desire an error if key does not exist.
func (m *Map) Delete(key interface{}) error {
	err := m.DeleteStrict(key)
	if err == syscall.ENOENT {
		return nil
	}
	return err
}

// DeleteStrict removes a key and returns an error if the
// key doesn't exist.
func (m *Map) DeleteStrict(key interface{}) error {
	keyBytes, err := marshalBytes(key, int(m.meta.KeySize))
	if err != nil {
		return err
	}
	attr := mapOpAttr{
		mapFd: m.fd,
		key:   newPtr(unsafe.Pointer(&keyBytes[0])),
	}
	_, err = bpfCall(_MapDeleteElem, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
	return err
}

// NextKey finds the key following an initial key.
//
// See NextKeyBytes for details.
func (m *Map) NextKey(key, nextKeyOut interface{}) (bool, error) {
	nextKeyBytes, err := m.NextKeyBytes(key)
	if err != nil {
		return false, err
	}
	if nextKeyBytes == nil {
		return false, nil
	}
	err = unmarshalBytes(nextKeyOut, nextKeyBytes)
	if err != nil {
		return false, err
	}
	return true, nil
}

// NextKeyBytes returns the key following an initial key as a byte slice.
//
// Passing nil will return the first key.
//
// Use Iterate if you want to traverse all entries in the map.
func (m *Map) NextKeyBytes(key interface{}) ([]byte, error) {
	var keyPtr syscallPtr
	if key != nil {
		keyBytes, err := marshalBytes(key, int(m.meta.KeySize))
		if err != nil {
			return nil, err
		}
		keyPtr = newPtr(unsafe.Pointer(&keyBytes[0]))
	}

	nextKey := make([]byte, m.meta.KeySize)
	attr := mapOpAttr{
		mapFd: m.fd,
		key:   keyPtr,
		value: newPtr(unsafe.Pointer(&nextKey[0])),
	}
	_, err := bpfCall(_MapGetNextKey, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
	if errors.Cause(err) == syscall.ENOENT {
		return nil, nil
	}
	return nextKey, err
}

// Iterate traverses a map.
//
// It's not possible to guarantee that all keys in a map will be
// returned if there are concurrent modifications to the map. If a
// map is modified too heavily iteration may abort.
func (m *Map) Iterate() *MapIterator {
	return &MapIterator{
		target: m,
	}
}

// Close removes a Map
func (m *Map) Close() error {
	runtime.SetFinalizer(m, nil)
	return syscall.Close(int(m.fd))
}

// FD gets the raw fd value of Map
func (m *Map) FD() int {
	return int(m.fd)
}

// Pin persists the map past the lifetime of the process that created it.
//
// This requires bpffs to be mounted above fileName. See http://cilium.readthedocs.io/en/doc-1.0/kubernetes/install/#mounting-the-bpf-fs-optional
func (m *Map) Pin(fileName string) error {
	return pinObject(fileName, m.fd)
}

// LoadPinnedMap load a Map from a BPF file.
//
// Requires at least Linux 4.13, use LoadPinnedMapExplicit on
// earlier versions.
func LoadPinnedMap(fileName string) (*Map, error) {
	fd, err := getObject(fileName)
	if err != nil {
		return nil, err
	}
	spec, err := getMapSpecByFD(fd)
	if err != nil {
		return nil, err
	}
	return newMap(fd, spec)
}

// LoadPinnedMapExplicit loads a map with explicit parameters.
func LoadPinnedMapExplicit(fileName string, spec *MapSpec) (*Map, error) {
	fd, err := getObject(fileName)
	if err != nil {
		return nil, err
	}
	return newMap(fd, spec)
}

func (m *Map) put(key, value interface{}, putType uint64) error {
	keyBytes, err := marshalBytes(key, int(m.meta.KeySize))
	if err != nil {
		return err
	}

	var valueBytes []byte
	if m.meta.Type.hasPerCPUValue() {
		valueBytes, err = marshalPerCPUValue(value, int(m.meta.ValueSize))
	} else {
		valueBytes, err = marshalBytes(value, int(m.meta.ValueSize))
	}
	if err != nil {
		return err
	}

	attr := mapOpAttr{
		mapFd: m.fd,
		key:   newPtr(unsafe.Pointer(&keyBytes[0])),
		value: newPtr(unsafe.Pointer(&valueBytes[0])),
		flags: putType,
	}
	_, err = bpfCall(_MapUpdateElem, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
	return err
}

// UnmarshalBinary implements BinaryUnmarshaler.
func (m *Map) UnmarshalBinary(buf []byte) error {
	if m.fd != 0 {
		return errors.New("ebpf: can't unmarshal into existing map")
	}
	if len(buf) != 4 {
		return errors.New("ebpf: map id requires uint32")
	}
	// Looking up an entry in a nested map or prog array returns an id,
	// not an fd.
	id := nativeEndian.Uint32(buf)
	fd, err := getMapFDByID(id)
	if err != nil {
		return err
	}
	spec, err := getMapSpecByFD(fd)
	if err != nil {
		return err
	}

	nm, err := newMap(fd, spec)
	if err != nil {
		return err
	}

	*m = *nm
	return nil
}

// MarshalBinary implements BinaryMarshaler.
func (m *Map) MarshalBinary() ([]byte, error) {
	buf := make([]byte, 4)
	nativeEndian.PutUint32(buf, m.fd)
	return buf, nil
}

// MapIterator iterates a Map.
//
// See Map.Iterate.
type MapIterator struct {
	target *Map
	prev   interface{}
	done   bool
	err    error
}

// Next decodes the next key and value.
//
// Returns false if there are no more entries.
func (mi *MapIterator) Next(keyOut, valueOut interface{}) bool {
	if mi.err != nil || mi.done {
		return false
	}

	var nextBytes []byte
	for i := 0; i < 3; i++ {
		nextBytes, mi.err = mi.target.NextKeyBytes(mi.prev)
		if mi.err != nil {
			return false
		}

		if nextBytes == nil {
			mi.done = true
			return false
		}

		var ok bool
		ok, mi.err = mi.target.Get(nextBytes, valueOut)
		if mi.err != nil {
			return false
		}

		if ok {
			break
		}

		// The next key was deleted before we could retrieve
		// it's value. As of Linux 4.16 there is no safe API which
		// prevents this race.
		nextBytes = nil
	}

	if nextBytes == nil {
		// We still hit the race condition even though we retried.
		mi.err = errors.New("ebpf: can't retrieve next entry, map mutated too quickly")
		return false
	}

	mi.err = unmarshalBytes(keyOut, nextBytes)
	if mi.err != nil {
		return false
	}

	// The user can get access to nextBytes since marshalBytes
	// does not copy when unmarshaling into a []byte.
	// Make a copy to prevent accidental corruption of
	// iterator state.
	prevBytes := make([]byte, len(nextBytes))
	copy(prevBytes, nextBytes)
	mi.prev = prevBytes
	return true
}

// Err returns any encountered error.
//
// The method must be called after Next returns nil.
func (mi *MapIterator) Err() error {
	return mi.err
}
<|endoftext|>


<commit_before>package capn

import (
	"bytes"
	"errors"
	"io"
	"math"
)

var (
	errBufferCall     = errors.New("capn: can't call on a memory buffer")
	ErrInvalidSegment = errors.New("capn: invalid segment id")
	ErrTooMuchData    = errors.New("capn: too much data in stream")
)

type buffer Segment

// NewBuffer creates an expanding single segment buffer. Creating new objects
// will expand the buffer. Data can be nil (or length 0 with some capacity) if
// creating a new session. If parsing an existing segment than data should be
// the segment contents and will not be copied.
func NewBuffer(data []byte) *Segment {
	if uint64(len(data)) > uint64(math.MaxUint32) {
		return nil
	}

	b := &buffer{}
	b.Message = b
	b.Data = data
	return (*Segment)(b)
}

func (b *buffer) NewSegment(minsz int) (*Segment, error) {
	if uint64(len(b.Data)) > uint64(math.MaxUint32-minsz) {
		return nil, ErrOverlarge
	}
	b.Data = append(b.Data, make([]byte, minsz)...)
	b.Data = b.Data[:len(b.Data)-minsz]
	return (*Segment)(b), nil
}

func (b *buffer) Lookup(segid uint32) (*Segment, error) {
	if segid == 0 {
		return (*Segment)(b), nil
	} else {
		return nil, ErrInvalidSegment
	}
}

type multiBuffer struct {
	segments []*Segment
}

// NewMultiBuffer creates a new multi segment message. Creating new objects
// will try and reuse the buffers available, but will create new ones if there
// is insufficient capacity. When parsing an existing message data should be
// the list of segments. The data buffers will not be copied.
func NewMultiBuffer(data [][]byte) *Segment {
	m := &multiBuffer{make([]*Segment, len(data))}
	for i, d := range data {
		m.segments[i] = &Segment{m, d, uint32(i)}
	}
	if len(data) > 0 {
		return m.segments[0]
	}
	return &Segment{m, nil, 0xFFFFFFFF}
}

var (
	MaxSegmentNumber = 1024
	MaxTotalSize     = 1024 * 1024 * 1024
)

func (m *multiBuffer) NewSegment(minsz int) (*Segment, error) {
	for _, s := range m.segments {
		if len(s.Data)+minsz <= cap(s.Data) {
			return s, nil
		}
	}

	if minsz < 4096 {
		minsz = 4096
	}
	s := &Segment{m, make([]byte, 0, minsz), uint32(len(m.segments))}
	m.segments = append(m.segments, s)
	return s, nil
}

func (m *multiBuffer) Lookup(segid uint32) (*Segment, error) {
	if uint(segid) < uint(len(m.segments)) {
		return m.segments[segid], nil
	} else {
		return nil, ErrInvalidSegment
	}
}

// ReadFromStream reads a non-packed serialized stream from r. buf is used to
// buffer the read contents, can be nil, and is provided so that the buffer
// can be reused between messages. The returned segment is the first segment
// read, which contains the root pointer.
func ReadFromStream(r io.Reader, buf *bytes.Buffer) (*Segment, error) {
	if buf == nil {
		buf = new(bytes.Buffer)
	} else {
		buf.Reset()
	}

	if _, err := io.CopyN(buf, r, 4); err != nil {
		return nil, err
	}

	if little32(buf.Bytes()[:]) >= uint32(MaxSegmentNumber) {
		return nil, ErrTooMuchData
	}

	segnum := int(little32(buf.Bytes()[:]) + 1)
	hdrsz := 8*(segnum/2) + 4

	if _, err := io.CopyN(buf, r, int64(hdrsz)); err != nil {
		return nil, err
	}

	total := 0
	for i := 0; i < segnum; i++ {
		sz := little32(buf.Bytes()[4*i+4:])
		if uint64(total)+uint64(sz)*8 > uint64(MaxTotalSize) {
			return nil, ErrTooMuchData
		}
		total += int(sz) * 8
	}

	if _, err := io.CopyN(buf, r, int64(total)); err != nil {
		return nil, err
	}

	hdrv := buf.Bytes()[4 : hdrsz+4]
	datav := buf.Bytes()[hdrsz+4:]
	m := &multiBuffer{make([]*Segment, segnum)}
	for i := 0; i < segnum; i++ {
		sz := int(little32(hdrv[4*i:])) * 8
		m.segments[i] = &Segment{m, datav[:sz], uint32(i)}
		datav = datav[sz:]
	}

	return m.segments[0], nil
}

// ReadFromMemoryZeroCopy: like ReadFromStream, but reads a non-packed
// serialized stream that already resides in memory in the argument data.
// The returned segment is the first segment read, which contains
// the root pointer. The returned bytesRead says how many bytes were
// consumed from data in making seg. The caller should advance the
// data slice by doing data = data[bytesRead:] between successive calls
// to ReadFromMemoryZeroCopy().
func ReadFromMemoryZeroCopy(data []byte) (seg *Segment, bytesRead int64, err error) {

	if little32(data[0:4]) >= uint32(MaxSegmentNumber) {
		return nil, 0, ErrTooMuchData
	}

	segnum := int(little32(data[0:4]) + 1)
	hdrsz := 8*(segnum/2) + 4

	b := data[0 : (hdrsz+4)]

	total := 0
	for i := 0; i < segnum; i++ {
		sz := little32(b[4*i+4:])
		if uint64(total)+uint64(sz)*8 > uint64(MaxTotalSize) {
			return nil, 0, ErrTooMuchData
		}
		total += int(sz) * 8
	}

	hdrv := data[4 : (hdrsz+4)]
	datav := data[hdrsz+4:]
	m := &multiBuffer{make([]*Segment, segnum)}
	for i := 0; i < segnum; i++ {
		sz := int(little32(hdrv[4*i:])) * 8
		m.segments[i] = &Segment{m, datav[:sz], uint32(i)}
		datav = datav[sz:]
	}

	return m.segments[0], int64(4+hdrsz+total), nil
}

// WriteTo writes the message that the segment is part of to the
// provided stream in serialized form.
func (s *Segment) WriteTo(w io.Writer) (int64, error) {
	segnum := uint32(1)
	for {
		if seg, _ := s.Message.Lookup(segnum); seg == nil {
			break
		}
		segnum++
	}

	hdrv := make([]uint8, 8*(segnum/2)+8)
	putLittle32(hdrv, segnum-1)
	for i := uint32(0); i < segnum; i++ {
		seg, _ := s.Message.Lookup(i)
		putLittle32(hdrv[4*i+4:], uint32(len(seg.Data)/8))
	}

	if n, err := w.Write(hdrv); err != nil {
		return int64(n), err
	}
	written := int64(len(hdrv))

	for i := uint32(0); i < segnum; i++ {
		seg, _ := s.Message.Lookup(i)
		if n, err := w.Write(seg.Data); err != nil {
			return written + int64(n), err
		} else {
			written += int64(n)
		}
	}

	return written, nil
}
<commit_msg>handle empty segment more gracefully<commit_after>package capn

import (
	"bytes"
	"errors"
	"io"
	"math"
)

var (
	errBufferCall     = errors.New("capn: can't call on a memory buffer")
	ErrInvalidSegment = errors.New("capn: invalid segment id")
	ErrTooMuchData    = errors.New("capn: too much data in stream")
)

type buffer Segment

// NewBuffer creates an expanding single segment buffer. Creating new objects
// will expand the buffer. Data can be nil (or length 0 with some capacity) if
// creating a new session. If parsing an existing segment than data should be
// the segment contents and will not be copied.
func NewBuffer(data []byte) *Segment {
	if uint64(len(data)) > uint64(math.MaxUint32) {
		return nil
	}

	b := &buffer{}
	b.Message = b
	b.Data = data
	return (*Segment)(b)
}

func (b *buffer) NewSegment(minsz int) (*Segment, error) {
	if uint64(len(b.Data)) > uint64(math.MaxUint32-minsz) {
		return nil, ErrOverlarge
	}
	b.Data = append(b.Data, make([]byte, minsz)...)
	b.Data = b.Data[:len(b.Data)-minsz]
	return (*Segment)(b), nil
}

func (b *buffer) Lookup(segid uint32) (*Segment, error) {
	if segid == 0 {
		return (*Segment)(b), nil
	} else {
		return nil, ErrInvalidSegment
	}
}

type multiBuffer struct {
	segments []*Segment
}

// NewMultiBuffer creates a new multi segment message. Creating new objects
// will try and reuse the buffers available, but will create new ones if there
// is insufficient capacity. When parsing an existing message data should be
// the list of segments. The data buffers will not be copied.
func NewMultiBuffer(data [][]byte) *Segment {
	m := &multiBuffer{make([]*Segment, len(data))}
	for i, d := range data {
		m.segments[i] = &Segment{m, d, uint32(i)}
	}
	if len(data) > 0 {
		return m.segments[0]
	}
	return &Segment{m, nil, 0xFFFFFFFF}
}

var (
	MaxSegmentNumber = 1024
	MaxTotalSize     = 1024 * 1024 * 1024
)

func (m *multiBuffer) NewSegment(minsz int) (*Segment, error) {
	for _, s := range m.segments {
		if len(s.Data)+minsz <= cap(s.Data) {
			return s, nil
		}
	}

	if minsz < 4096 {
		minsz = 4096
	}
	s := &Segment{m, make([]byte, 0, minsz), uint32(len(m.segments))}
	m.segments = append(m.segments, s)
	return s, nil
}

func (m *multiBuffer) Lookup(segid uint32) (*Segment, error) {
	if uint(segid) < uint(len(m.segments)) {
		return m.segments[segid], nil
	} else {
		return nil, ErrInvalidSegment
	}
}

// ReadFromStream reads a non-packed serialized stream from r. buf is used to
// buffer the read contents, can be nil, and is provided so that the buffer
// can be reused between messages. The returned segment is the first segment
// read, which contains the root pointer.
func ReadFromStream(r io.Reader, buf *bytes.Buffer) (*Segment, error) {
	if buf == nil {
		buf = new(bytes.Buffer)
	} else {
		buf.Reset()
	}

	if _, err := io.CopyN(buf, r, 4); err != nil {
		return nil, err
	}

	if little32(buf.Bytes()[:]) >= uint32(MaxSegmentNumber) {
		return nil, ErrTooMuchData
	}

	segnum := int(little32(buf.Bytes()[:]) + 1)
	hdrsz := 8*(segnum/2) + 4

	if _, err := io.CopyN(buf, r, int64(hdrsz)); err != nil {
		return nil, err
	}

	total := 0
	for i := 0; i < segnum; i++ {
		sz := little32(buf.Bytes()[4*i+4:])
		if uint64(total)+uint64(sz)*8 > uint64(MaxTotalSize) {
			return nil, ErrTooMuchData
		}
		total += int(sz) * 8
	}

	if _, err := io.CopyN(buf, r, int64(total)); err != nil {
		return nil, err
	}

	hdrv := buf.Bytes()[4 : hdrsz+4]
	datav := buf.Bytes()[hdrsz+4:]
	m := &multiBuffer{make([]*Segment, segnum)}
	for i := 0; i < segnum; i++ {
		sz := int(little32(hdrv[4*i:])) * 8
		m.segments[i] = &Segment{m, datav[:sz], uint32(i)}
		datav = datav[sz:]
	}

	return m.segments[0], nil
}

// ReadFromMemoryZeroCopy: like ReadFromStream, but reads a non-packed
// serialized stream that already resides in memory in the argument data.
// The returned segment is the first segment read, which contains
// the root pointer. The returned bytesRead says how many bytes were
// consumed from data in making seg. The caller should advance the
// data slice by doing data = data[bytesRead:] between successive calls
// to ReadFromMemoryZeroCopy().
func ReadFromMemoryZeroCopy(data []byte) (seg *Segment, bytesRead int64, err error) {

	if little32(data[0:4]) >= uint32(MaxSegmentNumber) {
		return nil, 0, ErrTooMuchData
	}

	segnum := int(little32(data[0:4]) + 1)
	hdrsz := 8*(segnum/2) + 4

	b := data[0 : (hdrsz+4)]

	total := 0
	for i := 0; i < segnum; i++ {
		sz := little32(b[4*i+4:])
		if uint64(total)+uint64(sz)*8 > uint64(MaxTotalSize) {
			return nil, 0, ErrTooMuchData
		}
		total += int(sz) * 8
	}
	if total == 0 { return nil, 0, io.EOF }

	hdrv := data[4 : (hdrsz+4)]
	datav := data[hdrsz+4:]
	m := &multiBuffer{make([]*Segment, segnum)}
	for i := 0; i < segnum; i++ {
		sz := int(little32(hdrv[4*i:])) * 8
		m.segments[i] = &Segment{m, datav[:sz], uint32(i)}
		datav = datav[sz:]
	}

	return m.segments[0], int64(4+hdrsz+total), nil
}

// WriteTo writes the message that the segment is part of to the
// provided stream in serialized form.
func (s *Segment) WriteTo(w io.Writer) (int64, error) {
	segnum := uint32(1)
	for {
		if seg, _ := s.Message.Lookup(segnum); seg == nil {
			break
		}
		segnum++
	}

	hdrv := make([]uint8, 8*(segnum/2)+8)
	putLittle32(hdrv, segnum-1)
	for i := uint32(0); i < segnum; i++ {
		seg, _ := s.Message.Lookup(i)
		putLittle32(hdrv[4*i+4:], uint32(len(seg.Data)/8))
	}

	if n, err := w.Write(hdrv); err != nil {
		return int64(n), err
	}
	written := int64(len(hdrv))

	for i := uint32(0); i < segnum; i++ {
		seg, _ := s.Message.Lookup(i)
		if n, err := w.Write(seg.Data); err != nil {
			return written + int64(n), err
		} else {
			written += int64(n)
		}
	}

	return written, nil
}
<|endoftext|>


<commit_before>package tchannel

import (
	"errors"
	"sync"

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

var (
	errDuplicateMex   = errors.New("multiple attempts to use the message id")
	errMexChannelFull = NewSystemError(ErrCodeBusy, "cannot send frame to message exchange channel")
)

const (
	messageExchangeSetInbound  = "inbound"
	messageExchangeSetOutbound = "outbound"
)

// A messageExchange tracks this channel's side of a message exchange with a peer.  Each
// message exchange has a channel that can be used to receive frames from the peer, and
// a Context that can controls when the exchange has timed out or been cancelled.
type messageExchange struct {
	recvCh  chan *Frame
	ctx     context.Context
	msgID   uint32
	msgType messageType
}

// forwardPeerFrame forwards a frame from a peer to the message exchange, where it can
// be pulled by whatever application thread is handling the exchange
func (mex *messageExchange) forwardPeerFrame(frame *Frame) error {
	select {
	case mex.recvCh <- frame:
		return nil
	default:
		return errMexChannelFull
	}
}

// messageExchangeSet manages a set of active message exchanges.  It is mainly
// used to route frames from a peer to the appropriate messageExchange, or to cancel
// or mark a messageExchange as being in error.  Each Connection maintains two
// messageExchangeSets, one to manage exchanges that it has initiated (outgoing), and another
// to manage exchanges that the peer has initiated (incoming).  The message-type specific handlers
// are responsible for ensuring that their message exchanges are properly registered and removed
// from the corresponding exchange set, but a background garbage collector also runs to find
// exchanges that have timed out or been cancelled without having been removed, to ensure that
// even buggy code doesn't result in out of memory situations
type messageExchangeSet struct {
	log       Logger
	name      string
	exchanges map[uint32]*messageExchange
	mut       sync.Mutex
}

// newExchange creates and adds a new message exchange to this set
func (mexset *messageExchangeSet) newExchange(ctx context.Context,
	msgType messageType, msgID uint32, bufferSize int) (*messageExchange, error) {
	mexset.log.Debugf("Creating new %s message exchange for [%s:%d]", mexset.name, msgType, msgID)

	mex := &messageExchange{
		msgType: msgType,
		msgID:   msgID,
		ctx:     ctx,
		recvCh:  make(chan *Frame, bufferSize),
	}

	mexset.mut.Lock()
	defer mexset.mut.Unlock()

	if existingMex := mexset.exchanges[mex.msgID]; existingMex != nil {
		if existingMex == mex {
			mexset.log.Warnf("%s mex for %s, %d registered multiple times",
				mexset.name, mex.msgType, mex.msgID)
		} else {
			mexset.log.Warnf("msg id %d used for both active mex %s and new mex %s",
				mex.msgID, existingMex.msgType, mex.msgType)
		}

		return nil, errDuplicateMex
	}

	mexset.exchanges[mex.msgID] = mex

	// TODO(mmihic): Put into a deadline ordered heap so we can garbage collected expired exchanges
	return mex, nil
}

// removeExchange removes a messge exchange from the set, if it exists.  It's perfectly
// fine to try and remove an exchange that has already completed
func (mexset *messageExchangeSet) removeExchange(msgID uint32) {
	mexset.log.Debugf("Removing %s message exchange %d", mexset.name, msgID)

	mexset.mut.Lock()
	defer mexset.mut.Unlock()

	delete(mexset.exchanges, msgID)
}

// forwardPeerFrame forwards a frame from the peer to the appropriate message exchange
// TODO(mmihic): We need to take the messageID here due to the weird originalMessageID field in error
// frame.  If we instead made the error frame message ID match the ID of the message in error, we could
// drop this additional parameters
func (mexset *messageExchangeSet) forwardPeerFrame(messageID uint32, frame *Frame) error {
	mexset.log.Debugf("forwarding %s %s", mexset.name, frame.Header)

	mexset.mut.Lock()
	mex := mexset.exchanges[messageID]
	mexset.mut.Unlock()

	if mex == nil {
		// This is ok since the exchange might have expired or been cancelled
		mexset.log.Warnf("received frame %s for message exchange that no longer exists", frame.Header)
		return nil
	}

	return mex.forwardPeerFrame(frame)
}
<commit_msg>Better docs and logging<commit_after>package tchannel

import (
	"errors"
	"sync"

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

var (
	errDuplicateMex   = errors.New("multiple attempts to use the message id")
	errMexChannelFull = NewSystemError(ErrCodeBusy, "cannot send frame to message exchange channel")
)

const (
	messageExchangeSetInbound  = "inbound"
	messageExchangeSetOutbound = "outbound"
)

// A messageExchange tracks this channel's side of a message exchange with a
// peer.  Each message exchange has a channel that can be used to receive
// frames from the peer, and a Context that can controls when the exchange has
// timed out or been cancelled.
type messageExchange struct {
	recvCh  chan *Frame
	ctx     context.Context
	msgID   uint32
	msgType messageType
}

// forwardPeerFrame forwards a frame from a peer to the message exchange, where
// it can be pulled by whatever application thread is handling the exchange
func (mex *messageExchange) forwardPeerFrame(frame *Frame) error {
	select {
	case mex.recvCh <- frame:
		return nil
	default:
		return errMexChannelFull
	}
}

// messageExchangeSet manages a set of active message exchanges.  It is mainly
// used to route frames from a peer to the appropriate messageExchange, or to
// cancel or mark a messageExchange as being in error.  Each Connection
// maintains two messageExchangeSets, one to manage exchanges that it has
// initiated (outgoing), and another to manage exchanges that the peer has
// initiated (incoming).  The message-type specific handlers are responsible
// for ensuring that their message exchanges are properly registered and
// removed from the corresponding exchange set, but a background garbage
// collector also runs to find exchanges that have timed out or been cancelled
// without having been removed, to ensure that even buggy code doesn't result
// in out of memory situations
type messageExchangeSet struct {
	log       Logger
	name      string
	exchanges map[uint32]*messageExchange
	mut       sync.Mutex
}

// newExchange creates and adds a new message exchange to this set
func (mexset *messageExchangeSet) newExchange(ctx context.Context,
	msgType messageType, msgID uint32, bufferSize int) (*messageExchange, error) {
	mexset.log.Debugf("Creating new %s message exchange for [%s:%d]", mexset.name, msgType, msgID)

	mex := &messageExchange{
		msgType: msgType,
		msgID:   msgID,
		ctx:     ctx,
		recvCh:  make(chan *Frame, bufferSize),
	}

	mexset.mut.Lock()
	defer mexset.mut.Unlock()

	if existingMex := mexset.exchanges[mex.msgID]; existingMex != nil {
		if existingMex == mex {
			mexset.log.Warnf("%s mex for %s, %d registered multiple times",
				mexset.name, mex.msgType, mex.msgID)
		} else {
			mexset.log.Warnf("msg id %d used for both active mex %s and new mex %s",
				mex.msgID, existingMex.msgType, mex.msgType)
		}

		return nil, errDuplicateMex
	}

	mexset.exchanges[mex.msgID] = mex

	// TODO(mmihic): Put into a deadline ordered heap so we can garbage collected expired exchanges
	return mex, nil
}

// removeExchange removes a message exchange from the set, if it exists.  It's
// perfectly fine to try and remove an exchange that has already completed
func (mexset *messageExchangeSet) removeExchange(msgID uint32) {
	mexset.log.Debugf("Removing %s message exchange %d", mexset.name, msgID)

	mexset.mut.Lock()
	defer mexset.mut.Unlock()

	delete(mexset.exchanges, msgID)
}

// forwardPeerFrame forwards a frame from the peer to the appropriate message
// exchange
// TODO(mmihic): We need to take the messageID here due to the weird
// originalMessageID field in error frame.  If we instead made the error frame
// message ID match the ID of the message in error, we could drop this
// additional parameters
func (mexset *messageExchangeSet) forwardPeerFrame(messageID uint32, frame *Frame) error {
	mexset.log.Debugf("forwarding %s %s", mexset.name, frame.Header)

	mexset.mut.Lock()
	mex := mexset.exchanges[messageID]
	mexset.mut.Unlock()

	if mex == nil {
		// This is ok since the exchange might have expired or been cancelled
		mexset.log.Warnf("received frame %s for message exchange that no longer exists", frame.Header)
		return nil
	}

	if err := mex.forwardPeerFrame(frame); err != nil {
		mexset.log.Warnf("Unable to forward %s to peer: %v", frame, err)
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package mig

import (
	"database/sql"
	"fmt"
)

// migration is a combination of a SQL snippet and a version number which
// represents it.
type migration struct {
	version int
	script  string
}

var migrations = make([]migration, 0, 100)

// Define will define and register a migration for running the next time
// Mig is run. A version is automatically chosen based on the sequence of
// previous calls to Define (aka version=len(migrations)+1).
//
//     mig.Define(`
//       create table foo ( id integer );
//     `)
func Define(script string) {
	DefineVersion(len(migrations)+1, script)
}

// DefineVersion will define and register a migration for running the next
// time Mig is run with a specific version number.
//
//     mig.DefineVersion(2, `
//       alter table foo add column bar varchar(10);
//     `)
func DefineVersion(version int, script string) {
	migrations = append(migrations, migration{version, script})
}

// Migrate will execute the defined migrations against the sql.DB. For
// each migration Mig will look in a db_version table to see if it has
// already been run. If there's a row in db_version for a migration it
// wont be run again, otherwise the migration will be run and a row stored
// in db_version.
//
// If this is the first time Mig has been run, Migrate will first create
// a db_version table and store a zero row.
func Migrate(db *sql.DB) error {
	fmt.Println("Running migrations")

	if err := createVersionTable(db); err != nil {
		return MigrationError{err, 0}
	}

	for _, m := range migrations {
		ok, err := shouldRun(m.version, db)
		if err != nil {
			return MigrationError{err, m.version}
		}

		if ok {
			fmt.Printf("Executing #%d...", m.version)

			if err := execute(m.script, db); err != nil {
				return MigrationError{err, m.version}
			}

			if err := insertVersion(m.version, db); err != nil {
				return MigrationError{err, m.version}
			}

			fmt.Println(" done.")
		} else {
			fmt.Printf("Skipping #%d, already run\n", m.version)
		}
	}

	fmt.Println("Migration complete")
	return nil
}

// createVersionTable creates the db_version table if it doesn't exist
func createVersionTable(db *sql.DB) error {
	return execute(`
		create table if not exists db_version (
			version     integer not null unique,
			updated_at  timestamp not null default(current_timestamp)
		);

		if not exists(select 1 from db_version where version = 0) then
			insert into db_version (version) values (0);
		end if;
	`, db)
}

// shouldRun tests whether a migration should be executed
func shouldRun(version int, db *sql.DB) (bool, error) {
	sql := fmt.Sprintf("select 1 from db_version where version = %d;", version)

	result, err := db.Exec(sql)
	if err != nil {
		return false, err
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return false, err
	}

	return rows == 0, nil
}

// execute runs a migration. This assumes the migration hasn't already
// been run.
func execute(script string, db *sql.DB) error {
	_, err := db.Exec(fmt.Sprintf(`DO $$
		BEGIN
		%s
		END;
		$$;`, script))

	return err
}

// insertVersion inserts a version into the db_version table. Assumes it
// doesn't already exist in there.
func insertVersion(version int, db *sql.DB) error {
	sql := fmt.Sprintf("insert into db_version (version) values (%d);", version)
	_, err := db.Exec(sql)
	return err
}
<commit_msg>Add Quiet flag for supressing normal output.<commit_after>package mig

import (
	"database/sql"
	"fmt"
)

// migration is a combination of a SQL snippet and a version number which
// represents it.
type migration struct {
	version int
	script  string
}

var migrations = make([]migration, 0, 100)

// Quiet will supress status output on STDOUT if true
var Quiet = false

// Define will define and register a migration for running the next time
// Mig is run. A version is automatically chosen based on the sequence of
// previous calls to Define (aka version=len(migrations)+1).
//
//     mig.Define(`
//       create table foo ( id integer );
//     `)
func Define(script string) {
	DefineVersion(len(migrations)+1, script)
}

// DefineVersion will define and register a migration for running the next
// time Mig is run with a specific version number.
//
//     mig.DefineVersion(2, `
//       alter table foo add column bar varchar(10);
//     `)
func DefineVersion(version int, script string) {
	migrations = append(migrations, migration{version, script})
}

// Migrate will execute the defined migrations against the sql.DB. For
// each migration Mig will look in a db_version table to see if it has
// already been run. If there's a row in db_version for a migration it
// wont be run again, otherwise the migration will be run and a row stored
// in db_version.
//
// If this is the first time Mig has been run, Migrate will first create
// a db_version table and store a zero row.
func Migrate(db *sql.DB) error {
	logln("Running migrations")

	if err := createVersionTable(db); err != nil {
		return MigrationError{err, 0}
	}

	for _, m := range migrations {
		ok, err := shouldRun(m.version, db)
		if err != nil {
			return MigrationError{err, m.version}
		}

		if ok {
			logf("Executing #%d...", m.version)

			if err := execute(m.script, db); err != nil {
				return MigrationError{err, m.version}
			}

			if err := insertVersion(m.version, db); err != nil {
				return MigrationError{err, m.version}
			}

			logln(" done.")
		} else {
			logf("Skipping #%d, already run\n", m.version)
		}
	}

	logln("Migration complete")
	return nil
}

// createVersionTable creates the db_version table if it doesn't exist
func createVersionTable(db *sql.DB) error {
	return execute(`
		create table if not exists db_version (
			version     integer not null unique,
			updated_at  timestamp not null default(current_timestamp)
		);

		if not exists(select 1 from db_version where version = 0) then
			insert into db_version (version) values (0);
		end if;
	`, db)
}

// shouldRun tests whether a migration should be executed
func shouldRun(version int, db *sql.DB) (bool, error) {
	sql := fmt.Sprintf("select 1 from db_version where version = %d;", version)

	result, err := db.Exec(sql)
	if err != nil {
		return false, err
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return false, err
	}

	return rows == 0, nil
}

// execute runs a migration. This assumes the migration hasn't already
// been run.
func execute(script string, db *sql.DB) error {
	_, err := db.Exec(fmt.Sprintf(`DO $$
		BEGIN
		%s
		END;
		$$;`, script))

	return err
}

// insertVersion inserts a version into the db_version table. Assumes it
// doesn't already exist in there.
func insertVersion(version int, db *sql.DB) error {
	sql := fmt.Sprintf("insert into db_version (version) values (%d);", version)
	_, err := db.Exec(sql)
	return err
}

func logf(format string, a ...interface{}) {
	if !Quiet {
		fmt.Printf(format, a)
	}
}

func logln(a ...interface{}) {
	if !Quiet {
		fmt.Println(a)
	}
}
<|endoftext|>


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

// +build windows

package mgr

import (
	"syscall"
	"unsafe"

	"golang.org/x/sys/windows"
)

const (
	// Service start types.
	StartManual    = windows.SERVICE_DEMAND_START // the service must be started manually
	StartAutomatic = windows.SERVICE_AUTO_START   // the service will start by itself whenever the computer reboots
	StartDisabled  = windows.SERVICE_DISABLED     // the service cannot be started

	// The severity of the error, and action taken,
	// if this service fails to start.
	ErrorCritical = windows.SERVICE_ERROR_CRITICAL
	ErrorIgnore   = windows.SERVICE_ERROR_IGNORE
	ErrorNormal   = windows.SERVICE_ERROR_NORMAL
	ErrorSevere   = windows.SERVICE_ERROR_SEVERE
)

// TODO(brainman): Password is not returned by windows.QueryServiceConfig, not sure how to get it.

type Config struct {
	ServiceType      uint32
	StartType        uint32
	ErrorControl     uint32
	BinaryPathName   string // fully qualified path to the service binary file, can also include arguments for an auto-start service
	LoadOrderGroup   string
	TagId            uint32
	Dependencies     []string
	ServiceStartName string // name of the account under which the service should run
	DisplayName      string
	Password         string
	Description      string
	SidType          uint32 // one of SERVICE_SID_TYPE, the type of sid to use for the service
	DelayedAutoStart bool   // the service is started after other auto-start services are started plus a short delay
}

func toStringSlice(ps *uint16) []string {
	r := make([]string, 0)
	p := unsafe.Pointer(ps)

	for {
		s := windows.UTF16PtrToString((*uint16)(p))
		if len(s) == 0 {
			break
		}

		r = append(r, s)
		offset := unsafe.Sizeof(uint16(0)) * (uintptr)(len(s)+1)
		p = unsafe.Pointer(uintptr(p) + offset)
	}

	return r
}

// Config retrieves service s configuration paramteres.
func (s *Service) Config() (Config, error) {
	var p *windows.QUERY_SERVICE_CONFIG
	n := uint32(1024)
	for {
		b := make([]byte, n)
		p = (*windows.QUERY_SERVICE_CONFIG)(unsafe.Pointer(&b[0]))
		err := windows.QueryServiceConfig(s.Handle, p, n, &n)
		if err == nil {
			break
		}
		if err.(syscall.Errno) != syscall.ERROR_INSUFFICIENT_BUFFER {
			return Config{}, err
		}
		if n <= uint32(len(b)) {
			return Config{}, err
		}
	}

	b, err := s.queryServiceConfig2(windows.SERVICE_CONFIG_DESCRIPTION)
	if err != nil {
		return Config{}, err
	}
	p2 := (*windows.SERVICE_DESCRIPTION)(unsafe.Pointer(&b[0]))

	b, err = s.queryServiceConfig2(windows.SERVICE_CONFIG_DELAYED_AUTO_START_INFO)
	if err != nil {
		return Config{}, err
	}
	p3 := (*windows.SERVICE_DELAYED_AUTO_START_INFO)(unsafe.Pointer(&b[0]))
	delayedStart := false
	if p3.IsDelayedAutoStartUp != 0 {
		delayedStart = true
	}

	return Config{
		ServiceType:      p.ServiceType,
		StartType:        p.StartType,
		ErrorControl:     p.ErrorControl,
		BinaryPathName:   windows.UTF16PtrToString(p.BinaryPathName),
		LoadOrderGroup:   windows.UTF16PtrToString(p.LoadOrderGroup),
		TagId:            p.TagId,
		Dependencies:     toStringSlice(p.Dependencies),
		ServiceStartName: windows.UTF16PtrToString(p.ServiceStartName),
		DisplayName:      windows.UTF16PtrToString(p.DisplayName),
		Description:      windows.UTF16PtrToString(p2.Description),
		DelayedAutoStart: delayedStart,
	}, nil
}

func updateDescription(handle windows.Handle, desc string) error {
	d := windows.SERVICE_DESCRIPTION{Description: toPtr(desc)}
	return windows.ChangeServiceConfig2(handle,
		windows.SERVICE_CONFIG_DESCRIPTION, (*byte)(unsafe.Pointer(&d)))
}

func updateSidType(handle windows.Handle, sidType uint32) error {
	return windows.ChangeServiceConfig2(handle, windows.SERVICE_CONFIG_SERVICE_SID_INFO, (*byte)(unsafe.Pointer(&sidType)))
}

func updateStartUp(handle windows.Handle, isDelayed bool) error {
	var d windows.SERVICE_DELAYED_AUTO_START_INFO
	if isDelayed {
		d.IsDelayedAutoStartUp = 1
	}
	return windows.ChangeServiceConfig2(handle,
		windows.SERVICE_CONFIG_DELAYED_AUTO_START_INFO, (*byte)(unsafe.Pointer(&d)))
}

// UpdateConfig updates service s configuration parameters.
func (s *Service) UpdateConfig(c Config) error {
	err := windows.ChangeServiceConfig(s.Handle, c.ServiceType, c.StartType,
		c.ErrorControl, toPtr(c.BinaryPathName), toPtr(c.LoadOrderGroup),
		nil, toStringBlock(c.Dependencies), toPtr(c.ServiceStartName),
		toPtr(c.Password), toPtr(c.DisplayName))
	if err != nil {
		return err
	}
	err = updateSidType(s.Handle, c.SidType)
	if err != nil {
		return err
	}

	err = updateStartUp(s.Handle, c.DelayedAutoStart)
	if err != nil {
		return err
	}

	return updateDescription(s.Handle, c.Description)
}

// queryServiceConfig2 calls Windows QueryServiceConfig2 with infoLevel parameter and returns retrieved service configuration information.
func (s *Service) queryServiceConfig2(infoLevel uint32) ([]byte, error) {
	n := uint32(1024)
	for {
		b := make([]byte, n)
		err := windows.QueryServiceConfig2(s.Handle, infoLevel, &b[0], n, &n)
		if err == nil {
			return b, nil
		}
		if err.(syscall.Errno) != syscall.ERROR_INSUFFICIENT_BUFFER {
			return nil, err
		}
		if n <= uint32(len(b)) {
			return nil, err
		}
	}
}
<commit_msg>windows/svc/mgr: fetch service sid type in Config()<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build windows

package mgr

import (
	"syscall"
	"unsafe"

	"golang.org/x/sys/windows"
)

const (
	// Service start types.
	StartManual    = windows.SERVICE_DEMAND_START // the service must be started manually
	StartAutomatic = windows.SERVICE_AUTO_START   // the service will start by itself whenever the computer reboots
	StartDisabled  = windows.SERVICE_DISABLED     // the service cannot be started

	// The severity of the error, and action taken,
	// if this service fails to start.
	ErrorCritical = windows.SERVICE_ERROR_CRITICAL
	ErrorIgnore   = windows.SERVICE_ERROR_IGNORE
	ErrorNormal   = windows.SERVICE_ERROR_NORMAL
	ErrorSevere   = windows.SERVICE_ERROR_SEVERE
)

// TODO(brainman): Password is not returned by windows.QueryServiceConfig, not sure how to get it.

type Config struct {
	ServiceType      uint32
	StartType        uint32
	ErrorControl     uint32
	BinaryPathName   string // fully qualified path to the service binary file, can also include arguments for an auto-start service
	LoadOrderGroup   string
	TagId            uint32
	Dependencies     []string
	ServiceStartName string // name of the account under which the service should run
	DisplayName      string
	Password         string
	Description      string
	SidType          uint32 // one of SERVICE_SID_TYPE, the type of sid to use for the service
	DelayedAutoStart bool   // the service is started after other auto-start services are started plus a short delay
}

func toStringSlice(ps *uint16) []string {
	r := make([]string, 0)
	p := unsafe.Pointer(ps)

	for {
		s := windows.UTF16PtrToString((*uint16)(p))
		if len(s) == 0 {
			break
		}

		r = append(r, s)
		offset := unsafe.Sizeof(uint16(0)) * (uintptr)(len(s)+1)
		p = unsafe.Pointer(uintptr(p) + offset)
	}

	return r
}

// Config retrieves service s configuration paramteres.
func (s *Service) Config() (Config, error) {
	var p *windows.QUERY_SERVICE_CONFIG
	n := uint32(1024)
	for {
		b := make([]byte, n)
		p = (*windows.QUERY_SERVICE_CONFIG)(unsafe.Pointer(&b[0]))
		err := windows.QueryServiceConfig(s.Handle, p, n, &n)
		if err == nil {
			break
		}
		if err.(syscall.Errno) != syscall.ERROR_INSUFFICIENT_BUFFER {
			return Config{}, err
		}
		if n <= uint32(len(b)) {
			return Config{}, err
		}
	}

	b, err := s.queryServiceConfig2(windows.SERVICE_CONFIG_DESCRIPTION)
	if err != nil {
		return Config{}, err
	}
	p2 := (*windows.SERVICE_DESCRIPTION)(unsafe.Pointer(&b[0]))

	b, err = s.queryServiceConfig2(windows.SERVICE_CONFIG_DELAYED_AUTO_START_INFO)
	if err != nil {
		return Config{}, err
	}
	p3 := (*windows.SERVICE_DELAYED_AUTO_START_INFO)(unsafe.Pointer(&b[0]))
	delayedStart := false
	if p3.IsDelayedAutoStartUp != 0 {
		delayedStart = true
	}

	b, err = s.queryServiceConfig2(windows.SERVICE_CONFIG_SERVICE_SID_INFO)
	if err != nil {
		return Config{}, err
	}
	sidType := *(*uint32)(unsafe.Pointer(&b[0]))

	return Config{
		ServiceType:      p.ServiceType,
		StartType:        p.StartType,
		ErrorControl:     p.ErrorControl,
		BinaryPathName:   windows.UTF16PtrToString(p.BinaryPathName),
		LoadOrderGroup:   windows.UTF16PtrToString(p.LoadOrderGroup),
		TagId:            p.TagId,
		Dependencies:     toStringSlice(p.Dependencies),
		ServiceStartName: windows.UTF16PtrToString(p.ServiceStartName),
		DisplayName:      windows.UTF16PtrToString(p.DisplayName),
		Description:      windows.UTF16PtrToString(p2.Description),
		DelayedAutoStart: delayedStart,
		SidType:          sidType,
	}, nil
}

func updateDescription(handle windows.Handle, desc string) error {
	d := windows.SERVICE_DESCRIPTION{Description: toPtr(desc)}
	return windows.ChangeServiceConfig2(handle,
		windows.SERVICE_CONFIG_DESCRIPTION, (*byte)(unsafe.Pointer(&d)))
}

func updateSidType(handle windows.Handle, sidType uint32) error {
	return windows.ChangeServiceConfig2(handle, windows.SERVICE_CONFIG_SERVICE_SID_INFO, (*byte)(unsafe.Pointer(&sidType)))
}

func updateStartUp(handle windows.Handle, isDelayed bool) error {
	var d windows.SERVICE_DELAYED_AUTO_START_INFO
	if isDelayed {
		d.IsDelayedAutoStartUp = 1
	}
	return windows.ChangeServiceConfig2(handle,
		windows.SERVICE_CONFIG_DELAYED_AUTO_START_INFO, (*byte)(unsafe.Pointer(&d)))
}

// UpdateConfig updates service s configuration parameters.
func (s *Service) UpdateConfig(c Config) error {
	err := windows.ChangeServiceConfig(s.Handle, c.ServiceType, c.StartType,
		c.ErrorControl, toPtr(c.BinaryPathName), toPtr(c.LoadOrderGroup),
		nil, toStringBlock(c.Dependencies), toPtr(c.ServiceStartName),
		toPtr(c.Password), toPtr(c.DisplayName))
	if err != nil {
		return err
	}
	err = updateSidType(s.Handle, c.SidType)
	if err != nil {
		return err
	}

	err = updateStartUp(s.Handle, c.DelayedAutoStart)
	if err != nil {
		return err
	}

	return updateDescription(s.Handle, c.Description)
}

// queryServiceConfig2 calls Windows QueryServiceConfig2 with infoLevel parameter and returns retrieved service configuration information.
func (s *Service) queryServiceConfig2(infoLevel uint32) ([]byte, error) {
	n := uint32(1024)
	for {
		b := make([]byte, n)
		err := windows.QueryServiceConfig2(s.Handle, infoLevel, &b[0], n, &n)
		if err == nil {
			return b, nil
		}
		if err.(syscall.Errno) != syscall.ERROR_INSUFFICIENT_BUFFER {
			return nil, err
		}
		if n <= uint32(len(b)) {
			return nil, err
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"

	"github.com/pdxgo/whispering-gophers/util"
)

// Open launches server from the media file path.
func Open(path string) (string, error) {
	path, err := filepath.Abs(path)
	if err != nil {
		return "", err
	}

	f, err := os.Open(path)
	defer f.Close()
	if err != nil {
		return "", err
	}

	stat, err := f.Stat()
	if err != nil {
		return "", err
	}

	switch mode := stat.Mode(); {
	case mode.IsDir():
		return "", errors.New("directory is unsupported")
	case mode.IsRegular():
		return serve(path), nil
	}

	return "", nil
}

// serve listens the file f from the local file system as local network address
func serve(f string) string {
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Cache-Control", "no-cache")
		http.ServeFile(w, r, f)
	})

	laddr := availableAddr()
	go http.ListenAndServe(laddr, nil)
	return fmt.Sprintf("http://%s", laddr)
}

// availableAddr returns an address that available non-lookback IPv4 network
// interface and port.
func availableAddr() string {
	ln, err := util.Listen()
	defer ln.Close()
	if err != nil {
		log.Fatal(err)
	}
	return ln.Addr().String()
}
<commit_msg>Use Listner to http.Serve<commit_after>package main

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"

	"github.com/pdxgo/whispering-gophers/util"
)

// Open launches server from the media file.
func Open(name string) (string, error) {
	path, err := filepath.Abs(name)
	if err != nil {
		return "", err
	}

	f, err := os.Open(path)
	defer f.Close()
	if err != nil {
		return "", err
	}

	stat, err := f.Stat()
	if err != nil {
		return "", err
	}

	switch mode := stat.Mode(); {
	case mode.IsDir():
		return "", errors.New("directory is unsupported")
	case mode.IsRegular():
		return serve(path), nil
	}

	return "", nil
}

func serve(name string) string {
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Set `no-cache` to flush the cache on Apple TV.
		w.Header().Set("Cache-Control", "no-cache")
		http.ServeFile(w, r, name)
	})

	l, err := util.Listen()
	if err != nil {
		log.Fatal(err)
	}
	go http.Serve(l, http.DefaultServeMux)
	return fmt.Sprintf("http://%s", l.Addr().String())
}
<|endoftext|>


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

/*
Package pbkdf2 implements the key derivation function PBKDF2 as defined in RFC
2898 / PKCS #5 v2.0.

A key derivation function is useful when encrypting data based on a password
or any other not-fully-random data. It uses a pseudorandom function to derive
a secure encryption key based on the password.

While v2.0 of the standard defines only one pseudorandom function to use,
HMAC-SHA1, the drafted v2.1 specification allows use of all five FIPS Approved
Hash Functions SHA-1, SHA-224, SHA-256, SHA-384 and SHA-512 for HMAC. To
choose, you can pass the `New` functions from the different SHA packages to
pbkdf2.Key.
*/
package pbkdf2

import (
	"crypto/hmac"
	"hash"
	"math/big"
)

// Key derives a key from the password, salt and iteration count, returning a
// []byte of length keylen that can be used as cryptographic key. The key is
// derived based on the method described as PBKDF2 with the HMAC variant using
// the supplied hash function.
//
// For example, to use a HMAC-SHA-1 based PBKDF2 key derivation function, you
// can get a derived key for e.g. AES-256 (which needs a 32-byte key) by
// doing:
//
// 	dk := pbkdf2.Key([]byte("some password"), salt, 4096, 32, sha1.New)
//
// Remember to get a good random salt. At least 8 bytes is recommended by the
// RFC.
//
// Using a higher iteration count will increase the cost of an exhaustive
// search but will also make derivation proportionally slower.
func Key(password, salt []byte, iter int, keyLen int, h func() hash.Hash) []byte {
	i := big.NewInt(int64(iter))
	return KeyBigIter(password, salt, i, keyLen, h)
}

// KeyBigIter derives a key from the password, salt and iteration count
// (as an arbitrary-precision arithmetic *big.Int), returning a []byte of
// length keylen that can be used as cryptographic key.
// The key is derived based on the method described as PBKDF2 with the HMAC
// variant using the supplied hash function.
//
// For example, to use a HMAC-SHA-1 based PBKDF2 key derivation function, you
// can get a derived key for e.g. AES-256 (which needs a 32-byte key) by
// doing:
//
// 	dk := pbkdf2.Key([]byte("some password"), salt, 4096, 32, sha1.New)
//
// Remember to get a good random salt. At least 8 bytes is recommended by the
// RFC.
//
// Using a higher iteration count will increase the cost of an exhaustive
// search but will also make derivation proportionally slower.
func KeyBigIter(password, salt []byte, iter *big.Int, keyLen int, h func() hash.Hash) []byte {
	prf := hmac.New(h, password)
	hashLen := prf.Size()
	numBlocks := (keyLen + hashLen - 1) / hashLen

	var buf [4]byte
	dk := make([]byte, 0, numBlocks*hashLen)
	U := make([]byte, hashLen)
	for block := 1; block <= numBlocks; block++ {
		// N.B.: || means concatenation, ^ means XOR
		// for each block T_i = U_1 ^ U_2 ^ ... ^ U_iter
		// U_1 = PRF(password, salt || uint(i))
		prf.Reset()
		prf.Write(salt)
		buf[0] = byte(block >> 24)
		buf[1] = byte(block >> 16)
		buf[2] = byte(block >> 8)
		buf[3] = byte(block)
		prf.Write(buf[:4])
		dk = prf.Sum(dk)
		T := dk[len(dk)-hashLen:]
		copy(U, T)

		// U_n = PRF(password, U_(n-1))
		i := big.NewInt(1)
		n := big.NewInt(1)
		for iter.Cmp(n) == 1 {
			//for n := 2; n <= iter; n++ {
			prf.Reset()
			prf.Write(U)
			U = U[:0]
			U = prf.Sum(U)
			for x := range U {
				T[x] ^= U[x]
			}
			n.Add(n, i)
		}
	}
	return dk[:keyLen]
}
<commit_msg>pbkdf2 key with int64 types<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
Package pbkdf2 implements the key derivation function PBKDF2 as defined in RFC
2898 / PKCS #5 v2.0.

A key derivation function is useful when encrypting data based on a password
or any other not-fully-random data. It uses a pseudorandom function to derive
a secure encryption key based on the password.

While v2.0 of the standard defines only one pseudorandom function to use,
HMAC-SHA1, the drafted v2.1 specification allows use of all five FIPS Approved
Hash Functions SHA-1, SHA-224, SHA-256, SHA-384 and SHA-512 for HMAC. To
choose, you can pass the `New` functions from the different SHA packages to
pbkdf2.Key.
*/
package pbkdf2

import (
	"crypto/hmac"
	"hash"
)

// Key derives a key from the password, salt and iteration count, returning a
// []byte of length keylen that can be used as cryptographic key. The key is
// derived based on the method described as PBKDF2 with the HMAC variant using
// the supplied hash function.
//
// For example, to use a HMAC-SHA-1 based PBKDF2 key derivation function, you
// can get a derived key for e.g. AES-256 (which needs a 32-byte key) by
// doing:
//
// 	dk := pbkdf2.Key([]byte("some password"), salt, 4096, 32, sha1.New)
//
// Remember to get a good random salt. At least 8 bytes is recommended by the
// RFC.
//
// Using a higher iteration count will increase the cost of an exhaustive
// search but will also make derivation proportionally slower.
func Key(password, salt []byte, iter, keyLen int, h func() hash.Hash) []byte {
	return Key64(password, salt, int64(iter), int64(keyLen), h)
}

// Key64 derives a key from the password, salt and iteration count, returning a
// []byte of length keylen that can be used as cryptographic key. Key64 uses
// int64 for the iteration count and key length to allow larger values.
// The key is derived based on the method described as PBKDF2 with the HMAC
// variant using the supplied hash function.
//
// For example, to use a HMAC-SHA-1 based PBKDF2 key derivation function, you
// can get a derived key for e.g. AES-256 (which needs a 32-byte key) by
// doing:
//
// 	dk := pbkdf2.Key([]byte("some password"), salt, 4096, 32, sha1.New)
//
// Remember to get a good random salt. At least 8 bytes is recommended by the
// RFC.
//
// Using a higher iteration count will increase the cost of an exhaustive
// search but will also make derivation proportionally slower.
func Key64(password, salt []byte, iter, keyLen int64, h func() hash.Hash) []byte {
	prf := hmac.New(h, password)
	hashLen := int64(prf.Size())
	numBlocks := (keyLen + hashLen - 1) / hashLen

	var buf [4]byte
	dk := make([]byte, 0, numBlocks*hashLen)
	U := make([]byte, hashLen)
	for block := int64(1); block <= numBlocks; block++ {
		// N.B.: || means concatenation, ^ means XOR
		// for each block T_i = U_1 ^ U_2 ^ ... ^ U_iter
		// U_1 = PRF(password, salt || uint(i))
		prf.Reset()
		prf.Write(salt)
		buf[0] = byte(block >> 24)
		buf[1] = byte(block >> 16)
		buf[2] = byte(block >> 8)
		buf[3] = byte(block)
		prf.Write(buf[:4])
		dk = prf.Sum(dk)
		T := dk[int64(len(dk))-hashLen:]
		copy(U, T)

		// U_n = PRF(password, U_(n-1))
		for n := int64(2); n <= iter; n++ {
			prf.Reset()
			prf.Write(U)
			U = U[:0]
			U = prf.Sum(U)
			for x := range U {
				T[x] ^= U[x]
			}
		}
	}
	return dk[:keyLen]
}
<|endoftext|>


<commit_before>//
// Copyright 2021, Sander van Harmelen
//
// 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 gitlab

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

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

func TestListProtectedEnvironments(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/projects/1/protected_environments", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, `[{"name":"1.0.0", "deploy_access_levels": [{"access_level": 40, "access_level_description": "Maintainers"}], "required_approval_count": 1},{"name":"*-release", "deploy_access_levels": [{"access_level": 30, "access_level_description": "Developers + Maintainers"}]}]`)
	})

	expected := []*ProtectedEnvironment{
		{
			Name:                  "1.0.0",
			RequiredApprovalCount: 1,
			DeployAccessLevels: []*EnvironmentAccessDescription{
				{
					AccessLevel:            40,
					AccessLevelDescription: "Maintainers",
				},
			},
		},
		{
			Name: "*-release",
			DeployAccessLevels: []*EnvironmentAccessDescription{
				{
					AccessLevel:            30,
					AccessLevelDescription: "Developers + Maintainers",
				},
			},
		},
	}

	opt := &ListProtectedEnvironmentsOptions{}
	environments, _, err := client.ProtectedEnvironments.ListProtectedEnvironments(1, opt)
	assert.NoError(t, err, "failed to get response")
	assert.Equal(t, expected, environments)
}

func TestGetProtectedEnvironment(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	// Test with RequiredApprovalCount
	environmentName := "my-awesome-environment"

	mux.HandleFunc(fmt.Sprintf("/api/v4/projects/1/protected_environments/%s", environmentName), func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, `{"name":"my-awesome-environment", "required_approval_count": 1, "deploy_access_levels": [{"access_level": 30, "access_level_description": "Developers + Maintainers"}]}`)
	})

	expected := &ProtectedEnvironment{
		Name:                  environmentName,
		RequiredApprovalCount: 1,
		DeployAccessLevels: []*EnvironmentAccessDescription{
			{
				AccessLevel:            30,
				AccessLevelDescription: "Developers + Maintainers",
			},
		},
	}

	environment, _, err := client.ProtectedEnvironments.GetProtectedEnvironment(1, environmentName)

	assert.NoError(t, err, "failed to get response")
	assert.Equal(t, expected, environment)

	// Test without RequiredApprovalCount
	environmentName = "my-awesome-environment2"

	mux.HandleFunc(fmt.Sprintf("/api/v4/projects/1/protected_environments/%s", environmentName), func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, `{"name":"my-awesome-environment2", "deploy_access_levels": [{"access_level": 30, "access_level_description": "Developers + Maintainers"}]}`)
	})

	expected = &ProtectedEnvironment{
		Name: environmentName,
		DeployAccessLevels: []*EnvironmentAccessDescription{
			{
				AccessLevel:            30,
				AccessLevelDescription: "Developers + Maintainers",
			},
		},
	}

	environment, _, err = client.ProtectedEnvironments.GetProtectedEnvironment(1, environmentName)

	assert.NoError(t, err, "failed to get response")
	assert.Equal(t, expected, environment)
}

func TestProtectRepositoryEnvironments(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/projects/1/protected_environments", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		fmt.Fprint(w, `{"name":"my-awesome-environment", "required_approval_count": 2, "deploy_access_levels": [{"access_level": 30, "access_level_description": "Developers + Maintainers"}]}`)
	})

	expected := &ProtectedEnvironment{
		Name:                  "my-awesome-environment",
		RequiredApprovalCount: 2,
		DeployAccessLevels: []*EnvironmentAccessDescription{
			{
				AccessLevel:            30,
				AccessLevelDescription: "Developers + Maintainers",
			},
		},
	}

	opt := &ProtectRepositoryEnvironmentsOptions{Name: String("my-awesome-environment"), DeployAccessLevels: &[]*EnvironmentAccessOptions{{AccessLevel: AccessLevel(30)}}}
	environment, _, err := client.ProtectedEnvironments.ProtectRepositoryEnvironments(1, opt)

	assert.NoError(t, err, "failed to get response")
	assert.Equal(t, expected, environment)
}

func TestUnprotectRepositoryEnvironments(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/projects/1/protected_environments/my-awesome-environment", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
	})

	resp, err := client.ProtectedEnvironments.UnprotectEnvironment(1, "my-awesome-environment")
	assert.NoError(t, err, "failed to get response")
	assert.Equal(t, http.StatusOK, resp.StatusCode)
}
<commit_msg>Updated tests to catch the new ProtectRepositoryEnvironmentsOptions change<commit_after>//
// Copyright 2021, Sander van Harmelen
//
// 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 gitlab

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

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

func TestListProtectedEnvironments(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/projects/1/protected_environments", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, `[{"name":"1.0.0", "deploy_access_levels": [{"access_level": 40, "access_level_description": "Maintainers"}], "required_approval_count": 1},{"name":"*-release", "deploy_access_levels": [{"access_level": 30, "access_level_description": "Developers + Maintainers"}]}]`)
	})

	expected := []*ProtectedEnvironment{
		{
			Name:                  "1.0.0",
			RequiredApprovalCount: 1,
			DeployAccessLevels: []*EnvironmentAccessDescription{
				{
					AccessLevel:            40,
					AccessLevelDescription: "Maintainers",
				},
			},
		},
		{
			Name: "*-release",
			DeployAccessLevels: []*EnvironmentAccessDescription{
				{
					AccessLevel:            30,
					AccessLevelDescription: "Developers + Maintainers",
				},
			},
		},
	}

	opt := &ListProtectedEnvironmentsOptions{}
	environments, _, err := client.ProtectedEnvironments.ListProtectedEnvironments(1, opt)
	assert.NoError(t, err, "failed to get response")
	assert.Equal(t, expected, environments)
}

func TestGetProtectedEnvironment(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	// Test with RequiredApprovalCount
	environmentName := "my-awesome-environment"

	mux.HandleFunc(fmt.Sprintf("/api/v4/projects/1/protected_environments/%s", environmentName), func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, `{"name":"my-awesome-environment", "required_approval_count": 1, "deploy_access_levels": [{"access_level": 30, "access_level_description": "Developers + Maintainers"}]}`)
	})

	expected := &ProtectedEnvironment{
		Name:                  environmentName,
		RequiredApprovalCount: 1,
		DeployAccessLevels: []*EnvironmentAccessDescription{
			{
				AccessLevel:            30,
				AccessLevelDescription: "Developers + Maintainers",
			},
		},
	}

	environment, _, err := client.ProtectedEnvironments.GetProtectedEnvironment(1, environmentName)

	assert.NoError(t, err, "failed to get response")
	assert.Equal(t, expected, environment)

	// Test without RequiredApprovalCount
	environmentName = "my-awesome-environment2"

	mux.HandleFunc(fmt.Sprintf("/api/v4/projects/2/protected_environments/%s", environmentName), func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		fmt.Fprint(w, `{"name":"my-awesome-environment2", "deploy_access_levels": [{"access_level": 30, "access_level_description": "Developers + Maintainers"}]}`)
	})

	expected = &ProtectedEnvironment{
		Name: environmentName,
		DeployAccessLevels: []*EnvironmentAccessDescription{
			{
				AccessLevel:            30,
				AccessLevelDescription: "Developers + Maintainers",
			},
		},
	}

	environment, _, err = client.ProtectedEnvironments.GetProtectedEnvironment(2, environmentName)

	assert.NoError(t, err, "failed to get response")
	assert.Equal(t, expected, environment)
}

func TestProtectRepositoryEnvironments(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	// Test with RequiredApprovalCount
	mux.HandleFunc("/api/v4/projects/1/protected_environments", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		fmt.Fprint(w, `{"name":"my-awesome-environment", "required_approval_count": 2, "deploy_access_levels": [{"access_level": 30, "access_level_description": "Developers + Maintainers"}]}`)
	})

	expected := &ProtectedEnvironment{
		Name:                  "my-awesome-environment",
		RequiredApprovalCount: 2,
		DeployAccessLevels: []*EnvironmentAccessDescription{
			{
				AccessLevel:            30,
				AccessLevelDescription: "Developers + Maintainers",
			},
		},
	}

	opt := &ProtectRepositoryEnvironmentsOptions{Name: String("my-awesome-environment"), RequiredApprovalCount: Int(2), DeployAccessLevels: &[]*EnvironmentAccessOptions{{AccessLevel: AccessLevel(30)}}}
	environment, _, err := client.ProtectedEnvironments.ProtectRepositoryEnvironments(1, opt)

	assert.NoError(t, err, "failed to get response")
	assert.Equal(t, expected, environment)

	// Test without RequiredApprovalCount
	mux.HandleFunc("/api/v4/projects/2/protected_environments", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodPost)
		fmt.Fprint(w, `{"name":"my-awesome-environment2", "deploy_access_levels": [{"access_level": 30, "access_level_description": "Developers + Maintainers"}]}`)
	})

	expected = &ProtectedEnvironment{
		Name: "my-awesome-environment2",
		DeployAccessLevels: []*EnvironmentAccessDescription{
			{
				AccessLevel:            30,
				AccessLevelDescription: "Developers + Maintainers",
			},
		},
	}

	opt = &ProtectRepositoryEnvironmentsOptions{Name: String("my-awesome-environment2"), DeployAccessLevels: &[]*EnvironmentAccessOptions{{AccessLevel: AccessLevel(30)}}}
	environment, _, err = client.ProtectedEnvironments.ProtectRepositoryEnvironments(2, opt)

	assert.NoError(t, err, "failed to get response")
	assert.Equal(t, expected, environment)
}

func TestUnprotectRepositoryEnvironments(t *testing.T) {
	mux, server, client := setup(t)
	defer teardown(server)

	mux.HandleFunc("/api/v4/projects/1/protected_environments/my-awesome-environment", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodDelete)
	})

	resp, err := client.ProtectedEnvironments.UnprotectEnvironment(1, "my-awesome-environment")
	assert.NoError(t, err, "failed to get response")
	assert.Equal(t, http.StatusOK, resp.StatusCode)
}
<|endoftext|>