<commit_before>package cl11

import (
	"unsafe"

	clw "github.com/rdwilliamson/clw11"
)

type Event struct {
	id           clw.Event
	Context      *Context
	CommandType  CommandType
	CommandQueue *CommandQueue

	Queued int64
	Submit int64
	Start  int64
	End    int64
}

func (c *Context) CreateUserEvent() (*Event, error) {

	event, err := clw.CreateUserEvent(c.id)
	if err != nil {
		return nil, err
	}

	return &Event{id: event, Context: c, CommandType: CommandUser}, nil
}

type CommandType int

const (
	CommandNDRangeKernel        = CommandType(clw.CommandNdrangeKernel)
	CommandTask                 = CommandType(clw.CommandTask)
	CommandNativeKernel         = CommandType(clw.CommandNativeKernel)
	CommandReadBuffer           = CommandType(clw.CommandReadBuffer)
	CommandWriteBuffer          = CommandType(clw.CommandWriteBuffer)
	CommandCopyBuffer           = CommandType(clw.CommandCopyBuffer)
	CommandReadImage            = CommandType(clw.CommandReadImage)
	CommandWriteImage           = CommandType(clw.CommandWriteImage)
	CommandCopyImage            = CommandType(clw.CommandCopyImage)
	CommandCopyImageToBuffer    = CommandType(clw.CommandCopyImageToBuffer)
	CommandCopyBufferToImage    = CommandType(clw.CommandCopyBufferToImage)
	CommandMapBuffer            = CommandType(clw.CommandMapBuffer)
	CommandMapImage             = CommandType(clw.CommandMapImage)
	CommandUnmapMemoryObject    = CommandType(clw.CommandUnmapMemoryObject)
	CommandMarker               = CommandType(clw.CommandMarker)
	CommandAcquireGlObjects     = CommandType(clw.CommandAcquireGlObjects)
	CommandReleaseGlObjects     = CommandType(clw.CommandReleaseGlObjects)
	CommandReadBufferRectangle  = CommandType(clw.CommandReadBufferRectangle)
	CommandWriteBufferRectangle = CommandType(clw.CommandWriteBufferRectangle)
	CommandCopyBufferRectangle  = CommandType(clw.CommandCopyBufferRectangle)
	CommandUser                 = CommandType(clw.CommandUser)
)

var commandTypeMap = map[CommandType]string{
	CommandNDRangeKernel:        "ND range kernel",
	CommandTask:                 "task",
	CommandNativeKernel:         "native kernel",
	CommandReadBuffer:           "read buffer",
	CommandWriteBuffer:          "write buffer",
	CommandCopyBuffer:           "copy buffer",
	CommandReadImage:            "read image",
	CommandWriteImage:           "write image",
	CommandCopyImage:            "copy image",
	CommandCopyImageToBuffer:    "copy image to buffer",
	CommandCopyBufferToImage:    "copy buffer to image",
	CommandMapBuffer:            "map buffer",
	CommandMapImage:             "map image",
	CommandUnmapMemoryObject:    "unmap memory object",
	CommandMarker:               "marker",
	CommandAcquireGlObjects:     "acquire GL objects",
	CommandReleaseGlObjects:     "release GL objects",
	CommandReadBufferRectangle:  "read buffer rectangle",
	CommandWriteBufferRectangle: "write buffer rectangle",
	CommandCopyBufferRectangle:  "copy buffer rectangle",
	CommandUser:                 "user",
}

func (ct CommandType) String() string {
	return commandTypeMap[ct]
}

type CommandExecutionStatus int

const (
	Complete  = CommandExecutionStatus(clw.Complete)
	Running   = CommandExecutionStatus(clw.Running)
	Submitted = CommandExecutionStatus(clw.Submitted)
	Queued    = CommandExecutionStatus(clw.Queued)
)

func (ces CommandExecutionStatus) String() string {
	switch ces {
	case Complete:
		return "complete"
	case Running:
		return "running"
	case Submitted:
		return "submitted"
	case Queued:
		return "queued"
	}
	return ""
}

// Returns the events status, an error that caused the event to terminate, or an
// error that occurred trying to retrieve the event status.
func (e *Event) Status() (CommandExecutionStatus, error, error) {

	var status clw.CommandExecutionStatus
	err := clw.GetEventInfo(e.id, clw.EventCommandExecutionStatus, clw.Size(unsafe.Sizeof(status)),
		unsafe.Pointer(&status), nil)
	if err != nil {
		return 0, nil, err
	}

	if status < 0 {
		return 0, clw.CodeToError(clw.Int(status)), nil
	}

	return CommandExecutionStatus(status), nil, nil
}

func (e *Event) GetProfilingInfo() error {

	var value clw.Ulong
	err := clw.GetEventProfilingInfo(e.id, clw.ProfilingCommandQueued, clw.Size(unsafe.Sizeof(value)),
		unsafe.Pointer(&value), nil)
	if err != nil {
		return err
	}
	e.Queued = int64(value)

	err = clw.GetEventProfilingInfo(e.id, clw.ProfilingCommandSubmit, clw.Size(unsafe.Sizeof(value)),
		unsafe.Pointer(&value), nil)
	if err != nil {
		return err
	}
	e.Submit = int64(value)

	err = clw.GetEventProfilingInfo(e.id, clw.ProfilingCommandStart, clw.Size(unsafe.Sizeof(value)),
		unsafe.Pointer(&value), nil)
	if err != nil {
		return err
	}
	e.Start = int64(value)

	err = clw.GetEventProfilingInfo(e.id, clw.ProfilingCommandEnd, clw.Size(unsafe.Sizeof(value)),
		unsafe.Pointer(&value), nil)
	if err != nil {
		return err
	}
	e.End = int64(value)

	return nil
}

func (e *Event) SetCallback(callback func(e *Event, userData interface{}), userData interface{}) error {

	return clw.SetEventCallback(e.id, clw.Complete,
		func(event clw.Event, ces clw.CommandExecutionStatus, _userData interface{}) {
			callback(e, _userData)
		},
		userData)
}
<commit_msg>Added wait for events.<commit_after>package cl11

import (
	"unsafe"

	clw "github.com/rdwilliamson/clw11"
)

type Event struct {
	id           clw.Event
	Context      *Context
	CommandType  CommandType
	CommandQueue *CommandQueue

	Queued int64
	Submit int64
	Start  int64
	End    int64
}

func (c *Context) CreateUserEvent() (*Event, error) {

	event, err := clw.CreateUserEvent(c.id)
	if err != nil {
		return nil, err
	}

	return &Event{id: event, Context: c, CommandType: CommandUser}, nil
}

type CommandType int

const (
	CommandNDRangeKernel        = CommandType(clw.CommandNdrangeKernel)
	CommandTask                 = CommandType(clw.CommandTask)
	CommandNativeKernel         = CommandType(clw.CommandNativeKernel)
	CommandReadBuffer           = CommandType(clw.CommandReadBuffer)
	CommandWriteBuffer          = CommandType(clw.CommandWriteBuffer)
	CommandCopyBuffer           = CommandType(clw.CommandCopyBuffer)
	CommandReadImage            = CommandType(clw.CommandReadImage)
	CommandWriteImage           = CommandType(clw.CommandWriteImage)
	CommandCopyImage            = CommandType(clw.CommandCopyImage)
	CommandCopyImageToBuffer    = CommandType(clw.CommandCopyImageToBuffer)
	CommandCopyBufferToImage    = CommandType(clw.CommandCopyBufferToImage)
	CommandMapBuffer            = CommandType(clw.CommandMapBuffer)
	CommandMapImage             = CommandType(clw.CommandMapImage)
	CommandUnmapMemoryObject    = CommandType(clw.CommandUnmapMemoryObject)
	CommandMarker               = CommandType(clw.CommandMarker)
	CommandAcquireGlObjects     = CommandType(clw.CommandAcquireGlObjects)
	CommandReleaseGlObjects     = CommandType(clw.CommandReleaseGlObjects)
	CommandReadBufferRectangle  = CommandType(clw.CommandReadBufferRectangle)
	CommandWriteBufferRectangle = CommandType(clw.CommandWriteBufferRectangle)
	CommandCopyBufferRectangle  = CommandType(clw.CommandCopyBufferRectangle)
	CommandUser                 = CommandType(clw.CommandUser)
)

var commandTypeMap = map[CommandType]string{
	CommandNDRangeKernel:        "ND range kernel",
	CommandTask:                 "task",
	CommandNativeKernel:         "native kernel",
	CommandReadBuffer:           "read buffer",
	CommandWriteBuffer:          "write buffer",
	CommandCopyBuffer:           "copy buffer",
	CommandReadImage:            "read image",
	CommandWriteImage:           "write image",
	CommandCopyImage:            "copy image",
	CommandCopyImageToBuffer:    "copy image to buffer",
	CommandCopyBufferToImage:    "copy buffer to image",
	CommandMapBuffer:            "map buffer",
	CommandMapImage:             "map image",
	CommandUnmapMemoryObject:    "unmap memory object",
	CommandMarker:               "marker",
	CommandAcquireGlObjects:     "acquire GL objects",
	CommandReleaseGlObjects:     "release GL objects",
	CommandReadBufferRectangle:  "read buffer rectangle",
	CommandWriteBufferRectangle: "write buffer rectangle",
	CommandCopyBufferRectangle:  "copy buffer rectangle",
	CommandUser:                 "user",
}

func (ct CommandType) String() string {
	return commandTypeMap[ct]
}

type CommandExecutionStatus int

const (
	Complete  = CommandExecutionStatus(clw.Complete)
	Running   = CommandExecutionStatus(clw.Running)
	Submitted = CommandExecutionStatus(clw.Submitted)
	Queued    = CommandExecutionStatus(clw.Queued)
)

func (ces CommandExecutionStatus) String() string {
	switch ces {
	case Complete:
		return "complete"
	case Running:
		return "running"
	case Submitted:
		return "submitted"
	case Queued:
		return "queued"
	}
	return ""
}

// Returns the events status, an error that caused the event to terminate, or an
// error that occurred trying to retrieve the event status.
func (e *Event) Status() (CommandExecutionStatus, error, error) {

	var status clw.CommandExecutionStatus
	err := clw.GetEventInfo(e.id, clw.EventCommandExecutionStatus, clw.Size(unsafe.Sizeof(status)),
		unsafe.Pointer(&status), nil)
	if err != nil {
		return 0, nil, err
	}

	if status < 0 {
		return 0, clw.CodeToError(clw.Int(status)), nil
	}

	return CommandExecutionStatus(status), nil, nil
}

func (e *Event) GetProfilingInfo() error {

	var value clw.Ulong
	err := clw.GetEventProfilingInfo(e.id, clw.ProfilingCommandQueued, clw.Size(unsafe.Sizeof(value)),
		unsafe.Pointer(&value), nil)
	if err != nil {
		return err
	}
	e.Queued = int64(value)

	err = clw.GetEventProfilingInfo(e.id, clw.ProfilingCommandSubmit, clw.Size(unsafe.Sizeof(value)),
		unsafe.Pointer(&value), nil)
	if err != nil {
		return err
	}
	e.Submit = int64(value)

	err = clw.GetEventProfilingInfo(e.id, clw.ProfilingCommandStart, clw.Size(unsafe.Sizeof(value)),
		unsafe.Pointer(&value), nil)
	if err != nil {
		return err
	}
	e.Start = int64(value)

	err = clw.GetEventProfilingInfo(e.id, clw.ProfilingCommandEnd, clw.Size(unsafe.Sizeof(value)),
		unsafe.Pointer(&value), nil)
	if err != nil {
		return err
	}
	e.End = int64(value)

	return nil
}

func (e *Event) SetCallback(callback func(e *Event, userData interface{}), userData interface{}) error {

	return clw.SetEventCallback(e.id, clw.Complete,
		func(event clw.Event, ces clw.CommandExecutionStatus, _userData interface{}) {
			callback(e, _userData)
		},
		userData)
}

func WaitForEvents(events ...*Event) error {
	e := make([]clw.Event, len(events))
	for i := range events {
		e[i] = events[i].id
	}
	return clw.WaitForEvents(e)
}
<|endoftext|>


<commit_before>package agent

import (
	"fmt"
	"github.com/golang/glog"
	"github.com/romana/core/common"
	"github.com/romana/core/pkg/util/firewall"
)

// addPolicy is a placeholder. TODO
func (a *Agent) addPolicy(input interface{}, ctx common.RestContext) (interface{}, error) {
	//	policy := input.(*common.Policy)
	return nil, nil
}

// deletePolicy is a placeholder. TODO
func (a *Agent) deletePolicy(input interface{}, ctx common.RestContext) (interface{}, error) {
	//	policyId := ctx.PathVariables["policyID"]
	return nil, nil
}

// listPolicies is a placeholder. TODO.
func (a *Agent) listPolicies(input interface{}, ctx common.RestContext) (interface{}, error) {
	return nil, nil
}

// Status is a structure containing statistics returned by statusHandler
type Status struct {
	Rules      []firewall.IPtablesRule `json:"rules"`
	Interfaces []NetIf                 `json:"interfaces"`
}

// statusHandler reports operational statistics.
func (a *Agent) statusHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.V(1).Infoln("Agent: Entering statusHandler()")
	fw, err := firewall.NewFirewall(firewall.ShellexProvider)
	if err != nil {
		return nil, err
	}

	err = fw.Init(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	rules, err := fw.ListRules()
	if err != nil {
		return nil, err
	}
	ifaces, err := a.store.listNetIfs()
	if err != nil {
		return nil, err
	}
	status := Status{Rules: rules, Interfaces: ifaces}
	return status, nil
}

// podDownHandler cleans up after pod deleted.
func (a *Agent) podDownHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.V(1).Infoln("Agent: Entering podDownHandler()")
	netReq := input.(*NetworkRequest)
	netif := netReq.NetIf

	// We need new firewall instance here to use it's Cleanup()
	// to uninstall firewall rules related to the endpoint.
	fw, err := firewall.NewFirewall(firewall.ShellexProvider)
	if err != nil {
		return nil, err
	}

	err = fw.Init(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	err = fw.Cleanup(netif)
	if err != nil {
		return nil, err
	}

	// Spawn new thread to process the request
	glog.Infof("Agent: Got request for pod teardown %v\n", netReq)

	return "OK", nil
}

// podUpHandler handles HTTP requests for endpoints provisioning.
func (a *Agent) podUpHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.Infof("Agent: Entering podUpHandler()")
	netReq := input.(*NetworkRequest)

	glog.Infof("Agent: Got request for network configuration: %v\n", netReq)
	// Spawn new thread to process the request

	// TODO don't know if fork-bombs are possible in go but if they are this
	// need to be refactored as buffered channel with fixed pool of workers
	go a.podUpHandlerAsync(*netReq)

	// TODO I wonder if this should actually return something like a
	// link to a status of this request which will later get updated
	// with success or failure -- Greg.
	return "OK", nil
}

// vmDownHandler handles HTTP requests for endpoints teardown.
func (a *Agent) vmDownHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.Infof("In vmDownHandler() with %T %v", input, input)
	netif := input.(*NetIf)
	if netif.Name == "" {
		// This is a request from OpenStack Mech driver who does not have a name, let's find it.
		err := a.store.findNetIf(netif)
		if err != nil {
			return nil, err
		}
	}
	glog.Infof("In vmDownHandler() with Name %s, IP %s Mac %s\n", netif.Name, netif.IP, netif.Mac)

	glog.Info("Agent: provisioning DHCP")
	if err := a.leaseFile.provisionLease(netif, leaseRemove); err != nil {
		glog.Error(agentError(err))
		return "Error removing DHCP lease", agentError(err)
	}

	// We need new firewall instance here to use it's Cleanup()
	// to uninstall firewall rules related to the endpoint.
	fw, err := firewall.NewFirewall(firewall.ShellexProvider)
	if err != nil {
		return nil, err
	}

	err = fw.Init(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	err = fw.Cleanup(netif)
	if err != nil {
		return nil, err
	}
	err = a.store.deleteNetIf(netif)
	if err != nil {
		return nil, err
	}
	return "OK", nil
}

// vmUpHandler handles HTTP requests for endpoints provisioning.
// Currently tested with Romana ML2 driver.
func (a *Agent) vmUpHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	// Parse out NetIf form the request
	netif := input.(*NetIf)

	glog.Infof("Got interface: Name %s, IP %s Mac %s\n", netif.Name, netif.IP, netif.Mac)

	// Spawn new thread to process the request

	// TODO don't know if fork-bombs are possible in go but if they are this
	// need to be refactored as buffered channel with fixed pool of workers
	go a.vmUpHandlerAsync(*netif)

	// TODO I wonder if this should actually return something like a
	// link to a status of this request which will later get updated
	// with success or failure -- Greg.
	return "OK", nil
}

// podUpHandlerAsync does a number of operations on given endpoint to ensure
// it's connected:
// 1. Ensures interface is ready
// 2. Creates ip route pointing new interface
// 3. Provisions firewall rules
func (a *Agent) podUpHandlerAsync(netReq NetworkRequest) error {
	glog.V(1).Info("Agent: Entering podUpHandlerAsync()")
	currentProvider := firewall.IPTsaveProvider

	netif := netReq.NetIf
	if netif.Name == "" {
		return agentErrorString("Agent: Interface name required")
	}
	if !a.Helper.waitForIface(netif.Name) {
		// TODO should we resubmit failed interface in queue for later
		// retry ? ... considering openstack will give up as well after
		// timeout
		msg := fmt.Sprintf("Requested interface not available in time - %s", netif.Name)
		glog.Infoln("Agent: ", msg)
		return agentErrorString(msg)
	}
	glog.Info("Agent: creating endpoint routes")
	if err := a.Helper.ensureRouteToEndpoint(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: provisioning firewall")
	fw, err := firewall.NewFirewall(currentProvider)
	if err != nil {
		return err
	}

	err = fw.Init(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return err
	}

	if err1 := fw.SetEndpoint(netif); err1 != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	var rules RuleSet
	switch currentProvider {
	case firewall.ShellexProvider:
		rules = KubeShellRules
	case firewall.IPTsaveProvider:
		rules = KubeSaveRestoreRules
	default:
		err := fmt.Errorf("Unkown firewall provider in podUpHandler")
		glog.Error(agentError(err))
		return agentError(err)
	}

	if err := prepareFirewallRules(fw, rules, currentProvider); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: All good", netif)
	return nil
}

func prepareFirewallRules(fw firewall.Firewall, rules RuleSet, firewallProvider firewall.Provider) error {
	var defaultRules []firewall.FirewallRule
	//	var u32filter string
	var chainNames []string
	var formatBody string

	switch firewallProvider {
	case firewall.ShellexProvider:
		metadata := fw.Metadata()
		chainNames = metadata["chains"].([]string)
		//		u32filter = metadata["u32filter"].(string)

		for _, rule := range rules {
			glog.V(2).Infof("In prepareFirewallRules(), with %v", rule)

			var currentChain string
			switch rule.Direction {
			case EgressLocalDirection:
				currentChain = chainNames[firewall.InputChainIndex]
			case EgressGlobalDirection:
				currentChain = chainNames[firewall.ForwardOutChainIndex]
			case IngressGlobalDirection:
				currentChain = chainNames[firewall.ForwardInChainIndex]
			default:
				return fmt.Errorf("Error, unsupported rule direction type with firewall provider %s", firewallProvider)
			}

			switch rule.Format {
			case FormatChain:
				formatBody = fmt.Sprintf(rule.Body, currentChain)
			default:
				return fmt.Errorf("Error, unsupported rule format type with firewall provider %s", firewallProvider)
			}

			r := firewall.NewFirewallRule()
			r.SetBody(formatBody)

			switch rule.Position {
			case DefaultPosition:
				defaultRules = append(defaultRules, r)
			default:
				return fmt.Errorf("Error, unsupported rule position with firewall provider %s", firewallProvider)
			}
		}
	case firewall.IPTsaveProvider:
		for _, rule := range rules {
			glog.V(2).Infof("In prepareFirewallRules(), with %v", rule)

			var currentChain string
			switch rule.Direction {
			case EgressLocalDirection:
				currentChain = firewall.ChainNameEndpointToHost
			case EgressGlobalDirection:
				currentChain = firewall.ChainNameEndpointEgress
			case IngressGlobalDirection:
				currentChain = firewall.ChainNameEndpointIngress
			default:
				return fmt.Errorf("Error, unsupported rule direction type with firewall provider %s", firewallProvider)
			}

			switch rule.Format {
			case FormatChain:
				formatBody = fmt.Sprintf(rule.Body, currentChain)
			default:
				return fmt.Errorf("Error, unsupported rule format type with firewall provider %s", firewallProvider)
			}

			r := firewall.NewFirewallRule()
			r.SetBody(formatBody)

			switch rule.Position {
			case TopPosition:
				fw.EnsureRule(r, firewall.EnsureFirst)
			case BottomPosition:
				fw.EnsureRule(r, firewall.EnsureLast)
			default:
				return fmt.Errorf("Error, unsupported rule position with firewall provider %s", firewallProvider)
			}
		}
	default:
		return fmt.Errorf("Error, unsupported firewall provider type when preparing firewall rules")
	}

	return nil
}

// vmUpHandlerAsync does a number of operations on given endpoint to ensure
// it's connected:
// 1. Ensures interface is ready
// 2. Checks if DHCP is running
// 3. Creates ip route pointing new interface
// 4. Provisions static DHCP lease for new interface
// 5. Provisions firewall rules
func (a *Agent) vmUpHandlerAsync(netif NetIf) error {
	glog.V(1).Info("Agent: Entering interfaceHandle()")
	if !a.Helper.waitForIface(netif.Name) {
		// TODO should we resubmit failed interface in queue for later
		// retry ? ... considering oenstack will give up as well after
		// timeout
		return agentErrorString(fmt.Sprintf("Requested interface not available in time - %s", netif.Name))
	}

	// dhcpPid is only needed here for fail fast check
	// will try to poll the pid again in provisionLease
	glog.Info("Agent: checking if DHCP is running")
	_, err := a.Helper.DhcpPid()
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	err = a.store.addNetIf(&netif)
	if err != nil {
		return agentError(err)
	}
	glog.Info("Agent: creating endpoint routes")
	if err := a.Helper.ensureRouteToEndpoint(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	glog.Info("Agent: provisioning DHCP")
	if err := a.leaseFile.provisionLease(&netif, leaseAdd); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: provisioning firewall")
	fw, err := firewall.NewFirewall(firewall.ShellexProvider)
	if err != nil {
		return err
	}

	err = fw.Init(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return err
	}

	if err1 := fw.SetEndpoint(netif); err1 != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	metadata := fw.Metadata()
	chainNames := metadata["chains"].([]string)
	u32filter := metadata["u32filter"]
	hostAddr := a.networkConfig.RomanaGW()

	// Default firewall rules for OpenStack
	inboundChain := chainNames[firewall.InputChainIndex]
	var defaultRules []firewall.FirewallRule

	inboundRule := firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, "-m comment --comment DefaultDrop -j DROP"))
	defaultRules = append(defaultRules, inboundRule)

	inboundRule = firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, "-m state --state ESTABLISHED -j ACCEPT"))
	defaultRules = append(defaultRules, inboundRule)

	forwardOutChain := chainNames[firewall.ForwardOutChainIndex]
	forwardOutRule := firewall.NewFirewallRule()
	forwardOutRule.SetBody(fmt.Sprintf("%s %s", forwardOutChain, "-m comment --comment Outgoing -j RETURN"))
	defaultRules = append(defaultRules, forwardOutRule)

	forwardInChain := chainNames[firewall.ForwardInChainIndex]
	forwardInRule := firewall.NewFirewallRule()
	forwardInRule.SetBody(fmt.Sprintf("%s %s", forwardInChain, "-m state --state ESTABLISHED -j ACCEPT"))
	defaultRules = append(defaultRules, forwardInRule)

	forwardInRule = firewall.NewFirewallRule()
	forwardInRule.SetBody(fmt.Sprintf("%s ! -s %s -m u32 --u32 %s %s", forwardInChain, hostAddr, u32filter, "-j ACCEPT"))
	defaultRules = append(defaultRules, forwardInRule)

	fw.SetDefaultRules(defaultRules)

	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("All good", netif)
	return nil
}
<commit_msg>Allow choosing firewall provider with romana agent configuration file.<commit_after>package agent

import (
	"fmt"
	"github.com/golang/glog"
	"github.com/romana/core/common"
	"github.com/romana/core/pkg/util/firewall"
)

// addPolicy is a placeholder. TODO
func (a *Agent) addPolicy(input interface{}, ctx common.RestContext) (interface{}, error) {
	//	policy := input.(*common.Policy)
	return nil, nil
}

// deletePolicy is a placeholder. TODO
func (a *Agent) deletePolicy(input interface{}, ctx common.RestContext) (interface{}, error) {
	//	policyId := ctx.PathVariables["policyID"]
	return nil, nil
}

// listPolicies is a placeholder. TODO.
func (a *Agent) listPolicies(input interface{}, ctx common.RestContext) (interface{}, error) {
	return nil, nil
}

// Status is a structure containing statistics returned by statusHandler
type Status struct {
	Rules      []firewall.IPtablesRule `json:"rules"`
	Interfaces []NetIf                 `json:"interfaces"`
}

// statusHandler reports operational statistics.
func (a *Agent) statusHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.V(1).Infoln("Agent: Entering statusHandler()")
	fw, err := firewall.NewFirewall(a.getFirewallType())
	if err != nil {
		return nil, err
	}

	err = fw.Init(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	rules, err := fw.ListRules()
	if err != nil {
		return nil, err
	}
	ifaces, err := a.store.listNetIfs()
	if err != nil {
		return nil, err
	}
	status := Status{Rules: rules, Interfaces: ifaces}
	return status, nil
}

// podDownHandler cleans up after pod deleted.
func (a *Agent) podDownHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.V(1).Infoln("Agent: Entering podDownHandler()")
	netReq := input.(*NetworkRequest)
	netif := netReq.NetIf

	// We need new firewall instance here to use it's Cleanup()
	// to uninstall firewall rules related to the endpoint.
	fw, err := firewall.NewFirewall(a.getFirewallType())
	if err != nil {
		return nil, err
	}

	err = fw.Init(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	err = fw.Cleanup(netif)
	if err != nil {
		return nil, err
	}

	// Spawn new thread to process the request
	glog.Infof("Agent: Got request for pod teardown %v\n", netReq)

	return "OK", nil
}

// podUpHandler handles HTTP requests for endpoints provisioning.
func (a *Agent) podUpHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.Infof("Agent: Entering podUpHandler()")
	netReq := input.(*NetworkRequest)

	glog.Infof("Agent: Got request for network configuration: %v\n", netReq)
	// Spawn new thread to process the request

	// TODO don't know if fork-bombs are possible in go but if they are this
	// need to be refactored as buffered channel with fixed pool of workers
	go a.podUpHandlerAsync(*netReq)

	// TODO I wonder if this should actually return something like a
	// link to a status of this request which will later get updated
	// with success or failure -- Greg.
	return "OK", nil
}

// vmDownHandler handles HTTP requests for endpoints teardown.
func (a *Agent) vmDownHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.Infof("In vmDownHandler() with %T %v", input, input)
	netif := input.(*NetIf)
	if netif.Name == "" {
		// This is a request from OpenStack Mech driver who does not have a name, let's find it.
		err := a.store.findNetIf(netif)
		if err != nil {
			return nil, err
		}
	}
	glog.Infof("In vmDownHandler() with Name %s, IP %s Mac %s\n", netif.Name, netif.IP, netif.Mac)

	glog.Info("Agent: provisioning DHCP")
	if err := a.leaseFile.provisionLease(netif, leaseRemove); err != nil {
		glog.Error(agentError(err))
		return "Error removing DHCP lease", agentError(err)
	}

	// We need new firewall instance here to use it's Cleanup()
	// to uninstall firewall rules related to the endpoint.
	fw, err := firewall.NewFirewall(a.getFirewallType())
	if err != nil {
		return nil, err
	}

	err = fw.Init(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	err = fw.Cleanup(netif)
	if err != nil {
		return nil, err
	}
	err = a.store.deleteNetIf(netif)
	if err != nil {
		return nil, err
	}
	return "OK", nil
}

// vmUpHandler handles HTTP requests for endpoints provisioning.
// Currently tested with Romana ML2 driver.
func (a *Agent) vmUpHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	// Parse out NetIf form the request
	netif := input.(*NetIf)

	glog.Infof("Got interface: Name %s, IP %s Mac %s\n", netif.Name, netif.IP, netif.Mac)

	// Spawn new thread to process the request

	// TODO don't know if fork-bombs are possible in go but if they are this
	// need to be refactored as buffered channel with fixed pool of workers
	go a.vmUpHandlerAsync(*netif)

	// TODO I wonder if this should actually return something like a
	// link to a status of this request which will later get updated
	// with success or failure -- Greg.
	return "OK", nil
}

// podUpHandlerAsync does a number of operations on given endpoint to ensure
// it's connected:
// 1. Ensures interface is ready
// 2. Creates ip route pointing new interface
// 3. Provisions firewall rules
func (a *Agent) podUpHandlerAsync(netReq NetworkRequest) error {
	glog.V(1).Info("Agent: Entering podUpHandlerAsync()")
	currentProvider := a.getFirewallType()

	netif := netReq.NetIf
	if netif.Name == "" {
		return agentErrorString("Agent: Interface name required")
	}
	if !a.Helper.waitForIface(netif.Name) {
		// TODO should we resubmit failed interface in queue for later
		// retry ? ... considering openstack will give up as well after
		// timeout
		msg := fmt.Sprintf("Requested interface not available in time - %s", netif.Name)
		glog.Infoln("Agent: ", msg)
		return agentErrorString(msg)
	}
	glog.Info("Agent: creating endpoint routes")
	if err := a.Helper.ensureRouteToEndpoint(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: provisioning firewall")
	fw, err := firewall.NewFirewall(currentProvider)
	if err != nil {
		return err
	}

	err = fw.Init(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return err
	}

	if err1 := fw.SetEndpoint(netif); err1 != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	var rules RuleSet
	switch currentProvider {
	case firewall.ShellexProvider:
		rules = KubeShellRules
	case firewall.IPTsaveProvider:
		rules = KubeSaveRestoreRules
	default:
		err := fmt.Errorf("Unkown firewall provider in podUpHandler")
		glog.Error(agentError(err))
		return agentError(err)
	}

	if err := prepareFirewallRules(fw, rules, currentProvider); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: All good", netif)
	return nil
}

func prepareFirewallRules(fw firewall.Firewall, rules RuleSet, firewallProvider firewall.Provider) error {
	var defaultRules []firewall.FirewallRule
	//	var u32filter string
	var chainNames []string
	var formatBody string

	switch firewallProvider {
	case firewall.ShellexProvider:
		metadata := fw.Metadata()
		chainNames = metadata["chains"].([]string)
		//		u32filter = metadata["u32filter"].(string)

		for _, rule := range rules {
			glog.V(2).Infof("In prepareFirewallRules(), with %v", rule)

			var currentChain string
			switch rule.Direction {
			case EgressLocalDirection:
				currentChain = chainNames[firewall.InputChainIndex]
			case EgressGlobalDirection:
				currentChain = chainNames[firewall.ForwardOutChainIndex]
			case IngressGlobalDirection:
				currentChain = chainNames[firewall.ForwardInChainIndex]
			default:
				return fmt.Errorf("Error, unsupported rule direction type with firewall provider %s", firewallProvider)
			}

			switch rule.Format {
			case FormatChain:
				formatBody = fmt.Sprintf(rule.Body, currentChain)
			default:
				return fmt.Errorf("Error, unsupported rule format type with firewall provider %s", firewallProvider)
			}

			r := firewall.NewFirewallRule()
			r.SetBody(formatBody)

			switch rule.Position {
			case DefaultPosition:
				defaultRules = append(defaultRules, r)
			default:
				return fmt.Errorf("Error, unsupported rule position with firewall provider %s", firewallProvider)
			}
		}
	case firewall.IPTsaveProvider:
		for _, rule := range rules {
			glog.V(2).Infof("In prepareFirewallRules(), with %v", rule)

			var currentChain string
			switch rule.Direction {
			case EgressLocalDirection:
				currentChain = firewall.ChainNameEndpointToHost
			case EgressGlobalDirection:
				currentChain = firewall.ChainNameEndpointEgress
			case IngressGlobalDirection:
				currentChain = firewall.ChainNameEndpointIngress
			default:
				return fmt.Errorf("Error, unsupported rule direction type with firewall provider %s", firewallProvider)
			}

			switch rule.Format {
			case FormatChain:
				formatBody = fmt.Sprintf(rule.Body, currentChain)
			default:
				return fmt.Errorf("Error, unsupported rule format type with firewall provider %s", firewallProvider)
			}

			r := firewall.NewFirewallRule()
			r.SetBody(formatBody)

			switch rule.Position {
			case TopPosition:
				fw.EnsureRule(r, firewall.EnsureFirst)
			case BottomPosition:
				fw.EnsureRule(r, firewall.EnsureLast)
			default:
				return fmt.Errorf("Error, unsupported rule position with firewall provider %s", firewallProvider)
			}
		}
	default:
		return fmt.Errorf("Error, unsupported firewall provider type when preparing firewall rules")
	}

	return nil
}

// vmUpHandlerAsync does a number of operations on given endpoint to ensure
// it's connected:
// 1. Ensures interface is ready
// 2. Checks if DHCP is running
// 3. Creates ip route pointing new interface
// 4. Provisions static DHCP lease for new interface
// 5. Provisions firewall rules
func (a *Agent) vmUpHandlerAsync(netif NetIf) error {
	glog.V(1).Info("Agent: Entering interfaceHandle()")
	if !a.Helper.waitForIface(netif.Name) {
		// TODO should we resubmit failed interface in queue for later
		// retry ? ... considering oenstack will give up as well after
		// timeout
		return agentErrorString(fmt.Sprintf("Requested interface not available in time - %s", netif.Name))
	}

	// dhcpPid is only needed here for fail fast check
	// will try to poll the pid again in provisionLease
	glog.Info("Agent: checking if DHCP is running")
	_, err := a.Helper.DhcpPid()
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	err = a.store.addNetIf(&netif)
	if err != nil {
		return agentError(err)
	}
	glog.Info("Agent: creating endpoint routes")
	if err := a.Helper.ensureRouteToEndpoint(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	glog.Info("Agent: provisioning DHCP")
	if err := a.leaseFile.provisionLease(&netif, leaseAdd); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: provisioning firewall")
	fw, err := firewall.NewFirewall(firewall.ShellexProvider)
	if err != nil {
		return err
	}

	err = fw.Init(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return err
	}

	if err1 := fw.SetEndpoint(netif); err1 != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	metadata := fw.Metadata()
	chainNames := metadata["chains"].([]string)
	u32filter := metadata["u32filter"]
	hostAddr := a.networkConfig.RomanaGW()

	// Default firewall rules for OpenStack
	inboundChain := chainNames[firewall.InputChainIndex]
	var defaultRules []firewall.FirewallRule

	inboundRule := firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, "-m comment --comment DefaultDrop -j DROP"))
	defaultRules = append(defaultRules, inboundRule)

	inboundRule = firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, "-m state --state ESTABLISHED -j ACCEPT"))
	defaultRules = append(defaultRules, inboundRule)

	forwardOutChain := chainNames[firewall.ForwardOutChainIndex]
	forwardOutRule := firewall.NewFirewallRule()
	forwardOutRule.SetBody(fmt.Sprintf("%s %s", forwardOutChain, "-m comment --comment Outgoing -j RETURN"))
	defaultRules = append(defaultRules, forwardOutRule)

	forwardInChain := chainNames[firewall.ForwardInChainIndex]
	forwardInRule := firewall.NewFirewallRule()
	forwardInRule.SetBody(fmt.Sprintf("%s %s", forwardInChain, "-m state --state ESTABLISHED -j ACCEPT"))
	defaultRules = append(defaultRules, forwardInRule)

	forwardInRule = firewall.NewFirewallRule()
	forwardInRule.SetBody(fmt.Sprintf("%s ! -s %s -m u32 --u32 %s %s", forwardInChain, hostAddr, u32filter, "-j ACCEPT"))
	defaultRules = append(defaultRules, forwardInRule)

	fw.SetDefaultRules(defaultRules)

	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("All good", netif)
	return nil
}

func (a Agent) getFirewallType() firewall.Provider {
	provider, ok := a.config.ServiceSpecific["firewall_provider"].(string)
	if !ok {
		panic("Unable to read firewall_provider from config")
	}

	switch provider {
	case "shellex":
		glog.Infoln("Agent: using ShellexProvider firewall provider")
		return firewall.ShellexProvider
	case "save-restore":
		glog.Infoln("Agent: using IPTsaveProvider firewall provider")
		return firewall.IPTsaveProvider
	default:
		panic(fmt.Sprintf("Unsupported firewall type value %s, supported values are 'shellex' and 'save-restore'", provider))
	}

}
<|endoftext|>


<commit_before>package broker

import (
	"fmt"
	"github.com/siddontang/golib/log"
	"github.com/siddontang/moonmq/proto"
	"net"
	"runtime"
	"sync"
	"time"
)

type conn struct {
	sync.Mutex

	app *App

	c net.Conn

	decoder *proto.Decoder

	lastUpdate int64

	handshaked bool

	routes map[string][]string
}

func newConn(app *App, co net.Conn) *conn {
	c := new(conn)

	c.app = app
	c.c = co

	c.handshaked = false

	c.decoder = proto.NewDecoder(co)

	c.checkKeepAlive()

	c.routes = make(map[string][]string)

	return c
}

func (c *conn) run() {
	c.onRead()
}

func (c *conn) onRead() {
	defer func() {
		if err := recover(); err != nil {
			buf := make([]byte, 1024)
			buf = buf[:runtime.Stack(buf, false)]
			log.Fatal("crash %v:%v", err, buf)
		}

		c.c.Close()

	}()

	for {
		p, err := c.decoder.DecodeProto()
		if err != nil {
			log.Info("on read error %v", err)
			return
		}

		if p.Method == proto.Handshake {
			err = c.handleHandshake(p)
		} else {
			if !c.handshaked {
				err = fmt.Errorf("must handshake first")
			} else {
				switch p.Method {
				case proto.Publish:
					err = c.handlePublish(p)
				case proto.Bind:
				case proto.Unbind:
				case proto.Ack:
					err = c.handleAck(p)
				case proto.Heartbeat:
					c.lastUpdate = time.Now().Unix()
				default:
					log.Info("invalid proto method %d", p.Method)
					return
				}
			}
		}

		if err != nil {
			c.writeError(err)
		}
	}
}

func (c *conn) handleHandshake(p *proto.Proto) error {
	//later check authorization

	rp := proto.NewProto(proto.Handshake_OK, nil, nil)
	c.writeProto(rp)

	c.handshaked = true

	return nil
}

func (c *conn) writeError(err error) {
	var p *proto.Proto
	if pe, ok := err.(*proto.ProtoError); ok {
		p = pe.P
	} else {
		pe = proto.NewProtoError(500, err.Error())
		p = pe.P
	}

	c.writeProto(p)
}

func (c *conn) protoError(code int, message string) error {
	return proto.NewProtoError(code, message)
}

func (c *conn) writeProto(p *proto.Proto) error {
	buf, err := proto.Marshal(p)
	if err != nil {
		return err
	}

	var n int
	c.Lock()
	n, err = c.c.Write(buf)
	c.Unlock()

	if err != nil {
		return err
	} else if n != len(buf) {
		return fmt.Errorf("write incomplete, %d less than %d", n, len(buf))
	} else {
		return nil
	}
}

func (c *conn) checkKeepAlive() {
	var f func()
	f = func() {
		if time.Now().Unix()-c.lastUpdate > int64(1.5*float32(c.app.cfg.KeepAlive)) {
			log.Info("keepalive timeout")
			c.c.Close()
			return
		} else {
			c.app.wheel.AddTask(time.Duration(c.app.cfg.KeepAlive), f)
		}
	}

	c.app.wheel.AddTask(time.Duration(c.app.cfg.KeepAlive), f)
}
<commit_msg>close after write error<commit_after>package broker

import (
	"fmt"
	"github.com/siddontang/golib/log"
	"github.com/siddontang/moonmq/proto"
	"net"
	"runtime"
	"sync"
	"time"
)

type conn struct {
	sync.Mutex

	app *App

	c net.Conn

	decoder *proto.Decoder

	lastUpdate int64

	handshaked bool

	routes map[string][]string
}

func newConn(app *App, co net.Conn) *conn {
	c := new(conn)

	c.app = app
	c.c = co

	c.handshaked = false

	c.decoder = proto.NewDecoder(co)

	c.checkKeepAlive()

	c.routes = make(map[string][]string)

	return c
}

func (c *conn) run() {
	c.onRead()
}

func (c *conn) onRead() {
	defer func() {
		if err := recover(); err != nil {
			buf := make([]byte, 1024)
			buf = buf[:runtime.Stack(buf, false)]
			log.Fatal("crash %v:%v", err, buf)
		}

		c.c.Close()

	}()

	for {
		p, err := c.decoder.DecodeProto()
		if err != nil {
			log.Info("on read error %v", err)
			return
		}

		if p.Method == proto.Handshake {
			err = c.handleHandshake(p)
		} else {
			if !c.handshaked {
				err = fmt.Errorf("must handshake first")
			} else {
				switch p.Method {
				case proto.Publish:
					err = c.handlePublish(p)
				case proto.Bind:
				case proto.Unbind:
				case proto.Ack:
					err = c.handleAck(p)
				case proto.Heartbeat:
					c.lastUpdate = time.Now().Unix()
				default:
					log.Info("invalid proto method %d", p.Method)
					return
				}
			}
		}

		if err != nil {
			c.writeError(err)
		}
	}
}

func (c *conn) handleHandshake(p *proto.Proto) error {
	//later check authorization

	rp := proto.NewProto(proto.Handshake_OK, nil, nil)
	c.writeProto(rp)

	c.handshaked = true

	return nil
}

func (c *conn) writeError(err error) {
	var p *proto.Proto
	if pe, ok := err.(*proto.ProtoError); ok {
		p = pe.P
	} else {
		pe = proto.NewProtoError(500, err.Error())
		p = pe.P
	}

	c.writeProto(p)
}

func (c *conn) protoError(code int, message string) error {
	return proto.NewProtoError(code, message)
}

func (c *conn) writeProto(p *proto.Proto) error {
	buf, err := proto.Marshal(p)
	if err != nil {
		return err
	}

	var n int
	c.Lock()
	n, err = c.c.Write(buf)
	c.Unlock()

	if err != nil {
		c.c.Close()
		return err
	} else if n != len(buf) {
		c.c.Close()
		return fmt.Errorf("write incomplete, %d less than %d", n, len(buf))
	} else {
		return nil
	}
}

func (c *conn) checkKeepAlive() {
	var f func()
	f = func() {
		if time.Now().Unix()-c.lastUpdate > int64(1.5*float32(c.app.cfg.KeepAlive)) {
			log.Info("keepalive timeout")
			c.c.Close()
			return
		} else {
			c.app.wheel.AddTask(time.Duration(c.app.cfg.KeepAlive), f)
		}
	}

	c.app.wheel.AddTask(time.Duration(c.app.cfg.KeepAlive), f)
}
<|endoftext|>


<commit_before>package lib

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"github.com/docker/docker/pkg/sockets"
	"github.com/docker/docker/pkg/tlsconfig"
)

// Client is the API client that performs all operations
// against a docker server.
type Client struct {
	// proto holds the client protocol i.e. unix.
	proto string
	// addr holds the client address.
	addr string
	// basePath holds the path to prepend to the requests
	basePath string
	// scheme holds the scheme of the client i.e. https.
	scheme string
	// tlsConfig holds the tls configuration to use in hijacked requests.
	tlsConfig *tls.Config
	// httpClient holds the client transport instance. Exported to keep the old code running.
	httpClient *http.Client
	// version of the server to talk to.
	version string
	// custom http headers configured by users
	customHTTPHeaders map[string]string
}

// NewClient initializes a new API client for the given host and API version.
// It won't send any version information if the version number is empty.
// It uses the tlsOptions to decide whether to use a secure connection or not.
// It also initializes the custom http headers to add to each request.
func NewClient(host string, version string, tlsOptions *tlsconfig.Options, httpHeaders map[string]string) (*Client, error) {
	var (
		basePath       string
		tlsConfig      *tls.Config
		scheme         = "http"
		protoAddrParts = strings.SplitN(host, "://", 2)
		proto, addr    = protoAddrParts[0], protoAddrParts[1]
	)

	if proto == "tcp" {
		parsed, err := url.Parse("tcp://" + addr)
		if err != nil {
			return nil, err
		}
		addr = parsed.Host
		basePath = parsed.Path
	}

	if tlsOptions != nil {
		scheme = "https"
		var err error
		tlsConfig, err = tlsconfig.Client(*tlsOptions)
		if err != nil {
			return nil, err
		}
	}

	// The transport is created here for reuse during the client session.
	transport := &http.Transport{
		TLSClientConfig: tlsConfig,
	}
	sockets.ConfigureTCPTransport(transport, proto, addr)

	return &Client{
		proto:             proto,
		addr:              addr,
		basePath:          basePath,
		scheme:            scheme,
		tlsConfig:         tlsConfig,
		httpClient:        &http.Client{Transport: transport},
		version:           version,
		customHTTPHeaders: httpHeaders,
	}, nil
}

// getAPIPath returns the versioned request path to call the api.
// It appends the query parameters to the path if they are not empty.
func (cli *Client) getAPIPath(p string, query url.Values) string {
	var apiPath string
	if cli.version != "" {
		v := strings.TrimPrefix(cli.version, "v")
		apiPath = fmt.Sprintf("%s/v%s%s", cli.basePath, v, p)
	} else {
		apiPath = fmt.Sprintf("%s%s", cli.basePath, p)
	}
	if len(query) > 0 {
		apiPath += "?" + query.Encode()
	}
	return apiPath
}
<commit_msg>Add a DOCKER_API_VERSION env var<commit_after>package lib

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"github.com/docker/docker/pkg/sockets"
	"github.com/docker/docker/pkg/tlsconfig"
)

// Client is the API client that performs all operations
// against a docker server.
type Client struct {
	// proto holds the client protocol i.e. unix.
	proto string
	// addr holds the client address.
	addr string
	// basePath holds the path to prepend to the requests
	basePath string
	// scheme holds the scheme of the client i.e. https.
	scheme string
	// tlsConfig holds the tls configuration to use in hijacked requests.
	tlsConfig *tls.Config
	// httpClient holds the client transport instance. Exported to keep the old code running.
	httpClient *http.Client
	// version of the server to talk to.
	version string
	// custom http headers configured by users
	customHTTPHeaders map[string]string
}

// NewClient initializes a new API client for the given host and API version.
// It won't send any version information if the version number is empty.
// It uses the tlsOptions to decide whether to use a secure connection or not.
// It also initializes the custom http headers to add to each request.
func NewClient(host string, version string, tlsOptions *tlsconfig.Options, httpHeaders map[string]string) (*Client, error) {
	var (
		basePath       string
		tlsConfig      *tls.Config
		scheme         = "http"
		protoAddrParts = strings.SplitN(host, "://", 2)
		proto, addr    = protoAddrParts[0], protoAddrParts[1]
	)

	if proto == "tcp" {
		parsed, err := url.Parse("tcp://" + addr)
		if err != nil {
			return nil, err
		}
		addr = parsed.Host
		basePath = parsed.Path
	}

	if tlsOptions != nil {
		scheme = "https"
		var err error
		tlsConfig, err = tlsconfig.Client(*tlsOptions)
		if err != nil {
			return nil, err
		}
	}

	// The transport is created here for reuse during the client session.
	transport := &http.Transport{
		TLSClientConfig: tlsConfig,
	}
	sockets.ConfigureTCPTransport(transport, proto, addr)

	return &Client{
		proto:             proto,
		addr:              addr,
		basePath:          basePath,
		scheme:            scheme,
		tlsConfig:         tlsConfig,
		httpClient:        &http.Client{Transport: transport},
		version:           version,
		customHTTPHeaders: httpHeaders,
	}, nil
}

// getAPIPath returns the versioned request path to call the api.
// It appends the query parameters to the path if they are not empty.
func (cli *Client) getAPIPath(p string, query url.Values) string {
	var apiPath string
	if cli.version != "" {
		v := strings.TrimPrefix(cli.version, "v")
		apiPath = fmt.Sprintf("%s/v%s%s", cli.basePath, v, p)
	} else {
		apiPath = fmt.Sprintf("%s%s", cli.basePath, p)
	}
	if len(query) > 0 {
		apiPath += "?" + query.Encode()
	}
	return apiPath
}

// ClientVersion returns the version string associated with this
// instance of the Client. Note that this value can be changed
// via the DOCKER_API_VERSION env var.
func (cli *Client) ClientVersion() string {
	return cli.version
}
<|endoftext|>


<commit_before>// Copyright 2017 Microsoft. All rights reserved.
// MIT License

package dockerclient

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"os/exec"

	"github.com/Azure/azure-container-networking/cns/imdsclient"
	"github.com/Azure/azure-container-networking/log"
)

const (
	defaultDockerConnectionURL = "http://127.0.0.1:2375"
	defaultIpamPlugin          = "azure-vnet"
	networkMode                = "com.microsoft.azure.network.mode"
	bridgeMode                 = "bridge"
)

// DockerClient specifies a client to connect to docker.
type DockerClient struct {
	connectionURL string
	imdsClient    *imdsclient.ImdsClient
}

func executeShellCommand(command string) error {
	log.Debugf("[ebtables] %s", command)
	cmd := exec.Command("sh", "-c", command)
	err := cmd.Start()
	if err != nil {
		return err
	}
	return cmd.Wait()
}

// NewDockerClient create a new docker client.
func NewDockerClient(url string) (*DockerClient, error) {
	return &DockerClient{
		connectionURL: url,
		imdsClient:    &imdsclient.ImdsClient{},
	}, nil
}

// NewDefaultDockerClient create a new docker client.
func NewDefaultDockerClient(imdsClient *imdsclient.ImdsClient) (*DockerClient, error) {
	return &DockerClient{
		connectionURL: defaultDockerConnectionURL,
		imdsClient:    imdsClient,
	}, nil
}

// NetworkExists tries to retrieve a network from docker (if it exists).
func (dockerClient *DockerClient) NetworkExists(networkName string) error {
	log.Printf("[Azure CNS] NetworkExists")

	res, err := http.Get(
		dockerClient.connectionURL + inspectNetworkPath + networkName)

	if err != nil {
		log.Printf("[Azure CNS] Error received from http Post for docker network inspect %v %v", networkName, err.Error())
		return err
	}

	// network exists
	if res.StatusCode == 200 {
		log.Debugf("[Azure CNS] Network with name %v already exists. Docker return code: %v", networkName, res.StatusCode)
		return nil
	}

	// network not found
	if res.StatusCode == 404 {
		log.Debugf("[Azure CNS] Network with name %v does not exist. Docker return code: %v", networkName, res.StatusCode)
		return fmt.Errorf("Network not found")
	}

	return fmt.Errorf("Unknown return code from docker inspect %d", res.StatusCode)
}

// CreateNetwork creates a network using docker network create.
func (dockerClient *DockerClient) CreateNetwork(networkName string, options map[string]interface{}) error {
	log.Printf("[Azure CNS] CreateNetwork")

	enableSnat := true

	primaryNic, err := dockerClient.imdsClient.GetPrimaryInterfaceInfoFromHost()
	if err != nil {
		return err
	}

	config := &Config{
		Subnet: primaryNic.Subnet,
	}

	configs := make([]Config, 1)
	configs[0] = *config
	ipamConfig := &IPAM{
		Driver: defaultIpamPlugin,
		Config: configs,
	}

	netConfig := &NetworkConfiguration{
		Name:     networkName,
		Driver:   defaultNetworkPlugin,
		IPAM:     *ipamConfig,
		Internal: true,
	}

	if options != nil {
		if _, ok := options[OptDisableSnat]; ok {
			enableSnat = false
		}
	}

	if enableSnat {
		netConfig.Options = make(map[string]interface{})
		netConfig.Options[networkMode] = bridgeMode
	}

	log.Printf("[Azure CNS] Going to create network with config: %+v", netConfig)

	netConfigJSON := new(bytes.Buffer)
	err = json.NewEncoder(netConfigJSON).Encode(netConfig)
	if err != nil {
		return err
	}

	res, err := http.Post(
		dockerClient.connectionURL+createNetworkPath,
		"application/json; charset=utf-8",
		netConfigJSON)

	if err != nil {
		log.Printf("[Azure CNS] Error received from http Post for docker network create %v", networkName)
		return err
	}
	if res.StatusCode != 201 {
		var createNetworkResponse DockerErrorResponse
		err = json.NewDecoder(res.Body).Decode(&createNetworkResponse)
		var ermsg string
		ermsg = ""
		if err != nil {
			ermsg = err.Error()
		}
		return fmt.Errorf("[Azure CNS] Create docker network failed with error code %v - %v - %v",
			res.StatusCode, createNetworkResponse.message, ermsg)
	}

	if enableSnat {
		cmd := fmt.Sprintf("iptables -t nat -A POSTROUTING -m iprange ! --dst-range 168.63.129.16 -m addrtype ! --dst-type local ! -d %v -j MASQUERADE",
			primaryNic.Subnet)
		err = executeShellCommand(cmd)
		if err != nil {
			log.Printf("SNAT Iptable rule was not set")
			return err
		}
	}

	return nil
}

// DeleteNetwork creates a network using docker network create.
func (dockerClient *DockerClient) DeleteNetwork(networkName string) error {
	log.Printf("[Azure CNS] DeleteNetwork")

	url := dockerClient.connectionURL + inspectNetworkPath + networkName
	req, err := http.NewRequest("DELETE", url, nil)
	if err != nil {
		log.Printf("[Azure CNS] Error received while creating http DELETE request for network delete %v %v", networkName, err.Error())
		return err
	}

	req.Header.Set("Content-Type", "application/json; charset=utf-8")
	client := &http.Client{}
	res, err := client.Do(req)

	// network successfully deleted.
	if res.StatusCode == 204 {
		primaryNic, err := dockerClient.imdsClient.GetPrimaryInterfaceInfoFromHost()
		if err != nil {
			return err
		}

		cmd := fmt.Sprintf("iptables -t nat -D POSTROUTING -m iprange ! --dst-range 168.63.129.16 -m addrtype ! --dst-type local ! -d %v -j MASQUERADE",
			primaryNic.Subnet)
		executeShellCommand(cmd)
		return nil
	}

	// network not found.
	if res.StatusCode == 404 {
		return fmt.Errorf("[Azure CNS] Network not found %v", networkName)
	}

	return fmt.Errorf("[Azure CNS] Unknown return code from docker delete network %v: ret = %d",
		networkName, res.StatusCode)
}
<commit_msg>updated log<commit_after>// Copyright 2017 Microsoft. All rights reserved.
// MIT License

package dockerclient

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"os/exec"

	"github.com/Azure/azure-container-networking/cns/imdsclient"
	"github.com/Azure/azure-container-networking/log"
)

const (
	defaultDockerConnectionURL = "http://127.0.0.1:2375"
	defaultIpamPlugin          = "azure-vnet"
	networkMode                = "com.microsoft.azure.network.mode"
	bridgeMode                 = "bridge"
)

// DockerClient specifies a client to connect to docker.
type DockerClient struct {
	connectionURL string
	imdsClient    *imdsclient.ImdsClient
}

func executeShellCommand(command string) error {
	log.Debugf("[Azure-CNS] %s", command)
	cmd := exec.Command("sh", "-c", command)
	err := cmd.Start()
	if err != nil {
		return err
	}
	return cmd.Wait()
}

// NewDockerClient create a new docker client.
func NewDockerClient(url string) (*DockerClient, error) {
	return &DockerClient{
		connectionURL: url,
		imdsClient:    &imdsclient.ImdsClient{},
	}, nil
}

// NewDefaultDockerClient create a new docker client.
func NewDefaultDockerClient(imdsClient *imdsclient.ImdsClient) (*DockerClient, error) {
	return &DockerClient{
		connectionURL: defaultDockerConnectionURL,
		imdsClient:    imdsClient,
	}, nil
}

// NetworkExists tries to retrieve a network from docker (if it exists).
func (dockerClient *DockerClient) NetworkExists(networkName string) error {
	log.Printf("[Azure CNS] NetworkExists")

	res, err := http.Get(
		dockerClient.connectionURL + inspectNetworkPath + networkName)

	if err != nil {
		log.Printf("[Azure CNS] Error received from http Post for docker network inspect %v %v", networkName, err.Error())
		return err
	}

	// network exists
	if res.StatusCode == 200 {
		log.Debugf("[Azure CNS] Network with name %v already exists. Docker return code: %v", networkName, res.StatusCode)
		return nil
	}

	// network not found
	if res.StatusCode == 404 {
		log.Debugf("[Azure CNS] Network with name %v does not exist. Docker return code: %v", networkName, res.StatusCode)
		return fmt.Errorf("Network not found")
	}

	return fmt.Errorf("Unknown return code from docker inspect %d", res.StatusCode)
}

// CreateNetwork creates a network using docker network create.
func (dockerClient *DockerClient) CreateNetwork(networkName string, options map[string]interface{}) error {
	log.Printf("[Azure CNS] CreateNetwork")

	enableSnat := true

	primaryNic, err := dockerClient.imdsClient.GetPrimaryInterfaceInfoFromHost()
	if err != nil {
		return err
	}

	config := &Config{
		Subnet: primaryNic.Subnet,
	}

	configs := make([]Config, 1)
	configs[0] = *config
	ipamConfig := &IPAM{
		Driver: defaultIpamPlugin,
		Config: configs,
	}

	netConfig := &NetworkConfiguration{
		Name:     networkName,
		Driver:   defaultNetworkPlugin,
		IPAM:     *ipamConfig,
		Internal: true,
	}

	if options != nil {
		if _, ok := options[OptDisableSnat]; ok {
			enableSnat = false
		}
	}

	if enableSnat {
		netConfig.Options = make(map[string]interface{})
		netConfig.Options[networkMode] = bridgeMode
	}

	log.Printf("[Azure CNS] Going to create network with config: %+v", netConfig)

	netConfigJSON := new(bytes.Buffer)
	err = json.NewEncoder(netConfigJSON).Encode(netConfig)
	if err != nil {
		return err
	}

	res, err := http.Post(
		dockerClient.connectionURL+createNetworkPath,
		"application/json; charset=utf-8",
		netConfigJSON)

	if err != nil {
		log.Printf("[Azure CNS] Error received from http Post for docker network create %v", networkName)
		return err
	}
	if res.StatusCode != 201 {
		var createNetworkResponse DockerErrorResponse
		err = json.NewDecoder(res.Body).Decode(&createNetworkResponse)
		var ermsg string
		ermsg = ""
		if err != nil {
			ermsg = err.Error()
		}
		return fmt.Errorf("[Azure CNS] Create docker network failed with error code %v - %v - %v",
			res.StatusCode, createNetworkResponse.message, ermsg)
	}

	if enableSnat {
		cmd := fmt.Sprintf("iptables -t nat -A POSTROUTING -m iprange ! --dst-range 168.63.129.16 -m addrtype ! --dst-type local ! -d %v -j MASQUERADE",
			primaryNic.Subnet)
		err = executeShellCommand(cmd)
		if err != nil {
			log.Printf("SNAT Iptable rule was not set")
			return err
		}
	}

	return nil
}

// DeleteNetwork creates a network using docker network create.
func (dockerClient *DockerClient) DeleteNetwork(networkName string) error {
	log.Printf("[Azure CNS] DeleteNetwork")

	url := dockerClient.connectionURL + inspectNetworkPath + networkName
	req, err := http.NewRequest("DELETE", url, nil)
	if err != nil {
		log.Printf("[Azure CNS] Error received while creating http DELETE request for network delete %v %v", networkName, err.Error())
		return err
	}

	req.Header.Set("Content-Type", "application/json; charset=utf-8")
	client := &http.Client{}
	res, err := client.Do(req)

	// network successfully deleted.
	if res.StatusCode == 204 {
		primaryNic, err := dockerClient.imdsClient.GetPrimaryInterfaceInfoFromHost()
		if err != nil {
			return err
		}

		cmd := fmt.Sprintf("iptables -t nat -D POSTROUTING -m iprange ! --dst-range 168.63.129.16 -m addrtype ! --dst-type local ! -d %v -j MASQUERADE",
			primaryNic.Subnet)
		executeShellCommand(cmd)
		return nil
	}

	// network not found.
	if res.StatusCode == 404 {
		return fmt.Errorf("[Azure CNS] Network not found %v", networkName)
	}

	return fmt.Errorf("[Azure CNS] Unknown return code from docker delete network %v: ret = %d",
		networkName, res.StatusCode)
}
<|endoftext|>


<commit_before>// Copyright Piero de Salvia.
// All Rights Reserved

/*
moved test to this package because of https://github.com/golang/go/issues/17928
*/
package dynaroutes_integration

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/pierods/dynaroutes"
)

var testDataDir string

func init() {

	goPath := os.Getenv("GOPATH")
	testDataDir = goPath + "/src/github.com/pierods/dynaroutes/testdata"
}

func TestHeaders(t *testing.T) {

	client := http.Client{}

	req, err := http.NewRequest("GET", "http://localhost:30000", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Add("Must", "gothrough")
	req.Header.Add("This", "too")
	resp, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	if resp.Header.Get("Must") != "gothrough" || resp.Header.Get("This") != "too" {
		t.Fatal("Should be able to carry headers through")
	}
}

func TestBody(t *testing.T) {

	client := http.Client{}

	body := "Body"
	req, err := http.NewRequest("GET", "http://localhost:30000", strings.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	resp, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	buffer := new(bytes.Buffer)
	io.Copy(buffer, resp.Body)
	responseBody := buffer.String()
	t.Log(body)

	if responseBody != body+" - filtered" {
		t.Fatal("Should be able to carry body through and filter it")
	}
}

func TestTimeouts(t *testing.T) {

	client := http.Client{}

	req, err := http.NewRequest("GET", "http://localhost:30000/timeout", nil)
	if err != nil {
		t.Fatal(err)
	}
	resp, err := client.Do(req)
	if resp.StatusCode == 200 {
		t.Fatal("Should err out on a timeout")
	}

	client = http.Client{}

	req, err = http.NewRequest("GET", "http://localhost:30000/timeoutplugin", nil)
	if err != nil {
		t.Fatal(err)
	}
	resp, err = client.Do(req)
	if resp.StatusCode == 200 {
		t.Fatal("Should err out on a timeout")
	}

}

func TestMain(m *testing.M) {

	listener, err := net.Listen("tcp", ":50000")
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	server := http.Server{
		Handler: &EchoHandler{},
	}

	go func() {
		if err := server.Serve(listener); err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	}()

	defer server.Close()

	router, err := dynaroutes.NewRouterF("localhost", 30000, 5*time.Second, 5*time.Second, testDataDir, "localhost", 31000)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	go func() {
		router.Start()
	}()

	defer router.Shutdown()
	time.Sleep(40 * time.Second)
	retCode := m.Run()
	os.Exit(retCode)

}

type EchoHandler struct{}

func (e *EchoHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {

	if strings.Contains(r.URL.String(), "timeout") {
		time.Sleep(1000 * time.Minute)
		return
	}
	copyHeader(rw.Header(), r.Header)
	rw.Header().Set("Request URL", r.URL.String())
	rw.Header().Set("Request method", r.Method)
	rw.Header().Set("Request Content-Length", strconv.FormatInt(r.ContentLength, 10))
	rw.WriteHeader(200)
	if r.ContentLength > 0 {
		defer r.Body.Close()
		io.Copy(rw, r.Body)
	}
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}
<commit_msg>removed spurious go func in tests<commit_after>// Copyright Piero de Salvia.
// All Rights Reserved

/*
moved test to this package because of https://github.com/golang/go/issues/17928
*/
package dynaroutes_integration

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/pierods/dynaroutes"
)

var testDataDir string

func init() {

	goPath := os.Getenv("GOPATH")
	testDataDir = goPath + "/src/github.com/pierods/dynaroutes/testdata"
}

func TestHeaders(t *testing.T) {

	client := http.Client{}

	req, err := http.NewRequest("GET", "http://localhost:30000", nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Add("Must", "gothrough")
	req.Header.Add("This", "too")
	resp, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	if resp.Header.Get("Must") != "gothrough" || resp.Header.Get("This") != "too" {
		t.Fatal("Should be able to carry headers through")
	}
}

func TestBody(t *testing.T) {

	client := http.Client{}

	body := "Body"
	req, err := http.NewRequest("GET", "http://localhost:30000", strings.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	resp, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	buffer := new(bytes.Buffer)
	io.Copy(buffer, resp.Body)
	responseBody := buffer.String()
	t.Log(body)

	if responseBody != body+" - filtered" {
		t.Fatal("Should be able to carry body through and filter it")
	}
}

func TestTimeouts(t *testing.T) {

	client := http.Client{}

	req, err := http.NewRequest("GET", "http://localhost:30000/timeout", nil)
	if err != nil {
		t.Fatal(err)
	}
	resp, err := client.Do(req)
	if resp.StatusCode == 200 {
		t.Fatal("Should err out on a timeout")
	}

	client = http.Client{}

	req, err = http.NewRequest("GET", "http://localhost:30000/timeoutplugin", nil)
	if err != nil {
		t.Fatal(err)
	}
	resp, err = client.Do(req)
	if resp.StatusCode == 200 {
		t.Fatal("Should err out on a timeout")
	}

}

func TestMain(m *testing.M) {

	listener, err := net.Listen("tcp", ":50000")
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	server := http.Server{
		Handler: &EchoHandler{},
	}

	go func() {
		if err := server.Serve(listener); err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	}()

	defer server.Close()

	router, err := dynaroutes.NewRouterF("localhost", 30000, 5*time.Second, 5*time.Second, testDataDir, "localhost", 31000)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	go router.Start()

	defer router.Shutdown()
	time.Sleep(50 * time.Second)
	retCode := m.Run()
	os.Exit(retCode)

}

type EchoHandler struct{}

func (e *EchoHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {

	if strings.Contains(r.URL.String(), "timeout") {
		time.Sleep(1000 * time.Minute)
		return
	}
	copyHeader(rw.Header(), r.Header)
	rw.Header().Set("Request URL", r.URL.String())
	rw.Header().Set("Request method", r.Method)
	rw.Header().Set("Request Content-Length", strconv.FormatInt(r.ContentLength, 10))
	rw.WriteHeader(200)
	if r.ContentLength > 0 {
		defer r.Body.Close()
		io.Copy(rw, r.Body)
	}
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}
<|endoftext|>


<commit_before>package sphere

import "errors"

const (
	// BrokerErrorOverrideOnSubscribe to warn use to override OnSubsribe function
	BrokerErrorOverrideOnSubscribe = "please override OnSubscribe"
	// BrokerErrorOverrideOnUnsubscribe to warn use to override OnUnsubscribe function
	BrokerErrorOverrideOnUnsubscribe = "please override OnUnsubscribe"
	// BrokerErrorOverrideOnPublish to warn use to override OnPublish function
	BrokerErrorOverrideOnPublish = "please override OnPublish"
)

// Broker allows you to interact directly with Websocket internal data and pub/sub channels
type Broker struct {
	// The broker's id
	id string
	// List of channels
	channels map[string]*Channel
}

// OnSubscribe when websocket subscribes to a channel
func (broker *Broker) OnSubscribe() error {
	return errors.New(BrokerErrorOverrideOnSubscribe)
}

// OnUnsubscribe when websocket unsubscribes from a channel
func (broker *Broker) OnUnsubscribe() error {
	return errors.New(BrokerErrorOverrideOnUnsubscribe)
}

// OnPublish when websocket publishes data to a particular channel from the current broker
func (broker *Broker) OnPublish(channel *Channel, data interface{}) error {
	return errors.New(BrokerErrorOverrideOnPublish)
}
<commit_msg>create new broker func<commit_after>package sphere

import "errors"

const (
	// BrokerErrorOverrideOnSubscribe to warn use to override OnSubsribe function
	BrokerErrorOverrideOnSubscribe = "please override OnSubscribe"
	// BrokerErrorOverrideOnUnsubscribe to warn use to override OnUnsubscribe function
	BrokerErrorOverrideOnUnsubscribe = "please override OnUnsubscribe"
	// BrokerErrorOverrideOnPublish to warn use to override OnPublish function
	BrokerErrorOverrideOnPublish = "please override OnPublish"
)

// Agent represents Broker instance
type Agent interface {
	OnSubscribe()   // => Broker OnSubscribe
	OnUnsubscribe() // => Broker OnUnsubscribe
	OnPublish()     // => Broker OnPublish
}

// NewBroker creates a broker instance
func NewBroker() *Broker {
	return &Broker{
		id:       "",
		channels: make(map[string]*Channel),
	}
}

// Broker allows you to interact directly with Websocket internal data and pub/sub channels
type Broker struct {
	// The broker's id
	id string
	// List of channels
	channels map[string]*Channel
}

// OnSubscribe when websocket subscribes to a channel
func (broker *Broker) OnSubscribe() error {
	return errors.New(BrokerErrorOverrideOnSubscribe)
}

// OnUnsubscribe when websocket unsubscribes from a channel
func (broker *Broker) OnUnsubscribe() error {
	return errors.New(BrokerErrorOverrideOnUnsubscribe)
}

// OnPublish when websocket publishes data to a particular channel from the current broker
func (broker *Broker) OnPublish(channel *Channel, data interface{}) error {
	return errors.New(BrokerErrorOverrideOnPublish)
}
<|endoftext|>


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

// Package login1 provides integration with the systemd logind API.  See http://www.freedesktop.org/wiki/Software/systemd/logind/
package login1

import (
	"fmt"
	"os"
	"strconv"

	"github.com/godbus/dbus/v5"
)

const (
	dbusDest      = "org.freedesktop.login1"
	dbusInterface = "org.freedesktop.login1.Manager"
	dbusPath      = "/org/freedesktop/login1"
)

// Conn is a connection to systemds dbus endpoint.
type Conn struct {
	conn   *dbus.Conn
	object dbus.BusObject
}

// New establishes a connection to the system bus and authenticates.
func New() (*Conn, error) {
	c := new(Conn)

	if err := c.initConnection(); err != nil {
		return nil, err
	}

	return c, nil
}

// Close closes the dbus connection
func (c *Conn) Close() {
	if c == nil {
		return
	}

	if c.conn != nil {
		c.conn.Close()
	}
}

func (c *Conn) initConnection() error {
	var err error
	c.conn, err = dbus.SystemBusPrivate()
	if err != nil {
		return err
	}

	// Only use EXTERNAL method, and hardcode the uid (not username)
	// to avoid a username lookup (which requires a dynamically linked
	// libc)
	methods := []dbus.Auth{dbus.AuthExternal(strconv.Itoa(os.Getuid()))}

	err = c.conn.Auth(methods)
	if err != nil {
		c.conn.Close()
		return err
	}

	err = c.conn.Hello()
	if err != nil {
		c.conn.Close()
		return err
	}

	c.object = c.conn.Object("org.freedesktop.login1", dbus.ObjectPath(dbusPath))

	return nil
}

// Session object definition.
type Session struct {
	ID   string
	UID  uint32
	User string
	Seat string
	Path dbus.ObjectPath
}

// User object definition.
type User struct {
	UID  uint32
	Name string
	Path dbus.ObjectPath
}

func (s Session) toInterface() []interface{} {
	return []interface{}{s.ID, s.UID, s.User, s.Seat, s.Path}
}

func sessionFromInterfaces(session []interface{}) (*Session, error) {
	if len(session) < 5 {
		return nil, fmt.Errorf("invalid number of session fields: %d", len(session))
	}
	id, ok := session[0].(string)
	if !ok {
		return nil, fmt.Errorf("failed to typecast session field 0 to string")
	}
	uid, ok := session[1].(uint32)
	if !ok {
		return nil, fmt.Errorf("failed to typecast session field 1 to uint32")
	}
	user, ok := session[2].(string)
	if !ok {
		return nil, fmt.Errorf("failed to typecast session field 2 to string")
	}
	seat, ok := session[3].(string)
	if !ok {
		return nil, fmt.Errorf("failed to typecast session field 2 to string")
	}
	path, ok := session[4].(dbus.ObjectPath)
	if !ok {
		return nil, fmt.Errorf("failed to typecast session field 4 to ObjectPath")
	}

	ret := Session{ID: id, UID: uid, User: user, Seat: seat, Path: path}
	return &ret, nil
}

func userFromInterfaces(user []interface{}) (*User, error) {
	if len(user) < 3 {
		return nil, fmt.Errorf("invalid number of user fields: %d", len(user))
	}
	uid, ok := user[0].(uint32)
	if !ok {
		return nil, fmt.Errorf("failed to typecast user field 0 to uint32")
	}
	name, ok := user[1].(string)
	if !ok {
		return nil, fmt.Errorf("failed to typecast session field 1 to string")
	}
	path, ok := user[2].(dbus.ObjectPath)
	if !ok {
		return nil, fmt.Errorf("failed to typecast user field 2 to ObjectPath")
	}

	ret := User{UID: uid, Name: name, Path: path}
	return &ret, nil
}

// GetActiveSession may be used to get the session object path for the current active session
func (c *Conn) GetActiveSession() (dbus.ObjectPath, error) {
	var activeSessionPath dbus.ObjectPath
	var seat0Path dbus.ObjectPath
	if err := c.object.Call(dbusInterface+".GetSeat", 0, "seat0").Store(&seat0Path); err != nil {
		return "", err
	}

	seat0Obj := c.conn.Object(dbusDest, seat0Path)
	activeSession, err := seat0Obj.GetProperty(dbusDest + ".Seat.ActiveSession")
	if err != nil {
		return "", err
	}
	activeSessionMap, ok := activeSession.Value().([]interface{})
	if !ok || len(activeSessionMap) < 2 {
		return "", fmt.Errorf("failed to typecast active session map")
	}

	activeSessionPath, ok = activeSessionMap[1].(dbus.ObjectPath)
	if !ok {
		return "", fmt.Errorf("failed to typecast dbus active session Path")
	}
	return activeSessionPath, nil
}

// GetSessionUser may be used to get the user of specific session
func (c *Conn) GetSessionUser(sessionPath dbus.ObjectPath) (*User, error) {
	if len(sessionPath) == 0 {
		return nil, fmt.Errorf("Empty sessionPath")
	}

	var user User
	activeSessionObj := c.conn.Object(dbusDest, sessionPath)
	sessionUserName, err := activeSessionObj.GetProperty(dbusDest + ".Session.Name")
	if err != nil {
		return nil, err
	}

	sessionUser, err := activeSessionObj.GetProperty(dbusDest + ".Session.User")
	if err != nil {
		return nil, err
	}
	dbusUser, ok := sessionUser.Value().([]interface{})
	if !ok {
		return nil, fmt.Errorf("failed to typecast dbus session user")
	}

	if len(dbusUser) < 2 {
		return nil, fmt.Errorf("invalid number of user fields: %d", len(dbusUser))
	}
	uid, ok := dbusUser[0].(uint32)
	if !ok {
		return nil, fmt.Errorf("failed to typecast user field 0 to uint32")
	}
	path, ok := dbusUser[1].(dbus.ObjectPath)
	if !ok {
		return nil, fmt.Errorf("failed to typecast user field 1 to ObjectPath")
	}

	user = User{UID: uid, Name: sessionUserName.String(), Path: path}

	return &user, nil
}

// GetSessionDisplay may be used to get the display for specific session
func (c *Conn) GetSessionDisplay(sessionPath dbus.ObjectPath) (string, error) {
	if len(sessionPath) == 0 {
		return "", fmt.Errorf("Empty sessionPath")
	}
	sessionObj := c.conn.Object(dbusDest, sessionPath)
	display, err := sessionObj.GetProperty(dbusDest + ".Session.Display")
	if err != nil {
		return "", err
	}

	return display.String(), nil
}

// GetSession may be used to get the session object path for the session with the specified ID.
func (c *Conn) GetSession(id string) (dbus.ObjectPath, error) {
	var out interface{}
	if err := c.object.Call(dbusInterface+".GetSession", 0, id).Store(&out); err != nil {
		return "", err
	}

	ret, ok := out.(dbus.ObjectPath)
	if !ok {
		return "", fmt.Errorf("failed to typecast session to ObjectPath")
	}

	return ret, nil
}

// ListSessions returns an array with all current sessions.
func (c *Conn) ListSessions() ([]Session, error) {
	out := [][]interface{}{}
	if err := c.object.Call(dbusInterface+".ListSessions", 0).Store(&out); err != nil {
		return nil, err
	}

	ret := []Session{}
	for _, el := range out {
		session, err := sessionFromInterfaces(el)
		if err != nil {
			return nil, err
		}
		ret = append(ret, *session)
	}
	return ret, nil
}

// ListUsers returns an array with all currently logged in users.
func (c *Conn) ListUsers() ([]User, error) {
	out := [][]interface{}{}
	if err := c.object.Call(dbusInterface+".ListUsers", 0).Store(&out); err != nil {
		return nil, err
	}

	ret := []User{}
	for _, el := range out {
		user, err := userFromInterfaces(el)
		if err != nil {
			return nil, err
		}
		ret = append(ret, *user)
	}
	return ret, nil
}

// LockSession asks the session with the specified ID to activate the screen lock.
func (c *Conn) LockSession(id string) {
	c.object.Call(dbusInterface+".LockSession", 0, id)
}

// LockSessions asks all sessions to activate the screen locks. This may be used to lock any access to the machine in one action.
func (c *Conn) LockSessions() {
	c.object.Call(dbusInterface+".LockSessions", 0)
}

// TerminateSession forcibly terminate one specific session.
func (c *Conn) TerminateSession(id string) {
	c.object.Call(dbusInterface+".TerminateSession", 0, id)
}

// TerminateUser forcibly terminates all processes of a user.
func (c *Conn) TerminateUser(uid uint32) {
	c.object.Call(dbusInterface+".TerminateUser", 0, uid)
}

// Reboot asks logind for a reboot optionally asking for auth.
func (c *Conn) Reboot(askForAuth bool) {
	c.object.Call(dbusInterface+".Reboot", 0, askForAuth)
}

// Inhibit takes inhibition lock in logind.
func (c *Conn) Inhibit(what, who, why, mode string) (*os.File, error) {
	var fd dbus.UnixFD

	err := c.object.Call(dbusInterface+".Inhibit", 0, what, who, why, mode).Store(&fd)
	if err != nil {
		return nil, err
	}

	return os.NewFile(uintptr(fd), "inhibit"), nil
}

// Subscribe to signals on the logind dbus
func (c *Conn) Subscribe(members ...string) chan *dbus.Signal {
	for _, member := range members {
		c.conn.BusObject().Call("org.freedesktop.DBus.AddMatch", 0,
			fmt.Sprintf("type='signal',interface='org.freedesktop.login1.Manager',member='%s'", member))
	}
	ch := make(chan *dbus.Signal, 10)
	c.conn.Signal(ch)
	return ch
}

// PowerOff asks logind for a power off optionally asking for auth.
func (c *Conn) PowerOff(askForAuth bool) {
	c.object.Call(dbusInterface+".PowerOff", 0, askForAuth)
}
<commit_msg>removed quotes from username and display. Made suggested fixes for PR<commit_after>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package login1 provides integration with the systemd logind API.  See http://www.freedesktop.org/wiki/Software/systemd/logind/
package login1

import (
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/godbus/dbus/v5"
)

const (
	dbusDest      = "org.freedesktop.login1"
	dbusInterface = "org.freedesktop.login1.Manager"
	dbusPath      = "/org/freedesktop/login1"
)

// Conn is a connection to systemds dbus endpoint.
type Conn struct {
	conn   *dbus.Conn
	object dbus.BusObject
}

// New establishes a connection to the system bus and authenticates.
func New() (*Conn, error) {
	c := new(Conn)

	if err := c.initConnection(); err != nil {
		return nil, err
	}

	return c, nil
}

// Close closes the dbus connection
func (c *Conn) Close() {
	if c == nil {
		return
	}

	if c.conn != nil {
		c.conn.Close()
	}
}

func (c *Conn) initConnection() error {
	var err error
	c.conn, err = dbus.SystemBusPrivate()
	if err != nil {
		return err
	}

	// Only use EXTERNAL method, and hardcode the uid (not username)
	// to avoid a username lookup (which requires a dynamically linked
	// libc)
	methods := []dbus.Auth{dbus.AuthExternal(strconv.Itoa(os.Getuid()))}

	err = c.conn.Auth(methods)
	if err != nil {
		c.conn.Close()
		return err
	}

	err = c.conn.Hello()
	if err != nil {
		c.conn.Close()
		return err
	}

	c.object = c.conn.Object("org.freedesktop.login1", dbus.ObjectPath(dbusPath))

	return nil
}

// Session object definition.
type Session struct {
	ID   string
	UID  uint32
	User string
	Seat string
	Path dbus.ObjectPath
}

// User object definition.
type User struct {
	UID  uint32
	Name string
	Path dbus.ObjectPath
}

func (s Session) toInterface() []interface{} {
	return []interface{}{s.ID, s.UID, s.User, s.Seat, s.Path}
}

func sessionFromInterfaces(session []interface{}) (*Session, error) {
	if len(session) < 5 {
		return nil, fmt.Errorf("invalid number of session fields: %d", len(session))
	}
	id, ok := session[0].(string)
	if !ok {
		return nil, fmt.Errorf("failed to typecast session field 0 to string")
	}
	uid, ok := session[1].(uint32)
	if !ok {
		return nil, fmt.Errorf("failed to typecast session field 1 to uint32")
	}
	user, ok := session[2].(string)
	if !ok {
		return nil, fmt.Errorf("failed to typecast session field 2 to string")
	}
	seat, ok := session[3].(string)
	if !ok {
		return nil, fmt.Errorf("failed to typecast session field 2 to string")
	}
	path, ok := session[4].(dbus.ObjectPath)
	if !ok {
		return nil, fmt.Errorf("failed to typecast session field 4 to ObjectPath")
	}

	ret := Session{ID: id, UID: uid, User: user, Seat: seat, Path: path}
	return &ret, nil
}

func userFromInterfaces(user []interface{}) (*User, error) {
	if len(user) < 3 {
		return nil, fmt.Errorf("invalid number of user fields: %d", len(user))
	}
	uid, ok := user[0].(uint32)
	if !ok {
		return nil, fmt.Errorf("failed to typecast user field 0 to uint32")
	}
	name, ok := user[1].(string)
	if !ok {
		return nil, fmt.Errorf("failed to typecast session field 1 to string")
	}
	path, ok := user[2].(dbus.ObjectPath)
	if !ok {
		return nil, fmt.Errorf("failed to typecast user field 2 to ObjectPath")
	}

	ret := User{UID: uid, Name: name, Path: path}
	return &ret, nil
}

// GetActiveSession may be used to get the session object path for the current active session
func (c *Conn) GetActiveSession() (dbus.ObjectPath, error) {
	var seat0Path dbus.ObjectPath
	if err := c.object.Call(dbusInterface+".GetSeat", 0, "seat0").Store(&seat0Path); err != nil {
		return "", err
	}

	seat0Obj := c.conn.Object(dbusDest, seat0Path)
	activeSession, err := seat0Obj.GetProperty(dbusDest + ".Seat.ActiveSession")
	if err != nil {
		return "", err
	}
	activeSessionMap, ok := activeSession.Value().([]interface{})
	if !ok || len(activeSessionMap) < 2 {
		return "", fmt.Errorf("failed to typecast active session map")
	}

	activeSessionPath, ok := activeSessionMap[1].(dbus.ObjectPath)
	if !ok {
		return "", fmt.Errorf("failed to typecast dbus active session Path")
	}
	return activeSessionPath, nil
}

// GetSessionUser may be used to get the user of specific session
func (c *Conn) GetSessionUser(sessionPath dbus.ObjectPath) (*User, error) {
	if len(sessionPath) == 0 {
		return nil, fmt.Errorf("empty sessionPath")
	}

	activeSessionObj := c.conn.Object(dbusDest, sessionPath)
	sessionUserName, err := activeSessionObj.GetProperty(dbusDest + ".Session.Name")
	if err != nil {
		return nil, err
	}

	sessionUser, err := activeSessionObj.GetProperty(dbusDest + ".Session.User")
	if err != nil {
		return nil, err
	}
	dbusUser, ok := sessionUser.Value().([]interface{})
	if !ok {
		return nil, fmt.Errorf("failed to typecast dbus session user")
	}

	if len(dbusUser) < 2 {
		return nil, fmt.Errorf("invalid number of user fields: %d", len(dbusUser))
	}
	uid, ok := dbusUser[0].(uint32)
	if !ok {
		return nil, fmt.Errorf("failed to typecast user field 0 to uint32")
	}
	path, ok := dbusUser[1].(dbus.ObjectPath)
	if !ok {
		return nil, fmt.Errorf("failed to typecast user field 1 to ObjectPath")
	}

	user := User{UID: uid, Name: strings.Trim(sessionUserName.String(), "\""), Path: path}

	return &user, nil
}

// GetSessionDisplay may be used to get the display for specific session
func (c *Conn) GetSessionDisplay(sessionPath dbus.ObjectPath) (string, error) {
	if len(sessionPath) == 0 {
		return "", fmt.Errorf("empty sessionPath")
	}
	sessionObj := c.conn.Object(dbusDest, sessionPath)
	display, err := sessionObj.GetProperty(dbusDest + ".Session.Display")
	if err != nil {
		return "", err
	}

	return strings.Trim(display.String(), "\""), nil
}

// GetSession may be used to get the session object path for the session with the specified ID.
func (c *Conn) GetSession(id string) (dbus.ObjectPath, error) {
	var out interface{}
	if err := c.object.Call(dbusInterface+".GetSession", 0, id).Store(&out); err != nil {
		return "", err
	}

	ret, ok := out.(dbus.ObjectPath)
	if !ok {
		return "", fmt.Errorf("failed to typecast session to ObjectPath")
	}

	return ret, nil
}

// ListSessions returns an array with all current sessions.
func (c *Conn) ListSessions() ([]Session, error) {
	out := [][]interface{}{}
	if err := c.object.Call(dbusInterface+".ListSessions", 0).Store(&out); err != nil {
		return nil, err
	}

	ret := []Session{}
	for _, el := range out {
		session, err := sessionFromInterfaces(el)
		if err != nil {
			return nil, err
		}
		ret = append(ret, *session)
	}
	return ret, nil
}

// ListUsers returns an array with all currently logged in users.
func (c *Conn) ListUsers() ([]User, error) {
	out := [][]interface{}{}
	if err := c.object.Call(dbusInterface+".ListUsers", 0).Store(&out); err != nil {
		return nil, err
	}

	ret := []User{}
	for _, el := range out {
		user, err := userFromInterfaces(el)
		if err != nil {
			return nil, err
		}
		ret = append(ret, *user)
	}
	return ret, nil
}

// LockSession asks the session with the specified ID to activate the screen lock.
func (c *Conn) LockSession(id string) {
	c.object.Call(dbusInterface+".LockSession", 0, id)
}

// LockSessions asks all sessions to activate the screen locks. This may be used to lock any access to the machine in one action.
func (c *Conn) LockSessions() {
	c.object.Call(dbusInterface+".LockSessions", 0)
}

// TerminateSession forcibly terminate one specific session.
func (c *Conn) TerminateSession(id string) {
	c.object.Call(dbusInterface+".TerminateSession", 0, id)
}

// TerminateUser forcibly terminates all processes of a user.
func (c *Conn) TerminateUser(uid uint32) {
	c.object.Call(dbusInterface+".TerminateUser", 0, uid)
}

// Reboot asks logind for a reboot optionally asking for auth.
func (c *Conn) Reboot(askForAuth bool) {
	c.object.Call(dbusInterface+".Reboot", 0, askForAuth)
}

// Inhibit takes inhibition lock in logind.
func (c *Conn) Inhibit(what, who, why, mode string) (*os.File, error) {
	var fd dbus.UnixFD

	err := c.object.Call(dbusInterface+".Inhibit", 0, what, who, why, mode).Store(&fd)
	if err != nil {
		return nil, err
	}

	return os.NewFile(uintptr(fd), "inhibit"), nil
}

// Subscribe to signals on the logind dbus
func (c *Conn) Subscribe(members ...string) chan *dbus.Signal {
	for _, member := range members {
		c.conn.BusObject().Call("org.freedesktop.DBus.AddMatch", 0,
			fmt.Sprintf("type='signal',interface='org.freedesktop.login1.Manager',member='%s'", member))
	}
	ch := make(chan *dbus.Signal, 10)
	c.conn.Signal(ch)
	return ch
}

// PowerOff asks logind for a power off optionally asking for auth.
func (c *Conn) PowerOff(askForAuth bool) {
	c.object.Call(dbusInterface+".PowerOff", 0, askForAuth)
}
<|endoftext|>


<commit_before>package goridge

import (
	"io"
	"sync"
)

// SocketRelay communicates with underlying process using sockets (TPC or Unix).
type SocketRelay struct {
	// How many bytes to write/read at once.
	BufferSize uint64

	muw sync.Mutex // concurrent write
	mur sync.Mutex // concurrent read
	rwc io.ReadWriteCloser
}

// NewSocketRelay creates new socket based data relay.
func NewSocketRelay(rwc io.ReadWriteCloser) *SocketRelay {
	return &SocketRelay{BufferSize: BufferSize, rwc: rwc}
}

// Send signed (prefixed) data to PHP process.
func (rl *SocketRelay) Send(data []byte, flags byte) (err error) {
	rl.muw.Lock()
	defer rl.muw.Unlock()

	prefix := NewPrefix().WithFlags(flags).WithSize(uint64(len(data)))
	if _, err := rl.rwc.Write(prefix[:]); err != nil {
		return err
	}

	if _, err := rl.rwc.Write(data); err != nil {
		return err
	}

	return nil
}

// Receive data from the underlying process and returns associated prefix or error.
func (rl *SocketRelay) Receive() (data []byte, p Prefix, err error) {
	rl.mur.Lock()
	defer rl.mur.Unlock()

	if _, err := rl.rwc.Read(p[:]); err != nil {
		return nil, p, err
	}

	if !p.HasPayload() {
		return nil, p, nil
	}

	leftBytes := p.Size()
	data = make([]byte, 0, leftBytes)
	buffer := make([]byte, min(leftBytes, rl.BufferSize))

	for {
		if n, err := rl.rwc.Read(buffer); err == nil {
			data = append(data, buffer[:n]...)
			leftBytes -= uint64(n)
		} else {
			return nil, p, err
		}

		if leftBytes == 0 {
			break
		}
	}

	return
}

// Close the connection.
func (rl *SocketRelay) Close() error {
	rl.muw.Lock()
	rl.mur.Lock()
	defer rl.muw.Unlock()
	defer rl.mur.Unlock()

	return rl.rwc.Close()
}
<commit_msg>error handling in receive<commit_after>package goridge

import (
	"io"
	"sync"
)

// SocketRelay communicates with underlying process using sockets (TPC or Unix).
type SocketRelay struct {
	// How many bytes to write/read at once.
	BufferSize uint64

	muw sync.Mutex // concurrent write
	mur sync.Mutex // concurrent read
	rwc io.ReadWriteCloser
}

// NewSocketRelay creates new socket based data relay.
func NewSocketRelay(rwc io.ReadWriteCloser) *SocketRelay {
	return &SocketRelay{BufferSize: BufferSize, rwc: rwc}
}

// Send signed (prefixed) data to PHP process.
func (rl *SocketRelay) Send(data []byte, flags byte) (err error) {
	rl.muw.Lock()
	defer rl.muw.Unlock()

	prefix := NewPrefix().WithFlags(flags).WithSize(uint64(len(data)))
	if _, err := rl.rwc.Write(prefix[:]); err != nil {
		return err
	}

	if _, err := rl.rwc.Write(data); err != nil {
		return err
	}

	return nil
}

// Receive data from the underlying process and returns associated prefix or error.
func (rl *SocketRelay) Receive() (data []byte, p Prefix, err error) {
	rl.mur.Lock()
	defer rl.mur.Unlock()

	defer func() {
		if rErr, ok := recover().(error); ok {
			err = rErr
		}
	}()

	if _, err := rl.rwc.Read(p[:]); err != nil {
		return nil, p, err
	}

	if !p.HasPayload() {
		return nil, p, nil
	}

	leftBytes := p.Size()
	data = make([]byte, 0, leftBytes)
	buffer := make([]byte, min(leftBytes, rl.BufferSize))

	for {
		if n, err := rl.rwc.Read(buffer); err == nil {
			data = append(data, buffer[:n]...)
			leftBytes -= uint64(n)
		} else {
			return nil, p, err
		}

		if leftBytes == 0 {
			break
		}
	}

	return
}

// Close the connection.
func (rl *SocketRelay) Close() error {
	rl.muw.Lock()
	rl.mur.Lock()
	defer rl.muw.Unlock()
	defer rl.mur.Unlock()

	return rl.rwc.Close()
}
<|endoftext|>