<commit_before>// Based on net/rpc/jsonrpc by:
// Copyright 2010 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 jsonrpc2

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/rpc"
	"sync"
)

type serverCodec struct {
	encmutex sync.Mutex    // protects enc
	dec      *json.Decoder // for reading JSON values
	enc      *json.Encoder // for writing JSON values
	c        io.Closer
	srv      *rpc.Server
	ctx      context.Context

	// temporary work space
	req serverRequest

	// JSON-RPC clients can use arbitrary json values as request IDs.
	// Package rpc expects uint64 request IDs.
	// We assign uint64 sequence numbers to incoming requests
	// but save the original request ID in the pending map.
	// When rpc responds, we use the sequence number in
	// the response to find the original request ID.
	mutex   sync.Mutex // protects seq, pending
	seq     uint64
	pending map[uint64]*json.RawMessage
}

// NewServerCodec returns a new rpc.ServerCodec using JSON-RPC 2.0 on conn,
// which will use srv to execute batch requests.
//
// If srv is nil then rpc.DefaultServer will be used.
//
// For most use cases NewServerCodec is too low-level and you should use
// ServeConn instead. You'll need NewServerCodec if you wanna register
// your own object of type named "JSONRPC2" (same as used internally to
// process batch requests) or you wanna use custom rpc server object
// instead of rpc.DefaultServer to process requests on conn.
func NewServerCodec(conn io.ReadWriteCloser, srv *rpc.Server) rpc.ServerCodec {
	if srv == nil {
		srv = rpc.DefaultServer
	}
	srv.Register(JSONRPC2{})
	return &serverCodec{
		dec:     json.NewDecoder(conn),
		enc:     json.NewEncoder(conn),
		c:       conn,
		srv:     srv,
		ctx:     context.Background(),
		pending: make(map[uint64]*json.RawMessage),
	}
}

// NewServerCodecContext is NewServerCodec with given context provided
// within parameters for compatible RPC methods.
func NewServerCodecContext(ctx context.Context, conn io.ReadWriteCloser, srv *rpc.Server) rpc.ServerCodec {
	codec := NewServerCodec(conn, srv)
	codec.(*serverCodec).ctx = ctx
	return codec
}

type serverRequest struct {
	Version string           `json:"jsonrpc"`
	Method  string           `json:"method"`
	Params  *json.RawMessage `json:"params"`
	ID      *json.RawMessage `json:"id"`
}

func (r *serverRequest) reset() {
	r.Version = ""
	r.Method = ""
	r.Params = nil
	r.ID = nil
}

func (r *serverRequest) UnmarshalJSON(raw []byte) error {
	r.reset()
	type req *serverRequest
	if err := json.Unmarshal(raw, req(r)); err != nil {
		return errors.New("bad request")
	}

	var o = make(map[string]*json.RawMessage)
	if err := json.Unmarshal(raw, &o); err != nil {
		return errors.New("bad request")
	}
	if o["jsonrpc"] == nil || o["method"] == nil {
		return errors.New("bad request")
	}
	_, okID := o["id"]
	_, okParams := o["params"]
	if len(o) == 3 && !(okID || okParams) || len(o) == 4 && !(okID && okParams) || len(o) > 4 {
		return errors.New("bad request")
	}
	if r.Version != "2.0" {
		return errors.New("bad request")
	}
	if okParams {
		if r.Params == nil || len(*r.Params) == 0 {
			return errors.New("bad request")
		}
		switch []byte(*r.Params)[0] {
		case '[', '{':
		default:
			return errors.New("bad request")
		}
	}
	if okID && r.ID == nil {
		r.ID = &null
	}
	if okID {
		if len(*r.ID) == 0 {
			return errors.New("bad request")
		}
		switch []byte(*r.ID)[0] {
		case 't', 'f', '{', '[':
			return errors.New("bad request")
		}
	}

	return nil
}

type serverResponse struct {
	Version string           `json:"jsonrpc"`
	ID      *json.RawMessage `json:"id"`
	Result  interface{}      `json:"result,omitempty"`
	Error   interface{}      `json:"error,omitempty"`
}

func (c *serverCodec) ReadRequestHeader(r *rpc.Request) (err error) {
	// If return error:
	// - codec will be closed
	// So, try to send error reply to client before returning error.
	var raw json.RawMessage
	if err := c.dec.Decode(&raw); err != nil {
		c.encmutex.Lock()
		c.enc.Encode(serverResponse{Version: "2.0", ID: &null, Error: errParse})
		c.encmutex.Unlock()
		return err
	}

	if len(raw) > 0 && raw[0] == '[' {
		c.req.Version = "2.0"
		c.req.Method = "JSONRPC2.Batch"
		c.req.Params = &raw
		c.req.ID = &null
	} else if err := json.Unmarshal(raw, &c.req); err != nil {
		if err.Error() == "bad request" {
			c.encmutex.Lock()
			c.enc.Encode(serverResponse{Version: "2.0", ID: &null, Error: errRequest})
			c.encmutex.Unlock()
		}
		return err
	}

	r.ServiceMethod = c.req.Method

	// JSON request id can be any JSON value;
	// RPC package expects uint64.  Translate to
	// internal uint64 and save JSON on the side.
	c.mutex.Lock()
	c.seq++
	c.pending[c.seq] = c.req.ID
	c.req.ID = nil
	r.Seq = c.seq
	c.mutex.Unlock()

	return nil
}

func (c *serverCodec) ReadRequestBody(x interface{}) error {
	// If x!=nil and return error e:
	// - WriteResponse() will be called with e.Error() in r.Error
	if x == nil {
		return nil
	}
	if c.req.Params == nil {
		return nil
	}
	if c.req.Method == "JSONRPC2.Batch" {
		arg := x.(*BatchArg)
		arg.srv = c.srv
		if err := json.Unmarshal(*c.req.Params, &arg.reqs); err != nil {
			return NewError(errParams.Code, err.Error())
		}
		if len(arg.reqs) == 0 {
			return errRequest
		}
	} else if err := json.Unmarshal(*c.req.Params, x); err != nil {
		return NewError(errParams.Code, err.Error())
	}
	if x, ok := x.(WithContext); ok {
		x.SetContext(c.ctx)
	}
	return nil
}

var null = json.RawMessage([]byte("null"))

func (c *serverCodec) WriteResponse(r *rpc.Response, x interface{}) error {
	// If return error: nothing happens.
	// In r.Error will be "" or .Error() of error returned by:
	// - ReadRequestBody()
	// - called RPC method
	c.mutex.Lock()
	b, ok := c.pending[r.Seq]
	if !ok {
		c.mutex.Unlock()
		return errors.New("invalid sequence number in response")
	}
	delete(c.pending, r.Seq)
	c.mutex.Unlock()

	if replies, ok := x.(*[]*json.RawMessage); r.ServiceMethod == "JSONRPC2.Batch" && ok {
		if len(*replies) == 0 {
			return nil
		}
		c.encmutex.Lock()
		defer c.encmutex.Unlock()
		return c.enc.Encode(replies)
	}

	if b == nil {
		// Notification. Do not respond.
		return nil
	}
	resp := serverResponse{Version: "2.0", ID: b}
	if r.Error == "" {
		if x == nil {
			resp.Result = &null
		} else {
			resp.Result = x
		}
	} else if r.Error[0] == '{' && r.Error[len(r.Error)-1] == '}' {
		// Well… this check for '{'…'}' isn't too strict, but I
		// suppose we're trusting our own RPC methods (this way they
		// can force sending wrong reply or many replies instead
		// of one) and normal errors won't be formatted this way.
		raw := json.RawMessage(r.Error)
		resp.Error = &raw
	} else {
		raw := json.RawMessage(newError(r.Error).Error())
		resp.Error = &raw
	}
	c.encmutex.Lock()
	defer c.encmutex.Unlock()
	return c.enc.Encode(resp)
}

func (c *serverCodec) Close() error {
	return c.c.Close()
}

// ServeConn runs the JSON-RPC 2.0 server on a single connection.
// ServeConn blocks, serving the connection until the client hangs up.
// The caller typically invokes ServeConn in a go statement.
func ServeConn(conn io.ReadWriteCloser) {
	rpc.ServeCodec(NewServerCodec(conn, nil))
}

// ServeConnContext is ServeConn with given context provided
// within parameters for compatible RPC methods.
func ServeConnContext(ctx context.Context, conn io.ReadWriteCloser) {
	rpc.ServeCodec(NewServerCodecContext(ctx, conn, nil))
}
<commit_msg>fix absent context for method without params<commit_after>// Based on net/rpc/jsonrpc by:
// Copyright 2010 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 jsonrpc2

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/rpc"
	"sync"
)

type serverCodec struct {
	encmutex sync.Mutex    // protects enc
	dec      *json.Decoder // for reading JSON values
	enc      *json.Encoder // for writing JSON values
	c        io.Closer
	srv      *rpc.Server
	ctx      context.Context

	// temporary work space
	req serverRequest

	// JSON-RPC clients can use arbitrary json values as request IDs.
	// Package rpc expects uint64 request IDs.
	// We assign uint64 sequence numbers to incoming requests
	// but save the original request ID in the pending map.
	// When rpc responds, we use the sequence number in
	// the response to find the original request ID.
	mutex   sync.Mutex // protects seq, pending
	seq     uint64
	pending map[uint64]*json.RawMessage
}

// NewServerCodec returns a new rpc.ServerCodec using JSON-RPC 2.0 on conn,
// which will use srv to execute batch requests.
//
// If srv is nil then rpc.DefaultServer will be used.
//
// For most use cases NewServerCodec is too low-level and you should use
// ServeConn instead. You'll need NewServerCodec if you wanna register
// your own object of type named "JSONRPC2" (same as used internally to
// process batch requests) or you wanna use custom rpc server object
// instead of rpc.DefaultServer to process requests on conn.
func NewServerCodec(conn io.ReadWriteCloser, srv *rpc.Server) rpc.ServerCodec {
	if srv == nil {
		srv = rpc.DefaultServer
	}
	srv.Register(JSONRPC2{})
	return &serverCodec{
		dec:     json.NewDecoder(conn),
		enc:     json.NewEncoder(conn),
		c:       conn,
		srv:     srv,
		ctx:     context.Background(),
		pending: make(map[uint64]*json.RawMessage),
	}
}

// NewServerCodecContext is NewServerCodec with given context provided
// within parameters for compatible RPC methods.
func NewServerCodecContext(ctx context.Context, conn io.ReadWriteCloser, srv *rpc.Server) rpc.ServerCodec {
	codec := NewServerCodec(conn, srv)
	codec.(*serverCodec).ctx = ctx
	return codec
}

type serverRequest struct {
	Version string           `json:"jsonrpc"`
	Method  string           `json:"method"`
	Params  *json.RawMessage `json:"params"`
	ID      *json.RawMessage `json:"id"`
}

func (r *serverRequest) reset() {
	r.Version = ""
	r.Method = ""
	r.Params = nil
	r.ID = nil
}

func (r *serverRequest) UnmarshalJSON(raw []byte) error {
	r.reset()
	type req *serverRequest
	if err := json.Unmarshal(raw, req(r)); err != nil {
		return errors.New("bad request")
	}

	var o = make(map[string]*json.RawMessage)
	if err := json.Unmarshal(raw, &o); err != nil {
		return errors.New("bad request")
	}
	if o["jsonrpc"] == nil || o["method"] == nil {
		return errors.New("bad request")
	}
	_, okID := o["id"]
	_, okParams := o["params"]
	if len(o) == 3 && !(okID || okParams) || len(o) == 4 && !(okID && okParams) || len(o) > 4 {
		return errors.New("bad request")
	}
	if r.Version != "2.0" {
		return errors.New("bad request")
	}
	if okParams {
		if r.Params == nil || len(*r.Params) == 0 {
			return errors.New("bad request")
		}
		switch []byte(*r.Params)[0] {
		case '[', '{':
		default:
			return errors.New("bad request")
		}
	}
	if okID && r.ID == nil {
		r.ID = &null
	}
	if okID {
		if len(*r.ID) == 0 {
			return errors.New("bad request")
		}
		switch []byte(*r.ID)[0] {
		case 't', 'f', '{', '[':
			return errors.New("bad request")
		}
	}

	return nil
}

type serverResponse struct {
	Version string           `json:"jsonrpc"`
	ID      *json.RawMessage `json:"id"`
	Result  interface{}      `json:"result,omitempty"`
	Error   interface{}      `json:"error,omitempty"`
}

func (c *serverCodec) ReadRequestHeader(r *rpc.Request) (err error) {
	// If return error:
	// - codec will be closed
	// So, try to send error reply to client before returning error.
	var raw json.RawMessage
	if err := c.dec.Decode(&raw); err != nil {
		c.encmutex.Lock()
		c.enc.Encode(serverResponse{Version: "2.0", ID: &null, Error: errParse})
		c.encmutex.Unlock()
		return err
	}

	if len(raw) > 0 && raw[0] == '[' {
		c.req.Version = "2.0"
		c.req.Method = "JSONRPC2.Batch"
		c.req.Params = &raw
		c.req.ID = &null
	} else if err := json.Unmarshal(raw, &c.req); err != nil {
		if err.Error() == "bad request" {
			c.encmutex.Lock()
			c.enc.Encode(serverResponse{Version: "2.0", ID: &null, Error: errRequest})
			c.encmutex.Unlock()
		}
		return err
	}

	r.ServiceMethod = c.req.Method

	// JSON request id can be any JSON value;
	// RPC package expects uint64.  Translate to
	// internal uint64 and save JSON on the side.
	c.mutex.Lock()
	c.seq++
	c.pending[c.seq] = c.req.ID
	c.req.ID = nil
	r.Seq = c.seq
	c.mutex.Unlock()

	return nil
}

func (c *serverCodec) ReadRequestBody(x interface{}) error {
	// If x!=nil and return error e:
	// - WriteResponse() will be called with e.Error() in r.Error
	if x == nil {
		return nil
	}
	if x, ok := x.(WithContext); ok {
		x.SetContext(c.ctx)
	}
	if c.req.Params == nil {
		return nil
	}
	if c.req.Method == "JSONRPC2.Batch" {
		arg := x.(*BatchArg)
		arg.srv = c.srv
		if err := json.Unmarshal(*c.req.Params, &arg.reqs); err != nil {
			return NewError(errParams.Code, err.Error())
		}
		if len(arg.reqs) == 0 {
			return errRequest
		}
	} else if err := json.Unmarshal(*c.req.Params, x); err != nil {
		return NewError(errParams.Code, err.Error())
	}
	return nil
}

var null = json.RawMessage([]byte("null"))

func (c *serverCodec) WriteResponse(r *rpc.Response, x interface{}) error {
	// If return error: nothing happens.
	// In r.Error will be "" or .Error() of error returned by:
	// - ReadRequestBody()
	// - called RPC method
	c.mutex.Lock()
	b, ok := c.pending[r.Seq]
	if !ok {
		c.mutex.Unlock()
		return errors.New("invalid sequence number in response")
	}
	delete(c.pending, r.Seq)
	c.mutex.Unlock()

	if replies, ok := x.(*[]*json.RawMessage); r.ServiceMethod == "JSONRPC2.Batch" && ok {
		if len(*replies) == 0 {
			return nil
		}
		c.encmutex.Lock()
		defer c.encmutex.Unlock()
		return c.enc.Encode(replies)
	}

	if b == nil {
		// Notification. Do not respond.
		return nil
	}
	resp := serverResponse{Version: "2.0", ID: b}
	if r.Error == "" {
		if x == nil {
			resp.Result = &null
		} else {
			resp.Result = x
		}
	} else if r.Error[0] == '{' && r.Error[len(r.Error)-1] == '}' {
		// Well… this check for '{'…'}' isn't too strict, but I
		// suppose we're trusting our own RPC methods (this way they
		// can force sending wrong reply or many replies instead
		// of one) and normal errors won't be formatted this way.
		raw := json.RawMessage(r.Error)
		resp.Error = &raw
	} else {
		raw := json.RawMessage(newError(r.Error).Error())
		resp.Error = &raw
	}
	c.encmutex.Lock()
	defer c.encmutex.Unlock()
	return c.enc.Encode(resp)
}

func (c *serverCodec) Close() error {
	return c.c.Close()
}

// ServeConn runs the JSON-RPC 2.0 server on a single connection.
// ServeConn blocks, serving the connection until the client hangs up.
// The caller typically invokes ServeConn in a go statement.
func ServeConn(conn io.ReadWriteCloser) {
	rpc.ServeCodec(NewServerCodec(conn, nil))
}

// ServeConnContext is ServeConn with given context provided
// within parameters for compatible RPC methods.
func ServeConnContext(ctx context.Context, conn io.ReadWriteCloser) {
	rpc.ServeCodec(NewServerCodecContext(ctx, conn, nil))
}
<|endoftext|>


<commit_before>package device_manager

import (
	"bufio"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/golang/mock/gomock"
	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"

	v1 "kubevirt.io/client-go/api/v1"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/yaml"

	"kubevirt.io/kubevirt/pkg/testutils"
	virtconfig "kubevirt.io/kubevirt/pkg/virt-config"
)

const (
	fakeMdevNameSelector      = "FAKE 123"
	fakeIntelMdevNameSelector = "i915-GVTg_V5_4"
	fakeMdevResourceName      = "example.org/fake123"
	fakeMdevUUID              = "53764d0e-85a0-42b4-af5c-2046b460b1dc"
	fakeIntelMdevUUID         = "54444d0e-85a0-42b4-af5c-2046b4bbb1aa"
)

var _ = Describe("Mediated Device", func() {
	var mockPCI *MockDeviceHandler
	var fakePermittedHostDevicesConfig string
	var fakePermittedHostDevices v1.PermittedHostDevices
	var ctrl *gomock.Controller
	var fakeSupportedTypesPath string
	reourceNameToTypeName := func(rawName string) string {
		typeNameStr := strings.Replace(string(rawName), " ", "_", -1)
		typeNameStr = strings.TrimSpace(typeNameStr)
		return typeNameStr
	}
	BeforeEach(func() {
		By("creating a temporary fake mdev directory tree")
		// create base mdev dir instead of /sys/bus/mdev/devices
		fakeMdevBasePath, err := ioutil.TempDir("/tmp", "mdevs")
		Expect(err).ToNot(HaveOccurred())
		// create an alternative mdev_supported_types dir instead of /sys/class/mdev_bus/[pciAddress]/
		fakeSupportedTypesPath, err = ioutil.TempDir("/tmp", "mdev_supported_types")
		Expect(err).ToNot(HaveOccurred())
		// create a fake path to nvidia mdev type
		fakeNvidiaTypePath := filepath.Join(fakeSupportedTypesPath, "nvidia-222")
		err = os.MkdirAll(fakeNvidiaTypePath, 0700)
		Expect(err).ToNot(HaveOccurred())
		// create a fake path to Intel mdev type
		fakeIntelTypePath := filepath.Join(fakeSupportedTypesPath, fakeIntelMdevNameSelector)
		err = os.MkdirAll(fakeIntelTypePath, 0700)
		Expect(err).ToNot(HaveOccurred())
		mdevBasePath = fakeMdevBasePath
		// create mdev directories and symlinks
		for _, uuid := range []string{fakeMdevUUID, fakeIntelMdevUUID} {
			mdevTypePath := filepath.Join(fakeMdevBasePath, uuid+"real")
			err = os.MkdirAll(mdevTypePath, 0700)
			Expect(err).ToNot(HaveOccurred())
			err = os.Symlink(filepath.Join(fakeMdevBasePath, uuid+"real"), filepath.Join(fakeMdevBasePath, uuid))
			Expect(err).ToNot(HaveOccurred())
		}
		// link nvidia type directory
		err = os.Symlink(fakeNvidiaTypePath, filepath.Join(fakeMdevBasePath, fakeMdevUUID+"real", "mdev_type"))
		Expect(err).ToNot(HaveOccurred())
		err = os.Symlink(fakeIntelTypePath, filepath.Join(fakeMdevBasePath, fakeIntelMdevUUID+"real", "mdev_type"))
		Expect(err).ToNot(HaveOccurred())
		// create a name file in the nvidia type directory
		mdevName, err := os.Create(filepath.Join(fakeNvidiaTypePath, "name"))
		Expect(err).ToNot(HaveOccurred())
		mdevNameWriter := bufio.NewWriter(mdevName)
		n, err := mdevNameWriter.WriteString(fakeMdevNameSelector + "\n")
		Expect(err).ToNot(HaveOccurred())
		Expect(n).To(Equal(len(fakeMdevNameSelector) + 1))
		mdevNameWriter.Flush()

	})

	AfterEach(func() {
		os.RemoveAll(mdevBasePath)
		os.RemoveAll(fakeSupportedTypesPath)
	})
	table.DescribeTable("should get correct file type name", func(namePathExist bool) {
		if namePathExist {
			mdevName, err := getMdevTypeName(fakeMdevUUID)
			Expect(err).ToNot(HaveOccurred())
			Expect(mdevName).To(Equal(reourceNameToTypeName(fakeMdevNameSelector)))
		} else {
			mdevName, err := getMdevTypeName(fakeIntelMdevUUID)
			Expect(err).ToNot(HaveOccurred())
			Expect(mdevName).To(Equal(reourceNameToTypeName(fakeIntelMdevNameSelector)))
		}
	},
		table.Entry("Nvidia name file exist", true),
		table.Entry("Intel name file doesn't exist", false),
	)
	Context("discover devices", func() {
		BeforeEach(func() {
			By("mocking PCI and MDEV functions to simulate an mdev an its parent PCI device")
			ctrl = gomock.NewController(GinkgoT())
			mockPCI = NewMockDeviceHandler(ctrl)
			Handler = mockPCI
			// Force pre-defined returned values and ensure the function only get called exacly once each on 0000:00:00.0
			mockPCI.EXPECT().GetMdevParentPCIAddr(fakeMdevUUID).Return(fakeAddress, nil).Times(1)
			mockPCI.EXPECT().GetDeviceIOMMUGroup(mdevBasePath, fakeMdevUUID).Return(fakeIommuGroup, nil).Times(1)
			mockPCI.EXPECT().GetDeviceNumaNode(pciBasePath, fakeAddress).Return(fakeNumaNode).Times(1)

			By("creating a list of fake device using the yaml decoder")
			fakePermittedHostDevicesConfig = `
    mediatedDevices:
    - mdevNameSelector: "` + fakeMdevNameSelector + `"
      resourceName: "` + fakeMdevResourceName + `"
    `
			err := yaml.NewYAMLOrJSONDecoder(strings.NewReader(fakePermittedHostDevicesConfig), 1024).Decode(&fakePermittedHostDevices)
			Expect(err).ToNot(HaveOccurred())
			Expect(len(fakePermittedHostDevices.MediatedDevices)).To(Equal(1))
			Expect(fakePermittedHostDevices.MediatedDevices[0].MDEVNameSelector).To(Equal(fakeMdevNameSelector))
			Expect(fakePermittedHostDevices.MediatedDevices[0].ResourceName).To(Equal(fakeMdevResourceName))
		})
		AfterEach(func() {
			ctrl.Finish()
		})
		It("Should parse the permitted devices and find 1 matching mediated device", func() {
			supportedMdevsMap := make(map[string]string)
			for _, supportedMdev := range fakePermittedHostDevices.MediatedDevices {
				// do not add a device plugin for this resource if it's being provided via an external device plugin
				if !supportedMdev.ExternalResourceProvider {
					selector := removeSelectorSpaces(supportedMdev.MDEVNameSelector)
					supportedMdevsMap[selector] = supportedMdev.ResourceName
				}
			}
			// discoverPermittedHostMediatedDevices() will walk real mdev devices wherever the tests are running
			devices := discoverPermittedHostMediatedDevices(supportedMdevsMap)
			Expect(len(devices)).To(Equal(1))
			selector := removeSelectorSpaces(fakeMdevNameSelector)
			Expect(len(devices[selector])).To(Equal(1))
			Expect(devices[selector][0].UUID).To(Equal(fakeMdevUUID))
			Expect(devices[selector][0].typeName).To(Equal(selector))
			Expect(devices[selector][0].parentPciAddress).To(Equal(fakeAddress))
			Expect(devices[selector][0].iommuGroup).To(Equal(fakeIommuGroup))
			Expect(devices[selector][0].numaNode).To(Equal(fakeNumaNode))
		})

		It("Should validate DPI devices", func() {
			iommuToMDEVMap := make(map[string]string)
			supportedMdevsMap := make(map[string]string)
			for _, supportedMdev := range fakePermittedHostDevices.MediatedDevices {
				// do not add a device plugin for this resource if it's being provided via an external device plugin
				if !supportedMdev.ExternalResourceProvider {
					selector := removeSelectorSpaces(supportedMdev.MDEVNameSelector)
					supportedMdevsMap[selector] = supportedMdev.ResourceName
				}
			}
			// discoverPermittedHostMediatedDevices() will walk real mdev devices wherever the tests are running
			mDevices := discoverPermittedHostMediatedDevices(supportedMdevsMap)
			selector := removeSelectorSpaces(fakeMdevNameSelector)
			devs := constructDPIdevicesFromMdev(mDevices[selector], iommuToMDEVMap)
			Expect(devs[0].ID).To(Equal(fakeIommuGroup))
			Expect(devs[0].Topology.Nodes[0].ID).To(Equal(int64(fakeNumaNode)))
		})

		It("Should update the device list according to the configmap", func() {
			By("creating a cluster config")
			kv := &v1.KubeVirt{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "kubevirt",
					Namespace: "kubevirt",
				},
				Spec: v1.KubeVirtSpec{
					Configuration: v1.KubeVirtConfiguration{
						DeveloperConfiguration: &v1.DeveloperConfiguration{},
					},
				},
				Status: v1.KubeVirtStatus{
					Phase: v1.KubeVirtPhaseDeploying,
				},
			}
			fakeClusterConfig, _, _, kvInformer := testutils.NewFakeClusterConfigUsingKV(kv)

			By("creating an empty device controller")
			deviceController := NewDeviceController("master", 10, fakeClusterConfig)
			deviceController.devicePlugins = make(map[string]ControlledDevice)

			By("adding a host device to the cluster config")
			kvConfig := kv.DeepCopy()
			kvConfig.Spec.Configuration.DeveloperConfiguration.FeatureGates = []string{virtconfig.HostDevicesGate}
			kvConfig.Spec.Configuration.PermittedHostDevices = &v1.PermittedHostDevices{
				MediatedDevices: []v1.MediatedHostDevice{
					{
						MDEVNameSelector: fakeMdevNameSelector,
						ResourceName:     fakeMdevResourceName,
					},
				},
			}
			testutils.UpdateFakeKubeVirtClusterConfig(kvInformer, kvConfig)
			permittedDevices := fakeClusterConfig.GetPermittedHostDevices()
			Expect(permittedDevices).ToNot(BeNil(), "something went wrong while parsing the configmap(s)")
			Expect(len(permittedDevices.MediatedDevices)).To(Equal(1), "the fake device was not found")

			By("ensuring a device plugin gets created for our fake device")
			enabledDevicePlugins, disabledDevicePlugins := deviceController.updatePermittedHostDevicePlugins()
			Expect(len(enabledDevicePlugins)).To(Equal(1), "a device plugin wasn't created for the fake device")
			Expect(len(disabledDevicePlugins)).To(Equal(0))
			Ω(enabledDevicePlugins).Should(HaveKey(fakeMdevResourceName))
			// Manually adding the enabled plugin, since the device controller is not actually running
			deviceController.devicePlugins[fakeMdevResourceName] = enabledDevicePlugins[fakeMdevResourceName]

			By("deletting the device from the configmap")
			kvConfig.Spec.Configuration.PermittedHostDevices = &v1.PermittedHostDevices{}
			testutils.UpdateFakeKubeVirtClusterConfig(kvInformer, kvConfig)
			permittedDevices = fakeClusterConfig.GetPermittedHostDevices()
			Expect(permittedDevices).ToNot(BeNil(), "something went wrong while parsing the configmap(s)")
			Expect(len(permittedDevices.MediatedDevices)).To(Equal(0), "the fake device was not deleted")

			By("ensuring the device plugin gets stopped")
			enabledDevicePlugins, disabledDevicePlugins = deviceController.updatePermittedHostDevicePlugins()
			Expect(len(enabledDevicePlugins)).To(Equal(0))
			Expect(len(disabledDevicePlugins)).To(Equal(1), "the fake device plugin did not get disabled")
			Ω(disabledDevicePlugins).Should(HaveKey(fakeMdevResourceName))
		})
	})
})
<commit_msg>fixing spelling error<commit_after>package device_manager

import (
	"bufio"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/golang/mock/gomock"
	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"

	v1 "kubevirt.io/client-go/api/v1"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/yaml"

	"kubevirt.io/kubevirt/pkg/testutils"
	virtconfig "kubevirt.io/kubevirt/pkg/virt-config"
)

const (
	fakeMdevNameSelector      = "FAKE 123"
	fakeIntelMdevNameSelector = "i915-GVTg_V5_4"
	fakeMdevResourceName      = "example.org/fake123"
	fakeMdevUUID              = "53764d0e-85a0-42b4-af5c-2046b460b1dc"
	fakeIntelMdevUUID         = "54444d0e-85a0-42b4-af5c-2046b4bbb1aa"
)

var _ = Describe("Mediated Device", func() {
	var mockPCI *MockDeviceHandler
	var fakePermittedHostDevicesConfig string
	var fakePermittedHostDevices v1.PermittedHostDevices
	var ctrl *gomock.Controller
	var fakeSupportedTypesPath string
	resourceNameToTypeName := func(rawName string) string {
		typeNameStr := strings.Replace(string(rawName), " ", "_", -1)
		typeNameStr = strings.TrimSpace(typeNameStr)
		return typeNameStr
	}
	BeforeEach(func() {
		By("creating a temporary fake mdev directory tree")
		// create base mdev dir instead of /sys/bus/mdev/devices
		fakeMdevBasePath, err := ioutil.TempDir("/tmp", "mdevs")
		Expect(err).ToNot(HaveOccurred())
		// create an alternative mdev_supported_types dir instead of /sys/class/mdev_bus/[pciAddress]/
		fakeSupportedTypesPath, err = ioutil.TempDir("/tmp", "mdev_supported_types")
		Expect(err).ToNot(HaveOccurred())
		// create a fake path to nvidia mdev type
		fakeNvidiaTypePath := filepath.Join(fakeSupportedTypesPath, "nvidia-222")
		err = os.MkdirAll(fakeNvidiaTypePath, 0700)
		Expect(err).ToNot(HaveOccurred())
		// create a fake path to Intel mdev type
		fakeIntelTypePath := filepath.Join(fakeSupportedTypesPath, fakeIntelMdevNameSelector)
		err = os.MkdirAll(fakeIntelTypePath, 0700)
		Expect(err).ToNot(HaveOccurred())
		mdevBasePath = fakeMdevBasePath
		// create mdev directories and symlinks
		for _, uuid := range []string{fakeMdevUUID, fakeIntelMdevUUID} {
			mdevTypePath := filepath.Join(fakeMdevBasePath, uuid+"real")
			err = os.MkdirAll(mdevTypePath, 0700)
			Expect(err).ToNot(HaveOccurred())
			err = os.Symlink(filepath.Join(fakeMdevBasePath, uuid+"real"), filepath.Join(fakeMdevBasePath, uuid))
			Expect(err).ToNot(HaveOccurred())
		}
		// link nvidia type directory
		err = os.Symlink(fakeNvidiaTypePath, filepath.Join(fakeMdevBasePath, fakeMdevUUID+"real", "mdev_type"))
		Expect(err).ToNot(HaveOccurred())
		err = os.Symlink(fakeIntelTypePath, filepath.Join(fakeMdevBasePath, fakeIntelMdevUUID+"real", "mdev_type"))
		Expect(err).ToNot(HaveOccurred())
		// create a name file in the nvidia type directory
		mdevName, err := os.Create(filepath.Join(fakeNvidiaTypePath, "name"))
		Expect(err).ToNot(HaveOccurred())
		mdevNameWriter := bufio.NewWriter(mdevName)
		n, err := mdevNameWriter.WriteString(fakeMdevNameSelector + "\n")
		Expect(err).ToNot(HaveOccurred())
		Expect(n).To(Equal(len(fakeMdevNameSelector) + 1))
		mdevNameWriter.Flush()

	})

	AfterEach(func() {
		os.RemoveAll(mdevBasePath)
		os.RemoveAll(fakeSupportedTypesPath)
	})
	table.DescribeTable("should get correct file type name", func(namePathExist bool) {
		if namePathExist {
			mdevName, err := getMdevTypeName(fakeMdevUUID)
			Expect(err).ToNot(HaveOccurred())
			Expect(mdevName).To(Equal(resourceNameToTypeName(fakeMdevNameSelector)))
		} else {
			mdevName, err := getMdevTypeName(fakeIntelMdevUUID)
			Expect(err).ToNot(HaveOccurred())
			Expect(mdevName).To(Equal(resourceNameToTypeName(fakeIntelMdevNameSelector)))
		}
	},
		table.Entry("Nvidia name file exist", true),
		table.Entry("Intel name file doesn't exist", false),
	)
	Context("discover devices", func() {
		BeforeEach(func() {
			By("mocking PCI and MDEV functions to simulate an mdev an its parent PCI device")
			ctrl = gomock.NewController(GinkgoT())
			mockPCI = NewMockDeviceHandler(ctrl)
			Handler = mockPCI
			// Force pre-defined returned values and ensure the function only get called exacly once each on 0000:00:00.0
			mockPCI.EXPECT().GetMdevParentPCIAddr(fakeMdevUUID).Return(fakeAddress, nil).Times(1)
			mockPCI.EXPECT().GetDeviceIOMMUGroup(mdevBasePath, fakeMdevUUID).Return(fakeIommuGroup, nil).Times(1)
			mockPCI.EXPECT().GetDeviceNumaNode(pciBasePath, fakeAddress).Return(fakeNumaNode).Times(1)

			By("creating a list of fake device using the yaml decoder")
			fakePermittedHostDevicesConfig = `
    mediatedDevices:
    - mdevNameSelector: "` + fakeMdevNameSelector + `"
      resourceName: "` + fakeMdevResourceName + `"
    `
			err := yaml.NewYAMLOrJSONDecoder(strings.NewReader(fakePermittedHostDevicesConfig), 1024).Decode(&fakePermittedHostDevices)
			Expect(err).ToNot(HaveOccurred())
			Expect(len(fakePermittedHostDevices.MediatedDevices)).To(Equal(1))
			Expect(fakePermittedHostDevices.MediatedDevices[0].MDEVNameSelector).To(Equal(fakeMdevNameSelector))
			Expect(fakePermittedHostDevices.MediatedDevices[0].ResourceName).To(Equal(fakeMdevResourceName))
		})
		AfterEach(func() {
			ctrl.Finish()
		})
		It("Should parse the permitted devices and find 1 matching mediated device", func() {
			supportedMdevsMap := make(map[string]string)
			for _, supportedMdev := range fakePermittedHostDevices.MediatedDevices {
				// do not add a device plugin for this resource if it's being provided via an external device plugin
				if !supportedMdev.ExternalResourceProvider {
					selector := removeSelectorSpaces(supportedMdev.MDEVNameSelector)
					supportedMdevsMap[selector] = supportedMdev.ResourceName
				}
			}
			// discoverPermittedHostMediatedDevices() will walk real mdev devices wherever the tests are running
			devices := discoverPermittedHostMediatedDevices(supportedMdevsMap)
			Expect(len(devices)).To(Equal(1))
			selector := removeSelectorSpaces(fakeMdevNameSelector)
			Expect(len(devices[selector])).To(Equal(1))
			Expect(devices[selector][0].UUID).To(Equal(fakeMdevUUID))
			Expect(devices[selector][0].typeName).To(Equal(selector))
			Expect(devices[selector][0].parentPciAddress).To(Equal(fakeAddress))
			Expect(devices[selector][0].iommuGroup).To(Equal(fakeIommuGroup))
			Expect(devices[selector][0].numaNode).To(Equal(fakeNumaNode))
		})

		It("Should validate DPI devices", func() {
			iommuToMDEVMap := make(map[string]string)
			supportedMdevsMap := make(map[string]string)
			for _, supportedMdev := range fakePermittedHostDevices.MediatedDevices {
				// do not add a device plugin for this resource if it's being provided via an external device plugin
				if !supportedMdev.ExternalResourceProvider {
					selector := removeSelectorSpaces(supportedMdev.MDEVNameSelector)
					supportedMdevsMap[selector] = supportedMdev.ResourceName
				}
			}
			// discoverPermittedHostMediatedDevices() will walk real mdev devices wherever the tests are running
			mDevices := discoverPermittedHostMediatedDevices(supportedMdevsMap)
			selector := removeSelectorSpaces(fakeMdevNameSelector)
			devs := constructDPIdevicesFromMdev(mDevices[selector], iommuToMDEVMap)
			Expect(devs[0].ID).To(Equal(fakeIommuGroup))
			Expect(devs[0].Topology.Nodes[0].ID).To(Equal(int64(fakeNumaNode)))
		})

		It("Should update the device list according to the configmap", func() {
			By("creating a cluster config")
			kv := &v1.KubeVirt{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "kubevirt",
					Namespace: "kubevirt",
				},
				Spec: v1.KubeVirtSpec{
					Configuration: v1.KubeVirtConfiguration{
						DeveloperConfiguration: &v1.DeveloperConfiguration{},
					},
				},
				Status: v1.KubeVirtStatus{
					Phase: v1.KubeVirtPhaseDeploying,
				},
			}
			fakeClusterConfig, _, _, kvInformer := testutils.NewFakeClusterConfigUsingKV(kv)

			By("creating an empty device controller")
			deviceController := NewDeviceController("master", 10, fakeClusterConfig)
			deviceController.devicePlugins = make(map[string]ControlledDevice)

			By("adding a host device to the cluster config")
			kvConfig := kv.DeepCopy()
			kvConfig.Spec.Configuration.DeveloperConfiguration.FeatureGates = []string{virtconfig.HostDevicesGate}
			kvConfig.Spec.Configuration.PermittedHostDevices = &v1.PermittedHostDevices{
				MediatedDevices: []v1.MediatedHostDevice{
					{
						MDEVNameSelector: fakeMdevNameSelector,
						ResourceName:     fakeMdevResourceName,
					},
				},
			}
			testutils.UpdateFakeKubeVirtClusterConfig(kvInformer, kvConfig)
			permittedDevices := fakeClusterConfig.GetPermittedHostDevices()
			Expect(permittedDevices).ToNot(BeNil(), "something went wrong while parsing the configmap(s)")
			Expect(len(permittedDevices.MediatedDevices)).To(Equal(1), "the fake device was not found")

			By("ensuring a device plugin gets created for our fake device")
			enabledDevicePlugins, disabledDevicePlugins := deviceController.updatePermittedHostDevicePlugins()
			Expect(len(enabledDevicePlugins)).To(Equal(1), "a device plugin wasn't created for the fake device")
			Expect(len(disabledDevicePlugins)).To(Equal(0))
			Ω(enabledDevicePlugins).Should(HaveKey(fakeMdevResourceName))
			// Manually adding the enabled plugin, since the device controller is not actually running
			deviceController.devicePlugins[fakeMdevResourceName] = enabledDevicePlugins[fakeMdevResourceName]

			By("deletting the device from the configmap")
			kvConfig.Spec.Configuration.PermittedHostDevices = &v1.PermittedHostDevices{}
			testutils.UpdateFakeKubeVirtClusterConfig(kvInformer, kvConfig)
			permittedDevices = fakeClusterConfig.GetPermittedHostDevices()
			Expect(permittedDevices).ToNot(BeNil(), "something went wrong while parsing the configmap(s)")
			Expect(len(permittedDevices.MediatedDevices)).To(Equal(0), "the fake device was not deleted")

			By("ensuring the device plugin gets stopped")
			enabledDevicePlugins, disabledDevicePlugins = deviceController.updatePermittedHostDevicePlugins()
			Expect(len(enabledDevicePlugins)).To(Equal(0))
			Expect(len(disabledDevicePlugins)).To(Equal(1), "the fake device plugin did not get disabled")
			Ω(disabledDevicePlugins).Should(HaveKey(fakeMdevResourceName))
		})
	})
})
<|endoftext|>


<commit_before>package server

import (
	"container/list"
	"sync"
	"time"

	"frp/models/consts"
	"frp/utils/conn"
	"frp/utils/log"
)

type ProxyServer struct {
	Name       string
	Passwd     string
	BindAddr   string
	ListenPort int64
	Status     int64

	listener     *conn.Listener  // accept new connection from remote users
	ctlMsgChan   chan int64      // every time accept a new user conn, put "1" to the channel
	cliConnChan  chan *conn.Conn // get client conns from control goroutine
	userConnList *list.List      // store user conns
	mutex        sync.Mutex
}

func (p *ProxyServer) Init() {
	p.Status = consts.Idle
	p.cliConnChan = make(chan *conn.Conn)
	p.ctlMsgChan = make(chan int64)
	p.userConnList = list.New()
}

func (p *ProxyServer) Lock() {
	p.mutex.Lock()
}

func (p *ProxyServer) Unlock() {
	p.mutex.Unlock()
}

// start listening for user conns
func (p *ProxyServer) Start() (err error) {
	p.Init()
	p.listener, err = conn.Listen(p.BindAddr, p.ListenPort)
	if err != nil {
		return err
	}

	p.Status = consts.Working

	// start a goroutine for listener to accept user connection
	go func() {
		for {
			// block
			// if listener is closed, err returned
			c, err := p.listener.GetConn()
			if err != nil {
				log.Info("ProxyName [%s], listener is closed", p.Name)
				return
			}
			log.Debug("ProxyName [%s], get one new user conn [%s]", p.Name, c.GetRemoteAddr())

			// insert into list
			p.Lock()
			if p.Status != consts.Working {
				log.Debug("ProxyName [%s] is not working, new user conn close", p.Name)
				c.Close()
				p.Unlock()
				return
			}
			p.userConnList.PushBack(c)
			p.Unlock()

			// put msg to control conn
			p.ctlMsgChan <- 1

			// set timeout
			time.AfterFunc(time.Duration(UserConnTimeout)*time.Second, func() {
				p.Lock()
				defer p.Unlock()
				element := p.userConnList.Front()
				if element == nil {
					return
				}

				userConn := element.Value.(*conn.Conn)
				if userConn == c {
					log.Warn("ProxyName [%s], user conn [%s] timeout", p.Name, c.GetRemoteAddr())
				}
			})
		}
	}()

	// start another goroutine for join two conns from client and user
	go func() {
		for {
			cliConn, ok := <-p.cliConnChan
			if !ok {
				return
			}

			p.Lock()
			element := p.userConnList.Front()

			var userConn *conn.Conn
			if element != nil {
				userConn = element.Value.(*conn.Conn)
				p.userConnList.Remove(element)
			} else {
				cliConn.Close()
				p.Unlock()
				continue
			}
			p.Unlock()

			// msg will transfer to another without modifying
			// l means local, r means remote
			log.Debug("Join two conns, (l[%s] r[%s]) (l[%s] r[%s])", cliConn.GetLocalAddr(), cliConn.GetRemoteAddr(),
				userConn.GetLocalAddr(), userConn.GetRemoteAddr())
			go conn.Join(cliConn, userConn)
		}
	}()

	return nil
}

func (p *ProxyServer) Close() {
	p.Lock()
	p.Status = consts.Idle
	p.listener.Close()
	close(p.ctlMsgChan)
	close(p.cliConnChan)
	p.userConnList = list.New()
	p.Unlock()
}

func (p *ProxyServer) WaitUserConn() (closeFlag bool) {
	closeFlag = false

	_, ok := <-p.ctlMsgChan
	if !ok {
		closeFlag = true
	}
	return
}

func (p *ProxyServer) GetNewCliConn(c *conn.Conn) {
	p.cliConnChan <- c
}
<commit_msg>models/server: fix bug, program will core if listener is nil<commit_after>package server

import (
	"container/list"
	"sync"
	"time"

	"frp/models/consts"
	"frp/utils/conn"
	"frp/utils/log"
)

type ProxyServer struct {
	Name       string
	Passwd     string
	BindAddr   string
	ListenPort int64
	Status     int64

	listener     *conn.Listener  // accept new connection from remote users
	ctlMsgChan   chan int64      // every time accept a new user conn, put "1" to the channel
	cliConnChan  chan *conn.Conn // get client conns from control goroutine
	userConnList *list.List      // store user conns
	mutex        sync.Mutex
}

func (p *ProxyServer) Init() {
	p.Status = consts.Idle
	p.cliConnChan = make(chan *conn.Conn)
	p.ctlMsgChan = make(chan int64)
	p.userConnList = list.New()
}

func (p *ProxyServer) Lock() {
	p.mutex.Lock()
}

func (p *ProxyServer) Unlock() {
	p.mutex.Unlock()
}

// start listening for user conns
func (p *ProxyServer) Start() (err error) {
	p.Init()
	p.listener, err = conn.Listen(p.BindAddr, p.ListenPort)
	if err != nil {
		return err
	}

	p.Status = consts.Working

	// start a goroutine for listener to accept user connection
	go func() {
		for {
			// block
			// if listener is closed, err returned
			c, err := p.listener.GetConn()
			if err != nil {
				log.Info("ProxyName [%s], listener is closed", p.Name)
				return
			}
			log.Debug("ProxyName [%s], get one new user conn [%s]", p.Name, c.GetRemoteAddr())

			// insert into list
			p.Lock()
			if p.Status != consts.Working {
				log.Debug("ProxyName [%s] is not working, new user conn close", p.Name)
				c.Close()
				p.Unlock()
				return
			}
			p.userConnList.PushBack(c)
			p.Unlock()

			// put msg to control conn
			p.ctlMsgChan <- 1

			// set timeout
			time.AfterFunc(time.Duration(UserConnTimeout)*time.Second, func() {
				p.Lock()
				defer p.Unlock()
				element := p.userConnList.Front()
				if element == nil {
					return
				}

				userConn := element.Value.(*conn.Conn)
				if userConn == c {
					log.Warn("ProxyName [%s], user conn [%s] timeout", p.Name, c.GetRemoteAddr())
				}
			})
		}
	}()

	// start another goroutine for join two conns from client and user
	go func() {
		for {
			cliConn, ok := <-p.cliConnChan
			if !ok {
				return
			}

			p.Lock()
			element := p.userConnList.Front()

			var userConn *conn.Conn
			if element != nil {
				userConn = element.Value.(*conn.Conn)
				p.userConnList.Remove(element)
			} else {
				cliConn.Close()
				p.Unlock()
				continue
			}
			p.Unlock()

			// msg will transfer to another without modifying
			// l means local, r means remote
			log.Debug("Join two conns, (l[%s] r[%s]) (l[%s] r[%s])", cliConn.GetLocalAddr(), cliConn.GetRemoteAddr(),
				userConn.GetLocalAddr(), userConn.GetRemoteAddr())
			go conn.Join(cliConn, userConn)
		}
	}()

	return nil
}

func (p *ProxyServer) Close() {
	p.Lock()
	p.Status = consts.Idle
	if p.listener != nil {
		p.listener.Close()
	}
	close(p.ctlMsgChan)
	close(p.cliConnChan)
	p.userConnList = list.New()
	p.Unlock()
}

func (p *ProxyServer) WaitUserConn() (closeFlag bool) {
	closeFlag = false

	_, ok := <-p.ctlMsgChan
	if !ok {
		closeFlag = true
	}
	return
}

func (p *ProxyServer) GetNewCliConn(c *conn.Conn) {
	p.cliConnChan <- c
}
<|endoftext|>


<commit_before><commit_msg>Clean up integration tests for main logic<commit_after><|endoftext|>


<commit_before>// Package token : Token definitions used when tokenizing program input
package token

// TokenType : String representation of each type of Token
type TokenType string

// TokenTypes
const (
	ILLEGAL = "ILLEGAL"
	EOF     = "EOF"

	// Literals
	IDENT  = "IDENT"
	INT    = "INT"
	FLOAT  = "FLOAT"
	STRING = "STRING"

	// Operators
	ASSIGN   = "="
	PLUS     = "+"
	MINUS    = "-"
	BANG     = "!"
	ASTERISK = "*"
	SLASH    = "/"
	PERCENT  = "%"

	LT  = "<"
	GT  = ">"
	LTE = "<="
	GTE = ">="

	EQ     = "=="
	NOT_EQ = "!="

	AND = "&&"
	OR  = "||"

	INCREMENT = "++"
	DECREMENT = "--"

	// Delimiters
	COMMA     = ","
	SEMICOLON = ";"
	COLON     = ":"

	LPAREN   = "("
	RPAREN   = ")"
	LBRACE   = "{"
	RBRACE   = "}"
	LBRACKET = "["
	RBRACKET = "]"

	// Keywords
	FUNCTION = "FUNCTION"
	LET      = "LET"
	TRUE     = "TRUE"
	FALSE    = "FALSE"
	IF       = "IF"
	ELSE     = "ELSE"
	RETURN   = "RETURN"
	WHILE    = "WHILE"
)

// Token : Defines the type and literal representation for the tokens to be used in program analysis
type Token struct {
	Type    TokenType
	Literal string
}

var keywords = map[string]TokenType{
	"fn":     FUNCTION,
	"let":    LET,
	"true":   TRUE,
	"false":  FALSE,
	"if":     IF,
	"else":   ELSE,
	"return": RETURN,
	"while":  WHILE,
}

// LookupIdent : Compares identifier input against list of keywords to return proper TokenType
func LookupIdent(ident string) TokenType {
	if tok, ok := keywords[ident]; ok {
		return tok
	}
	return IDENT
}
<commit_msg>add power token<commit_after>// Package token : Token definitions used when tokenizing program input
package token

// TokenType : String representation of each type of Token
type TokenType string

// TokenTypes
const (
	ILLEGAL = "ILLEGAL"
	EOF     = "EOF"

	// Literals
	IDENT  = "IDENT"
	INT    = "INT"
	FLOAT  = "FLOAT"
	STRING = "STRING"

	// Operators
	ASSIGN   = "="
	PLUS     = "+"
	MINUS    = "-"
	BANG     = "!"
	ASTERISK = "*"
	SLASH    = "/"
	PERCENT  = "%"
	POWER    = "**"

	LT  = "<"
	GT  = ">"
	LTE = "<="
	GTE = ">="

	EQ     = "=="
	NOT_EQ = "!="

	AND = "&&"
	OR  = "||"

	INCREMENT = "++"
	DECREMENT = "--"

	// Delimiters
	COMMA     = ","
	SEMICOLON = ";"
	COLON     = ":"

	LPAREN   = "("
	RPAREN   = ")"
	LBRACE   = "{"
	RBRACE   = "}"
	LBRACKET = "["
	RBRACKET = "]"

	// Keywords
	FUNCTION = "FUNCTION"
	LET      = "LET"
	TRUE     = "TRUE"
	FALSE    = "FALSE"
	IF       = "IF"
	ELSE     = "ELSE"
	RETURN   = "RETURN"
	WHILE    = "WHILE"
)

// Token : Defines the type and literal representation for the tokens to be used in program analysis
type Token struct {
	Type    TokenType
	Literal string
}

var keywords = map[string]TokenType{
	"fn":     FUNCTION,
	"let":    LET,
	"true":   TRUE,
	"false":  FALSE,
	"if":     IF,
	"else":   ELSE,
	"return": RETURN,
	"while":  WHILE,
}

// LookupIdent : Compares identifier input against list of keywords to return proper TokenType
func LookupIdent(ident string) TokenType {
	if tok, ok := keywords[ident]; ok {
		return tok
	}
	return IDENT
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Tigera, Inc. All rights reserved.

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

package ipam

import (
	"errors"
	"fmt"
	"math/big"
	"net"
	"reflect"

	"github.com/projectcalico/libcalico-go/lib/apis/v3"
	log "github.com/sirupsen/logrus"

	"github.com/projectcalico/libcalico-go/lib/backend/model"
	cnet "github.com/projectcalico/libcalico-go/lib/net"
)

// Wrap the backend AllocationBlock struct so that we can
// attach methods to it.
type allocationBlock struct {
	*model.AllocationBlock
}

func newBlock(cidr cnet.IPNet) allocationBlock {
	ones, size := cidr.Mask.Size()
	numAddresses := 2 << uint(size-ones-1)
	if numAddresses == 0 {
		// Shifting doesn't handle 2^0 properly. Override the value.
		numAddresses = 1
	}
	b := model.AllocationBlock{}
	b.Allocations = make([]*int, numAddresses)
	b.Unallocated = make([]int, numAddresses)
	b.StrictAffinity = false
	b.CIDR = cidr

	// Initialize unallocated ordinals.
	for i := 0; i < numAddresses; i++ {
		b.Unallocated[i] = i
	}

	return allocationBlock{&b}
}

func (b *allocationBlock) autoAssign(
	num int, handleID *string, host string, attrs map[string]string, affinityCheck bool) ([]cnet.IP, error) {

	// Determine if we need to check for affinity.
	checkAffinity := b.StrictAffinity || affinityCheck
	if checkAffinity && b.Affinity != nil && !hostAffinityMatches(host, b.AllocationBlock) {
		// Affinity check is enabled but the host does not match - error.
		s := fmt.Sprintf("Block affinity (%s) does not match provided (%s)", *b.Affinity, host)
		return nil, errors.New(s)
	} else if b.Affinity == nil {
		log.Warnf("Attempting to assign IPs from block with no affinity: %v", b)
		if checkAffinity {
			// If we're checking strict affinity, we can't assign from a block with no affinity.
			return nil, fmt.Errorf("Attempt to assign from block %v with no affinity", b.CIDR)
		}
	}

	// Walk the allocations until we find enough addresses.
	ordinals := []int{}
	for len(b.Unallocated) > 0 && len(ordinals) < num {
		ordinals = append(ordinals, b.Unallocated[0])
		b.Unallocated = b.Unallocated[1:]
	}

	// Create slice of IPs and perform the allocations.
	ips := []cnet.IP{}
	for _, o := range ordinals {
		attrIndex := b.findOrAddAttribute(handleID, attrs)
		b.Allocations[o] = &attrIndex
		ips = append(ips, incrementIP(cnet.IP{b.CIDR.IP}, big.NewInt(int64(o))))
	}

	log.Debugf("Block %s returned ips: %v", b.CIDR.String(), ips)
	return ips, nil
}

func (b *allocationBlock) assign(address cnet.IP, handleID *string, attrs map[string]string, host string) error {
	if b.StrictAffinity && b.Affinity != nil && !hostAffinityMatches(host, b.AllocationBlock) {
		// Affinity check is enabled but the host does not match - error.
		return errors.New("Block host affinity does not match")
	} else if b.Affinity == nil {
		log.Warnf("Attempting to assign IP from block with no affinity: %v", b)
		if b.StrictAffinity {
			// If we're checking strict affinity, we can't assign from a block with no affinity.
			return fmt.Errorf("Attempt to assign from block %v with no affinity", b.CIDR)
		}
	}

	// Convert to an ordinal.
	ordinal, err := ipToOrdinal(address, *b)
	if err != nil {
		return err
	}

	// Check if already allocated.
	if b.Allocations[ordinal] != nil {
		return errors.New("Address already assigned in block")
	}

	// Set up attributes.
	attrIndex := b.findOrAddAttribute(handleID, attrs)
	b.Allocations[ordinal] = &attrIndex

	// Remove from unallocated.
	for i, unallocated := range b.Unallocated {
		if unallocated == ordinal {
			b.Unallocated = append(b.Unallocated[:i], b.Unallocated[i+1:]...)
			break
		}
	}
	return nil
}

// hostAffinityMatches checks if the provided host matches the provided affinity.
func hostAffinityMatches(host string, block *model.AllocationBlock) bool {
	return *block.Affinity == "host:"+host
}

func (b allocationBlock) numFreeAddresses() int {
	return len(b.Unallocated)
}

func (b allocationBlock) empty() bool {
	return b.numFreeAddresses() == b.numAddresses()
}

func (b *allocationBlock) release(addresses []cnet.IP) ([]cnet.IP, map[string]int, error) {
	// Store return values.
	unallocated := []cnet.IP{}
	countByHandle := map[string]int{}

	// Used internally.
	var ordinals []int
	delRefCounts := map[int]int{}
	attrsToDelete := []int{}

	// Determine the ordinals that need to be released and the
	// attributes that need to be cleaned up.
	for _, ip := range addresses {
		// Convert to an ordinal.
		ordinal, err := ipToOrdinal(ip, *b)
		if err != nil {
			return nil, nil, err
		}

		// Check if allocated.
		attrIdx := b.Allocations[ordinal]
		if attrIdx == nil {
			log.Debugf("Asked to release address that was not allocated")
			unallocated = append(unallocated, ip)
			continue
		}
		ordinals = append(ordinals, ordinal)

		// Increment referece counting for attributes.
		cnt := 1
		if cur, exists := delRefCounts[*attrIdx]; exists {
			cnt = cur + 1
		}
		delRefCounts[*attrIdx] = cnt

		// Increment count of addresses by handle if a handle
		// exists.
		handleID := b.Attributes[*attrIdx].AttrPrimary
		if handleID != nil {
			handleCount := 0
			if count, ok := countByHandle[*handleID]; !ok {
				handleCount = count
			}
			handleCount += 1
			countByHandle[*handleID] = handleCount
		}
	}

	// Handle cleaning up of attributes.  We do this by
	// reference counting.  If we're deleting the last reference to
	// a given attribute, then it needs to be cleaned up.
	refCounts := b.attributeRefCounts()
	for idx, refs := range delRefCounts {
		if refCounts[idx] == refs {
			attrsToDelete = append(attrsToDelete, idx)
		}
	}
	if len(attrsToDelete) != 0 {
		log.Debugf("Deleting attributes: %v", attrsToDelete)
		b.deleteAttributes(attrsToDelete, ordinals)
	}

	// Release requested addresses.
	for _, ordinal := range ordinals {
		b.Allocations[ordinal] = nil
		b.Unallocated = append(b.Unallocated, ordinal)
	}
	return unallocated, countByHandle, nil
}

func (b *allocationBlock) deleteAttributes(delIndexes, ordinals []int) {
	newIndexes := make([]*int, len(b.Attributes))
	newAttrs := []model.AllocationAttribute{}
	y := 0 // Next free slot in the new attributes list.
	for x := range b.Attributes {
		if !intInSlice(x, delIndexes) {
			// Attribute at x is not being deleted.  Build a mapping
			// of old attribute index (x) to new attribute index (y).
			log.Debugf("%d in %v", x, delIndexes)
			newIndex := y
			newIndexes[x] = &newIndex
			y += 1
			newAttrs = append(newAttrs, b.Attributes[x])
		}
	}
	b.Attributes = newAttrs

	// Update attribute indexes for all allocations in this block.
	for i := 0; i < b.numAddresses(); i++ {
		if b.Allocations[i] != nil {
			// Get the new index that corresponds to the old index
			// and update the allocation.
			newIndex := newIndexes[*b.Allocations[i]]
			b.Allocations[i] = newIndex
		}
	}
}

func (b allocationBlock) attributeRefCounts() map[int]int {
	refCounts := map[int]int{}
	for _, a := range b.Allocations {
		if a == nil {
			continue
		}

		if count, ok := refCounts[*a]; !ok {
			// No entry for given attribute index.
			refCounts[*a] = 1
		} else {
			refCounts[*a] = count + 1
		}
	}
	return refCounts
}

func (b allocationBlock) attributeIndexesByHandle(handleID string) []int {
	indexes := []int{}
	for i, attr := range b.Attributes {
		if attr.AttrPrimary != nil && *attr.AttrPrimary == handleID {
			indexes = append(indexes, i)
		}
	}
	return indexes
}

func (b *allocationBlock) releaseByHandle(handleID string) int {
	attrIndexes := b.attributeIndexesByHandle(handleID)
	log.Debugf("Attribute indexes to release: %v", attrIndexes)
	if len(attrIndexes) == 0 {
		// Nothing to release.
		log.Debugf("No addresses assigned to handle '%s'", handleID)
		return 0
	}

	// There are addresses to release.
	ordinals := []int{}
	var o int
	for o = 0; o < b.numAddresses(); o++ {
		// Only check allocated ordinals.
		if b.Allocations[o] != nil && intInSlice(*b.Allocations[o], attrIndexes) {
			// Release this ordinal.
			ordinals = append(ordinals, o)
		}
	}

	// Clean and reorder attributes.
	b.deleteAttributes(attrIndexes, ordinals)

	// Release the addresses.
	for _, o := range ordinals {
		b.Allocations[o] = nil
		b.Unallocated = append(b.Unallocated, o)
	}
	return len(ordinals)
}

func (b allocationBlock) ipsByHandle(handleID string) []cnet.IP {
	ips := []cnet.IP{}
	attrIndexes := b.attributeIndexesByHandle(handleID)
	var o int
	for o = 0; o < b.numAddresses(); o++ {
		if b.Allocations[o] != nil && intInSlice(*b.Allocations[o], attrIndexes) {
			ip := ordinalToIP(o, b)
			ips = append(ips, ip)
		}
	}
	return ips
}

func (b allocationBlock) attributesForIP(ip cnet.IP) (map[string]string, error) {
	// Convert to an ordinal.
	ordinal, err := ipToOrdinal(ip, b)
	if err != nil {
		return nil, err
	}

	// Check if allocated.
	attrIndex := b.Allocations[ordinal]
	if attrIndex == nil {
		return nil, errors.New(fmt.Sprintf("IP %s is not currently assigned in block", ip))
	}
	return b.Attributes[*attrIndex].AttrSecondary, nil
}

func (b *allocationBlock) findOrAddAttribute(handleID *string, attrs map[string]string) int {
	logCtx := log.WithField("attrs", attrs)
	if handleID != nil {
		logCtx = log.WithField("handle", *handleID)
	}
	attr := model.AllocationAttribute{handleID, attrs}
	for idx, existing := range b.Attributes {
		if reflect.DeepEqual(attr, existing) {
			log.Debugf("Attribute '%+v' already exists", attr)
			return idx
		}
	}

	// Does not exist - add it.
	logCtx.Debugf("New allocation attribute: %#v", attr)
	attrIndex := len(b.Attributes)
	b.Attributes = append(b.Attributes, attr)
	return attrIndex
}

// Get number of addresses covered by the block
func (b allocationBlock) numAddresses() int {
	ones, size := b.CIDR.Mask.Size()
	numAddresses := 2 << uint(size-ones-1)
	return numAddresses
}

func getBlockCIDRForAddress(addr cnet.IP, pool *v3.IPPool) cnet.IPNet {
	var mask net.IPMask
	if addr.Version() == 6 {
		// This is an IPv6 address.
		mask = net.CIDRMask(pool.Spec.BlockSize, 128)
	} else {
		// This is an IPv4 address.
		mask = net.CIDRMask(pool.Spec.BlockSize, 32)
	}
	masked := addr.Mask(mask)
	return cnet.IPNet{IPNet: net.IPNet{IP: masked, Mask: mask}}
}

func getIPVersion(ip cnet.IP) int {
	if ip.To4() == nil {
		return 6
	}
	return 4
}

func largerThanOrEqualToBlock(blockCIDR cnet.IPNet, pool *v3.IPPool) bool {
	ones, _ := blockCIDR.Mask.Size()
	return ones <= pool.Spec.BlockSize
}

func intInSlice(searchInt int, slice []int) bool {
	for _, v := range slice {
		if v == searchInt {
			return true
		}
	}
	return false
}

func ipToInt(ip cnet.IP) *big.Int {
	if ip.To4() != nil {
		return big.NewInt(0).SetBytes(ip.To4())
	} else {
		return big.NewInt(0).SetBytes(ip.To16())
	}
}

func intToIP(ipInt *big.Int) cnet.IP {
	ip := cnet.IP{net.IP(ipInt.Bytes())}
	return ip
}

func incrementIP(ip cnet.IP, increment *big.Int) cnet.IP {
	sum := big.NewInt(0).Add(ipToInt(ip), increment)
	return intToIP(sum)
}

func ipToOrdinal(ip cnet.IP, b allocationBlock) (int, error) {
	ip_int := ipToInt(ip)
	base_int := ipToInt(cnet.IP{b.CIDR.IP})
	ord := big.NewInt(0).Sub(ip_int, base_int).Int64()
	if ord < 0 || ord >= int64(b.numAddresses()) {
		return 0, fmt.Errorf("IP %s not in block %s", ip, b.CIDR)
	}
	return int(ord), nil
}

func ordinalToIP(ord int, b allocationBlock) cnet.IP {
	sum := big.NewInt(0).Add(ipToInt(cnet.IP{b.CIDR.IP}), big.NewInt(int64(ord)))
	return intToIP(sum)
}
<commit_msg>Fix shifts in block size calculation.<commit_after>// Copyright (c) 2016-2018 Tigera, Inc. All rights reserved.

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

package ipam

import (
	"errors"
	"fmt"
	"math/big"
	"net"
	"reflect"

	"github.com/projectcalico/libcalico-go/lib/apis/v3"
	log "github.com/sirupsen/logrus"

	"github.com/projectcalico/libcalico-go/lib/backend/model"
	cnet "github.com/projectcalico/libcalico-go/lib/net"
)

// Wrap the backend AllocationBlock struct so that we can
// attach methods to it.
type allocationBlock struct {
	*model.AllocationBlock
}

func newBlock(cidr cnet.IPNet) allocationBlock {
	ones, size := cidr.Mask.Size()
	numAddresses := 1 << uint(size-ones)
	b := model.AllocationBlock{}
	b.Allocations = make([]*int, numAddresses)
	b.Unallocated = make([]int, numAddresses)
	b.StrictAffinity = false
	b.CIDR = cidr

	// Initialize unallocated ordinals.
	for i := 0; i < numAddresses; i++ {
		b.Unallocated[i] = i
	}

	return allocationBlock{&b}
}

func (b *allocationBlock) autoAssign(
	num int, handleID *string, host string, attrs map[string]string, affinityCheck bool) ([]cnet.IP, error) {

	// Determine if we need to check for affinity.
	checkAffinity := b.StrictAffinity || affinityCheck
	if checkAffinity && b.Affinity != nil && !hostAffinityMatches(host, b.AllocationBlock) {
		// Affinity check is enabled but the host does not match - error.
		s := fmt.Sprintf("Block affinity (%s) does not match provided (%s)", *b.Affinity, host)
		return nil, errors.New(s)
	} else if b.Affinity == nil {
		log.Warnf("Attempting to assign IPs from block with no affinity: %v", b)
		if checkAffinity {
			// If we're checking strict affinity, we can't assign from a block with no affinity.
			return nil, fmt.Errorf("Attempt to assign from block %v with no affinity", b.CIDR)
		}
	}

	// Walk the allocations until we find enough addresses.
	ordinals := []int{}
	for len(b.Unallocated) > 0 && len(ordinals) < num {
		ordinals = append(ordinals, b.Unallocated[0])
		b.Unallocated = b.Unallocated[1:]
	}

	// Create slice of IPs and perform the allocations.
	ips := []cnet.IP{}
	for _, o := range ordinals {
		attrIndex := b.findOrAddAttribute(handleID, attrs)
		b.Allocations[o] = &attrIndex
		ips = append(ips, incrementIP(cnet.IP{b.CIDR.IP}, big.NewInt(int64(o))))
	}

	log.Debugf("Block %s returned ips: %v", b.CIDR.String(), ips)
	return ips, nil
}

func (b *allocationBlock) assign(address cnet.IP, handleID *string, attrs map[string]string, host string) error {
	if b.StrictAffinity && b.Affinity != nil && !hostAffinityMatches(host, b.AllocationBlock) {
		// Affinity check is enabled but the host does not match - error.
		return errors.New("Block host affinity does not match")
	} else if b.Affinity == nil {
		log.Warnf("Attempting to assign IP from block with no affinity: %v", b)
		if b.StrictAffinity {
			// If we're checking strict affinity, we can't assign from a block with no affinity.
			return fmt.Errorf("Attempt to assign from block %v with no affinity", b.CIDR)
		}
	}

	// Convert to an ordinal.
	ordinal, err := ipToOrdinal(address, *b)
	if err != nil {
		return err
	}

	// Check if already allocated.
	if b.Allocations[ordinal] != nil {
		return errors.New("Address already assigned in block")
	}

	// Set up attributes.
	attrIndex := b.findOrAddAttribute(handleID, attrs)
	b.Allocations[ordinal] = &attrIndex

	// Remove from unallocated.
	for i, unallocated := range b.Unallocated {
		if unallocated == ordinal {
			b.Unallocated = append(b.Unallocated[:i], b.Unallocated[i+1:]...)
			break
		}
	}
	return nil
}

// hostAffinityMatches checks if the provided host matches the provided affinity.
func hostAffinityMatches(host string, block *model.AllocationBlock) bool {
	return *block.Affinity == "host:"+host
}

func (b allocationBlock) numFreeAddresses() int {
	return len(b.Unallocated)
}

func (b allocationBlock) empty() bool {
	return b.numFreeAddresses() == b.numAddresses()
}

func (b *allocationBlock) release(addresses []cnet.IP) ([]cnet.IP, map[string]int, error) {
	// Store return values.
	unallocated := []cnet.IP{}
	countByHandle := map[string]int{}

	// Used internally.
	var ordinals []int
	delRefCounts := map[int]int{}
	attrsToDelete := []int{}

	// Determine the ordinals that need to be released and the
	// attributes that need to be cleaned up.
	for _, ip := range addresses {
		// Convert to an ordinal.
		ordinal, err := ipToOrdinal(ip, *b)
		if err != nil {
			return nil, nil, err
		}

		// Check if allocated.
		attrIdx := b.Allocations[ordinal]
		if attrIdx == nil {
			log.Debugf("Asked to release address that was not allocated")
			unallocated = append(unallocated, ip)
			continue
		}
		ordinals = append(ordinals, ordinal)

		// Increment referece counting for attributes.
		cnt := 1
		if cur, exists := delRefCounts[*attrIdx]; exists {
			cnt = cur + 1
		}
		delRefCounts[*attrIdx] = cnt

		// Increment count of addresses by handle if a handle
		// exists.
		handleID := b.Attributes[*attrIdx].AttrPrimary
		if handleID != nil {
			handleCount := 0
			if count, ok := countByHandle[*handleID]; !ok {
				handleCount = count
			}
			handleCount += 1
			countByHandle[*handleID] = handleCount
		}
	}

	// Handle cleaning up of attributes.  We do this by
	// reference counting.  If we're deleting the last reference to
	// a given attribute, then it needs to be cleaned up.
	refCounts := b.attributeRefCounts()
	for idx, refs := range delRefCounts {
		if refCounts[idx] == refs {
			attrsToDelete = append(attrsToDelete, idx)
		}
	}
	if len(attrsToDelete) != 0 {
		log.Debugf("Deleting attributes: %v", attrsToDelete)
		b.deleteAttributes(attrsToDelete, ordinals)
	}

	// Release requested addresses.
	for _, ordinal := range ordinals {
		b.Allocations[ordinal] = nil
		b.Unallocated = append(b.Unallocated, ordinal)
	}
	return unallocated, countByHandle, nil
}

func (b *allocationBlock) deleteAttributes(delIndexes, ordinals []int) {
	newIndexes := make([]*int, len(b.Attributes))
	newAttrs := []model.AllocationAttribute{}
	y := 0 // Next free slot in the new attributes list.
	for x := range b.Attributes {
		if !intInSlice(x, delIndexes) {
			// Attribute at x is not being deleted.  Build a mapping
			// of old attribute index (x) to new attribute index (y).
			log.Debugf("%d in %v", x, delIndexes)
			newIndex := y
			newIndexes[x] = &newIndex
			y += 1
			newAttrs = append(newAttrs, b.Attributes[x])
		}
	}
	b.Attributes = newAttrs

	// Update attribute indexes for all allocations in this block.
	for i := 0; i < b.numAddresses(); i++ {
		if b.Allocations[i] != nil {
			// Get the new index that corresponds to the old index
			// and update the allocation.
			newIndex := newIndexes[*b.Allocations[i]]
			b.Allocations[i] = newIndex
		}
	}
}

func (b allocationBlock) attributeRefCounts() map[int]int {
	refCounts := map[int]int{}
	for _, a := range b.Allocations {
		if a == nil {
			continue
		}

		if count, ok := refCounts[*a]; !ok {
			// No entry for given attribute index.
			refCounts[*a] = 1
		} else {
			refCounts[*a] = count + 1
		}
	}
	return refCounts
}

func (b allocationBlock) attributeIndexesByHandle(handleID string) []int {
	indexes := []int{}
	for i, attr := range b.Attributes {
		if attr.AttrPrimary != nil && *attr.AttrPrimary == handleID {
			indexes = append(indexes, i)
		}
	}
	return indexes
}

func (b *allocationBlock) releaseByHandle(handleID string) int {
	attrIndexes := b.attributeIndexesByHandle(handleID)
	log.Debugf("Attribute indexes to release: %v", attrIndexes)
	if len(attrIndexes) == 0 {
		// Nothing to release.
		log.Debugf("No addresses assigned to handle '%s'", handleID)
		return 0
	}

	// There are addresses to release.
	ordinals := []int{}
	var o int
	for o = 0; o < b.numAddresses(); o++ {
		// Only check allocated ordinals.
		if b.Allocations[o] != nil && intInSlice(*b.Allocations[o], attrIndexes) {
			// Release this ordinal.
			ordinals = append(ordinals, o)
		}
	}

	// Clean and reorder attributes.
	b.deleteAttributes(attrIndexes, ordinals)

	// Release the addresses.
	for _, o := range ordinals {
		b.Allocations[o] = nil
		b.Unallocated = append(b.Unallocated, o)
	}
	return len(ordinals)
}

func (b allocationBlock) ipsByHandle(handleID string) []cnet.IP {
	ips := []cnet.IP{}
	attrIndexes := b.attributeIndexesByHandle(handleID)
	var o int
	for o = 0; o < b.numAddresses(); o++ {
		if b.Allocations[o] != nil && intInSlice(*b.Allocations[o], attrIndexes) {
			ip := ordinalToIP(o, b)
			ips = append(ips, ip)
		}
	}
	return ips
}

func (b allocationBlock) attributesForIP(ip cnet.IP) (map[string]string, error) {
	// Convert to an ordinal.
	ordinal, err := ipToOrdinal(ip, b)
	if err != nil {
		return nil, err
	}

	// Check if allocated.
	attrIndex := b.Allocations[ordinal]
	if attrIndex == nil {
		return nil, errors.New(fmt.Sprintf("IP %s is not currently assigned in block", ip))
	}
	return b.Attributes[*attrIndex].AttrSecondary, nil
}

func (b *allocationBlock) findOrAddAttribute(handleID *string, attrs map[string]string) int {
	logCtx := log.WithField("attrs", attrs)
	if handleID != nil {
		logCtx = log.WithField("handle", *handleID)
	}
	attr := model.AllocationAttribute{handleID, attrs}
	for idx, existing := range b.Attributes {
		if reflect.DeepEqual(attr, existing) {
			log.Debugf("Attribute '%+v' already exists", attr)
			return idx
		}
	}

	// Does not exist - add it.
	logCtx.Debugf("New allocation attribute: %#v", attr)
	attrIndex := len(b.Attributes)
	b.Attributes = append(b.Attributes, attr)
	return attrIndex
}

// Get number of addresses covered by the block
func (b allocationBlock) numAddresses() int {
	ones, size := b.CIDR.Mask.Size()
	numAddresses := 1 << uint(size-ones)
	return numAddresses
}

func getBlockCIDRForAddress(addr cnet.IP, pool *v3.IPPool) cnet.IPNet {
	var mask net.IPMask
	if addr.Version() == 6 {
		// This is an IPv6 address.
		mask = net.CIDRMask(pool.Spec.BlockSize, 128)
	} else {
		// This is an IPv4 address.
		mask = net.CIDRMask(pool.Spec.BlockSize, 32)
	}
	masked := addr.Mask(mask)
	return cnet.IPNet{IPNet: net.IPNet{IP: masked, Mask: mask}}
}

func getIPVersion(ip cnet.IP) int {
	if ip.To4() == nil {
		return 6
	}
	return 4
}

func largerThanOrEqualToBlock(blockCIDR cnet.IPNet, pool *v3.IPPool) bool {
	ones, _ := blockCIDR.Mask.Size()
	return ones <= pool.Spec.BlockSize
}

func intInSlice(searchInt int, slice []int) bool {
	for _, v := range slice {
		if v == searchInt {
			return true
		}
	}
	return false
}

func ipToInt(ip cnet.IP) *big.Int {
	if ip.To4() != nil {
		return big.NewInt(0).SetBytes(ip.To4())
	} else {
		return big.NewInt(0).SetBytes(ip.To16())
	}
}

func intToIP(ipInt *big.Int) cnet.IP {
	ip := cnet.IP{net.IP(ipInt.Bytes())}
	return ip
}

func incrementIP(ip cnet.IP, increment *big.Int) cnet.IP {
	sum := big.NewInt(0).Add(ipToInt(ip), increment)
	return intToIP(sum)
}

func ipToOrdinal(ip cnet.IP, b allocationBlock) (int, error) {
	ip_int := ipToInt(ip)
	base_int := ipToInt(cnet.IP{b.CIDR.IP})
	ord := big.NewInt(0).Sub(ip_int, base_int).Int64()
	if ord < 0 || ord >= int64(b.numAddresses()) {
		return 0, fmt.Errorf("IP %s not in block %s", ip, b.CIDR)
	}
	return int(ord), nil
}

func ordinalToIP(ord int, b allocationBlock) cnet.IP {
	sum := big.NewInt(0).Add(ipToInt(cnet.IP{b.CIDR.IP}), big.NewInt(int64(ord)))
	return intToIP(sum)
}
<|endoftext|>


<commit_before>package certificate_test

import (
	"testing"

	"github.com/RobotsAndPencils/buford/certificate"
)

func TestValidCert(t *testing.T) {
	// TODO: figure out how to test certificate loading and validation in CI
	const name = "../fixtures/cert.p12"

	_, _, err := certificate.Load(name, "")
	if err != nil {
		t.Fatal(err)
	}
}

func TestExpiredCert(t *testing.T) {
	// TODO: figure out how to test certificate loading and validation in CI
	const name = "../cert-expired.p12"

	_, _, err := certificate.Load(name, "")
	if err != certificate.ErrExpired {
		t.Fatal("Expected expired cert error, got", err)
	}
}

func TestMissingFile(t *testing.T) {
	_, _, err := certificate.Load("hide-and-seek.p12", "")
	if err == nil {
		t.Fatal("Expected file not found, got", err)
	}
}
<commit_msg>using a self-signed cert for testing<commit_after>package certificate_test

import (
	"testing"

	"github.com/RobotsAndPencils/buford/certificate"
)

func TestValidCert(t *testing.T) {
	const name = "../fixtures/cert.p12"

	_, _, err := certificate.Load(name, "")
	if err != nil {
		t.Fatal(err)
	}
}

func TestExpiredCert(t *testing.T) {
	// TODO: figure out how to test certificate loading and validation in CI
	const name = "../cert-expired.p12"

	_, _, err := certificate.Load(name, "")
	if err != certificate.ErrExpired {
		t.Fatal("Expected expired cert error, got", err)
	}
}

func TestMissingFile(t *testing.T) {
	_, _, err := certificate.Load("hide-and-seek.p12", "")
	if err == nil {
		t.Fatal("Expected file not found, got", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"github.com/dchest/uniuri"
	"log"
	"net"
	"sync"
	"github.com/sysr-q/kyubu/packets"
	"time"
)

type Kurafuto struct {
	Players []*Player
	mutex   sync.Mutex

	salt string
	Name string
	Motd string

	Hub    *Server
	Config *Config

	Listener net.Listener
	Done     chan bool
	Running  bool

	rMut sync.Mutex
}

func (ku *Kurafuto) Quit() {
	ku.rMut.Lock()
	if !ku.Running {
		ku.rMut.Unlock()
		return
	}

	ku.Running = false
	ku.rMut.Unlock()

	// So we don't take on any new players.
	ku.Listener.Close()

	for _, p := range ku.Players {
		disc, _ := packets.NewDisconnectPlayer("Server shutting down.")
		p.toClient <- disc
		p.Quit()
	}

	go func() {
		// TODO: `while len(ku.Players) > 0 {}`?
		time.Sleep(2 * time.Second)
		ku.Done <- true
	}()
}

func (ku *Kurafuto) Run() {
	ku.rMut.Lock()
	ku.Running = true
	ku.rMut.Unlock()

	for {
		ku.rMut.Lock()
		if !ku.Running {
			ku.rMut.Unlock()
			break
		}
		ku.rMut.Unlock()

		c, err := ku.Listener.Accept()
		if err != nil && !ku.Running {
			break
		} else if err != nil {
			log.Fatal(err)
		}

		p, err := NewPlayer(c, ku)
		if err != nil {
			c.Close()
			continue
		}
		ku.Players = append(ku.Players, p)

		Infof("New connection from %s (%d clients)", c.RemoteAddr().String(), len(ku.Players))
		Debugf("(%s) New connection from %s", p.Id, c.RemoteAddr().String())

		go p.Parse()
	}
}

func (ku *Kurafuto) Remove(p *Player) bool {
	ku.mutex.Lock()
	defer ku.mutex.Unlock()
	for i, player := range ku.Players {
		if player != p {
			continue
		}
		p.Quit() // just in case
		// Remove and zero player to allow GC to collect it.
		copy(ku.Players[i:], ku.Players[i+1:])
		ku.Players[len(ku.Players)-1] = nil
		ku.Players = ku.Players[:len(ku.Players)-1]
		f := "%s (%s) disconnected"
		if p.Name == "" {
			f = "%s(%s) disconnected"
		}
		Infof(f, p.Name, p.Client.RemoteAddr().String())
		Debugf("(%s) %s disconnected from slot %d", p.Id, p.Client.RemoteAddr().String(), i)
		return true
	}
	return false
}

func NewKurafuto(config *Config) (ku *Kurafuto, err error) {
	if len(config.Servers) < 1 {
		err = errors.New("kurafuto: Need at least 1 server in config.")
		return
	}

	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", config.Address, config.Port))
	if err != nil {
		return
	}

	ku = &Kurafuto{
		Players:  []*Player{},
		mutex:    sync.Mutex{},
		salt:     uniuri.New(),
		Hub:      &config.Servers[0],
		Config:   config,
		Listener: listener,
		Done:     make(chan bool, 1),

		rMut: sync.Mutex{},
	}
	return
}
<commit_msg>Done after all players are gone<commit_after>package main

import (
	"errors"
	"fmt"
	"github.com/dchest/uniuri"
	"log"
	"net"
	"sync"
	"github.com/sysr-q/kyubu/packets"
)

type Kurafuto struct {
	Players []*Player
	mutex   sync.Mutex

	salt string
	Name string
	Motd string

	Hub    *Server
	Config *Config

	Listener net.Listener
	Done     chan bool
	Running  bool

	rMut sync.Mutex
}

func (ku *Kurafuto) Quit() {
	ku.rMut.Lock()
	if !ku.Running {
		ku.rMut.Unlock()
		return
	}

	ku.Running = false
	ku.rMut.Unlock()

	// So we don't take on any new players.
	ku.Listener.Close()
	for len(ku.Players) <= 0 {
		for _, p := range ku.Players {
			disc, _ := packets.NewDisconnectPlayer("Server shutting down.")
			p.toClient <- disc
			p.Quit()
		}
	}
	ku.Done <- true
}

func (ku *Kurafuto) Run() {
	ku.rMut.Lock()
	ku.Running = true
	ku.rMut.Unlock()

	for {
		ku.rMut.Lock()
		if !ku.Running {
			ku.rMut.Unlock()
			break
		}
		ku.rMut.Unlock()

		c, err := ku.Listener.Accept()
		if err != nil && !ku.Running {
			break
		} else if err != nil {
			log.Fatal(err)
		}

		p, err := NewPlayer(c, ku)
		if err != nil {
			c.Close()
			continue
		}
		ku.Players = append(ku.Players, p)

		Infof("New connection from %s (%d clients)", c.RemoteAddr().String(), len(ku.Players))
		Debugf("(%s) New connection from %s", p.Id, c.RemoteAddr().String())

		go p.Parse()
	}
}

func (ku *Kurafuto) Remove(p *Player) bool {
	ku.mutex.Lock()
	defer ku.mutex.Unlock()
	for i, player := range ku.Players {
		if player != p {
			continue
		}
		p.Quit() // just in case
		// Remove and zero player to allow GC to collect it.
		copy(ku.Players[i:], ku.Players[i+1:])
		ku.Players[len(ku.Players)-1] = nil
		ku.Players = ku.Players[:len(ku.Players)-1]
		f := "%s (%s) disconnected"
		if p.Name == "" {
			f = "%s(%s) disconnected"
		}
		Infof(f, p.Name, p.Client.RemoteAddr().String())
		Debugf("(%s) %s disconnected from slot %d", p.Id, p.Client.RemoteAddr().String(), i)
		return true
	}
	return false
}

func NewKurafuto(config *Config) (ku *Kurafuto, err error) {
	if len(config.Servers) < 1 {
		err = errors.New("kurafuto: Need at least 1 server in config.")
		return
	}

	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", config.Address, config.Port))
	if err != nil {
		return
	}

	ku = &Kurafuto{
		Players:  []*Player{},
		mutex:    sync.Mutex{},
		salt:     uniuri.New(),
		Hub:      &config.Servers[0],
		Config:   config,
		Listener: listener,
		Done:     make(chan bool, 1),

		rMut: sync.Mutex{},
	}
	return
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 s3 implements a key/value store in an Amazon S3 bucket.
package s3

import (
	"fmt"
	"github.com/jacobsa/aws/s3"
	"github.com/jacobsa/comeback/kv"
	"sync"
)

// Create a key/value store that stores data in the supplied S3 bucket. Keys
// supplied to its methods must be valid S3 keys. It is assumed that no keys in
// the bucket are ever removed.
//
// This function blocks while listing keys in the bucket.
func NewS3KvStore(bucket s3.Bucket) (kv.Store, error) {
	// List the keys in the bucket.
	keys, err := getAllKeys(bucket)
	if err != nil {
		return nil, err
	}

	// Create an appropriate map for efficient lookups.
	keyMap := make(map[string]bool)
	for _, key := range keys {
		keyMap[key] = true
	}

	store := &kvStore{
		bucket: bucket,
		knownKeys: keyMap,
	}

	return store, nil
}

func getAllKeys(bucket s3.Bucket) ([]string, error) {
	keys := []string{}
	for {
		var prevKey string
		if len(keys) > 0 {
			prevKey = keys[len(keys)-1]
		}

		partialKeys, err := bucket.ListKeys(prevKey)
		if err != nil {
			return nil, fmt.Errorf("ListKeys: %v", err)
		}

		if len(partialKeys) == 0 {
			break
		}

		keys = append(keys, partialKeys...)
	}

	return keys, nil
}

type kvStore struct {
	bucket s3.Bucket

	mutex sync.RWMutex
	knownKeys map[string]bool  // Protected by mutex
}

func (s *kvStore) Set(key []byte, val []byte) error {
	// Call the bucket.
	if err := s.bucket.StoreObject(string(key), val); err != nil {
		return fmt.Errorf("StoreObject: %v", err)
	}

	// Record the fact that the key is now known.
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.knownKeys[string(key)] = true

	return nil
}

func (s *kvStore) Get(key []byte) (val []byte, err error) {
	return nil, fmt.Errorf("TODO")
}

func (s *kvStore) Contains(key []byte) (res bool, err error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	_, ok := s.knownKeys[string(key)]
	return ok, nil
}
<commit_msg>Implemented Get.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 s3 implements a key/value store in an Amazon S3 bucket.
package s3

import (
	"fmt"
	"github.com/jacobsa/aws/s3"
	"github.com/jacobsa/comeback/kv"
	"sync"
)

// Create a key/value store that stores data in the supplied S3 bucket. Keys
// supplied to its methods must be valid S3 keys. It is assumed that no keys in
// the bucket are ever removed.
//
// This function blocks while listing keys in the bucket.
func NewS3KvStore(bucket s3.Bucket) (kv.Store, error) {
	// List the keys in the bucket.
	keys, err := getAllKeys(bucket)
	if err != nil {
		return nil, err
	}

	// Create an appropriate map for efficient lookups.
	keyMap := make(map[string]bool)
	for _, key := range keys {
		keyMap[key] = true
	}

	store := &kvStore{
		bucket: bucket,
		knownKeys: keyMap,
	}

	return store, nil
}

func getAllKeys(bucket s3.Bucket) ([]string, error) {
	keys := []string{}
	for {
		var prevKey string
		if len(keys) > 0 {
			prevKey = keys[len(keys)-1]
		}

		partialKeys, err := bucket.ListKeys(prevKey)
		if err != nil {
			return nil, fmt.Errorf("ListKeys: %v", err)
		}

		if len(partialKeys) == 0 {
			break
		}

		keys = append(keys, partialKeys...)
	}

	return keys, nil
}

type kvStore struct {
	bucket s3.Bucket

	mutex sync.RWMutex
	knownKeys map[string]bool  // Protected by mutex
}

func (s *kvStore) Set(key []byte, val []byte) error {
	// Call the bucket.
	if err := s.bucket.StoreObject(string(key), val); err != nil {
		return fmt.Errorf("StoreObject: %v", err)
	}

	// Record the fact that the key is now known.
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.knownKeys[string(key)] = true

	return nil
}

func (s *kvStore) Get(key []byte) (val []byte, err error) {
	if val, err = s.bucket.GetObject(string(key)); err != nil {
		err = fmt.Errorf("GetObject: %v", err)
		return
	}

	return
}

func (s *kvStore) Contains(key []byte) (res bool, err error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	_, ok := s.knownKeys[string(key)]
	return ok, nil
}
<|endoftext|>


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

package usb_test

import (
	"bytes"
	"log"
	"os"
	"testing"

	. "github.com/kylelemons/gousb/usb"
	"github.com/kylelemons/gousb/usbid"
)

func TestNoop(t *testing.T) {
	c := NewContext()
	defer c.Close()
	c.Debug(0)
}

func TestEnum(t *testing.T) {
	c := NewContext()
	defer c.Close()
	c.Debug(0)

	logDevice := func(t *testing.T, desc *Descriptor) {
		t.Logf("%03d.%03d %s", desc.Bus, desc.Address, usbid.Describe(desc))
		t.Logf("- Protocol: %s", usbid.Classify(desc))

		for _, cfg := range desc.Configs {
			t.Logf("- %s:", cfg)
			for _, alt := range cfg.Interfaces {
				t.Logf("  --------------")
				for _, iface := range alt.Setups {
					t.Logf("  - %s", iface)
					t.Logf("    - %s", usbid.Classify(iface))
					for _, end := range iface.Endpoints {
						t.Logf("    - %s (packet size: %d bytes)", end, end.MaxPacketSize)
					}
				}
			}
			t.Logf("  --------------")
		}
	}

	descs := []*Descriptor{}
	devs, err := c.ListDevices(func(desc *Descriptor) bool {
		logDevice(t, desc)
		descs = append(descs, desc)
		return true
	})
	defer func() {
		for _, d := range devs {
			d.Close()
		}
	}()
	if err != nil {
		t.Fatalf("list: %s", err)
	}

	if got, want := len(devs), len(descs); got != want {
		t.Fatalf("len(devs) = %d, want %d", got, want)
	}

	for i := range devs {
		if got, want := devs[i].Descriptor, descs[i]; got != want {
			t.Errorf("dev[%d].Descriptor = %p, want %p", i, got, want)
		}
	}
}

func TestOpenDeviceWithVidPid(t *testing.T) {
	c := NewContext()
	defer c.Close()
	c.Debug(0)

	// Accept for all device
	devs, err := c.ListDevices(func(desc *Descriptor) bool {
		return true
	})
	defer func() {
		for _, d := range devs {
			d.Close()
		}
	}()

	if err != nil {
		t.Fatalf("list: %s", err)
	}

	for i := range devs {
		vid := devs[i].Vendor
		pid := devs[i].Product
		device, err := c.OpenDeviceWithVidPid((int)(vid), (int)(pid))

		// if the context failed to open device
		if err != nil {
			t.Fail()
		}

		// if opened device was not valid
		if device.Descriptor.Bus != devs[i].Bus ||
			device.Descriptor.Address != devs[i].Address ||
			device.Vendor != devs[i].Vendor ||
			device.Product != devs[i].Product {
			t.Fail()
		}

	}
}

func TestMultipleContexts(t *testing.T) {
	var buf bytes.Buffer
	log.SetOutput(&buf)
	for i := 0; i < 2; i++ {
		ctx := NewContext()
		_, err := ctx.ListDevices(func(desc *Descriptor) bool {
			return true
		})
		if err != nil {
			t.Fatal(err)
		}
		ctx.Close()
	}
	log.SetOutput(os.Stderr)
	if buf.Len() > 0 {
		t.Errorf("Non zero output to log, while testing:  %s", buf.String())
	}
}
<commit_msg>undo the edited usb_test.go in TestMultipleContexts to origin<commit_after>// Copyright 2013 Google Inc.  All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package usb_test

import (
	"bytes"
	"log"
	"os"
	"testing"

	. "github.com/kylelemons/gousb/usb"
	"github.com/kylelemons/gousb/usbid"
)

func TestNoop(t *testing.T) {
	c := NewContext()
	defer c.Close()
	c.Debug(0)
}

func TestEnum(t *testing.T) {
	c := NewContext()
	defer c.Close()
	c.Debug(0)

	logDevice := func(t *testing.T, desc *Descriptor) {
		t.Logf("%03d.%03d %s", desc.Bus, desc.Address, usbid.Describe(desc))
		t.Logf("- Protocol: %s", usbid.Classify(desc))

		for _, cfg := range desc.Configs {
			t.Logf("- %s:", cfg)
			for _, alt := range cfg.Interfaces {
				t.Logf("  --------------")
				for _, iface := range alt.Setups {
					t.Logf("  - %s", iface)
					t.Logf("    - %s", usbid.Classify(iface))
					for _, end := range iface.Endpoints {
						t.Logf("    - %s (packet size: %d bytes)", end, end.MaxPacketSize)
					}
				}
			}
			t.Logf("  --------------")
		}
	}

	descs := []*Descriptor{}
	devs, err := c.ListDevices(func(desc *Descriptor) bool {
		logDevice(t, desc)
		descs = append(descs, desc)
		return true
	})
	defer func() {
		for _, d := range devs {
			d.Close()
		}
	}()
	if err != nil {
		t.Fatalf("list: %s", err)
	}

	if got, want := len(devs), len(descs); got != want {
		t.Fatalf("len(devs) = %d, want %d", got, want)
	}

	for i := range devs {
		if got, want := devs[i].Descriptor, descs[i]; got != want {
			t.Errorf("dev[%d].Descriptor = %p, want %p", i, got, want)
		}
	}
}

func TestOpenDeviceWithVidPid(t *testing.T) {
	c := NewContext()
	defer c.Close()
	c.Debug(0)

	// Accept for all device
	devs, err := c.ListDevices(func(desc *Descriptor) bool {
		return true
	})
	defer func() {
		for _, d := range devs {
			d.Close()
		}
	}()

	if err != nil {
		t.Fatalf("list: %s", err)
	}

	for i := range devs {
		vid := devs[i].Vendor
		pid := devs[i].Product
		device, err := c.OpenDeviceWithVidPid((int)(vid), (int)(pid))

		// if the context failed to open device
		if err != nil {
			t.Fail()
		}

		// if opened device was not valid
		if device.Descriptor.Bus != devs[i].Bus ||
			device.Descriptor.Address != devs[i].Address ||
			device.Vendor != devs[i].Vendor ||
			device.Product != devs[i].Product {
			t.Fail()
		}

	}
}

func TestMultipleContexts(t *testing.T) {
	var buf bytes.Buffer
	log.SetOutput(&buf)
	for i := 0; i < 2; i++ {
		ctx := NewContext()
		_, err := ctx.ListDevices(func(desc *Descriptor) bool {
			return false
		})
		if err != nil {
			t.Fatal(err)
		}
		ctx.Close()
	}
	log.SetOutput(os.Stderr)
	if buf.Len() > 0 {
		t.Errorf("Non zero output to log, while testing:  %s", buf.String())
	}
}
<|endoftext|>


<commit_before>/*
 * ZAnnotate Copyright 2017 Regents of the University of Michigan
 *
 * 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 zrouting

import (
	"encoding/json"
	"io"
	"net"

	"github.com/osrg/gobgp/packet/bgp"
	log "github.com/sirupsen/logrus"
	"github.com/zmap/go-iptree/iptree"
	"github.com/zmap/zannotate/zmrt"
)

type ASNameNode struct {
	ASN          uint32 `json:"asn,omitempty"`
	Description  string `json:"description,omitempty"`
	Name         string `json:"name,omitempty"`
	Organization string `json:"organization,omitempty"`
	CountryCode  string `json:"country_code,omitempty"`
}

type ASTreeNode struct {
	Prefix string
	ASN    uint32
	Path   []uint32
}

type RoutingOutput struct {
	Prefix string       `json:"prefix"`
	ASN    uint32       `json:"asn,omitempty"`
	Path   []uint32     `json:"path,omitempty"`
	Origin *ASNameNode  `json:"as,omitempty"`
	Data   *interface{} `json:"data,omitempty"`
}

type RoutingLookupTree struct {
	ASNames map[uint32]ASNameNode
	ASData  map[uint32]interface{}
	IPTree  *iptree.IPTree
}

//	RoutingTablePath string
//	ASNamesPath      string

func (t *RoutingLookupTree) PopulateFromMRT(raw io.Reader) {
	t.IPTree = iptree.New()
	zmrt.MrtPathIterate(raw, func(e *zmrt.RIBEntry) {
		if e.AFI == bgp.AFI_IP {
			var n ASTreeNode
			n.Prefix = e.Prefix
			n.Path = e.Attributes.ASPath
			if len(n.Path) > 0 {
				n.ASN = n.Path[len(n.Path)-1]
			}
			t.IPTree.AddByString(e.Prefix, n)
		}
	})
}

func (t *RoutingLookupTree) SetASName(asn uint32, m ASNameNode) {
	if t.ASNames == nil {
		t.ASNames = make(map[uint32]ASNameNode)
	}
	t.ASNames[asn] = m
}

func (t *RoutingLookupTree) SetASData(asn uint32, m interface{}) {
	if t.ASData == nil {
		t.ASData = make(map[uint32]interface{})
	}
	t.ASData[asn] = m
}

func (t *RoutingLookupTree) PopulateASnames(raw io.Reader) {
	d := json.NewDecoder(raw)
	for {
		var m ASNameNode
		if err := d.Decode(&m); err == io.EOF {
			break
		} else if err != nil {
			log.Fatalf("%s", err)
		}
		t.SetASName(m.ASN, m)
	}
}

func (t *RoutingLookupTree) Get(ip net.IP) (*RoutingOutput, error) {
	var out RoutingOutput
	if n, ok, err := t.IPTree.Get(ip); ok && err == nil {
		node := n.(ASTreeNode)
		out.Prefix = node.Prefix
		out.Path = node.Path
		out.ASN = node.ASN
		if t.ASNames != nil {
			var n ASNameNode
			if name, ok := t.ASNames[out.ASN]; ok {
				n.Description = name.Description
				n.Organization = name.Organization
				n.Name = name.Name
				n.CountryCode = name.CountryCode
				n.ASN = node.ASN
				out.Origin = &n
			}
		}
		return &out, nil
	} else {
		return nil, err
	}
}
<commit_msg>return errors instead of fatal'ing in library<commit_after>/*
 * ZAnnotate Copyright 2017 Regents of the University of Michigan
 *
 * 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 zrouting

import (
	"encoding/json"
	"io"
	"net"

	"github.com/osrg/gobgp/packet/bgp"
	"github.com/zmap/go-iptree/iptree"
	"github.com/zmap/zannotate/zmrt"
)

type ASNameNode struct {
	ASN          uint32 `json:"asn,omitempty"`
	Description  string `json:"description,omitempty"`
	Name         string `json:"name,omitempty"`
	Organization string `json:"organization,omitempty"`
	CountryCode  string `json:"country_code,omitempty"`
}

type ASTreeNode struct {
	Prefix string
	ASN    uint32
	Path   []uint32
}

type RoutingOutput struct {
	Prefix string       `json:"prefix"`
	ASN    uint32       `json:"asn,omitempty"`
	Path   []uint32     `json:"path,omitempty"`
	Origin *ASNameNode  `json:"as,omitempty"`
	Data   *interface{} `json:"data,omitempty"`
}

type RoutingLookupTree struct {
	ASNames map[uint32]ASNameNode
	ASData  map[uint32]interface{}
	IPTree  *iptree.IPTree
}

func (t *RoutingLookupTree) PopulateFromMRT(raw io.Reader) {
	t.IPTree = iptree.New()
	zmrt.MrtPathIterate(raw, func(e *zmrt.RIBEntry) {
		if e.AFI == bgp.AFI_IP {
			var n ASTreeNode
			n.Prefix = e.Prefix
			n.Path = e.Attributes.ASPath
			if len(n.Path) > 0 {
				n.ASN = n.Path[len(n.Path)-1]
			}
			t.IPTree.AddByString(e.Prefix, n)
		}
	})
}

func (t *RoutingLookupTree) SetASName(asn uint32, m ASNameNode) {
	if t.ASNames == nil {
		t.ASNames = make(map[uint32]ASNameNode)
	}
	t.ASNames[asn] = m
}

func (t *RoutingLookupTree) SetASData(asn uint32, m interface{}) {
	if t.ASData == nil {
		t.ASData = make(map[uint32]interface{})
	}
	t.ASData[asn] = m
}

func (t *RoutingLookupTree) PopulateASnames(raw io.Reader) error {
	d := json.NewDecoder(raw)
	for {
		var m ASNameNode
		if err := d.Decode(&m); err == io.EOF {
			break
		} else if err != nil {
			return err
		}
		t.SetASName(m.ASN, m)
	}
	return nil
}

func (t *RoutingLookupTree) Get(ip net.IP) (*RoutingOutput, error) {
	var out RoutingOutput
	if n, ok, err := t.IPTree.Get(ip); ok && err == nil {
		node := n.(ASTreeNode)
		out.Prefix = node.Prefix
		out.Path = node.Path
		out.ASN = node.ASN
		if t.ASNames != nil {
			var n ASNameNode
			if name, ok := t.ASNames[out.ASN]; ok {
				n.Description = name.Description
				n.Organization = name.Organization
				n.Name = name.Name
				n.CountryCode = name.CountryCode
				n.ASN = node.ASN
				out.Origin = &n
			}
		}
		return &out, nil
	} else {
		return nil, err
	}
}
<|endoftext|>