<commit_before>package hypervisor

import (
	"fmt"
	"strings"
	"sync"

	"github.com/hyperhq/runv/api"
	hyperstartapi "github.com/hyperhq/runv/hyperstart/api/json"
	"github.com/hyperhq/runv/hypervisor/network"
)

const (
	MAX_NIC = int(^uint(0) >> 1) // Eth is network card, while lo is alias, what's the maximum for each? same?
	// let upper level logic care about the restriction. here is just an upbond.
	DEFAULT_LO_DEVICE_NAME = "lo"
)

type NetworkContext struct {
	*api.SandboxConfig

	sandbox *VmContext

	ports []*api.PortDescription
	eth   map[int]*InterfaceCreated
	lo    map[string]*InterfaceCreated

	idMap map[string]*InterfaceCreated // a secondary index for both eth and lo, for lo, the hostdevice is empty

	slotLock *sync.RWMutex
}

func NewNetworkContext() *NetworkContext {
	return &NetworkContext{
		ports:    []*api.PortDescription{},
		eth:      make(map[int]*InterfaceCreated),
		lo:       make(map[string]*InterfaceCreated),
		idMap:    make(map[string]*InterfaceCreated),
		slotLock: &sync.RWMutex{},
	}
}

func (nc *NetworkContext) sandboxInfo() *hyperstartapi.Pod {

	vmSpec := NewVmSpec()

	vmSpec.Hostname = nc.Hostname
	vmSpec.Dns = nc.Dns
	vmSpec.DnsSearch = nc.DnsSearch
	vmSpec.DnsOptions = nc.DnsOptions
	if nc.Neighbors != nil {
		vmSpec.PortmappingWhiteLists = &hyperstartapi.PortmappingWhiteList{
			InternalNetworks: nc.Neighbors.InternalNetworks,
			ExternalNetworks: nc.Neighbors.ExternalNetworks,
		}
	}

	return vmSpec
}

func (nc *NetworkContext) applySlot() int {
	for i := 0; i <= MAX_NIC; i++ {
		if _, ok := nc.eth[i]; !ok {
			nc.eth[i] = nil
			return i
		}
	}

	return -1
}

func (nc *NetworkContext) freeSlot(slot int) {
	if inf, ok := nc.eth[slot]; !ok {
		nc.sandbox.Log(WARNING, "Freeing an unoccupied eth slot %d", slot)
		return
	} else if inf != nil {
		if _, ok := nc.idMap[inf.Id]; ok {
			delete(nc.idMap, inf.Id)
		}
	}
	nc.sandbox.Log(DEBUG, "Free slot %d of eth", slot)
	delete(nc.eth, slot)
}

// nextAvailableDevName find the initial device name in guest when add a new tap device
// then rename it to some new name.
// For example: user want to insert a new nic named "eth5" into container, which already owns
// "eth0", "eth3" and "eth4". After tap device is added to VM, guest will detect a new device
// named "eth1" which is first available "ethX" device, then guest will try to rename "eth1" to
// "eth5". Then container will have "eth0", "eth3", "eth4" and "eth5" in the last.
// This function try to find the first available "ethX" as said above. @WeiZhang555
func (nc *NetworkContext) nextAvailableDevName() string {
	for i := 0; i <= MAX_NIC; i++ {
		find := false
		for _, inf := range nc.eth {
			if inf != nil && inf.NewName == fmt.Sprintf("eth%d", i) {
				find = true
				break
			}
		}
		if !find {
			return fmt.Sprintf("eth%d", i)
		}
	}

	return ""
}

func (nc *NetworkContext) addInterface(inf *api.InterfaceDescription, result chan api.Result) {
	if inf.Lo {
		if len(inf.Ip) == 0 {
			estr := fmt.Sprintf("creating an interface without an IP address: %#v", inf)
			nc.sandbox.Log(ERROR, estr)
			result <- NewSpecError(inf.Id, estr)
			return
		}
		if inf.Id == "" {
			inf.Id = "lo"
		}
		i := &InterfaceCreated{
			Id:         inf.Id,
			DeviceName: DEFAULT_LO_DEVICE_NAME,
			IpAddr:     inf.Ip,
			Mtu:        inf.Mtu,
		}
		nc.lo[inf.Ip] = i
		nc.idMap[inf.Id] = i

		result <- &api.ResultBase{
			Id:      inf.Id,
			Success: true,
		}
		return
	}

	var devChan chan VmEvent = make(chan VmEvent, 1)

	go func() {
		nc.slotLock.Lock()
		defer nc.slotLock.Unlock()

		idx := nc.applySlot()
		initialDevName := nc.nextAvailableDevName()
		if inf.Id == "" {
			inf.Id = fmt.Sprintf("%d", idx)
		}
		if idx < 0 || initialDevName == "" {
			estr := fmt.Sprintf("no available ethernet slot/name for interface %#v", inf)
			nc.sandbox.Log(ERROR, estr)
			result <- NewBusyError(inf.Id, estr)
			close(devChan)
			return
		}

		nc.configureInterface(idx, nc.sandbox.NextPciAddr(), initialDevName, inf, devChan)
	}()

	go func() {
		ev, ok := <-devChan
		if !ok {
			nc.sandbox.Log(ERROR, "chan closed while waiting network inserted event: %#v", ev)
			return
		}
		// ev might be DeviceInsert failed, or inserted
		if fe, ok := ev.(*DeviceFailed); ok {
			if inf, ok := fe.Session.(*InterfaceCreated); ok {
				nc.netdevInsertFailed(inf.Index, inf.DeviceName)
				nc.sandbox.Log(ERROR, "interface creation failed: %#v", inf)
			} else if inf, ok := fe.Session.(*NetDevInsertedEvent); ok {
				nc.netdevInsertFailed(inf.Index, inf.DeviceName)
				nc.sandbox.Log(ERROR, "interface creation failed: %#v", inf)
			}
			result <- fe
			return
		} else if ni, ok := ev.(*NetDevInsertedEvent); ok {
			created := nc.idMap[inf.Id]
			created.TapFd = ni.TapFd
			nc.sandbox.Log(DEBUG, "nic insert success: %s", ni.Id)
			result <- ni
			return
		}
		nc.sandbox.Log(ERROR, "got unknown event while waiting network inserted event: %#v", ev)
		result <- NewDeviceError(inf.Id, "unknown event")
	}()
}

func (nc *NetworkContext) removeInterface(id string, result chan api.Result) {
	if inf, ok := nc.idMap[id]; !ok {
		nc.sandbox.Log(WARNING, "trying remove a non-exist interface %s", id)
		result <- api.NewResultBase(id, true, "not exist")
		return
	} else if inf.HostDevice == "" { // a virtual interface
		delete(nc.idMap, id)
		delete(nc.lo, inf.IpAddr)
		result <- api.NewResultBase(id, true, "")
		return
	} else {
		nc.slotLock.Lock()
		defer nc.slotLock.Unlock()

		if _, ok := nc.eth[inf.Index]; !ok {
			delete(nc.idMap, id)
			nc.sandbox.Log(INFO, "non-configured network device %d remove failed", inf.Index)
			result <- api.NewResultBase(id, true, "not configured eth")
			return
		}

		var devChan chan VmEvent = make(chan VmEvent, 1)

		nc.sandbox.Log(DEBUG, "remove network card %d: %s", inf.Index, inf.IpAddr)
		nc.sandbox.DCtx.RemoveNic(nc.sandbox, inf, &NetDevRemovedEvent{Index: inf.Index}, devChan)

		go func() {
			ev, ok := <-devChan
			if !ok {
				return
			}

			success := true
			message := ""

			if fe, ok := ev.(*DeviceFailed); ok {
				success = false
				message = "unplug failed"
				if inf, ok := fe.Session.(*NetDevRemovedEvent); ok {
					nc.sandbox.Log(ERROR, "interface remove failed: %#v", inf)
				}
			}

			nc.slotLock.Lock()
			defer nc.slotLock.Unlock()
			nc.freeSlot(inf.Index)
			nc.cleanupInf(inf)

			result <- api.NewResultBase(id, success, message)
		}()
	}
}

// allInterfaces return all the network interfaces except loop
func (nc *NetworkContext) allInterfaces() (nics []*InterfaceCreated) {
	nc.slotLock.Lock()
	defer nc.slotLock.Unlock()

	for _, v := range nc.eth {
		if v != nil {
			nics = append(nics, v)
		}
	}
	return
}

func (nc *NetworkContext) updateInterface(inf *api.InterfaceDescription) error {
	oldInf, ok := nc.idMap[inf.Id]
	if !ok {
		nc.sandbox.Log(WARNING, "trying update a non-exist interface %s", inf.Id)
		return fmt.Errorf("interface %q not exists", inf.Id)
	}

	// only support update some fields: Name, ip addresses, mtu
	nc.slotLock.Lock()
	defer nc.slotLock.Unlock()

	if inf.Name != "" {
		oldInf.NewName = inf.Name
	}

	if inf.Mtu > 0 {
		oldInf.Mtu = inf.Mtu
	}

	if len(inf.Ip) > 0 {
		addrs := strings.Split(inf.Ip, ",")
		oldAddrs := strings.Split(oldInf.IpAddr, ",")
		for _, ip := range addrs {
			var found bool
			if ip[0] == '-' { // to delete
				ip = ip[1:]
				for k, i := range oldAddrs {
					if i == ip {
						oldAddrs = append(oldAddrs[:k], oldAddrs[k+1:]...)
						found = true
						break
					}
				}
				if !found {
					return fmt.Errorf("failed to delete %q: not found", ip)
				}
			} else { // to add
				oldAddrs = append(oldAddrs, ip)
			}
		}
		oldInf.IpAddr = strings.Join(oldAddrs, ",")
	}
	return nil
}

func (nc *NetworkContext) netdevInsertFailed(idx int, name string) {
	nc.slotLock.Lock()
	defer nc.slotLock.Unlock()

	if _, ok := nc.eth[idx]; !ok {
		nc.sandbox.Log(INFO, "network device %d (%s) insert failed before configured", idx, name)
		return
	}

	nc.sandbox.Log(INFO, "network device %d (%s) insert failed", idx, name)
	nc.freeSlot(idx)
}

func (nc *NetworkContext) configureInterface(index, pciAddr int, name string, inf *api.InterfaceDescription, result chan<- VmEvent) {
	if inf.TapName == "" {
		inf.TapName = network.NicName(nc.sandbox.Id, index)
	}

	settings, err := network.Configure(inf)
	if err != nil {
		nc.sandbox.Log(ERROR, "interface creating failed: %v", err.Error())
		session := &InterfaceCreated{Id: inf.Id, Index: index, PCIAddr: pciAddr, DeviceName: name, NewName: inf.Name, Mtu: inf.Mtu}
		result <- &DeviceFailed{Session: session}
		return
	}

	created, err := interfaceGot(inf.Id, index, pciAddr, name, inf.Name, settings)
	if err != nil {
		result <- &DeviceFailed{Session: created}
		return
	}

	h := &HostNicInfo{
		Id:      created.Id,
		Device:  created.HostDevice,
		Mac:     created.MacAddr,
		Bridge:  created.Bridge,
		Gateway: created.Bridge,
		Options: inf.Options,
	}

	// Note: Use created.NewName add tap name
	// this is because created.DeviceName isn't always uniq,
	// instead NewName is real nic name in VM which is certainly uniq
	g := &GuestNicInfo{
		Device:  created.NewName,
		Ipaddr:  created.IpAddr,
		Index:   created.Index,
		Busaddr: created.PCIAddr,
	}

	nc.eth[index] = created
	nc.idMap[created.Id] = created
	nc.sandbox.DCtx.AddNic(nc.sandbox, h, g, result)
}

func (nc *NetworkContext) cleanupInf(inf *InterfaceCreated) {
	network.ReleaseAddr(inf.IpAddr)
}

func (nc *NetworkContext) getInterface(id string) *InterfaceCreated {
	nc.slotLock.RLock()
	defer nc.slotLock.RUnlock()

	inf, ok := nc.idMap[id]
	if ok {
		return inf
	}
	return nil
}

func (nc *NetworkContext) getIPAddrs() []string {
	nc.slotLock.RLock()
	defer nc.slotLock.RUnlock()

	res := []string{}
	for _, inf := range nc.eth {
		if inf.IpAddr != "" {
			addrs := strings.Split(inf.IpAddr, ",")
			res = append(res, addrs...)
		}
	}

	return res
}

func (nc *NetworkContext) getRoutes() []hyperstartapi.Route {
	nc.slotLock.RLock()
	defer nc.slotLock.RUnlock()
	routes := []hyperstartapi.Route{}

	for _, inf := range nc.idMap {
		for _, r := range inf.RouteTable {
			routes = append(routes, hyperstartapi.Route{
				Dest:    r.Destination,
				Gateway: r.Gateway,
				Device:  inf.NewName,
			})
		}
	}

	return routes
}

func (nc *NetworkContext) close() {
	nc.slotLock.Lock()
	defer nc.slotLock.Unlock()

	for _, inf := range nc.eth {
		nc.cleanupInf(inf)
	}
	nc.eth = map[int]*InterfaceCreated{}
	nc.lo = map[string]*InterfaceCreated{}
	nc.idMap = map[string]*InterfaceCreated{}
}

func interfaceGot(id string, index int, pciAddr int, deviceName, newName string, inf *network.Settings) (*InterfaceCreated, error) {
	rt := []*RouteRule{}
	/* Route rule is generated automaticly on first interface,
	 * or generated on the gateway configured interface. */
	if (index == 0 && inf.Automatic) || (!inf.Automatic && inf.Gateway != "") {
		rt = append(rt, &RouteRule{
			Destination: "0.0.0.0/0",
			Gateway:     inf.Gateway, ViaThis: true,
		})
	}

	infc := &InterfaceCreated{
		Id:         id,
		Index:      index,
		PCIAddr:    pciAddr,
		Bridge:     inf.Bridge,
		HostDevice: inf.Device,
		DeviceName: deviceName,
		NewName:    newName,
		MacAddr:    inf.Mac,
		IpAddr:     inf.IPAddress,
		Mtu:        inf.Mtu,
		RouteTable: rt,
	}
	return infc, nil
}
<commit_msg>make sure Name of InterfaceDescription is not null<commit_after>package hypervisor

import (
	"fmt"
	"strings"
	"sync"

	"github.com/hyperhq/runv/api"
	hyperstartapi "github.com/hyperhq/runv/hyperstart/api/json"
	"github.com/hyperhq/runv/hypervisor/network"
)

const (
	MAX_NIC = int(^uint(0) >> 1) // Eth is network card, while lo is alias, what's the maximum for each? same?
	// let upper level logic care about the restriction. here is just an upbond.
	DEFAULT_LO_DEVICE_NAME = "lo"
)

type NetworkContext struct {
	*api.SandboxConfig

	sandbox *VmContext

	ports []*api.PortDescription
	eth   map[int]*InterfaceCreated
	lo    map[string]*InterfaceCreated

	idMap map[string]*InterfaceCreated // a secondary index for both eth and lo, for lo, the hostdevice is empty

	slotLock *sync.RWMutex
}

func NewNetworkContext() *NetworkContext {
	return &NetworkContext{
		ports:    []*api.PortDescription{},
		eth:      make(map[int]*InterfaceCreated),
		lo:       make(map[string]*InterfaceCreated),
		idMap:    make(map[string]*InterfaceCreated),
		slotLock: &sync.RWMutex{},
	}
}

func (nc *NetworkContext) sandboxInfo() *hyperstartapi.Pod {

	vmSpec := NewVmSpec()

	vmSpec.Hostname = nc.Hostname
	vmSpec.Dns = nc.Dns
	vmSpec.DnsSearch = nc.DnsSearch
	vmSpec.DnsOptions = nc.DnsOptions
	if nc.Neighbors != nil {
		vmSpec.PortmappingWhiteLists = &hyperstartapi.PortmappingWhiteList{
			InternalNetworks: nc.Neighbors.InternalNetworks,
			ExternalNetworks: nc.Neighbors.ExternalNetworks,
		}
	}

	return vmSpec
}

func (nc *NetworkContext) applySlot() int {
	for i := 0; i <= MAX_NIC; i++ {
		if _, ok := nc.eth[i]; !ok {
			nc.eth[i] = nil
			return i
		}
	}

	return -1
}

func (nc *NetworkContext) freeSlot(slot int) {
	if inf, ok := nc.eth[slot]; !ok {
		nc.sandbox.Log(WARNING, "Freeing an unoccupied eth slot %d", slot)
		return
	} else if inf != nil {
		if _, ok := nc.idMap[inf.Id]; ok {
			delete(nc.idMap, inf.Id)
		}
	}
	nc.sandbox.Log(DEBUG, "Free slot %d of eth", slot)
	delete(nc.eth, slot)
}

// nextAvailableDevName find the initial device name in guest when add a new tap device
// then rename it to some new name.
// For example: user want to insert a new nic named "eth5" into container, which already owns
// "eth0", "eth3" and "eth4". After tap device is added to VM, guest will detect a new device
// named "eth1" which is first available "ethX" device, then guest will try to rename "eth1" to
// "eth5". Then container will have "eth0", "eth3", "eth4" and "eth5" in the last.
// This function try to find the first available "ethX" as said above. @WeiZhang555
func (nc *NetworkContext) nextAvailableDevName() string {
	for i := 0; i <= MAX_NIC; i++ {
		find := false
		for _, inf := range nc.eth {
			if inf != nil && inf.NewName == fmt.Sprintf("eth%d", i) {
				find = true
				break
			}
		}
		if !find {
			return fmt.Sprintf("eth%d", i)
		}
	}

	return ""
}

func (nc *NetworkContext) addInterface(inf *api.InterfaceDescription, result chan api.Result) {
	if inf.Lo {
		if len(inf.Ip) == 0 {
			estr := fmt.Sprintf("creating an interface without an IP address: %#v", inf)
			nc.sandbox.Log(ERROR, estr)
			result <- NewSpecError(inf.Id, estr)
			return
		}
		if inf.Id == "" {
			inf.Id = "lo"
		}
		i := &InterfaceCreated{
			Id:         inf.Id,
			DeviceName: DEFAULT_LO_DEVICE_NAME,
			IpAddr:     inf.Ip,
			Mtu:        inf.Mtu,
		}
		nc.lo[inf.Ip] = i
		nc.idMap[inf.Id] = i

		result <- &api.ResultBase{
			Id:      inf.Id,
			Success: true,
		}
		return
	}

	var devChan chan VmEvent = make(chan VmEvent, 1)

	go func() {
		nc.slotLock.Lock()
		defer nc.slotLock.Unlock()

		idx := nc.applySlot()
		initialDevName := nc.nextAvailableDevName()
		if inf.Id == "" {
			inf.Id = fmt.Sprintf("%d", idx)
		}
		if idx < 0 || initialDevName == "" {
			estr := fmt.Sprintf("no available ethernet slot/name for interface %#v", inf)
			nc.sandbox.Log(ERROR, estr)
			result <- NewBusyError(inf.Id, estr)
			close(devChan)
			return
		}

		nc.configureInterface(idx, nc.sandbox.NextPciAddr(), initialDevName, inf, devChan)
	}()

	go func() {
		ev, ok := <-devChan
		if !ok {
			nc.sandbox.Log(ERROR, "chan closed while waiting network inserted event: %#v", ev)
			return
		}
		// ev might be DeviceInsert failed, or inserted
		if fe, ok := ev.(*DeviceFailed); ok {
			if inf, ok := fe.Session.(*InterfaceCreated); ok {
				nc.netdevInsertFailed(inf.Index, inf.DeviceName)
				nc.sandbox.Log(ERROR, "interface creation failed: %#v", inf)
			} else if inf, ok := fe.Session.(*NetDevInsertedEvent); ok {
				nc.netdevInsertFailed(inf.Index, inf.DeviceName)
				nc.sandbox.Log(ERROR, "interface creation failed: %#v", inf)
			}
			result <- fe
			return
		} else if ni, ok := ev.(*NetDevInsertedEvent); ok {
			created := nc.idMap[inf.Id]
			created.TapFd = ni.TapFd
			nc.sandbox.Log(DEBUG, "nic insert success: %s", ni.Id)
			result <- ni
			return
		}
		nc.sandbox.Log(ERROR, "got unknown event while waiting network inserted event: %#v", ev)
		result <- NewDeviceError(inf.Id, "unknown event")
	}()
}

func (nc *NetworkContext) removeInterface(id string, result chan api.Result) {
	if inf, ok := nc.idMap[id]; !ok {
		nc.sandbox.Log(WARNING, "trying remove a non-exist interface %s", id)
		result <- api.NewResultBase(id, true, "not exist")
		return
	} else if inf.HostDevice == "" { // a virtual interface
		delete(nc.idMap, id)
		delete(nc.lo, inf.IpAddr)
		result <- api.NewResultBase(id, true, "")
		return
	} else {
		nc.slotLock.Lock()
		defer nc.slotLock.Unlock()

		if _, ok := nc.eth[inf.Index]; !ok {
			delete(nc.idMap, id)
			nc.sandbox.Log(INFO, "non-configured network device %d remove failed", inf.Index)
			result <- api.NewResultBase(id, true, "not configured eth")
			return
		}

		var devChan chan VmEvent = make(chan VmEvent, 1)

		nc.sandbox.Log(DEBUG, "remove network card %d: %s", inf.Index, inf.IpAddr)
		nc.sandbox.DCtx.RemoveNic(nc.sandbox, inf, &NetDevRemovedEvent{Index: inf.Index}, devChan)

		go func() {
			ev, ok := <-devChan
			if !ok {
				return
			}

			success := true
			message := ""

			if fe, ok := ev.(*DeviceFailed); ok {
				success = false
				message = "unplug failed"
				if inf, ok := fe.Session.(*NetDevRemovedEvent); ok {
					nc.sandbox.Log(ERROR, "interface remove failed: %#v", inf)
				}
			}

			nc.slotLock.Lock()
			defer nc.slotLock.Unlock()
			nc.freeSlot(inf.Index)
			nc.cleanupInf(inf)

			result <- api.NewResultBase(id, success, message)
		}()
	}
}

// allInterfaces return all the network interfaces except loop
func (nc *NetworkContext) allInterfaces() (nics []*InterfaceCreated) {
	nc.slotLock.Lock()
	defer nc.slotLock.Unlock()

	for _, v := range nc.eth {
		if v != nil {
			nics = append(nics, v)
		}
	}
	return
}

func (nc *NetworkContext) updateInterface(inf *api.InterfaceDescription) error {
	oldInf, ok := nc.idMap[inf.Id]
	if !ok {
		nc.sandbox.Log(WARNING, "trying update a non-exist interface %s", inf.Id)
		return fmt.Errorf("interface %q not exists", inf.Id)
	}

	// only support update some fields: Name, ip addresses, mtu
	nc.slotLock.Lock()
	defer nc.slotLock.Unlock()

	if inf.Name != "" {
		oldInf.NewName = inf.Name
	}

	if inf.Mtu > 0 {
		oldInf.Mtu = inf.Mtu
	}

	if len(inf.Ip) > 0 {
		addrs := strings.Split(inf.Ip, ",")
		oldAddrs := strings.Split(oldInf.IpAddr, ",")
		for _, ip := range addrs {
			var found bool
			if ip[0] == '-' { // to delete
				ip = ip[1:]
				for k, i := range oldAddrs {
					if i == ip {
						oldAddrs = append(oldAddrs[:k], oldAddrs[k+1:]...)
						found = true
						break
					}
				}
				if !found {
					return fmt.Errorf("failed to delete %q: not found", ip)
				}
			} else { // to add
				oldAddrs = append(oldAddrs, ip)
			}
		}
		oldInf.IpAddr = strings.Join(oldAddrs, ",")
	}
	return nil
}

func (nc *NetworkContext) netdevInsertFailed(idx int, name string) {
	nc.slotLock.Lock()
	defer nc.slotLock.Unlock()

	if _, ok := nc.eth[idx]; !ok {
		nc.sandbox.Log(INFO, "network device %d (%s) insert failed before configured", idx, name)
		return
	}

	nc.sandbox.Log(INFO, "network device %d (%s) insert failed", idx, name)
	nc.freeSlot(idx)
}

func (nc *NetworkContext) configureInterface(index, pciAddr int, name string, inf *api.InterfaceDescription, result chan<- VmEvent) {
	if inf.TapName == "" {
		inf.TapName = network.NicName(nc.sandbox.Id, index)
	}
	if inf.Name == "" {
		inf.Name = name
	}
	settings, err := network.Configure(inf)
	if err != nil {
		nc.sandbox.Log(ERROR, "interface creating failed: %v", err.Error())
		session := &InterfaceCreated{Id: inf.Id, Index: index, PCIAddr: pciAddr, DeviceName: name, NewName: inf.Name, Mtu: inf.Mtu}
		result <- &DeviceFailed{Session: session}
		return
	}

	created, err := interfaceGot(inf.Id, index, pciAddr, name, inf.Name, settings)
	if err != nil {
		result <- &DeviceFailed{Session: created}
		return
	}

	h := &HostNicInfo{
		Id:      created.Id,
		Device:  created.HostDevice,
		Mac:     created.MacAddr,
		Bridge:  created.Bridge,
		Gateway: created.Bridge,
		Options: inf.Options,
	}

	// Note: Use created.NewName add tap name
	// this is because created.DeviceName isn't always uniq,
	// instead NewName is real nic name in VM which is certainly uniq
	g := &GuestNicInfo{
		Device:  created.NewName,
		Ipaddr:  created.IpAddr,
		Index:   created.Index,
		Busaddr: created.PCIAddr,
	}

	nc.eth[index] = created
	nc.idMap[created.Id] = created
	nc.sandbox.DCtx.AddNic(nc.sandbox, h, g, result)
}

func (nc *NetworkContext) cleanupInf(inf *InterfaceCreated) {
	network.ReleaseAddr(inf.IpAddr)
}

func (nc *NetworkContext) getInterface(id string) *InterfaceCreated {
	nc.slotLock.RLock()
	defer nc.slotLock.RUnlock()

	inf, ok := nc.idMap[id]
	if ok {
		return inf
	}
	return nil
}

func (nc *NetworkContext) getIPAddrs() []string {
	nc.slotLock.RLock()
	defer nc.slotLock.RUnlock()

	res := []string{}
	for _, inf := range nc.eth {
		if inf.IpAddr != "" {
			addrs := strings.Split(inf.IpAddr, ",")
			res = append(res, addrs...)
		}
	}

	return res
}

func (nc *NetworkContext) getRoutes() []hyperstartapi.Route {
	nc.slotLock.RLock()
	defer nc.slotLock.RUnlock()
	routes := []hyperstartapi.Route{}

	for _, inf := range nc.idMap {
		for _, r := range inf.RouteTable {
			routes = append(routes, hyperstartapi.Route{
				Dest:    r.Destination,
				Gateway: r.Gateway,
				Device:  inf.NewName,
			})
		}
	}

	return routes
}

func (nc *NetworkContext) close() {
	nc.slotLock.Lock()
	defer nc.slotLock.Unlock()

	for _, inf := range nc.eth {
		nc.cleanupInf(inf)
	}
	nc.eth = map[int]*InterfaceCreated{}
	nc.lo = map[string]*InterfaceCreated{}
	nc.idMap = map[string]*InterfaceCreated{}
}

func interfaceGot(id string, index int, pciAddr int, deviceName, newName string, inf *network.Settings) (*InterfaceCreated, error) {
	rt := []*RouteRule{}
	/* Route rule is generated automaticly on first interface,
	 * or generated on the gateway configured interface. */
	if (index == 0 && inf.Automatic) || (!inf.Automatic && inf.Gateway != "") {
		rt = append(rt, &RouteRule{
			Destination: "0.0.0.0/0",
			Gateway:     inf.Gateway, ViaThis: true,
		})
	}

	infc := &InterfaceCreated{
		Id:         id,
		Index:      index,
		PCIAddr:    pciAddr,
		Bridge:     inf.Bridge,
		HostDevice: inf.Device,
		DeviceName: deviceName,
		NewName:    newName,
		MacAddr:    inf.Mac,
		IpAddr:     inf.IPAddress,
		Mtu:        inf.Mtu,
		RouteTable: rt,
	}
	return infc, nil
}
<|endoftext|>


<commit_before>package conf

import (
	"fmt"
	"github.com/astaxie/beego"
	"github.com/astaxie/beego/orm"
	// "github.com/naokij/gotalk/models"
	"github.com/naokij/gotalk/setting"
	"runtime"
	// "time"
)

var DiscuzDb string
var Orm orm.Ormer
var OrmGotalk orm.Ormer
var Workers = runtime.NumCPU()
var WorkerLoad int = 20000
var AvatarPath string

func init() {
	runtime.GOMAXPROCS(Workers)
	beego.AppConfigPath = "../conf/app.conf"
	beego.ParseConfig()
	setting.ReadConfig()
	DiscuzDb = beego.AppConfig.String("importer::discuzdb")
	AvatarPath = beego.AppConfig.String("importer::avatarpath")
	if err := orm.RegisterDataBase("discuz", "mysql", fmt.Sprintf("%s:%s@tcp(%s:3306)/%s?charset=utf8", setting.MySQLUser, setting.MySQLPassword, setting.MySQLHost, DiscuzDb)+"&loc=Asia%2FShanghai", 30); err != nil {
		fmt.Println("MySQL error:", err.Error())
		return
	}
	OrmGotalk = orm.NewOrm()
	if db, err := orm.GetDB("discuz"); err == nil {
		Orm, err = orm.NewOrmWithDB("mysql", "discuz", db)
		if err != nil {
			fmt.Println("MySQL error:", err.Error())
			return
		}
	} else {
		fmt.Println("MySQL error:", err.Error())
		return
	}
	orm.RunSyncdb("default", true, false)
}
<commit_msg> 优化import 参数<commit_after>package conf

import (
	"fmt"
	"github.com/astaxie/beego"
	"github.com/astaxie/beego/orm"
	// "github.com/naokij/gotalk/models"
	"github.com/naokij/gotalk/setting"
	"runtime"
	// "time"
)

var DiscuzDb string
var Orm orm.Ormer
var OrmGotalk orm.Ormer
var Workers = runtime.NumCPU()
var WorkerLoad int = 30000
var AvatarPath string

func init() {
	runtime.GOMAXPROCS(Workers)
	beego.AppConfigPath = "../conf/app.conf"
	beego.ParseConfig()
	setting.ReadConfig()
	DiscuzDb = beego.AppConfig.String("importer::discuzdb")
	AvatarPath = beego.AppConfig.String("importer::avatarpath")
	if err := orm.RegisterDataBase("discuz", "mysql", fmt.Sprintf("%s:%s@tcp(%s:3306)/%s?charset=utf8", setting.MySQLUser, setting.MySQLPassword, setting.MySQLHost, DiscuzDb)+"&loc=Asia%2FShanghai", 30); err != nil {
		fmt.Println("MySQL error:", err.Error())
		return
	}
	OrmGotalk = orm.NewOrm()
	if db, err := orm.GetDB("discuz"); err == nil {
		Orm, err = orm.NewOrmWithDB("mysql", "discuz", db)
		if err != nil {
			fmt.Println("MySQL error:", err.Error())
			return
		}
	} else {
		fmt.Println("MySQL error:", err.Error())
		return
	}
	orm.RunSyncdb("default", true, false)
}
<|endoftext|>


<commit_before>package device

import (
	"encoding/json"
	"fmt"
	"github.com/hemtjanst/hemtjanst/messaging"
	"log"
	"sync"
)

type DeviceHandler interface {
	DeviceUpdated(*Device)
	DeviceLeave(*Device)
	DeviceRemoved(*Device)
}

type Manager struct {
	devices  map[string]*Device
	handlers []DeviceHandler
	client   messaging.PublishSubscriber
	sync.RWMutex
}

func NewManager(c messaging.PublishSubscriber) *Manager {
	return &Manager{
		client:   c,
		devices:  make(map[string]*Device, 10),
		handlers: []DeviceHandler{},
	}
}

func (m *Manager) Add(topic string) {
	if _, ok := m.devices[topic]; ok {
		log.Print("Got announce for existing device ", topic)
		return
	}
	log.Print("Going to add device ", topic)
	dev := &Device{Topic: topic, transport: m.client}
	m.client.Subscribe(fmt.Sprintf("%s/meta", topic), 1, func(msg messaging.Message) {
		err := json.Unmarshal(msg.Payload(), dev)
		if err != nil {
			log.Print(err)
			return
		}
		for name, ft := range dev.Features {
			if ft.SetTopic == "" {
				ft.SetTopic = fmt.Sprintf("%s/%s/set", topic, name)
			}
			if ft.GetTopic == "" {
				ft.GetTopic = fmt.Sprintf("%s/%s/get", topic, name)
			}
		}

		go m.forHandler(func(handler DeviceHandler) {
			handler.DeviceUpdated(dev)
		})
	})

	m.Lock()
	defer m.Unlock()
	m.devices[topic] = dev
}

func (m *Manager) Get(id string) (*Device, error) {
	log.Print("Looking for device ", id)
	m.RLock()
	defer m.RUnlock()
	if val, ok := m.devices[id]; ok {
		return val, nil
	}
	return nil, fmt.Errorf("Unknown device %s", id)
}

func (m *Manager) Remove(msg string) {
	log.Print("Attempting to remove device ", msg)
	m.Lock()
	defer m.Unlock()
	if val, ok := m.devices[msg]; ok {
		log.Print("Found device, unsubscribing and removing")
		m.client.Unsubscribe(fmt.Sprintf("%s/meta", msg))
		for _, ft := range val.Features {
			m.client.Unsubscribe(ft.GetTopic)
		}

		go m.forHandler(func(handler DeviceHandler) {
			handler.DeviceLeave(val)
		})

		delete(m.devices, msg)
		return
	}
	for _, d := range m.devices {
		if d.LastWillID == msg {
			log.Print("Found device match for LastWillUID, calling Remove")
			m.Remove(d.Topic)
		}
	}
}

func (m *Manager) forHandler(f func(handler DeviceHandler)) {
	for _, h := range m.handlers {
		f(h)
	}
}

func (m *Manager) AddHandler(handler DeviceHandler) {
	m.Lock()
	defer m.Unlock()
	m.handlers = append(m.handlers, handler)
	go func() {
		for _, device := range m.devices {
			handler.DeviceUpdated(device)
		}
	}()
}
<commit_msg>:ambulance: Fix infinite recursive loop in manager.Remove<commit_after>package device

import (
	"encoding/json"
	"fmt"
	"github.com/hemtjanst/hemtjanst/messaging"
	"log"
	"sync"
)

type DeviceHandler interface {
	DeviceUpdated(*Device)
	DeviceLeave(*Device)
	DeviceRemoved(*Device)
}

type Manager struct {
	devices  map[string]*Device
	handlers []DeviceHandler
	client   messaging.PublishSubscriber
	sync.RWMutex
}

func NewManager(c messaging.PublishSubscriber) *Manager {
	return &Manager{
		client:   c,
		devices:  make(map[string]*Device, 10),
		handlers: []DeviceHandler{},
	}
}

func (m *Manager) Add(topic string) {
	if _, ok := m.devices[topic]; ok {
		log.Print("Got announce for existing device ", topic)
		return
	}
	log.Print("Going to add device ", topic)
	dev := &Device{Topic: topic, transport: m.client}
	m.client.Subscribe(fmt.Sprintf("%s/meta", topic), 1, func(msg messaging.Message) {
		err := json.Unmarshal(msg.Payload(), dev)
		if err != nil {
			log.Print(err)
			return
		}
		for name, ft := range dev.Features {
			if ft.SetTopic == "" {
				ft.SetTopic = fmt.Sprintf("%s/%s/set", topic, name)
			}
			if ft.GetTopic == "" {
				ft.GetTopic = fmt.Sprintf("%s/%s/get", topic, name)
			}
		}

		go m.forHandler(func(handler DeviceHandler) {
			handler.DeviceUpdated(dev)
		})
	})

	m.Lock()
	defer m.Unlock()
	m.devices[topic] = dev
}

func (m *Manager) Get(id string) (*Device, error) {
	log.Print("Looking for device ", id)
	m.RLock()
	defer m.RUnlock()
	if val, ok := m.devices[id]; ok {
		return val, nil
	}
	return nil, fmt.Errorf("Unknown device %s", id)
}

func (m *Manager) Remove(msg string) {
	log.Print("Attempting to remove device ", msg)
	m.Lock()
	defer m.Unlock()
	if val, ok := m.devices[msg]; ok {
		log.Print("Found device, unsubscribing and removing")
		m.client.Unsubscribe(fmt.Sprintf("%s/meta", msg))
		for _, ft := range val.Features {
			m.client.Unsubscribe(ft.GetTopic)
		}

		go m.forHandler(func(handler DeviceHandler) {
			handler.DeviceLeave(val)
		})

		delete(m.devices, msg)
		return
	}
	for _, d := range m.devices {
		if d.LastWillID == msg {
			log.Print("Found device match for LastWillUID, calling Remove")
			delete(m.devices, d.Topic)
		}
	}
}

func (m *Manager) forHandler(f func(handler DeviceHandler)) {
	for _, h := range m.handlers {
		f(h)
	}
}

func (m *Manager) AddHandler(handler DeviceHandler) {
	m.Lock()
	defer m.Unlock()
	m.handlers = append(m.handlers, handler)
	go func() {
		for _, device := range m.devices {
			handler.DeviceUpdated(device)
		}
	}()
}
<|endoftext|>


<commit_before>package hashstructure

import (
	"testing"
)

func TestHash_equal(t *testing.T) {
	cases := []interface{}{
		nil,
		"foo",
		42,
		[]string{"foo", "bar"},
		[]interface{}{1, nil, "foo"},
		map[string]string{"foo": "bar"},
		map[interface{}]string{"foo": "bar"},
		map[interface{}]interface{}{"foo": "bar", "bar": 0},
		struct {
			Foo string
			Bar []interface{}
		}{
			Foo: "foo",
			Bar: []interface{}{nil, nil, nil},
		},
	}

	for _, tc := range cases {
		// We run the test 100 times to try to tease out variability
		// in the runtime in terms of ordering.
		valuelist := make([]uint64, 100)
		for i, _ := range valuelist {
			v, err := Hash(tc, nil)
			if err != nil {
				t.Fatalf("Error: %s\n\n%#v", err, tc)
			}

			valuelist[i] = v
		}

		// Zero is always wrong
		if valuelist[0] == 0 {
			t.Fatalf("zero hash: %#v", tc)
		}

		// Make sure all the values match
		t.Logf("%#v: %d", tc, valuelist[0])
		for i := 1; i < len(valuelist); i++ {
			if valuelist[i] != valuelist[0] {
				t.Fatalf("non-matching: %d, %d\n\n%#v", i, 0, tc)
			}
		}
	}
}
<commit_msg>more tests<commit_after>package hashstructure

import (
	"testing"
)

func TestHash_identity(t *testing.T) {
	cases := []interface{}{
		nil,
		"foo",
		42,
		[]string{"foo", "bar"},
		[]interface{}{1, nil, "foo"},
		map[string]string{"foo": "bar"},
		map[interface{}]string{"foo": "bar"},
		map[interface{}]interface{}{"foo": "bar", "bar": 0},
		struct {
			Foo string
			Bar []interface{}
		}{
			Foo: "foo",
			Bar: []interface{}{nil, nil, nil},
		},
	}

	for _, tc := range cases {
		// We run the test 100 times to try to tease out variability
		// in the runtime in terms of ordering.
		valuelist := make([]uint64, 100)
		for i, _ := range valuelist {
			v, err := Hash(tc, nil)
			if err != nil {
				t.Fatalf("Error: %s\n\n%#v", err, tc)
			}

			valuelist[i] = v
		}

		// Zero is always wrong
		if valuelist[0] == 0 {
			t.Fatalf("zero hash: %#v", tc)
		}

		// Make sure all the values match
		t.Logf("%#v: %d", tc, valuelist[0])
		for i := 1; i < len(valuelist); i++ {
			if valuelist[i] != valuelist[0] {
				t.Fatalf("non-matching: %d, %d\n\n%#v", i, 0, tc)
			}
		}
	}
}

func TestHash_equal(t *testing.T) {
	cases := []struct {
		One, Two interface{}
		Match    bool
	}{
		{
			map[string]string{"foo": "bar"},
			map[interface{}]string{"foo": "bar"},
			true,
		},
	}

	for _, tc := range cases {
		one, err := Hash(tc.One, nil)
		if err != nil {
			t.Fatalf("Failed to hash %#v: %s", tc.One, err)
		}
		two, err := Hash(tc.Two, nil)
		if err != nil {
			t.Fatalf("Failed to hash %#v: %s", tc.Two, err)
		}

		// Zero is always wrong
		if one == 0 {
			t.Fatalf("zero hash: %#v", tc.One)
		}

		// Compare
		if (one == two) != tc.Match {
			t.Fatalf("bad, expected: %#v\n\n%#v\n\n%#v", tc.Match, tc.One, tc.Two)
		}
	}
}
<|endoftext|>


<commit_before>package rabbitmq

import (
	"github.com/dmportella/docker-beat/plugin"
	_ "golang.org/x/net/websocket" // not needed at the moment
)

type consumer struct {
}

func (consumer *consumer) OnEvent(event plugin.DockerEvent) {

}

func init() {
	// do something here
}

/*package main

import (
	"fmt"
	"log"

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

var origin = "http://localhost/"
var url = "ws://localhost:8080/echo"

func main() {
	ws, err := websocket.Dial(url, "", origin)
	if err != nil {
		log.Fatal(err)
	}

	message := []byte("hello, world!")
	_, err = ws.Write(message)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("Send: %s\n", message)

	var msg = make([]byte, 512)
	_, err = ws.Read(msg)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("Receive: %s\n", msg)
}*/
<commit_msg>added support for websockets<commit_after>package rabbitmq

import (
	"encoding/json"
	"flag"
	"github.com/dmportella/docker-beat/logging"
	"github.com/dmportella/docker-beat/plugin"
	"golang.org/x/net/websocket"
)

var (
	WebsocketOrigin   string
	WebsocketProtocol string
	WebsocketEndpoint string
)

const (
	defaultWebsocketEndpoint = ""
	WebsocketEndpointUsage   = "websocket: The URL that events will be streamed too."

	defaultWebsocketProtocol = ""
	WebsocketProtocolUsage   = "websocket: The protocol to be used in the web socket stream."

	defaultWebsocketOrigin = ""
	WebsocketOriginUsage   = "websocket: The origin of the request to be used in the web socket stream."

	userAgent = "Docker-Beat (https://github.com/dmportella/docker-beat, 0.0.0)"
)

type consumer struct {
	socket *websocket.Conn
}

func (consumer *consumer) OnEvent(event plugin.DockerEvent) {
	ws, err := websocket.Dial(WebsocketEndpoint, WebsocketProtocol, WebsocketOrigin)
	if err != nil {
		logging.Error.Println(err.Error())
		return
	}

	data, _ := json.MarshalIndent(event, "", "    ")
	_, err = ws.Write(data)
	if err != nil {
		logging.Error.Printf(err.Error())
	}
}

func init() {
	flag.StringVar(&WebsocketEndpoint, "websocket-endpoint", defaultWebsocketEndpoint, WebsocketEndpointUsage)
	flag.StringVar(&WebsocketProtocol, "websocket-protocol", defaultWebsocketProtocol, WebsocketProtocolUsage)
	flag.StringVar(&WebsocketOrigin, "websocket-origin", defaultWebsocketOrigin, WebsocketOriginUsage)

	consumer := &consumer{}

	plugin.RegisterConsumer("websocket", consumer)
}
<|endoftext|>


<commit_before>package shell

import (
	"context"
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/erasure_coding"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/storage/super_block"
	"io"
	"sort"
)

func init() {
	Commands = append(Commands, &commandVolumeServerEvacuate{})
}

type commandVolumeServerEvacuate struct {
}

func (c *commandVolumeServerEvacuate) Name() string {
	return "volumeServer.evacuate"
}

func (c *commandVolumeServerEvacuate) Help() string {
	return `move out all data on a volume server

	volumeServer.evacuate -node <host:port>

	This command moves all data away from the volume server.
	The volumes on the volume servers will be redistributed.

	Usually this is used to prepare to shutdown or upgrade the volume server.

	Sometimes a volume can not be moved because there are no
	good destination to meet the replication requirement.
	E.g. a volume replication 001 in a cluster with 2 volume servers can not be moved.
	You can use "-skipNonMoveable" to move the rest volumes.

`
}

func (c *commandVolumeServerEvacuate) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	if err = commandEnv.confirmIsLocked(); err != nil {
		return
	}

	vsEvacuateCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	volumeServer := vsEvacuateCommand.String("node", "", "<host>:<port> of the volume server")
	skipNonMoveable := vsEvacuateCommand.Bool("skipNonMoveable", false, "skip volumes that can not be moved")
	applyChange := vsEvacuateCommand.Bool("force", false, "actually apply the changes")
	if err = vsEvacuateCommand.Parse(args); err != nil {
		return nil
	}

	if *volumeServer == "" {
		return fmt.Errorf("need to specify volume server by -node=<host>:<port>")
	}

	return volumeServerEvacuate(commandEnv, *volumeServer, *skipNonMoveable, *applyChange, writer)

}

func volumeServerEvacuate(commandEnv *CommandEnv, volumeServer string, skipNonMoveable, applyChange bool, writer io.Writer) (err error) {
	// 1. confirm the volume server is part of the cluster
	// 2. collect all other volume servers, sort by empty slots
	// 3. move to any other volume server as long as it satisfy the replication requirements

	// list all the volumes
	var resp *master_pb.VolumeListResponse
	err = commandEnv.MasterClient.WithClient(func(client master_pb.SeaweedClient) error {
		resp, err = client.VolumeList(context.Background(), &master_pb.VolumeListRequest{})
		return err
	})
	if err != nil {
		return err
	}

	if err := evacuateNormalVolumes(commandEnv, resp, volumeServer, skipNonMoveable, applyChange, writer); err != nil {
		return err
	}

	if err := evacuateEcVolumes(commandEnv, resp, volumeServer, skipNonMoveable, applyChange, writer); err != nil {
		return err
	}

	return nil
}

func evacuateNormalVolumes(commandEnv *CommandEnv, resp *master_pb.VolumeListResponse, volumeServer string, skipNonMoveable, applyChange bool, writer io.Writer) error {
	// find this volume server
	volumeServers := collectVolumeServersByDc(resp.TopologyInfo, "")
	thisNode, otherNodes := nodesOtherThan(volumeServers, volumeServer)
	if thisNode == nil {
		return fmt.Errorf("%s is not found in this cluster", volumeServer)
	}

	// move away normal volumes
	volumeReplicas, _ := collectVolumeReplicaLocations(resp)
	for _, vol := range thisNode.info.VolumeInfos {
		hasMoved, err := moveAwayOneNormalVolume(commandEnv, volumeReplicas, vol, thisNode, otherNodes, applyChange)
		if err != nil {
			return fmt.Errorf("move away volume %d from %s: %v", vol.Id, volumeServer, err)
		}
		if !hasMoved {
			if skipNonMoveable {
				replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(vol.ReplicaPlacement))
				fmt.Fprintf(writer, "skipping non moveable volume %d replication:%s\n", vol.Id, replicaPlacement.String())
			} else {
				return fmt.Errorf("failed to move volume %d from %s", vol.Id, volumeServer)
			}
		}
	}
	return nil
}

func evacuateEcVolumes(commandEnv *CommandEnv, resp *master_pb.VolumeListResponse, volumeServer string, skipNonMoveable, applyChange bool, writer io.Writer) error {
	// find this ec volume server
	ecNodes, _ := collectEcVolumeServersByDc(resp.TopologyInfo, "")
	thisNode, otherNodes := ecNodesOtherThan(ecNodes, volumeServer)
	if thisNode == nil {
		return fmt.Errorf("%s is not found in this cluster\n", volumeServer)
	}

	// move away ec volumes
	for _, ecShardInfo := range thisNode.info.EcShardInfos {
		hasMoved, err := moveAwayOneEcVolume(commandEnv, ecShardInfo, thisNode, otherNodes, applyChange)
		if err != nil {
			return fmt.Errorf("move away volume %d from %s: %v", ecShardInfo.Id, volumeServer, err)
		}
		if !hasMoved {
			if skipNonMoveable {
				fmt.Fprintf(writer, "failed to move away ec volume %d from %s\n", ecShardInfo.Id, volumeServer)
			} else {
				return fmt.Errorf("failed to move away ec volume %d from %s", ecShardInfo.Id, volumeServer)
			}
		}
	}
	return nil
}

func moveAwayOneEcVolume(commandEnv *CommandEnv, ecShardInfo *master_pb.VolumeEcShardInformationMessage, thisNode *EcNode, otherNodes []*EcNode, applyChange bool) (hasMoved bool, err error) {

	for _, shardId := range erasure_coding.ShardBits(ecShardInfo.EcIndexBits).ShardIds() {

		sort.Slice(otherNodes, func(i, j int) bool {
			return otherNodes[i].localShardIdCount(ecShardInfo.Id) < otherNodes[j].localShardIdCount(ecShardInfo.Id)
		})

		for i := 0; i < len(otherNodes); i++ {
			emptyNode := otherNodes[i]
			err = moveMountedShardToEcNode(commandEnv, thisNode, ecShardInfo.Collection, needle.VolumeId(ecShardInfo.Id), shardId, emptyNode, applyChange)
			if err != nil {
				return
			} else {
				hasMoved = true
				break
			}
		}
		if !hasMoved {
			return
		}
	}

	return
}

func moveAwayOneNormalVolume(commandEnv *CommandEnv, volumeReplicas map[uint32][]*VolumeReplica, vol *master_pb.VolumeInformationMessage, thisNode *Node, otherNodes []*Node, applyChange bool) (hasMoved bool, err error) {
	sort.Slice(otherNodes, func(i, j int) bool {
		return otherNodes[i].localVolumeRatio() < otherNodes[j].localVolumeRatio()
	})

	for i := 0; i < len(otherNodes); i++ {
		emptyNode := otherNodes[i]
		hasMoved, err = maybeMoveOneVolume(commandEnv, volumeReplicas, thisNode, vol, emptyNode, applyChange)
		if err != nil {
			return
		}
		if hasMoved {
			break
		}
	}
	return
}

func nodesOtherThan(volumeServers []*Node, thisServer string) (thisNode *Node, otherNodes []*Node) {
	for _, node := range volumeServers {
		if node.info.Id == thisServer {
			thisNode = node
			continue
		}
		otherNodes = append(otherNodes, node)
	}
	return
}

func ecNodesOtherThan(volumeServers []*EcNode, thisServer string) (thisNode *EcNode, otherNodes []*EcNode) {
	for _, node := range volumeServers {
		if node.info.Id == thisServer {
			thisNode = node
			continue
		}
		otherNodes = append(otherNodes, node)
	}
	return
}
<commit_msg>shell: volumeServer.evacuate adds printout for ec volumes<commit_after>package shell

import (
	"context"
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/erasure_coding"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/storage/super_block"
	"io"
	"os"
	"sort"
)

func init() {
	Commands = append(Commands, &commandVolumeServerEvacuate{})
}

type commandVolumeServerEvacuate struct {
}

func (c *commandVolumeServerEvacuate) Name() string {
	return "volumeServer.evacuate"
}

func (c *commandVolumeServerEvacuate) Help() string {
	return `move out all data on a volume server

	volumeServer.evacuate -node <host:port>

	This command moves all data away from the volume server.
	The volumes on the volume servers will be redistributed.

	Usually this is used to prepare to shutdown or upgrade the volume server.

	Sometimes a volume can not be moved because there are no
	good destination to meet the replication requirement.
	E.g. a volume replication 001 in a cluster with 2 volume servers can not be moved.
	You can use "-skipNonMoveable" to move the rest volumes.

`
}

func (c *commandVolumeServerEvacuate) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	if err = commandEnv.confirmIsLocked(); err != nil {
		return
	}

	vsEvacuateCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	volumeServer := vsEvacuateCommand.String("node", "", "<host>:<port> of the volume server")
	skipNonMoveable := vsEvacuateCommand.Bool("skipNonMoveable", false, "skip volumes that can not be moved")
	applyChange := vsEvacuateCommand.Bool("force", false, "actually apply the changes")
	if err = vsEvacuateCommand.Parse(args); err != nil {
		return nil
	}

	if *volumeServer == "" {
		return fmt.Errorf("need to specify volume server by -node=<host>:<port>")
	}

	return volumeServerEvacuate(commandEnv, *volumeServer, *skipNonMoveable, *applyChange, writer)

}

func volumeServerEvacuate(commandEnv *CommandEnv, volumeServer string, skipNonMoveable, applyChange bool, writer io.Writer) (err error) {
	// 1. confirm the volume server is part of the cluster
	// 2. collect all other volume servers, sort by empty slots
	// 3. move to any other volume server as long as it satisfy the replication requirements

	// list all the volumes
	var resp *master_pb.VolumeListResponse
	err = commandEnv.MasterClient.WithClient(func(client master_pb.SeaweedClient) error {
		resp, err = client.VolumeList(context.Background(), &master_pb.VolumeListRequest{})
		return err
	})
	if err != nil {
		return err
	}

	if err := evacuateNormalVolumes(commandEnv, resp, volumeServer, skipNonMoveable, applyChange, writer); err != nil {
		return err
	}

	if err := evacuateEcVolumes(commandEnv, resp, volumeServer, skipNonMoveable, applyChange, writer); err != nil {
		return err
	}

	return nil
}

func evacuateNormalVolumes(commandEnv *CommandEnv, resp *master_pb.VolumeListResponse, volumeServer string, skipNonMoveable, applyChange bool, writer io.Writer) error {
	// find this volume server
	volumeServers := collectVolumeServersByDc(resp.TopologyInfo, "")
	thisNode, otherNodes := nodesOtherThan(volumeServers, volumeServer)
	if thisNode == nil {
		return fmt.Errorf("%s is not found in this cluster", volumeServer)
	}

	// move away normal volumes
	volumeReplicas, _ := collectVolumeReplicaLocations(resp)
	for _, vol := range thisNode.info.VolumeInfos {
		hasMoved, err := moveAwayOneNormalVolume(commandEnv, volumeReplicas, vol, thisNode, otherNodes, applyChange)
		if err != nil {
			return fmt.Errorf("move away volume %d from %s: %v", vol.Id, volumeServer, err)
		}
		if !hasMoved {
			if skipNonMoveable {
				replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(vol.ReplicaPlacement))
				fmt.Fprintf(writer, "skipping non moveable volume %d replication:%s\n", vol.Id, replicaPlacement.String())
			} else {
				return fmt.Errorf("failed to move volume %d from %s", vol.Id, volumeServer)
			}
		}
	}
	return nil
}

func evacuateEcVolumes(commandEnv *CommandEnv, resp *master_pb.VolumeListResponse, volumeServer string, skipNonMoveable, applyChange bool, writer io.Writer) error {
	// find this ec volume server
	ecNodes, _ := collectEcVolumeServersByDc(resp.TopologyInfo, "")
	thisNode, otherNodes := ecNodesOtherThan(ecNodes, volumeServer)
	if thisNode == nil {
		return fmt.Errorf("%s is not found in this cluster\n", volumeServer)
	}

	// move away ec volumes
	for _, ecShardInfo := range thisNode.info.EcShardInfos {
		hasMoved, err := moveAwayOneEcVolume(commandEnv, ecShardInfo, thisNode, otherNodes, applyChange)
		if err != nil {
			return fmt.Errorf("move away volume %d from %s: %v", ecShardInfo.Id, volumeServer, err)
		}
		if !hasMoved {
			if skipNonMoveable {
				fmt.Fprintf(writer, "failed to move away ec volume %d from %s\n", ecShardInfo.Id, volumeServer)
			} else {
				return fmt.Errorf("failed to move away ec volume %d from %s", ecShardInfo.Id, volumeServer)
			}
		}
	}
	return nil
}

func moveAwayOneEcVolume(commandEnv *CommandEnv, ecShardInfo *master_pb.VolumeEcShardInformationMessage, thisNode *EcNode, otherNodes []*EcNode, applyChange bool) (hasMoved bool, err error) {

	for _, shardId := range erasure_coding.ShardBits(ecShardInfo.EcIndexBits).ShardIds() {

		sort.Slice(otherNodes, func(i, j int) bool {
			return otherNodes[i].localShardIdCount(ecShardInfo.Id) < otherNodes[j].localShardIdCount(ecShardInfo.Id)
		})

		for i := 0; i < len(otherNodes); i++ {
			emptyNode := otherNodes[i]
			collectionPrefix := ""
			if ecShardInfo.Collection != "" {
				collectionPrefix = ecShardInfo.Collection + "_"
			}
			fmt.Fprintf(os.Stdout, "moving ec volume %s%d.%d %s => %s\n", collectionPrefix, ecShardInfo.Id, shardId, thisNode.info.Id, emptyNode.info.Id)
			err = moveMountedShardToEcNode(commandEnv, thisNode, ecShardInfo.Collection, needle.VolumeId(ecShardInfo.Id), shardId, emptyNode, applyChange)
			if err != nil {
				return
			} else {
				hasMoved = true
				break
			}
		}
		if !hasMoved {
			return
		}
	}

	return
}

func moveAwayOneNormalVolume(commandEnv *CommandEnv, volumeReplicas map[uint32][]*VolumeReplica, vol *master_pb.VolumeInformationMessage, thisNode *Node, otherNodes []*Node, applyChange bool) (hasMoved bool, err error) {
	sort.Slice(otherNodes, func(i, j int) bool {
		return otherNodes[i].localVolumeRatio() < otherNodes[j].localVolumeRatio()
	})

	for i := 0; i < len(otherNodes); i++ {
		emptyNode := otherNodes[i]
		hasMoved, err = maybeMoveOneVolume(commandEnv, volumeReplicas, thisNode, vol, emptyNode, applyChange)
		if err != nil {
			return
		}
		if hasMoved {
			break
		}
	}
	return
}

func nodesOtherThan(volumeServers []*Node, thisServer string) (thisNode *Node, otherNodes []*Node) {
	for _, node := range volumeServers {
		if node.info.Id == thisServer {
			thisNode = node
			continue
		}
		otherNodes = append(otherNodes, node)
	}
	return
}

func ecNodesOtherThan(volumeServers []*EcNode, thisServer string) (thisNode *EcNode, otherNodes []*EcNode) {
	for _, node := range volumeServers {
		if node.info.Id == thisServer {
			thisNode = node
			continue
		}
		otherNodes = append(otherNodes, node)
	}
	return
}
<|endoftext|>


<commit_before>package disk

import (
	"fmt"
	"runtime"
	"testing"
)

func TestDisk_usage(t *testing.T) {
	path := "/"
	if runtime.GOOS == "windows" {
		path = "C:"
	}
	v, err := Usage(path)
	if err != nil {
		t.Errorf("error %v", err)
	}
	if v.Path != path {
		t.Errorf("error %v", err)
	}
}

func TestDisk_partitions(t *testing.T) {
	ret, err := Partitions(false)
	if err != nil || len(ret) == 0 {
		t.Errorf("error %v", err)
	}
	t.Log(ret)

	empty := PartitionStat{}
	if len(ret) == 0 {
		t.Errorf("ret is empty")
	}
	for _, disk := range ret {
		if disk == empty {
			t.Errorf("Could not get device info %v", disk)
		}
	}
}

func TestDisk_io_counters(t *testing.T) {
	ret, err := IOCounters()
	if err != nil {
		t.Errorf("error %v", err)
	}
	if len(ret) == 0 {
		t.Errorf("ret is empty")
	}
	empty := IOCountersStat{}
	for part, io := range ret {
		t.Log(part, io)
		if io == empty {
			t.Errorf("io_counter error %v, %v", part, io)
		}
	}
}

func TestDiskUsageStat_String(t *testing.T) {
	v := UsageStat{
		Path:              "/",
		Total:             1000,
		Free:              2000,
		Used:              3000,
		UsedPercent:       50.1,
		InodesTotal:       4000,
		InodesUsed:        5000,
		InodesFree:        6000,
		InodesUsedPercent: 49.1,
		Fstype:            "ext4",
	}
	e := `{"path":"/","fstype":"ext4","total":1000,"free":2000,"used":3000,"usedPercent":50.1,"inodesTotal":4000,"inodesUsed":5000,"inodesFree":6000,"inodesUsedPercent":49.1}`
	if e != fmt.Sprintf("%v", v) {
		t.Errorf("DiskUsageStat string is invalid: %v", v)
	}
}

func TestDiskPartitionStat_String(t *testing.T) {
	v := PartitionStat{
		Device:     "sd01",
		Mountpoint: "/",
		Fstype:     "ext4",
		Opts:       "ro",
	}
	e := `{"device":"sd01","mountpoint":"/","fstype":"ext4","opts":"ro"}`
	if e != fmt.Sprintf("%v", v) {
		t.Errorf("DiskUsageStat string is invalid: %v", v)
	}
}

func TestDiskIOCountersStat_String(t *testing.T) {
	v := IOCountersStat{
		Name:         "sd01",
		ReadCount:    100,
		WriteCount:   200,
		ReadBytes:    300,
		WriteBytes:   400,
		SerialNumber: "SERIAL",
	}
	e := `{"readCount":100,"mergedReadCount":0,"writeCount":200,"mergedWriteCount":0,"readBytes":300,"writeBytes":400,"readTime":0,"writeTime":0,"iopsInProgress":0,"ioTime":0,"weightedIO":0,"name":"sd01","serialNumber":"SERIAL"}`
	if e != fmt.Sprintf("%v", v) {
		t.Errorf("DiskUsageStat string is invalid: %v", v)
	}
}
<commit_msg>[disk]linux: fix test<commit_after>package disk

import (
	"fmt"
	"runtime"
	"testing"
)

func TestDisk_usage(t *testing.T) {
	path := "/"
	if runtime.GOOS == "windows" {
		path = "C:"
	}
	v, err := Usage(path)
	if err != nil {
		t.Errorf("error %v", err)
	}
	if v.Path != path {
		t.Errorf("error %v", err)
	}
}

func TestDisk_partitions(t *testing.T) {
	ret, err := Partitions(false)
	if err != nil || len(ret) == 0 {
		t.Errorf("error %v", err)
	}
	t.Log(ret)

	empty := PartitionStat{}
	if len(ret) == 0 {
		t.Errorf("ret is empty")
	}
	for _, disk := range ret {
		if disk == empty {
			t.Errorf("Could not get device info %v", disk)
		}
	}
}

func TestDisk_io_counters(t *testing.T) {
	ret, err := IOCounters()
	if err != nil {
		t.Errorf("error %v", err)
	}
	if len(ret) == 0 {
		t.Errorf("ret is empty")
	}
	empty := IOCountersStat{}
	for part, io := range ret {
		t.Log(part, io)
		if io == empty {
			t.Errorf("io_counter error %v, %v", part, io)
		}
	}
}

func TestDiskUsageStat_String(t *testing.T) {
	v := UsageStat{
		Path:              "/",
		Total:             1000,
		Free:              2000,
		Used:              3000,
		UsedPercent:       50.1,
		InodesTotal:       4000,
		InodesUsed:        5000,
		InodesFree:        6000,
		InodesUsedPercent: 49.1,
		Fstype:            "ext4",
	}
	e := `{"path":"/","fstype":"ext4","total":1000,"free":2000,"used":3000,"usedPercent":50.1,"inodesTotal":4000,"inodesUsed":5000,"inodesFree":6000,"inodesUsedPercent":49.1}`
	if e != fmt.Sprintf("%v", v) {
		t.Errorf("DiskUsageStat string is invalid: %v", v)
	}
}

func TestDiskPartitionStat_String(t *testing.T) {
	v := PartitionStat{
		Device:     "sd01",
		Mountpoint: "/",
		Fstype:     "ext4",
		Opts:       "ro",
	}
	e := `{"device":"sd01","mountpoint":"/","fstype":"ext4","opts":"ro"}`
	if e != fmt.Sprintf("%v", v) {
		t.Errorf("DiskUsageStat string is invalid: %v", v)
	}
}

func TestDiskIOCountersStat_String(t *testing.T) {
	v := IOCountersStat{
		Name:         "sd01",
		ReadCount:    100,
		WriteCount:   200,
		ReadBytes:    300,
		WriteBytes:   400,
		SerialNumber: "SERIAL",
	}
	e := `{"readCount":100,"mergedReadCount":0,"writeCount":200,"mergedWriteCount":0,"readBytes":300,"writeBytes":400,"readTime":0,"writeTime":0,"iopsInProgress":0,"ioTime":0,"weightedIO":0,"name":"sd01","serialNumber":"SERIAL","label":""}`
	if e != fmt.Sprintf("%v", v) {
		t.Errorf("DiskUsageStat string is invalid: %v", v)
	}
}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp. 2016 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 main

//WARNING - this chaincode's ID is hard-coded in chaincode_example04 to illustrate one way of
//calling chaincode from a chaincode. If this example is modified, chaincode_example04.go has
//to be modified as well with the new ID of chaincode_example02.
//chaincode_example05 show's how chaincode ID can be passed in as a parameter instead of
//hard-coding.

import (
	"errors"
	"fmt"
	"strconv"
	"bytes"

	"github.com/hyperledger/fabric/core/chaincode/shim"
)

// SimpleChaincode example simple Chaincode implementation
type SimpleChaincode struct {
}

func (t *SimpleChaincode) Init(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
	var A, B string    // Entities
	var Aval, Bval int // Asset holdings
	var err error

	if len(args) != 4 {
		return nil, errors.New("Incorrect number of arguments. Expecting 4")
	}

	// Initialize the chaincode
	A = args[0]
	Aval, err = strconv.Atoi(args[1])
	if err != nil {
		return nil, errors.New("Expecting integer value for asset holding")
	}
	B = args[2]
	Bval, err = strconv.Atoi(args[3])
	if err != nil {
		return nil, errors.New("Expecting integer value for asset holding")
	}
	fmt.Printf("Aval = %d, Bval = %d\n", Aval, Bval)

	// Write the state to the ledger
	err = stub.PutState(A, []byte(strconv.Itoa(Aval)))
	if err != nil {
		return nil, err
	}

	err = stub.PutState(B, []byte(strconv.Itoa(Bval)))
	if err != nil {
		return nil, err
	}

	return nil, nil
}

func (t *SimpleChaincode) write(stub *shim.ChaincodeStub, args []string) ([]byte, error) {
    var key, value string
    var err error
    fmt.Println("Storing the parameters in hyperledger fabric...")

    /*if len(args) != 2 {
        return nil, errors.New("Incorrect number of arguments. Expecting 2. name of the key and value to set")
    }*/

	if(len(args)%2 != 0) {
		  fmt.Printf("Incorrect number of arguments. One of the keys or values is missing.")
		  fmt.Println("")

    }else{
	     for i := 0; i < len(args); i++ {
	     if(i%2 == 0){
		     if args[i] != "" {
                  fmt.Printf("Key: %s", args[i])
				  fmt.Println("")
				  key = args[i]
				  i++
             }
		     if(i!=len(args)) {
			      fmt.Printf("Value: %s", args[i])
			      fmt.Println("")
				  value = args[i]
			 }
		     err = stub.PutState(key, []byte(value))  //write the variable into the chaincode state
			 if err != nil {
				return nil, err
			 }
		 }
    }
	}

	/*
    key = args[0]                            //rename for fun
    value = args[1]
    err = stub.PutState(key, []byte(value))  //write the variable into the chaincode state
    if err != nil {
        return nil, err
    }
	*/

    return nil, nil
}

// Transaction makes payment of X units from A to B
func (t *SimpleChaincode) Invoke(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
	if function == "delete" {
		// Deletes an entity from its state
		return t.delete(stub, args)
	}
	if function == "write" {
		fmt.Println("Calling write()")
        return t.write(stub, args)
    }

	var A, B string    // Entities
	var Aval, Bval int // Asset holdings
	var X int          // Transaction value
	var err error

	if len(args) != 3 {
		return nil, errors.New("Incorrect number of arguments. Expecting 3")
	}

	A = args[0]
	B = args[1]

	// Get the state from the ledger
	// TODO: will be nice to have a GetAllState call to ledger
	Avalbytes, err := stub.GetState(A)
	if err != nil {
		return nil, errors.New("Failed to get state")
	}
	if Avalbytes == nil {
		return nil, errors.New("Entity not found")
	}
	Aval, _ = strconv.Atoi(string(Avalbytes))

	Bvalbytes, err := stub.GetState(B)
	if err != nil {
		return nil, errors.New("Failed to get state")
	}
	if Bvalbytes == nil {
		return nil, errors.New("Entity not found")
	}
	Bval, _ = strconv.Atoi(string(Bvalbytes))

	// Perform the execution
	X, err = strconv.Atoi(args[2])
	Aval = Aval - X
	Bval = Bval + X
	fmt.Printf("Aval = %d, Bval = %d\n", Aval, Bval)

	// Write the state back to the ledger
	err = stub.PutState(A, []byte(strconv.Itoa(Aval)))
	if err != nil {
		return nil, err
	}

	err = stub.PutState(B, []byte(strconv.Itoa(Bval)))
	if err != nil {
		return nil, err
	}

	return nil, nil
}

// Deletes an entity from state
func (t *SimpleChaincode) delete(stub *shim.ChaincodeStub, args []string) ([]byte, error) {
	if len(args) != 1 {
		return nil, errors.New("Incorrect number of arguments. Expecting 1")
	}

	A := args[0]

	// Delete the key from the state in ledger
	err := stub.DelState(A)
	if err != nil {
		return nil, errors.New("Failed to delete state")
	}

	return nil, nil
}

// Query callback representing the query of a chaincode
func (t *SimpleChaincode) Query(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
	/*if function != "query" {
		return nil, errors.New("Invalid query function name. Expecting \"query\"")
	}*/

	if function == "queryAll" {
		fmt.Println("Calling QueryAll()")
        return t.queryAll(stub, args)
    }

	var A string // Entities
	var err error

	if len(args) != 1 {
		return nil, errors.New("Incorrect number of arguments. Expecting name of the person to query")
	}

	A = args[0]

	// Get the state from the ledger
	Avalbytes, err := stub.GetState(A)
	if err != nil {
		jsonResp := "{\"Error\":\"Failed to get state for " + A + "\"}"
		return nil, errors.New(jsonResp)
	}

	if Avalbytes == nil {
		jsonResp := "{\"Error\":\"Nil amount for " + A + "\"}"
		return nil, errors.New(jsonResp)
	}

	jsonResp := "{\"Name\":\"" + A + "\",\"Amount\":\"" + string(Avalbytes) + "\"}"
	fmt.Printf("Query Response:%s\n", jsonResp)
	return Avalbytes, nil
}

// Query callback representing the query of a chaincode
func (t *SimpleChaincode) queryAll(stub *shim.ChaincodeStub, args []string) ([]byte, error) {

	//var A string // Entities
	//var err error

	/*if len(args) != 1 {
		return nil, errors.New("Incorrect number of arguments. Expecting name of the person to query")
	}*/
    var RetValue []byte
	var buffer bytes.Buffer
		for i := 0; i < len(args); i++ {
		    Avalbytes, err := stub.GetState(args[i])
			if err != nil {
				jsonResp := "{\"Error\":\"Failed to get state for " + args[i] + "\"}"
				return nil, errors.New(jsonResp)
			}

			if Avalbytes == nil {
				jsonResp := "{\"Error\":\"Nil amount for " + args[i] + "\"}"
				return nil, errors.New(jsonResp)
			}
			RetValue = Avalbytes
			jsonRespString :=  "\"Name\":\"" + args[i] + "\",\"Value\":\"" + string(RetValue) + "\""
			buffer.WriteString(jsonRespString)

		}
		jsonResp := "{"+buffer.String()+"}"
		fmt.Printf("Query Response:%s\n", jsonResp)
		return RetValue, nil

	/*
	A = args[0]

	// Get the state from the ledger
	Avalbytes, err := stub.GetState(A)
	if err != nil {
		jsonResp := "{\"Error\":\"Failed to get state for " + A + "\"}"
		return nil, errors.New(jsonResp)
	}

	if Avalbytes == nil {
		jsonResp := "{\"Error\":\"Nil amount for " + A + "\"}"
		return nil, errors.New(jsonResp)
	}

	jsonResp := "{\"Name\":\"" + A + "\",\"Amount\":\"" + string(Avalbytes) + "\"}"
	fmt.Printf("Query Response:%s\n", jsonResp)
	return Avalbytes, nil
	*/
}

func main() {
	err := shim.Start(new(SimpleChaincode))
	if err != nil {
		fmt.Printf("Error starting Simple chaincode: %s", err)
	}
}
<commit_msg>Updated queryAll method - 1<commit_after>/*
Copyright IBM Corp. 2016 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 main

//WARNING - this chaincode's ID is hard-coded in chaincode_example04 to illustrate one way of
//calling chaincode from a chaincode. If this example is modified, chaincode_example04.go has
//to be modified as well with the new ID of chaincode_example02.
//chaincode_example05 show's how chaincode ID can be passed in as a parameter instead of
//hard-coding.

import (
	"errors"
	"fmt"
	"strconv"
	"bytes"

	"github.com/hyperledger/fabric/core/chaincode/shim"
)

// SimpleChaincode example simple Chaincode implementation
type SimpleChaincode struct {
}

func (t *SimpleChaincode) Init(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
	var A, B string    // Entities
	var Aval, Bval int // Asset holdings
	var err error

	if len(args) != 4 {
		return nil, errors.New("Incorrect number of arguments. Expecting 4")
	}

	// Initialize the chaincode
	A = args[0]
	Aval, err = strconv.Atoi(args[1])
	if err != nil {
		return nil, errors.New("Expecting integer value for asset holding")
	}
	B = args[2]
	Bval, err = strconv.Atoi(args[3])
	if err != nil {
		return nil, errors.New("Expecting integer value for asset holding")
	}
	fmt.Printf("Aval = %d, Bval = %d\n", Aval, Bval)

	// Write the state to the ledger
	err = stub.PutState(A, []byte(strconv.Itoa(Aval)))
	if err != nil {
		return nil, err
	}

	err = stub.PutState(B, []byte(strconv.Itoa(Bval)))
	if err != nil {
		return nil, err
	}

	return nil, nil
}

func (t *SimpleChaincode) write(stub *shim.ChaincodeStub, args []string) ([]byte, error) {
    var key, value string
    var err error
    fmt.Println("Storing the parameters in hyperledger fabric...")

    /*if len(args) != 2 {
        return nil, errors.New("Incorrect number of arguments. Expecting 2. name of the key and value to set")
    }*/

	if(len(args)%2 != 0) {
		  fmt.Printf("Incorrect number of arguments. One of the keys or values is missing.")
		  fmt.Println("")

    }else{
	     for i := 0; i < len(args); i++ {
	     if(i%2 == 0){
		     if args[i] != "" {
                  fmt.Printf("Key: %s", args[i])
				  fmt.Println("")
				  key = args[i]
				  i++
             }
		     if(i!=len(args)) {
			      fmt.Printf("Value: %s", args[i])
			      fmt.Println("")
				  value = args[i]
			 }
		     err = stub.PutState(key, []byte(value))  //write the variable into the chaincode state
			 if err != nil {
				return nil, err
			 }
		 }
    }
	}

	/*
    key = args[0]                            //rename for fun
    value = args[1]
    err = stub.PutState(key, []byte(value))  //write the variable into the chaincode state
    if err != nil {
        return nil, err
    }
	*/

    return nil, nil
}

// Transaction makes payment of X units from A to B
func (t *SimpleChaincode) Invoke(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
	if function == "delete" {
		// Deletes an entity from its state
		return t.delete(stub, args)
	}
	if function == "write" {
		fmt.Println("Calling write()")
        return t.write(stub, args)
    }

	var A, B string    // Entities
	var Aval, Bval int // Asset holdings
	var X int          // Transaction value
	var err error

	if len(args) != 3 {
		return nil, errors.New("Incorrect number of arguments. Expecting 3")
	}

	A = args[0]
	B = args[1]

	// Get the state from the ledger
	// TODO: will be nice to have a GetAllState call to ledger
	Avalbytes, err := stub.GetState(A)
	if err != nil {
		return nil, errors.New("Failed to get state")
	}
	if Avalbytes == nil {
		return nil, errors.New("Entity not found")
	}
	Aval, _ = strconv.Atoi(string(Avalbytes))

	Bvalbytes, err := stub.GetState(B)
	if err != nil {
		return nil, errors.New("Failed to get state")
	}
	if Bvalbytes == nil {
		return nil, errors.New("Entity not found")
	}
	Bval, _ = strconv.Atoi(string(Bvalbytes))

	// Perform the execution
	X, err = strconv.Atoi(args[2])
	Aval = Aval - X
	Bval = Bval + X
	fmt.Printf("Aval = %d, Bval = %d\n", Aval, Bval)

	// Write the state back to the ledger
	err = stub.PutState(A, []byte(strconv.Itoa(Aval)))
	if err != nil {
		return nil, err
	}

	err = stub.PutState(B, []byte(strconv.Itoa(Bval)))
	if err != nil {
		return nil, err
	}

	return nil, nil
}

// Deletes an entity from state
func (t *SimpleChaincode) delete(stub *shim.ChaincodeStub, args []string) ([]byte, error) {
	if len(args) != 1 {
		return nil, errors.New("Incorrect number of arguments. Expecting 1")
	}

	A := args[0]

	// Delete the key from the state in ledger
	err := stub.DelState(A)
	if err != nil {
		return nil, errors.New("Failed to delete state")
	}

	return nil, nil
}

// Query callback representing the query of a chaincode
func (t *SimpleChaincode) Query(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
	/*if function != "query" {
		return nil, errors.New("Invalid query function name. Expecting \"query\"")
	}*/

	if function == "queryAll" {
		fmt.Println("Calling QueryAll()")
        return t.queryAll(stub, args)
    }

	var A string // Entities
	var err error

	if len(args) != 1 {
		return nil, errors.New("Incorrect number of arguments. Expecting name of the person to query")
	}

	A = args[0]

	// Get the state from the ledger
	Avalbytes, err := stub.GetState(A)
	if err != nil {
		jsonResp := "{\"Error\":\"Failed to get state for " + A + "\"}"
		return nil, errors.New(jsonResp)
	}

	if Avalbytes == nil {
		jsonResp := "{\"Error\":\"Nil amount for " + A + "\"}"
		return nil, errors.New(jsonResp)
	}

	jsonResp := "{\"Name\":\"" + A + "\",\"Amount\":\"" + string(Avalbytes) + "\"}"
	fmt.Printf("Query Response:%s\n", jsonResp)
	return Avalbytes, nil
}

// Query callback representing the query of a chaincode
func (t *SimpleChaincode) queryAll(stub *shim.ChaincodeStub, args []string) ([]byte, error) {

	//var A string // Entities
	//var err error

	/*if len(args) != 1 {
		return nil, errors.New("Incorrect number of arguments. Expecting name of the person to query")
	}*/
    var RetValue []byte
	var buffer bytes.Buffer
		for i := 0; i < len(args); i++ {
		    Avalbytes, err := stub.GetState(args[i])
			if err != nil {
				jsonResp := "{\"Error\":\"Failed to get state for " + args[i] + "\"}"
				return nil, errors.New(jsonResp)
			}

			if Avalbytes == nil {
				jsonResp := "{\"Error\":\"Nil amount for " + args[i] + "\"}"
				return nil, errors.New(jsonResp)
			}

			jsonRespString :=  "\"Name\":\"" + args[i] + "\",\"Value\":\"" + string(Avalbytes) + "\""
			RetValue = []byte(jsonRespString)
			buffer.WriteString(jsonRespString)

		}
		jsonResp := "{"+buffer.String()+"}"
		fmt.Printf("Query Response:%s\n", jsonResp)
		return RetValue, nil

	/*
	A = args[0]

	// Get the state from the ledger
	Avalbytes, err := stub.GetState(A)
	if err != nil {
		jsonResp := "{\"Error\":\"Failed to get state for " + A + "\"}"
		return nil, errors.New(jsonResp)
	}

	if Avalbytes == nil {
		jsonResp := "{\"Error\":\"Nil amount for " + A + "\"}"
		return nil, errors.New(jsonResp)
	}

	jsonResp := "{\"Name\":\"" + A + "\",\"Amount\":\"" + string(Avalbytes) + "\"}"
	fmt.Printf("Query Response:%s\n", jsonResp)
	return Avalbytes, nil
	*/
}

func main() {
	err := shim.Start(new(SimpleChaincode))
	if err != nil {
		fmt.Printf("Error starting Simple chaincode: %s", err)
	}
}
<|endoftext|>


<commit_before>package interfaces

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/Tinker-Ware/gh-service/domain"
	"github.com/dvsekhvalnov/jose2go"
)

// Adapter is the signature of an HTTPHandler for middlewares
type Adapter func(http.Handler) http.Handler
type repository interface {
	SetToken(token string)
}

// Notify is a middleware to measure the time that a request takes
func Notify() Adapter {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			start := time.Now()
			defer log.Printf("%s on %s took %s\n", r.Method, r.URL.Path, time.Since(start))
			h.ServeHTTP(w, r)
		})
	}
}

// Adapt takes several Adapters and calls them in order
func Adapt(h http.Handler, adapters ...Adapter) http.Handler {
	for _, adapter := range adapters {
		h = adapter(h)
	}
	return h
}

// SetToken injects the token from the request
func SetToken(repo repository) Adapter {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			token := r.Header.Get(domain.TokenHeader)
			repo.SetToken(token)
			h.ServeHTTP(w, r)
		})
	}
}

type integration struct {
	UserID     int    `json:"user_id"`
	Token      string `json:"token"`
	Username   string `json:"username"`
	Provider   string `json:"provider"`
	ExpireDate int64  `json:"expire_date"`
}

type integrationsWrapper struct {
	Integrations []integration `json:"integrations"`
}

const integrationsURL string = "/api/v1/users/%s/integrations"

// GetToken gets the token from the users microservice
func GetToken(repo repository, apiURL string, salt string) Adapter {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			userToken := r.Header.Get("authorization")

			payload, _, _ := jose.Decode(userToken, []byte(salt))

			var objmap map[string]*json.RawMessage
			json.Unmarshal([]byte(payload), &objmap)

			var aud string
			json.Unmarshal(*objmap["aud"], &aud)

			vals := strings.Split(aud, ":")

			path := fmt.Sprintf(integrationsURL, vals[1])

			request, _ := http.NewRequest(http.MethodGet, apiURL+path, nil)
			request.Header.Set("authorization", userToken)

			client := &http.Client{}
			resp, _ := client.Do(request)
			defer resp.Body.Close()

			integrations := integrationsWrapper{}
			decoder := json.NewDecoder(resp.Body)
			decoder.Decode(&integrations)
			var token string

			for _, integ := range integrations.Integrations {
				if integ.Provider == "github" {
					token = integ.Token
				}
			}

			repo.SetToken(token)

			h.ServeHTTP(w, r)
		})
	}
}
<commit_msg>Split header to get token<commit_after>package interfaces

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/Tinker-Ware/gh-service/domain"
	"github.com/dvsekhvalnov/jose2go"
)

// Adapter is the signature of an HTTPHandler for middlewares
type Adapter func(http.Handler) http.Handler
type repository interface {
	SetToken(token string)
}

// Notify is a middleware to measure the time that a request takes
func Notify() Adapter {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			start := time.Now()
			defer log.Printf("%s on %s took %s\n", r.Method, r.URL.Path, time.Since(start))
			h.ServeHTTP(w, r)
		})
	}
}

// Adapt takes several Adapters and calls them in order
func Adapt(h http.Handler, adapters ...Adapter) http.Handler {
	for _, adapter := range adapters {
		h = adapter(h)
	}
	return h
}

// SetToken injects the token from the request
func SetToken(repo repository) Adapter {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			token := r.Header.Get(domain.TokenHeader)
			repo.SetToken(token)
			h.ServeHTTP(w, r)
		})
	}
}

type integration struct {
	UserID     int    `json:"user_id"`
	Token      string `json:"token"`
	Username   string `json:"username"`
	Provider   string `json:"provider"`
	ExpireDate int64  `json:"expire_date"`
}

type integrationsWrapper struct {
	Integrations []integration `json:"integrations"`
}

const integrationsURL string = "/api/v1/users/%s/integration"

// GetToken gets the token from the users microservice
func GetToken(repo repository, apiURL string, salt string) Adapter {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			userToken := r.Header.Get("authorization")

			sp := strings.Split(userToken, " ")
			payload, _, _ := jose.Decode(sp[1], []byte(salt))

			var objmap map[string]*json.RawMessage
			json.Unmarshal([]byte(payload), &objmap)

			var aud string
			json.Unmarshal(*objmap["aud"], &aud)

			vals := strings.Split(aud, ":")

			path := fmt.Sprintf(integrationsURL, vals[1])

			request, _ := http.NewRequest(http.MethodGet, apiURL+path, nil)
			request.Header.Set("authorization", userToken)

			client := &http.Client{}
			resp, _ := client.Do(request)
			defer resp.Body.Close()

			integrations := integrationsWrapper{}
			decoder := json.NewDecoder(resp.Body)
			decoder.Decode(&integrations)
			var token string

			for _, integ := range integrations.Integrations {
				if integ.Provider == "github" {
					token = integ.Token
				}
			}

			repo.SetToken(token)

			h.ServeHTTP(w, r)
		})
	}
}
<|endoftext|>


<commit_before>package core

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"

	"github.com/citihub/probr-sdk/config"
	"github.com/citihub/probr-sdk/plugin"
	"github.com/citihub/probr-sdk/probeengine"
	"github.com/citihub/probr-sdk/utils"

	hclog "github.com/hashicorp/go-hclog"
	hcplugin "github.com/hashicorp/go-plugin"
)

// BinariesPath represents the path where service pack binaries are installed
// Must be a pointer to accept the flag when it is set
var BinariesPath *string

// Verbose is a CLI option to increase output detail
var Verbose *bool

// AllPacks is a CLI option to target all installed packs, instead of just those specified in config.yml
var AllPacks *bool

// NewClient client handles the lifecycle of a plugin application
// Plugin hosts should use one Client for each plugin executable
// (this is different from the client that manages gRPC)
func NewClient(cmd *exec.Cmd) *hcplugin.Client {
	logger := hclog.New(&hclog.LoggerOptions{
		Name:   plugin.ServicePackPluginName,
		Output: os.Stdout,
		Level:  hclog.Debug,
	})
	var pluginMap = map[string]hcplugin.Plugin{
		plugin.ServicePackPluginName: &plugin.ServicePackPlugin{},
	}
	var handshakeConfig = plugin.GetHandshakeConfig()
	return hcplugin.NewClient(&hcplugin.ClientConfig{
		HandshakeConfig: handshakeConfig,
		Plugins:         pluginMap,
		Cmd:             cmd,
		Logger:          logger,
	})
}

// SetupCloseHandler creates a 'listener' on a new goroutine which will notify the
// program if it receives an interrupt from the OS. We then handle this by calling
// our clean up procedure and exiting the program.
// Ref: https://golangcode.com/handle-ctrl-c-exit-in-terminal/
func SetupCloseHandler() {
	c := make(chan os.Signal)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-c
		log.Printf("Execution aborted - %v", "SIGTERM")
		probeengine.CleanupTmp()
		os.Exit(0)
	}()
}

// GetCommands ...
func GetCommands() (cmdSet []*exec.Cmd, err error) {
	// TODO: give any exec errors a familiar format
	configPath, err := getConfigPath()
	if err != nil {
		return
	}
	packNames, err := GetPackNames()
	if err != nil {
		return
	}

	for _, pack := range packNames {
		binaryName, binErr := GetPackBinary(pack)
		if binErr != nil {
			err = binErr
			break
		}
		cmd := exec.Command(binaryName)
		cmd.Args = append(cmd.Args, fmt.Sprintf("--varsfile=%s", configPath))
		cmdSet = append(cmdSet, cmd)
	}
	if err == nil && len(cmdSet) == 0 {
		err = utils.ReformatError("No valid service packs specified")
	}
	return
}

// UserHomeDir provides the OS-aware user home directory
// TODO: move this to SDK
func UserHomeDir() string {
	dirname, err := os.UserHomeDir()
	if err != nil {
		log.Fatalf(err.Error())
	}
	return dirname
}

// GetPackBinary finds provided service pack in installation folder and return binary name
func GetPackBinary(name string) (binaryName string, err error) {
	name = filepath.Base(strings.ToLower(name)) // in some cases a filepath may arrive here instead of the base name
	if runtime.GOOS == "windows" && !strings.HasSuffix(name, ".exe") {
		name = fmt.Sprintf("%s.exe", name)
	}
	*BinariesPath = strings.Replace(*BinariesPath, "~", UserHomeDir(), 1)
	plugins, _ := hcplugin.Discover(name, *BinariesPath)
	if len(plugins) != 1 {
		err = fmt.Errorf("failed to locate requested plugin '%s'", name)
		return
	}
	binaryName = plugins[0]

	return
}

func getConfigPath() (string, error) {
	workDir, err := os.Getwd()
	if err != nil {
		return "", err
	}
	return filepath.Join(workDir, "config.yml"), nil
}

// GetPackNames returns all service packs declared in config file
func GetPackNames() (packNames []string, err error) {
	if AllPacks != nil && *AllPacks {
		return hcplugin.Discover("*", *BinariesPath)
	}
	packNames, err = getPackNamesFromConfig()
	return
}

func getPackNamesFromConfig() ([]string, error) {
	type simpleVars struct {
		Run []string `yaml:"Run"`
	}
	var vars simpleVars

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

	configDecoder, file, err := config.NewConfigDecoder(configPath)
	if err != nil {
		return nil, err
	}

	err = configDecoder.Decode(&vars)
	file.Close()
	return vars.Run, err

}
<commit_msg> will show all packs if no config file is present (#44)<commit_after>package core

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"

	"github.com/citihub/probr-sdk/config"
	"github.com/citihub/probr-sdk/plugin"
	"github.com/citihub/probr-sdk/probeengine"
	"github.com/citihub/probr-sdk/utils"

	hclog "github.com/hashicorp/go-hclog"
	hcplugin "github.com/hashicorp/go-plugin"
)

// BinariesPath represents the path where service pack binaries are installed
// Must be a pointer to accept the flag when it is set
var BinariesPath *string

// Verbose is a CLI option to increase output detail
var Verbose *bool

// AllPacks is a CLI option to target all installed packs, instead of just those specified in config.yml
var AllPacks *bool

// NewClient client handles the lifecycle of a plugin application
// Plugin hosts should use one Client for each plugin executable
// (this is different from the client that manages gRPC)
func NewClient(cmd *exec.Cmd) *hcplugin.Client {
	logger := hclog.New(&hclog.LoggerOptions{
		Name:   plugin.ServicePackPluginName,
		Output: os.Stdout,
		Level:  hclog.Debug,
	})
	var pluginMap = map[string]hcplugin.Plugin{
		plugin.ServicePackPluginName: &plugin.ServicePackPlugin{},
	}
	var handshakeConfig = plugin.GetHandshakeConfig()
	return hcplugin.NewClient(&hcplugin.ClientConfig{
		HandshakeConfig: handshakeConfig,
		Plugins:         pluginMap,
		Cmd:             cmd,
		Logger:          logger,
	})
}

// SetupCloseHandler creates a 'listener' on a new goroutine which will notify the
// program if it receives an interrupt from the OS. We then handle this by calling
// our clean up procedure and exiting the program.
// Ref: https://golangcode.com/handle-ctrl-c-exit-in-terminal/
func SetupCloseHandler() {
	c := make(chan os.Signal)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-c
		log.Printf("Execution aborted - %v", "SIGTERM")
		probeengine.CleanupTmp()
		os.Exit(0)
	}()
}

// GetCommands ...
func GetCommands() (cmdSet []*exec.Cmd, err error) {
	// TODO: give any exec errors a familiar format
	configPath, err := getConfigPath()
	if err != nil {
		return
	}
	packNames, err := GetPackNames()
	if err != nil {
		return
	}

	for _, pack := range packNames {
		binaryName, binErr := GetPackBinary(pack)
		if binErr != nil {
			err = binErr
			break
		}
		cmd := exec.Command(binaryName)
		cmd.Args = append(cmd.Args, fmt.Sprintf("--varsfile=%s", configPath))
		cmdSet = append(cmdSet, cmd)
	}
	if err == nil && len(cmdSet) == 0 {
		err = utils.ReformatError("No valid service packs specified")
	}
	return
}

// UserHomeDir provides the OS-aware user home directory
// TODO: move this to SDK
func UserHomeDir() string {
	dirname, err := os.UserHomeDir()
	if err != nil {
		log.Fatalf(err.Error())
	}
	return dirname
}

// GetPackBinary finds provided service pack in installation folder and return binary name
func GetPackBinary(name string) (binaryName string, err error) {
	name = filepath.Base(strings.ToLower(name)) // in some cases a filepath may arrive here instead of the base name
	if runtime.GOOS == "windows" && !strings.HasSuffix(name, ".exe") {
		name = fmt.Sprintf("%s.exe", name)
	}
	*BinariesPath = strings.Replace(*BinariesPath, "~", UserHomeDir(), 1)
	plugins, _ := hcplugin.Discover(name, *BinariesPath)
	if len(plugins) != 1 {
		err = fmt.Errorf("failed to locate requested plugin '%s'", name)
		return
	}
	binaryName = plugins[0]

	return
}

func getConfigPath() (string, error) {
	workDir, err := os.Getwd()
	if err != nil {
		return "", err
	}
	return filepath.Join(workDir, "config.yml"), nil
}

// GetPackNames returns all service packs declared in config file
func GetPackNames() (packNames []string, err error) {
	packNames, err = getPackNamesFromConfig()
	if err != nil || (AllPacks != nil && *AllPacks) {
		return hcplugin.Discover("*", *BinariesPath)
	}
	return
}

func getPackNamesFromConfig() ([]string, error) {
	type simpleVars struct {
		Run []string `yaml:"Run"`
	}
	var vars simpleVars

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

	configDecoder, file, err := config.NewConfigDecoder(configPath)
	if err != nil {
		return nil, err
	}

	err = configDecoder.Decode(&vars)
	file.Close()
	return vars.Run, err

}
<|endoftext|>


<commit_before>// Copyright 2019 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

// Package host 管理路由中域名的切换等操作
package host

import (
	"fmt"
	"net/http"
	"sort"
	"strings"

	"github.com/issue9/mux/v2/internal/handlers"
	"github.com/issue9/mux/v2/internal/tree"
	"github.com/issue9/mux/v2/params"
)

type host struct {
	raw      string // 域名的原始值，非通配符版本，与 domain 相同
	domain   string // 域名
	wildcard bool   // 是否带通配符

	tree *tree.Tree
}

// Hosts 域名管理
type Hosts struct {
	disableOptions bool
	disableHead    bool
	skipCleanPath  bool
	hosts          []*host
	tree           *tree.Tree // 非域名限定的路由项
}

// New 声明新的 Host 变量
func New(disableOptions, disableHead, skipCleanPath bool) *Hosts {
	return &Hosts{
		disableHead:    disableHead,
		disableOptions: disableOptions,
		skipCleanPath:  skipCleanPath,
		hosts:          make([]*host, 0, 10),
		tree:           tree.New(disableOptions, disableHead),
	}
}

// Add 添加路由项
func (hs *Hosts) Add(pattern string, h http.Handler, method ...string) error {
	tree, err := hs.getTree(pattern)
	if err != nil {
		return err
	}

	return tree.Add(pattern, h, method...)
}

// SetAllow 设置 Options 的 allow 报头值
func (hs *Hosts) SetAllow(pattern string, allow string) error {
	tree, err := hs.getTree(pattern)
	if err != nil {
		return err
	}

	return tree.SetAllow(pattern, allow)
}

// Remove 移除指定的路由项。
func (hs *Hosts) Remove(pattern string, method ...string) {
	tree, err := hs.getTree(pattern)
	if err != nil {
		panic(err)
	}

	tree.Remove(pattern, method...)
}

// URL 根据参数生成地址。
func (hs *Hosts) URL(pattern string, params map[string]string) (string, error) {
	tree, err := hs.getTree(pattern)
	if err != nil {
		return "", err
	}

	return tree.URL(pattern, params)
}

// CleanAll 清除所有的路由项
func (hs *Hosts) CleanAll() {
	for _, host := range hs.hosts {
		host.tree.Clean("")
	}
	hs.tree.Clean("")
}

// Clean 消除指定前缀的路由项
func (hs *Hosts) Clean(prefix string) {
	if prefix == "" {
		hs.CleanAll()
		return
	}

	if prefix[0] == '/' {
		hs.tree.Clean(prefix)
		return
	}

	index := strings.IndexByte(prefix, '/')
	if index < 0 {
		panic(fmt.Errorf("%s 不能只指定域名部分", prefix))
	}

	domain := prefix[:index]

	for _, host := range hs.hosts {
		if host.raw == domain {
			host.tree.Clean(prefix[index:])
			return
		}
	}
}

// Handler 获取匹配的路由项
func (hs *Hosts) Handler(r *http.Request) (*handlers.Handlers, params.Params) {
	p := r.URL.Path
	if !hs.skipCleanPath {
		p = cleanPath(p)
	}

	for _, host := range hs.hosts {
		if host.wildcard && strings.HasSuffix(r.Host, host.domain) {
			return host.tree.Handler(p)
		}

		if r.Host == host.domain {
			return host.tree.Handler(p)
		}
	}

	return hs.tree.Handler(p)
}

func (hs *Hosts) getTree(pattern string) (*tree.Tree, error) {
	if pattern == "" {
		panic("路由项地址不能为空")
	}

	if pattern[0] == '/' {
		return hs.tree, nil
	}

	index := strings.IndexByte(pattern, '/')
	if index < 0 {
		return nil, fmt.Errorf("%s 不能只指定域名部分", pattern)
	}

	domain := pattern[:index]

	for _, host := range hs.hosts {
		if host.raw == domain {
			return host.tree, nil
		}
	}

	host := &host{
		raw:    domain,
		domain: domain,
		tree:   tree.New(hs.disableOptions, hs.disableHead),
	}

	if strings.HasPrefix(host.domain, "*.") {
		host.wildcard = true
		host.domain = domain[1:] // 保留 . 符号
	}

	// 对域名列表进行排序，非通配符版本在前面
	hs.hosts = append(hs.hosts, host)
	sort.SliceStable(hs.hosts, func(i, j int) bool {
		ii := hs.hosts[i]
		jj := hs.hosts[j]

		switch {
		case ii.wildcard:
			return true
		case jj.wildcard:
			return true
		default:
			return ii.domain < jj.domain
		}
	})

	return host.tree, nil
}

// 清除路径中的重复的 / 字符
func cleanPath(p string) string {
	if p == "" {
		return "/"
	}

	if p[0] != '/' {
		p = "/" + p
	}

	index := strings.Index(p, "//")
	if index == -1 {
		return p
	}

	pp := make([]byte, index+1, len(p))
	copy(pp, p[:index+1])

	slash := true
	for i := index + 2; i < len(p); i++ {
		if p[i] == '/' {
			if slash {
				continue
			}
			slash = true
		} else {
			slash = false
		}
		pp = append(pp, p[i])
	}

	return string(pp)
}
<commit_msg>[internal/host] 添加 findTree<commit_after>// Copyright 2019 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

// Package host 管理路由中域名的切换等操作
package host

import (
	"fmt"
	"net/http"
	"sort"
	"strings"

	"github.com/issue9/mux/v2/internal/handlers"
	"github.com/issue9/mux/v2/internal/tree"
	"github.com/issue9/mux/v2/params"
)

type host struct {
	raw      string // 域名的原始值，非通配符版本，与 domain 相同
	domain   string // 域名
	wildcard bool   // 是否带通配符

	tree *tree.Tree
}

// Hosts 域名管理
type Hosts struct {
	disableOptions bool
	disableHead    bool
	skipCleanPath  bool
	hosts          []*host
	tree           *tree.Tree // 非域名限定的路由项
}

// New 声明新的 Host 变量
func New(disableOptions, disableHead, skipCleanPath bool) *Hosts {
	return &Hosts{
		disableHead:    disableHead,
		disableOptions: disableOptions,
		skipCleanPath:  skipCleanPath,
		hosts:          make([]*host, 0, 10),
		tree:           tree.New(disableOptions, disableHead),
	}
}

// Add 添加路由项
func (hs *Hosts) Add(pattern string, h http.Handler, method ...string) error {
	tree, err := hs.getTree(pattern)
	if err != nil {
		return err
	}

	return tree.Add(pattern, h, method...)
}

// SetAllow 设置 Options 的 allow 报头值
func (hs *Hosts) SetAllow(pattern string, allow string) error {
	tree, err := hs.getTree(pattern)
	if err != nil {
		return err
	}

	return tree.SetAllow(pattern, allow)
}

// Remove 移除指定的路由项。
func (hs *Hosts) Remove(pattern string, method ...string) {
	tree, err := hs.findTree(pattern)
	if err != nil {
		panic(err)
	}

	tree.Remove(pattern, method...)
}

// URL 根据参数生成地址。
func (hs *Hosts) URL(pattern string, params map[string]string) (string, error) {
	tree, err := hs.findTree(pattern)
	if err != nil {
		return "", err
	}

	return tree.URL(pattern, params)
}

// CleanAll 清除所有的路由项
func (hs *Hosts) CleanAll() {
	for _, host := range hs.hosts {
		host.tree.Clean("")
	}
	hs.tree.Clean("")
}

// Clean 消除指定前缀的路由项
func (hs *Hosts) Clean(prefix string) {
	if prefix == "" {
		hs.CleanAll()
		return
	}

	if prefix[0] == '/' {
		hs.tree.Clean(prefix)
		return
	}

	index := strings.IndexByte(prefix, '/')
	if index < 0 {
		panic(fmt.Errorf("%s 不能只指定域名部分", prefix))
	}

	domain := prefix[:index]

	for _, host := range hs.hosts {
		if host.raw == domain {
			host.tree.Clean(prefix[index:])
			return
		}
	}
}

// Handler 获取匹配的路由项
func (hs *Hosts) Handler(r *http.Request) (*handlers.Handlers, params.Params) {
	p := r.URL.Path
	if !hs.skipCleanPath {
		p = cleanPath(p)
	}

	for _, host := range hs.hosts {
		if host.wildcard && strings.HasSuffix(r.Host, host.domain) {
			return host.tree.Handler(p)
		}

		if r.Host == host.domain {
			return host.tree.Handler(p)
		}
	}

	return hs.tree.Handler(p)
}

// 获取指定路由项对应的 tree.Tree 实例，如果不存在，则返回空值。
func (hs *Hosts) findTree(pattern string) (*tree.Tree, error) {
	if pattern == "" {
		panic("路由项地址不能为空")
	}

	if pattern[0] == '/' {
		return hs.tree, nil
	}

	index := strings.IndexByte(pattern, '/')
	if index < 0 {
		return nil, fmt.Errorf("%s 不能只指定域名部分", pattern)
	}

	domain := pattern[:index]

	for _, host := range hs.hosts {
		if host.raw == domain {
			return host.tree, nil
		}
	}

	return nil, nil
}

// 获取指定路由项对应的 tree.Tree 实例，如果不存在，则添加并返回。
func (hs *Hosts) getTree(pattern string) (*tree.Tree, error) {
	if pattern == "" {
		panic("路由项地址不能为空")
	}

	if pattern[0] == '/' {
		return hs.tree, nil
	}

	index := strings.IndexByte(pattern, '/')
	if index < 0 {
		return nil, fmt.Errorf("%s 不能只指定域名部分", pattern)
	}

	domain := pattern[:index]

	for _, host := range hs.hosts {
		if host.raw == domain {
			return host.tree, nil
		}
	}

	host := &host{
		raw:    domain,
		domain: domain,
		tree:   tree.New(hs.disableOptions, hs.disableHead),
	}

	if strings.HasPrefix(host.domain, "*.") {
		host.wildcard = true
		host.domain = domain[1:] // 保留 . 符号
	}

	// 对域名列表进行排序，非通配符版本在前面
	hs.hosts = append(hs.hosts, host)
	sort.SliceStable(hs.hosts, func(i, j int) bool {
		ii := hs.hosts[i]
		jj := hs.hosts[j]

		switch {
		case ii.wildcard:
			return true
		case jj.wildcard:
			return true
		default:
			return ii.domain < jj.domain
		}
	})

	return host.tree, nil
}

// 清除路径中的重复的 / 字符
func cleanPath(p string) string {
	if p == "" {
		return "/"
	}

	if p[0] != '/' {
		p = "/" + p
	}

	index := strings.Index(p, "//")
	if index == -1 {
		return p
	}

	pp := make([]byte, index+1, len(p))
	copy(pp, p[:index+1])

	slash := true
	for i := index + 2; i < len(p); i++ {
		if p[i] == '/' {
			if slash {
				continue
			}
			slash = true
		} else {
			slash = false
		}
		pp = append(pp, p[i])
	}

	return string(pp)
}
<|endoftext|>