<commit_before>package drouter

import (
	"fmt"
	log "github.com/Sirupsen/logrus"
	dockertypes "github.com/docker/engine-api/types"
	"github.com/vishvananda/netlink"
	"golang.org/x/net/context"
	"net"
	"strings"
)

type container struct {
	handle *netlink.Handle
}

func newContainerFromID(id string) (*container, error) {
	cjson, err := dockerClient.ContainerInspect(context.Background(), id)
	if err != nil {
		return nil, err
	}
	ch, err := netlinkHandleFromPid(cjson.State.Pid)
	if err != nil {
		return nil, err
	}

	return &container{handle: ch}, nil
}

//adds all known routes for provided container
func (c *container) addAllRoutes() error {
	//Loop through all static routes, ensure each one is installed in the container
	log.Info("Syncing static routes.")
	//add routes for all the static routes
	for _, sr := range staticRoutes {
		go c.addRoute(sr)
	}

	//Loop through all discovered networks, ensure each one is installed in the container
	//Unless it is covered by a static route already
	log.Info("Syncing discovered routes.")
	routes, err := netlink.RouteList(nil, netlink.FAMILY_ALL)
	if err != nil {
		log.Error("Failed to get my routes.")
		log.Error(err)
	}

Subnets:
	for _, r := range routes {
		if r.Gw != nil {
			continue
		}
		if localShortcut {
			if subnetEqualSubnet(r.Dst, networkID(p2p.network)) {
				continue
			}
		}
		for _, sr := range staticRoutes {
			if subnetContainsSubnet(sr, r.Dst) {
				log.Debugf("Skipping route %v covered by %v.", r.Dst, sr)
				break Subnets
			}
		}
		go c.addRoute(r.Dst)
	}

	return nil
}

func (c *container) addRoute(prefix *net.IPNet) {
	gateway, err := c.getPathIP()
	if err != nil {
		log.Error(err)
		return
	}

	if (prefix.IP.To4() == nil) != (gateway.To4() == nil) {
		// Dst is a different IP family
		return
	}

	route := &netlink.Route{
		Dst: prefix,
		Gw:  gateway,
	}

	log.Infof("Adding route to %v via %v.", prefix, gateway)
	err = c.handle.RouteAdd(route)
	if err != nil {
		if !strings.Contains(err.Error(), "file exists") {
			log.Error(err)
			return
		}
	}
}

func (c *container) delRoutes(prefix *net.IPNet) {
	//get all container routes
	routes, err := c.handle.RouteList(nil, netlink.FAMILY_V4)
	if err != nil {
		log.Error("Failed to get container route table.")
		log.Error(err)
		return
	}

	//get all drouter ips
	ips, err := netlink.AddrList(nil, netlink.FAMILY_V4)
	if err != nil {
		log.Error("Failed to get drouter ip addresses.")
		log.Error(err)
		return
	}

	for _, r := range routes {
		if r.Dst == nil {
			continue
		}
		if !subnetContainsSubnet(prefix, r.Dst) {
			continue
		}

		for _, ipaddr := range ips {
			if r.Gw.Equal(ipaddr.IP) {
				err := c.handle.RouteDel(&r)
				if err != nil {
					log.Errorf("Failed to delete container route to %v via %v", r.Dst, r.Gw)
					continue
				}
			}
		}
	}
}

//sets the provided container's default route to the provided gateway
func (c *container) replaceGateway(gateway net.IP) error {
	log.Debugf("container.replaceGateway(%v)", gateway)

	var defr *netlink.Route
	//replace containers default gateway with drouter
	routes, err := c.handle.RouteList(nil, netlink.FAMILY_V4)
	if err != nil {
		return err
	}
	log.Debugf("container routes: %v", routes)
	for _, r := range routes {
		if r.Dst != nil {
			// Not the container gateway
			continue
		}

		defr = &r
	}

	//bail if the container gateway is already set to gateway
	if gateway.Equal(defr.Gw) {
		return nil
	}

	log.Debugf("Remove existing default route: %v", defr)
	err = c.handle.RouteDel(defr)
	if err != nil {
		return err
	}

	if gateway == nil || gateway.Equal(net.IP{}) {
		return nil
	}

	defr.Gw = gateway
	err = c.handle.RouteAdd(defr)
	if err != nil {
		return err
	}
	log.Debugf("Default route changed to: %v", defr)

	return nil
}

// called during a network connect event
func (c *container) networkConnectEvent(drn *network) error {
	if !drn.isConnected() {
		drn.connect()
	}
	//let's push our routes into this new container
	gateway, err := c.getPathIP()
	if err != nil {
		log.Error("Failed to get container path IP.")
		return err
	}

	if localGateway {
		go c.replaceGateway(gateway)
	} else {
		go c.addAllRoutes()
	}
	return nil
}

// called during a network disconnect event
func (c *container) networkDisconnectEvent(drn *network) error {

	all, _ := netlink.ParseIPNet("0.0.0.0/0")
	err := c.delRoutes(all)
	if err != nil {
		log.Errorf("Error removing routes from container")
	}

	if pIP, _ := c.getPathIP(); pIP != nil {
		c.addAllRoutes()
	}

	if aggressive {
		return nil
	}

	//if not aggressive mode, then we disconnect from the network if this is the last connected container

	//loop through all the containers
	dockerContainers, err := dockerClient.ContainerList(context.Background(), dockertypes.ContainerListOptions{})
	if err != nil {
		return err
	}
	for _, dc := range dockerContainers {
		if dc.HostConfig.NetworkMode == "host" {
			continue
		}
		if dc.ID == selfContainerID {
			continue
		}

		for n, _ := range dc.NetworkSettings.Networks {
			if n == drn.Name {
				// This netowork is still in use
				return nil
			}
		}
	}

	drn.disconnect()

	return nil
}

//returns a drouter IP that is on some same network as provided container
func (c *container) getPathIP() (net.IP, error) {
	addrs, err := c.handle.AddrList(nil, netlink.FAMILY_V4)
	if err != nil {
		log.Error("Failed to list container addresses.")
		return nil, err
	}

	for _, addr := range addrs {
		if addr.Label == "lo" {
			continue
		}

		log.Debugf("Getting my route to container IP: %v", addr.IP)
		srcRoutes, err := netlink.RouteGet(addr.IP)
		if err != nil {
			log.Error(err)
			continue
		}
		for _, srcRoute := range srcRoutes {
			if srcRoute.Gw == nil {
				return srcRoute.Src, nil
			}
		}
	}

	return nil, fmt.Errorf("No direct connection to container.")
}
<commit_msg>only delete routes via a specific drn<commit_after>package drouter

import (
	"fmt"
	log "github.com/Sirupsen/logrus"
	dockertypes "github.com/docker/engine-api/types"
	"github.com/vishvananda/netlink"
	"golang.org/x/net/context"
	"net"
	"strings"
)

type container struct {
	handle *netlink.Handle
}

func newContainerFromID(id string) (*container, error) {
	cjson, err := dockerClient.ContainerInspect(context.Background(), id)
	if err != nil {
		return nil, err
	}
	ch, err := netlinkHandleFromPid(cjson.State.Pid)
	if err != nil {
		return nil, err
	}

	return &container{handle: ch}, nil
}

//adds all known routes for provided container
func (c *container) addAllRoutes() error {
	//Loop through all static routes, ensure each one is installed in the container
	log.Info("Syncing static routes.")
	//add routes for all the static routes
	for _, sr := range staticRoutes {
		go c.addRoute(sr)
	}

	//Loop through all discovered networks, ensure each one is installed in the container
	//Unless it is covered by a static route already
	log.Info("Syncing discovered routes.")
	routes, err := netlink.RouteList(nil, netlink.FAMILY_ALL)
	if err != nil {
		log.Error("Failed to get my routes.")
		log.Error(err)
	}

Subnets:
	for _, r := range routes {
		if r.Gw != nil {
			continue
		}
		if localShortcut {
			if subnetEqualSubnet(r.Dst, networkID(p2p.network)) {
				continue
			}
		}
		for _, sr := range staticRoutes {
			if subnetContainsSubnet(sr, r.Dst) {
				log.Debugf("Skipping route %v covered by %v.", r.Dst, sr)
				break Subnets
			}
		}
		go c.addRoute(r.Dst)
	}

	return nil
}

func (c *container) addRoute(prefix *net.IPNet) {
	gateway, err := c.getPathIP()
	if err != nil {
		log.Error(err)
		return
	}

	if (prefix.IP.To4() == nil) != (gateway.To4() == nil) {
		// Dst is a different IP family
		return
	}

	route := &netlink.Route{
		Dst: prefix,
		Gw:  gateway,
	}

	log.Infof("Adding route to %v via %v.", prefix, gateway)
	err = c.handle.RouteAdd(route)
	if err != nil {
		if !strings.Contains(err.Error(), "file exists") {
			log.Error(err)
			return
		}
	}
}

func (c *container) delRoutesVia(drn *network) {
	routes, err := c.handle.RouteList(nil, netlink.FAMILY_V4)
	if err != nil {
		log.Error("Failed to get container route table.")
		log.Error(err)
		return
	}

	//get all drouter ips
	ips, err := netlink.AddrList(nil, netlink.FAMILY_V4)
	if err != nil {
		log.Error("Failed to get drouter ip addresses.")
		log.Error(err)
		return
	}

	for _, r := range routes {
		if r.Dst == nil {
			continue
		}

		for _, ipaddr := range ips {
			if !r.Gw.Equal(ipaddr.IP) {
				continue
			}
			for _, ic := range drn.IPAM.Config {
				sn, err := netlink.ParseIPNet(ic.Subnet)
				if err != nil {
					log.Errorf("Failed to parse IPNet %v", ic.Subnet)
					continue
				}
				if !sn.Contains(r.Gw) {
					continue
				}
				err = c.handle.RouteDel(&r)
				if err != nil {
					log.Errorf("Failed to delete container route to %v via %v", r.Dst, r.Gw)
					continue
				}
			}
		}
	}
}

func (c *container) delRoutes(prefix *net.IPNet) {
	//get all container routes
	routes, err := c.handle.RouteList(nil, netlink.FAMILY_V4)
	if err != nil {
		log.Error("Failed to get container route table.")
		log.Error(err)
		return
	}

	//get all drouter ips
	ips, err := netlink.AddrList(nil, netlink.FAMILY_V4)
	if err != nil {
		log.Error("Failed to get drouter ip addresses.")
		log.Error(err)
		return
	}

	for _, r := range routes {
		if r.Dst == nil {
			continue
		}
		if !subnetContainsSubnet(prefix, r.Dst) {
			continue
		}

		for _, ipaddr := range ips {
			if !r.Gw.Equal(ipaddr.IP) {
				continue
			}
			err := c.handle.RouteDel(&r)
			if err != nil {
				log.Errorf("Failed to delete container route to %v via %v", r.Dst, r.Gw)
				continue
			}
		}
	}
}

//sets the provided container's default route to the provided gateway
func (c *container) replaceGateway(gateway net.IP) error {
	log.Debugf("container.replaceGateway(%v)", gateway)

	var defr *netlink.Route
	//replace containers default gateway with drouter
	routes, err := c.handle.RouteList(nil, netlink.FAMILY_V4)
	if err != nil {
		return err
	}
	log.Debugf("container routes: %v", routes)
	for _, r := range routes {
		if r.Dst != nil {
			// Not the container gateway
			continue
		}

		defr = &r
	}

	//bail if the container gateway is already set to gateway
	if gateway.Equal(defr.Gw) {
		return nil
	}

	log.Debugf("Remove existing default route: %v", defr)
	err = c.handle.RouteDel(defr)
	if err != nil {
		return err
	}

	if gateway == nil || gateway.Equal(net.IP{}) {
		return nil
	}

	defr.Gw = gateway
	err = c.handle.RouteAdd(defr)
	if err != nil {
		return err
	}
	log.Debugf("Default route changed to: %v", defr)

	return nil
}

// called during a network connect event
func (c *container) networkConnectEvent(drn *network) error {
	if !drn.isConnected() {
		drn.connect()
	}
	//let's push our routes into this new container
	gateway, err := c.getPathIP()
	if err != nil {
		log.Error("Failed to get container path IP.")
		return err
	}

	if localGateway {
		go c.replaceGateway(gateway)
	} else {
		go c.addAllRoutes()
	}
	return nil
}

// called during a network disconnect event
func (c *container) networkDisconnectEvent(drn *network) error {

	c.delRoutesVia(drn)

	if pIP, _ := c.getPathIP(); pIP != nil {
		c.addAllRoutes()
	}

	if aggressive {
		return nil
	}

	//if not aggressive mode, then we disconnect from the network if this is the last connected container

	//loop through all the containers
	dockerContainers, err := dockerClient.ContainerList(context.Background(), dockertypes.ContainerListOptions{})
	if err != nil {
		return err
	}
	for _, dc := range dockerContainers {
		if dc.HostConfig.NetworkMode == "host" {
			continue
		}
		if dc.ID == selfContainerID {
			continue
		}

		for n, _ := range dc.NetworkSettings.Networks {
			if n == drn.Name {
				// This netowork is still in use
				return nil
			}
		}
	}

	drn.disconnect()

	return nil
}

//returns a drouter IP that is on some same network as provided container
func (c *container) getPathIP() (net.IP, error) {
	addrs, err := c.handle.AddrList(nil, netlink.FAMILY_V4)
	if err != nil {
		log.Error("Failed to list container addresses.")
		return nil, err
	}

	for _, addr := range addrs {
		if addr.Label == "lo" {
			continue
		}

		log.Debugf("Getting my route to container IP: %v", addr.IP)
		srcRoutes, err := netlink.RouteGet(addr.IP)
		if err != nil {
			log.Error(err)
			continue
		}
		for _, srcRoute := range srcRoutes {
			if srcRoute.Gw == nil {
				return srcRoute.Src, nil
			}
		}
	}

	return nil, fmt.Errorf("No direct connection to container.")
}
<|endoftext|>


<commit_before>package ngrok

import (
	"context"
	"errors"
	"fmt"
	"net/url"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"

	"github.com/efarrer/gmash/ptyutils"
	"github.com/kr/pty"
)

// Reason is the reason why executing ngrok failed
type Reason int

const (
	// 0 is unused to help find issues with errors that have a default value
	_ Reason = iota
	// MissingNgrok indicates ngrok executable can't be found
	MissingNgrok Reason = iota
	// UnexecutableNgrok indicates ngrok can't be executed
	UnexecutableNgrok Reason = iota
	// MissingAuthToken indicates ngrok must be authed
	MissingAuthToken Reason = iota
	// Canceled indicates that the user canceled the execution
	Canceled Reason = iota
	// CantReadFromPty indicates that there was a problem reading the stdout from ngrok
	CantReadFromPty Reason = iota
	// PortParsingError indicates that there was a problem parsing the forwarding url's port from ngrok
	PortParsingError Reason = iota
	// URLParsingError indicates that there was a problem parsing the forwarding url from ngrok
	URLParsingError Reason = iota
	// CantSetPtyWindowSize indicates that there was a problem setting the pty's window size
	CantSetPtyWindowSize Reason = iota
)

// ExecutionError is an error type returned by Execute
type ExecutionError struct {
	Reason Reason
	Err    error
}

// Error returns the error string
func (r *ExecutionError) Error() string {
	return fmt.Sprintf("Error(%d) %s", r.Reason, r.Err)
}

// Value is the Host and Port found by executing ngrok
type Value struct {
	Host string
	Port int
}

// A Response contains either an error from executing ngrok or the Value
type Response struct {
	Err   *ExecutionError
	Value *Value
}

// String returns a human friendly representation of a Response
func (rs Response) String() string {
	res := "ngrok.Response"
	if rs.Err != nil {
		res += fmt.Sprintf("%d %s", rs.Err.Reason, rs.Err.Err)
	}
	if rs.Value != nil {
		res += fmt.Sprintf("%s %s", rs.Value.Host, rs.Value.Port)
	}
	return res
}

func newErrorResponse(reason Reason, err error) Response {
	return Response{
		Err: &ExecutionError{
			Reason: reason,
			Err:    err,
		},
		Value: nil,
	}
}

// Execute executes ngrok forwarding to the given port
func Execute(ctx context.Context, port int) Response {
	return execute(ctx, port, "ngrok")
}

func execute(ctx context.Context, port int, bin string) Response {
	cmd := exec.CommandContext(ctx, bin, "tcp", strconv.FormatInt(int64(port), 10))
	_pty, err := pty.Start(cmd)
	if err != nil {
		var reason Reason
		switch err.(type) {
		case *exec.Error:
			reason = MissingNgrok
		case *os.PathError:
			reason = UnexecutableNgrok
		default:
			reason = Canceled
		}

		return newErrorResponse(reason, err)
	}

	err = ptyutils.SetWindowSize(_pty, 100, 100)
	if err != nil {
		return newErrorResponse(CantSetPtyWindowSize, err)
	}

	output := ""
	bytes := make([]byte, 1024)
	for {
		count, err := _pty.Read(bytes)
		if err != nil {
			if ctx.Err() == context.Canceled || ctx.Err() == context.DeadlineExceeded {
				return newErrorResponse(
					Canceled,
					errors.New("ngrok was canceled"),
				)
			}
			return newErrorResponse(CantReadFromPty, err)
		}
		output += string(bytes[0:count])

		if strings.Contains(output, "ERR_NGROK_302") {
			return newErrorResponse(
				MissingAuthToken,
				errors.New("Please signup at https://ngrok.com/signup or make sure your athtoken is installed https://dashboard.ngrok.com"),
			)
		}
		rx := regexp.MustCompile("Forwarding[ ]+(tcp://[^ ]+)[ ].*")

		match := rx.FindStringSubmatch(output)
		if len(match) == 2 {
			ngrokurl, err := url.Parse(match[1])
			if err != nil {
				return newErrorResponse(
					URLParsingError,
					errors.New("Unable to parse ngrok's forwarding url"),
				)
			}
			iport, err := strconv.Atoi(ngrokurl.Port())
			if err != nil {
				return newErrorResponse(
					PortParsingError,
					errors.New("Unable to parse ngrok's port"),
				)
			}

			return Response{
				Err:   nil,
				Value: &Value{Host: ngrokurl.Hostname(), Port: iport},
			}
		}
	}
}
<commit_msg>Fix Sprintf format bug<commit_after>package ngrok

import (
	"context"
	"errors"
	"fmt"
	"net/url"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"

	"github.com/efarrer/gmash/ptyutils"
	"github.com/kr/pty"
)

// Reason is the reason why executing ngrok failed
type Reason int

const (
	// 0 is unused to help find issues with errors that have a default value
	_ Reason = iota
	// MissingNgrok indicates ngrok executable can't be found
	MissingNgrok Reason = iota
	// UnexecutableNgrok indicates ngrok can't be executed
	UnexecutableNgrok Reason = iota
	// MissingAuthToken indicates ngrok must be authed
	MissingAuthToken Reason = iota
	// Canceled indicates that the user canceled the execution
	Canceled Reason = iota
	// CantReadFromPty indicates that there was a problem reading the stdout from ngrok
	CantReadFromPty Reason = iota
	// PortParsingError indicates that there was a problem parsing the forwarding url's port from ngrok
	PortParsingError Reason = iota
	// URLParsingError indicates that there was a problem parsing the forwarding url from ngrok
	URLParsingError Reason = iota
	// CantSetPtyWindowSize indicates that there was a problem setting the pty's window size
	CantSetPtyWindowSize Reason = iota
)

// ExecutionError is an error type returned by Execute
type ExecutionError struct {
	Reason Reason
	Err    error
}

// Error returns the error string
func (r *ExecutionError) Error() string {
	return fmt.Sprintf("Error(%d) %s", r.Reason, r.Err)
}

// Value is the Host and Port found by executing ngrok
type Value struct {
	Host string
	Port int
}

// A Response contains either an error from executing ngrok or the Value
type Response struct {
	Err   *ExecutionError
	Value *Value
}

// String returns a human friendly representation of a Response
func (rs Response) String() string {
	res := "ngrok.Response"
	if rs.Err != nil {
		res += fmt.Sprintf("%d %s", rs.Err.Reason, rs.Err.Err)
	}
	if rs.Value != nil {
		res += fmt.Sprintf("%s %d", rs.Value.Host, rs.Value.Port)
	}
	return res
}

func newErrorResponse(reason Reason, err error) Response {
	return Response{
		Err: &ExecutionError{
			Reason: reason,
			Err:    err,
		},
		Value: nil,
	}
}

// Execute executes ngrok forwarding to the given port
func Execute(ctx context.Context, port int) Response {
	return execute(ctx, port, "ngrok")
}

func execute(ctx context.Context, port int, bin string) Response {
	cmd := exec.CommandContext(ctx, bin, "tcp", strconv.FormatInt(int64(port), 10))
	_pty, err := pty.Start(cmd)
	if err != nil {
		var reason Reason
		switch err.(type) {
		case *exec.Error:
			reason = MissingNgrok
		case *os.PathError:
			reason = UnexecutableNgrok
		default:
			reason = Canceled
		}

		return newErrorResponse(reason, err)
	}

	err = ptyutils.SetWindowSize(_pty, 100, 100)
	if err != nil {
		return newErrorResponse(CantSetPtyWindowSize, err)
	}

	output := ""
	bytes := make([]byte, 1024)
	for {
		count, err := _pty.Read(bytes)
		if err != nil {
			if ctx.Err() == context.Canceled || ctx.Err() == context.DeadlineExceeded {
				return newErrorResponse(
					Canceled,
					errors.New("ngrok was canceled"),
				)
			}
			return newErrorResponse(CantReadFromPty, err)
		}
		output += string(bytes[0:count])

		if strings.Contains(output, "ERR_NGROK_302") {
			return newErrorResponse(
				MissingAuthToken,
				errors.New("Please signup at https://ngrok.com/signup or make sure your athtoken is installed https://dashboard.ngrok.com"),
			)
		}
		rx := regexp.MustCompile("Forwarding[ ]+(tcp://[^ ]+)[ ].*")

		match := rx.FindStringSubmatch(output)
		if len(match) == 2 {
			ngrokurl, err := url.Parse(match[1])
			if err != nil {
				return newErrorResponse(
					URLParsingError,
					errors.New("Unable to parse ngrok's forwarding url"),
				)
			}
			iport, err := strconv.Atoi(ngrokurl.Port())
			if err != nil {
				return newErrorResponse(
					PortParsingError,
					errors.New("Unable to parse ngrok's port"),
				)
			}

			return Response{
				Err:   nil,
				Value: &Value{Host: ngrokurl.Hostname(), Port: iport},
			}
		}
	}
}
<|endoftext|>


<commit_before>// Package nl has low level primitives for making Netlink calls.
package nl

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"net"
	"sync/atomic"
	"syscall"
	"unsafe"
)

const (
	// Family type definitions
	FAMILY_ALL = syscall.AF_UNSPEC
	FAMILY_V4  = syscall.AF_INET
	FAMILY_V6  = syscall.AF_INET6
)

var nextSeqNr uint32

// GetIPFamily returns the family type of a net.IP.
func GetIPFamily(ip net.IP) int {
	if len(ip) <= net.IPv4len {
		return FAMILY_V4
	}
	if ip.To4() != nil {
		return FAMILY_V4
	}
	return FAMILY_V6
}

var nativeEndian binary.ByteOrder

// Get native endianness for the system
func NativeEndian() binary.ByteOrder {
	if nativeEndian == nil {
		var x uint32 = 0x01020304
		if *(*byte)(unsafe.Pointer(&x)) == 0x01 {
			nativeEndian = binary.BigEndian
		}
		nativeEndian = binary.LittleEndian
	}
	return nativeEndian
}

// Byte swap a 16 bit value if we aren't big endian
func Swap16(i uint16) uint16 {
	if NativeEndian() == binary.BigEndian {
		return i
	}
	return (i&0xff00)>>8 | (i&0xff)<<8
}

// Byte swap a 32 bit value if aren't big endian
func Swap32(i uint32) uint32 {
	if NativeEndian() == binary.BigEndian {
		return i
	}
	return (i&0xff000000)>>24 | (i&0xff0000)>>8 | (i&0xff00)<<8 | (i&0xff)<<24
}

type NetlinkRequestData interface {
	Len() int
	Serialize() []byte
}

// IfInfomsg is related to links, but it is used for list requests as well
type IfInfomsg struct {
	syscall.IfInfomsg
}

// Create an IfInfomsg with family specified
func NewIfInfomsg(family int) *IfInfomsg {
	return &IfInfomsg{
		IfInfomsg: syscall.IfInfomsg{
			Family: uint8(family),
		},
	}
}

func DeserializeIfInfomsg(b []byte) *IfInfomsg {
	return (*IfInfomsg)(unsafe.Pointer(&b[0:syscall.SizeofIfInfomsg][0]))
}

func (msg *IfInfomsg) Serialize() []byte {
	return (*(*[syscall.SizeofIfInfomsg]byte)(unsafe.Pointer(msg)))[:]
}

func (msg *IfInfomsg) Len() int {
	return syscall.SizeofIfInfomsg
}

func rtaAlignOf(attrlen int) int {
	return (attrlen + syscall.RTA_ALIGNTO - 1) & ^(syscall.RTA_ALIGNTO - 1)
}

func NewIfInfomsgChild(parent *RtAttr, family int) *IfInfomsg {
	msg := NewIfInfomsg(family)
	parent.children = append(parent.children, msg)
	return msg
}

// Extend RtAttr to handle data and children
type RtAttr struct {
	syscall.RtAttr
	Data     []byte
	children []NetlinkRequestData
}

// Create a new Extended RtAttr object
func NewRtAttr(attrType int, data []byte) *RtAttr {
	return &RtAttr{
		RtAttr: syscall.RtAttr{
			Type: uint16(attrType),
		},
		children: []NetlinkRequestData{},
		Data:     data,
	}
}

// Create a new RtAttr obj anc add it as a child of an existing object
func NewRtAttrChild(parent *RtAttr, attrType int, data []byte) *RtAttr {
	attr := NewRtAttr(attrType, data)
	parent.children = append(parent.children, attr)
	return attr
}

func (a *RtAttr) Len() int {
	if len(a.children) == 0 {
		return (syscall.SizeofRtAttr + len(a.Data))
	}

	l := 0
	for _, child := range a.children {
		l += rtaAlignOf(child.Len())
	}
	l += syscall.SizeofRtAttr
	return rtaAlignOf(l + len(a.Data))
}

// Serialize the RtAttr into a byte array
// This can't ust unsafe.cast because it must iterate through children.
func (a *RtAttr) Serialize() []byte {
	native := NativeEndian()

	length := a.Len()
	buf := make([]byte, rtaAlignOf(length))

	if a.Data != nil {
		copy(buf[4:], a.Data)
	} else {
		next := 4
		for _, child := range a.children {
			childBuf := child.Serialize()
			copy(buf[next:], childBuf)
			next += rtaAlignOf(len(childBuf))
		}
	}

	if l := uint16(length); l != 0 {
		native.PutUint16(buf[0:2], l)
	}
	native.PutUint16(buf[2:4], a.Type)
	return buf
}

type NetlinkRequest struct {
	syscall.NlMsghdr
	Data []NetlinkRequestData
}

// Serialize the Netlink Request into a byte array
func (req *NetlinkRequest) Serialize() []byte {
	length := syscall.SizeofNlMsghdr
	dataBytes := make([][]byte, len(req.Data))
	for i, data := range req.Data {
		dataBytes[i] = data.Serialize()
		length = length + len(dataBytes[i])
	}
	req.Len = uint32(length)
	b := make([]byte, length)
	hdr := (*(*[syscall.SizeofNlMsghdr]byte)(unsafe.Pointer(req)))[:]
	next := syscall.SizeofNlMsghdr
	copy(b[0:next], hdr)
	for _, data := range dataBytes {
		for _, dataByte := range data {
			b[next] = dataByte
			next = next + 1
		}
	}
	return b
}

func (req *NetlinkRequest) AddData(data NetlinkRequestData) {
	if data != nil {
		req.Data = append(req.Data, data)
	}
}

// Execute the request against a the given sockType.
// Returns a list of netlink messages in seriaized format, optionally filtered
// by resType.
func (req *NetlinkRequest) Execute(sockType int, resType uint16) ([][]byte, error) {
	s, err := getNetlinkSocket(sockType)
	if err != nil {
		return nil, err
	}
	defer s.Close()

	if err := s.Send(req); err != nil {
		return nil, err
	}

	pid, err := s.GetPid()
	if err != nil {
		return nil, err
	}

	var res [][]byte

done:
	for {
		msgs, err := s.Receive()
		if err != nil {
			return nil, err
		}
		for _, m := range msgs {
			if m.Header.Seq != req.Seq {
				return nil, fmt.Errorf("Wrong Seq nr %d, expected 1", m.Header.Seq)
			}
			if m.Header.Pid != pid {
				return nil, fmt.Errorf("Wrong pid %d, expected %d", m.Header.Pid, pid)
			}
			if m.Header.Type == syscall.NLMSG_DONE {
				break done
			}
			if m.Header.Type == syscall.NLMSG_ERROR {
				native := NativeEndian()
				error := int32(native.Uint32(m.Data[0:4]))
				if error == 0 {
					break done
				}
				return nil, syscall.Errno(-error)
			}
			if resType != 0 && m.Header.Type != resType {
				continue
			}
			res = append(res, m.Data)
			if m.Header.Flags&syscall.NLM_F_MULTI == 0 {
				break done
			}
		}
	}
	return res, nil
}

// Create a new netlink request from proto and flags
// Note the Len value will be inaccurate once data is added until
// the message is serialized
func NewNetlinkRequest(proto, flags int) *NetlinkRequest {
	return &NetlinkRequest{
		NlMsghdr: syscall.NlMsghdr{
			Len:   uint32(syscall.SizeofNlMsghdr),
			Type:  uint16(proto),
			Flags: syscall.NLM_F_REQUEST | uint16(flags),
			Seq:   atomic.AddUint32(&nextSeqNr, 1),
		},
	}
}

type NetlinkSocket struct {
	fd  int
	lsa syscall.SockaddrNetlink
}

func getNetlinkSocket(protocol int) (*NetlinkSocket, error) {
	fd, err := syscall.Socket(syscall.AF_NETLINK, syscall.SOCK_RAW, protocol)
	if err != nil {
		return nil, err
	}
	s := &NetlinkSocket{
		fd: fd,
	}
	s.lsa.Family = syscall.AF_NETLINK
	if err := syscall.Bind(fd, &s.lsa); err != nil {
		syscall.Close(fd)
		return nil, err
	}

	return s, nil
}

// Create a netlink socket with a given protocol (e.g. NETLINK_ROUTE)
// and subscribe it to multicast groups passed in variable argument list.
// Returns the netlink socket on which Receive() method can be called
// to retrieve the messages from the kernel.
func Subscribe(protocol int, groups ...uint) (*NetlinkSocket, error) {
	fd, err := syscall.Socket(syscall.AF_NETLINK, syscall.SOCK_RAW, protocol)
	if err != nil {
		return nil, err
	}
	s := &NetlinkSocket{
		fd: fd,
	}
	s.lsa.Family = syscall.AF_NETLINK

	for _, g := range groups {
		s.lsa.Groups |= (1 << (g - 1))
	}

	if err := syscall.Bind(fd, &s.lsa); err != nil {
		syscall.Close(fd)
		return nil, err
	}

	return s, nil
}

func (s *NetlinkSocket) Close() {
	syscall.Close(s.fd)
}

func (s *NetlinkSocket) Send(request *NetlinkRequest) error {
	if err := syscall.Sendto(s.fd, request.Serialize(), 0, &s.lsa); err != nil {
		return err
	}
	return nil
}

func (s *NetlinkSocket) Receive() ([]syscall.NetlinkMessage, error) {
	rb := make([]byte, syscall.Getpagesize())
	nr, _, err := syscall.Recvfrom(s.fd, rb, 0)
	if err != nil {
		return nil, err
	}
	if nr < syscall.NLMSG_HDRLEN {
		return nil, fmt.Errorf("Got short response from netlink")
	}
	rb = rb[:nr]
	return syscall.ParseNetlinkMessage(rb)
}

func (s *NetlinkSocket) GetPid() (uint32, error) {
	lsa, err := syscall.Getsockname(s.fd)
	if err != nil {
		return 0, err
	}
	switch v := lsa.(type) {
	case *syscall.SockaddrNetlink:
		return v.Pid, nil
	}
	return 0, fmt.Errorf("Wrong socket type")
}

func ZeroTerminated(s string) []byte {
	bytes := make([]byte, len(s)+1)
	for i := 0; i < len(s); i++ {
		bytes[i] = s[i]
	}
	bytes[len(s)] = 0
	return bytes
}

func NonZeroTerminated(s string) []byte {
	bytes := make([]byte, len(s))
	for i := 0; i < len(s); i++ {
		bytes[i] = s[i]
	}
	return bytes
}

func BytesToString(b []byte) string {
	n := bytes.Index(b, []byte{0})
	return string(b[:n])
}

func Uint8Attr(v uint8) []byte {
	return []byte{byte(v)}
}

func Uint16Attr(v uint16) []byte {
	native := NativeEndian()
	bytes := make([]byte, 2)
	native.PutUint16(bytes, v)
	return bytes
}

func Uint32Attr(v uint32) []byte {
	native := NativeEndian()
	bytes := make([]byte, 4)
	native.PutUint32(bytes, v)
	return bytes
}

func ParseRouteAttr(b []byte) ([]syscall.NetlinkRouteAttr, error) {
	var attrs []syscall.NetlinkRouteAttr
	for len(b) >= syscall.SizeofRtAttr {
		a, vbuf, alen, err := netlinkRouteAttrAndValue(b)
		if err != nil {
			return nil, err
		}
		ra := syscall.NetlinkRouteAttr{Attr: *a, Value: vbuf[:int(a.Len)-syscall.SizeofRtAttr]}
		attrs = append(attrs, ra)
		b = b[alen:]
	}
	return attrs, nil
}

func netlinkRouteAttrAndValue(b []byte) (*syscall.RtAttr, []byte, int, error) {
	a := (*syscall.RtAttr)(unsafe.Pointer(&b[0]))
	if int(a.Len) < syscall.SizeofRtAttr || int(a.Len) > len(b) {
		return nil, nil, 0, syscall.EINVAL
	}
	return a, b[syscall.SizeofRtAttr:], rtaAlignOf(int(a.Len)), nil
}
<commit_msg>fixing issue with setting big endian<commit_after>// Package nl has low level primitives for making Netlink calls.
package nl

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"net"
	"sync/atomic"
	"syscall"
	"unsafe"
)

const (
	// Family type definitions
	FAMILY_ALL = syscall.AF_UNSPEC
	FAMILY_V4  = syscall.AF_INET
	FAMILY_V6  = syscall.AF_INET6
)

var nextSeqNr uint32

// GetIPFamily returns the family type of a net.IP.
func GetIPFamily(ip net.IP) int {
	if len(ip) <= net.IPv4len {
		return FAMILY_V4
	}
	if ip.To4() != nil {
		return FAMILY_V4
	}
	return FAMILY_V6
}

var nativeEndian binary.ByteOrder

// Get native endianness for the system
func NativeEndian() binary.ByteOrder {
	if nativeEndian == nil {
		var x uint32 = 0x01020304
		if *(*byte)(unsafe.Pointer(&x)) == 0x01 {
			nativeEndian = binary.BigEndian
		} else {
			nativeEndian = binary.LittleEndian
		}
	}
	return nativeEndian
}

// Byte swap a 16 bit value if we aren't big endian
func Swap16(i uint16) uint16 {
	if NativeEndian() == binary.BigEndian {
		return i
	}
	return (i&0xff00)>>8 | (i&0xff)<<8
}

// Byte swap a 32 bit value if aren't big endian
func Swap32(i uint32) uint32 {
	if NativeEndian() == binary.BigEndian {
		return i
	}
	return (i&0xff000000)>>24 | (i&0xff0000)>>8 | (i&0xff00)<<8 | (i&0xff)<<24
}

type NetlinkRequestData interface {
	Len() int
	Serialize() []byte
}

// IfInfomsg is related to links, but it is used for list requests as well
type IfInfomsg struct {
	syscall.IfInfomsg
}

// Create an IfInfomsg with family specified
func NewIfInfomsg(family int) *IfInfomsg {
	return &IfInfomsg{
		IfInfomsg: syscall.IfInfomsg{
			Family: uint8(family),
		},
	}
}

func DeserializeIfInfomsg(b []byte) *IfInfomsg {
	return (*IfInfomsg)(unsafe.Pointer(&b[0:syscall.SizeofIfInfomsg][0]))
}

func (msg *IfInfomsg) Serialize() []byte {
	return (*(*[syscall.SizeofIfInfomsg]byte)(unsafe.Pointer(msg)))[:]
}

func (msg *IfInfomsg) Len() int {
	return syscall.SizeofIfInfomsg
}

func rtaAlignOf(attrlen int) int {
	return (attrlen + syscall.RTA_ALIGNTO - 1) & ^(syscall.RTA_ALIGNTO - 1)
}

func NewIfInfomsgChild(parent *RtAttr, family int) *IfInfomsg {
	msg := NewIfInfomsg(family)
	parent.children = append(parent.children, msg)
	return msg
}

// Extend RtAttr to handle data and children
type RtAttr struct {
	syscall.RtAttr
	Data     []byte
	children []NetlinkRequestData
}

// Create a new Extended RtAttr object
func NewRtAttr(attrType int, data []byte) *RtAttr {
	return &RtAttr{
		RtAttr: syscall.RtAttr{
			Type: uint16(attrType),
		},
		children: []NetlinkRequestData{},
		Data:     data,
	}
}

// Create a new RtAttr obj anc add it as a child of an existing object
func NewRtAttrChild(parent *RtAttr, attrType int, data []byte) *RtAttr {
	attr := NewRtAttr(attrType, data)
	parent.children = append(parent.children, attr)
	return attr
}

func (a *RtAttr) Len() int {
	if len(a.children) == 0 {
		return (syscall.SizeofRtAttr + len(a.Data))
	}

	l := 0
	for _, child := range a.children {
		l += rtaAlignOf(child.Len())
	}
	l += syscall.SizeofRtAttr
	return rtaAlignOf(l + len(a.Data))
}

// Serialize the RtAttr into a byte array
// This can't ust unsafe.cast because it must iterate through children.
func (a *RtAttr) Serialize() []byte {
	native := NativeEndian()

	length := a.Len()
	buf := make([]byte, rtaAlignOf(length))

	if a.Data != nil {
		copy(buf[4:], a.Data)
	} else {
		next := 4
		for _, child := range a.children {
			childBuf := child.Serialize()
			copy(buf[next:], childBuf)
			next += rtaAlignOf(len(childBuf))
		}
	}

	if l := uint16(length); l != 0 {
		native.PutUint16(buf[0:2], l)
	}
	native.PutUint16(buf[2:4], a.Type)
	return buf
}

type NetlinkRequest struct {
	syscall.NlMsghdr
	Data []NetlinkRequestData
}

// Serialize the Netlink Request into a byte array
func (req *NetlinkRequest) Serialize() []byte {
	length := syscall.SizeofNlMsghdr
	dataBytes := make([][]byte, len(req.Data))
	for i, data := range req.Data {
		dataBytes[i] = data.Serialize()
		length = length + len(dataBytes[i])
	}
	req.Len = uint32(length)
	b := make([]byte, length)
	hdr := (*(*[syscall.SizeofNlMsghdr]byte)(unsafe.Pointer(req)))[:]
	next := syscall.SizeofNlMsghdr
	copy(b[0:next], hdr)
	for _, data := range dataBytes {
		for _, dataByte := range data {
			b[next] = dataByte
			next = next + 1
		}
	}
	return b
}

func (req *NetlinkRequest) AddData(data NetlinkRequestData) {
	if data != nil {
		req.Data = append(req.Data, data)
	}
}

// Execute the request against a the given sockType.
// Returns a list of netlink messages in seriaized format, optionally filtered
// by resType.
func (req *NetlinkRequest) Execute(sockType int, resType uint16) ([][]byte, error) {
	s, err := getNetlinkSocket(sockType)
	if err != nil {
		return nil, err
	}
	defer s.Close()

	if err := s.Send(req); err != nil {
		return nil, err
	}

	pid, err := s.GetPid()
	if err != nil {
		return nil, err
	}

	var res [][]byte

done:
	for {
		msgs, err := s.Receive()
		if err != nil {
			return nil, err
		}
		for _, m := range msgs {
			if m.Header.Seq != req.Seq {
				return nil, fmt.Errorf("Wrong Seq nr %d, expected 1", m.Header.Seq)
			}
			if m.Header.Pid != pid {
				return nil, fmt.Errorf("Wrong pid %d, expected %d", m.Header.Pid, pid)
			}
			if m.Header.Type == syscall.NLMSG_DONE {
				break done
			}
			if m.Header.Type == syscall.NLMSG_ERROR {
				native := NativeEndian()
				error := int32(native.Uint32(m.Data[0:4]))
				if error == 0 {
					break done
				}
				return nil, syscall.Errno(-error)
			}
			if resType != 0 && m.Header.Type != resType {
				continue
			}
			res = append(res, m.Data)
			if m.Header.Flags&syscall.NLM_F_MULTI == 0 {
				break done
			}
		}
	}
	return res, nil
}

// Create a new netlink request from proto and flags
// Note the Len value will be inaccurate once data is added until
// the message is serialized
func NewNetlinkRequest(proto, flags int) *NetlinkRequest {
	return &NetlinkRequest{
		NlMsghdr: syscall.NlMsghdr{
			Len:   uint32(syscall.SizeofNlMsghdr),
			Type:  uint16(proto),
			Flags: syscall.NLM_F_REQUEST | uint16(flags),
			Seq:   atomic.AddUint32(&nextSeqNr, 1),
		},
	}
}

type NetlinkSocket struct {
	fd  int
	lsa syscall.SockaddrNetlink
}

func getNetlinkSocket(protocol int) (*NetlinkSocket, error) {
	fd, err := syscall.Socket(syscall.AF_NETLINK, syscall.SOCK_RAW, protocol)
	if err != nil {
		return nil, err
	}
	s := &NetlinkSocket{
		fd: fd,
	}
	s.lsa.Family = syscall.AF_NETLINK
	if err := syscall.Bind(fd, &s.lsa); err != nil {
		syscall.Close(fd)
		return nil, err
	}

	return s, nil
}

// Create a netlink socket with a given protocol (e.g. NETLINK_ROUTE)
// and subscribe it to multicast groups passed in variable argument list.
// Returns the netlink socket on which Receive() method can be called
// to retrieve the messages from the kernel.
func Subscribe(protocol int, groups ...uint) (*NetlinkSocket, error) {
	fd, err := syscall.Socket(syscall.AF_NETLINK, syscall.SOCK_RAW, protocol)
	if err != nil {
		return nil, err
	}
	s := &NetlinkSocket{
		fd: fd,
	}
	s.lsa.Family = syscall.AF_NETLINK

	for _, g := range groups {
		s.lsa.Groups |= (1 << (g - 1))
	}

	if err := syscall.Bind(fd, &s.lsa); err != nil {
		syscall.Close(fd)
		return nil, err
	}

	return s, nil
}

func (s *NetlinkSocket) Close() {
	syscall.Close(s.fd)
}

func (s *NetlinkSocket) Send(request *NetlinkRequest) error {
	if err := syscall.Sendto(s.fd, request.Serialize(), 0, &s.lsa); err != nil {
		return err
	}
	return nil
}

func (s *NetlinkSocket) Receive() ([]syscall.NetlinkMessage, error) {
	rb := make([]byte, syscall.Getpagesize())
	nr, _, err := syscall.Recvfrom(s.fd, rb, 0)
	if err != nil {
		return nil, err
	}
	if nr < syscall.NLMSG_HDRLEN {
		return nil, fmt.Errorf("Got short response from netlink")
	}
	rb = rb[:nr]
	return syscall.ParseNetlinkMessage(rb)
}

func (s *NetlinkSocket) GetPid() (uint32, error) {
	lsa, err := syscall.Getsockname(s.fd)
	if err != nil {
		return 0, err
	}
	switch v := lsa.(type) {
	case *syscall.SockaddrNetlink:
		return v.Pid, nil
	}
	return 0, fmt.Errorf("Wrong socket type")
}

func ZeroTerminated(s string) []byte {
	bytes := make([]byte, len(s)+1)
	for i := 0; i < len(s); i++ {
		bytes[i] = s[i]
	}
	bytes[len(s)] = 0
	return bytes
}

func NonZeroTerminated(s string) []byte {
	bytes := make([]byte, len(s))
	for i := 0; i < len(s); i++ {
		bytes[i] = s[i]
	}
	return bytes
}

func BytesToString(b []byte) string {
	n := bytes.Index(b, []byte{0})
	return string(b[:n])
}

func Uint8Attr(v uint8) []byte {
	return []byte{byte(v)}
}

func Uint16Attr(v uint16) []byte {
	native := NativeEndian()
	bytes := make([]byte, 2)
	native.PutUint16(bytes, v)
	return bytes
}

func Uint32Attr(v uint32) []byte {
	native := NativeEndian()
	bytes := make([]byte, 4)
	native.PutUint32(bytes, v)
	return bytes
}

func ParseRouteAttr(b []byte) ([]syscall.NetlinkRouteAttr, error) {
	var attrs []syscall.NetlinkRouteAttr
	for len(b) >= syscall.SizeofRtAttr {
		a, vbuf, alen, err := netlinkRouteAttrAndValue(b)
		if err != nil {
			return nil, err
		}
		ra := syscall.NetlinkRouteAttr{Attr: *a, Value: vbuf[:int(a.Len)-syscall.SizeofRtAttr]}
		attrs = append(attrs, ra)
		b = b[alen:]
	}
	return attrs, nil
}

func netlinkRouteAttrAndValue(b []byte) (*syscall.RtAttr, []byte, int, error) {
	a := (*syscall.RtAttr)(unsafe.Pointer(&b[0]))
	if int(a.Len) < syscall.SizeofRtAttr || int(a.Len) > len(b) {
		return nil, nil, 0, syscall.EINVAL
	}
	return a, b[syscall.SizeofRtAttr:], rtaAlignOf(int(a.Len)), nil
}
<|endoftext|>


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

// Copyright 2015 PingCAP, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package stmts

import (
	"strings"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/context"
	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/plan"
	"github.com/pingcap/tidb/rset"
	"github.com/pingcap/tidb/rset/rsets"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/stmt"
	"github.com/pingcap/tidb/table"
	"github.com/pingcap/tidb/util"
	"github.com/pingcap/tidb/util/format"
)

var _ stmt.Statement = (*DeleteStmt)(nil) // TODO optimizer plan

// DeleteStmt is a statement to delete rows from table.
// See: https://dev.mysql.com/doc/refman/5.7/en/delete.html
type DeleteStmt struct {
	Where       expression.Expression
	Order       *rsets.OrderByRset
	Limit       *rsets.LimitRset
	LowPriority bool
	Ignore      bool
	Quick       bool
	MultiTable  bool
	BeforeFrom  bool
	TableIdents []table.Ident
	Refs        *rsets.JoinRset

	Text string
}

// Explain implements the stmt.Statement Explain interface.
func (s *DeleteStmt) Explain(ctx context.Context, w format.Formatter) {
	p, err := s.plan(ctx)
	if err != nil {
		log.Error(err)
		return
	}
	p.Explain(w)
	w.Format("└Delete row\n")
}

// IsDDL implements the stmt.Statement IsDDL interface.
func (s *DeleteStmt) IsDDL() bool {
	return false
}

// OriginText implements the stmt.Statement OriginText interface.
func (s *DeleteStmt) OriginText() string {
	return s.Text
}

// SetText implements the stmt.Statement SetText interface.
func (s *DeleteStmt) SetText(text string) {
	s.Text = text
}

func (s *DeleteStmt) plan(ctx context.Context) (plan.Plan, error) {
	var (
		r   plan.Plan
		err error
	)
	if s.Refs != nil {
		r, err = s.Refs.Plan(ctx)
		if err != nil {
			return nil, err
		}
	}
	if s.Where != nil {
		r, err = (&rsets.WhereRset{Expr: s.Where, Src: r}).Plan(ctx)
		if err != nil {
			return nil, err
		}
	}
	return r, nil
}

func (s *DeleteStmt) removeRow(ctx context.Context, t table.Table, h int64, data []interface{}) error {
	// remove row's all indexies
	if err := t.RemoveRowAllIndex(ctx, h, data); err != nil {
		return err
	}
	// remove row
	if err := t.RemoveRow(ctx, h); err != nil {
		return err
	}
	variable.GetSessionVars(ctx).AddAffectedRows(1)
	return nil
}

// Exec implements the stmt.Statement Exec interface.
func (s *DeleteStmt) Exec(ctx context.Context) (_ rset.Recordset, err error) {
	if s.MultiTable && len(s.TableIdents) == 0 {
		return nil, nil
	}
	p, err := s.plan(ctx)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if p == nil {
		return nil, nil
	}
	defer p.Close()
	tblIDMap := make(map[int64]bool, len(s.TableIdents))
	// Get table alias map.
	tblNames := make(map[string]string)
	if s.MultiTable {
		// Delete from multiple tables should consider table ident list.
		fs := p.GetFields()
		for _, f := range fs {
			if f.TableName != f.OrgTableName {
				tblNames[f.TableName] = f.OrgTableName
			} else {
				tblNames[f.TableName] = f.TableName
			}
		}
		for _, t := range s.TableIdents {
			// Consider DBName.
			oname, ok := tblNames[t.Name.O]
			if !ok {
				return nil, errors.Errorf("Unknown table '%s' in MULTI DELETE", t.Name.O)
			}

			t.Name.O = oname
			t.Name.L = strings.ToLower(oname)

			var tbl table.Table
			tbl, err = getTable(ctx, t)
			if err != nil {
				return nil, errors.Trace(err)
			}
			tblIDMap[tbl.TableID()] = true
		}
	}
	rowKeyMap := make(map[string]table.Table)
	for {
		row, err1 := p.Next(ctx)
		if err1 != nil {
			return nil, errors.Trace(err1)
		}
		if row == nil {
			break
		}

		for _, entry := range row.RowKeys {
			if s.MultiTable {
				tid := entry.Tbl.TableID()
				if _, ok := tblIDMap[tid]; !ok {
					continue
				}
			}
			rowKeyMap[entry.Key] = entry.Tbl
		}
	}

	for k, t := range rowKeyMap {
		handle, err := util.DecodeHandleFromRowKey(k)
		if err != nil {
			return nil, errors.Trace(err)
		}
		data, err := t.Row(ctx, handle)
		if err != nil {
			return nil, errors.Trace(err)
		}
		err = s.removeRow(ctx, t, handle, data)
		if err != nil {
			return nil, errors.Trace(err)
		}
	}
	return nil, nil
}
<commit_msg>stmt: change remove method to function<commit_after>// Copyright 2013 The ql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSES/QL-LICENSE file.

// Copyright 2015 PingCAP, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package stmts

import (
	"strings"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/context"
	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/plan"
	"github.com/pingcap/tidb/rset"
	"github.com/pingcap/tidb/rset/rsets"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/stmt"
	"github.com/pingcap/tidb/table"
	"github.com/pingcap/tidb/util"
	"github.com/pingcap/tidb/util/format"
)

var _ stmt.Statement = (*DeleteStmt)(nil) // TODO optimizer plan

// DeleteStmt is a statement to delete rows from table.
// See: https://dev.mysql.com/doc/refman/5.7/en/delete.html
type DeleteStmt struct {
	Where       expression.Expression
	Order       *rsets.OrderByRset
	Limit       *rsets.LimitRset
	LowPriority bool
	Ignore      bool
	Quick       bool
	MultiTable  bool
	BeforeFrom  bool
	TableIdents []table.Ident
	Refs        *rsets.JoinRset

	Text string
}

// Explain implements the stmt.Statement Explain interface.
func (s *DeleteStmt) Explain(ctx context.Context, w format.Formatter) {
	p, err := s.plan(ctx)
	if err != nil {
		log.Error(err)
		return
	}
	p.Explain(w)
	w.Format("└Delete row\n")
}

// IsDDL implements the stmt.Statement IsDDL interface.
func (s *DeleteStmt) IsDDL() bool {
	return false
}

// OriginText implements the stmt.Statement OriginText interface.
func (s *DeleteStmt) OriginText() string {
	return s.Text
}

// SetText implements the stmt.Statement SetText interface.
func (s *DeleteStmt) SetText(text string) {
	s.Text = text
}

func (s *DeleteStmt) plan(ctx context.Context) (plan.Plan, error) {
	var (
		r   plan.Plan
		err error
	)
	if s.Refs != nil {
		r, err = s.Refs.Plan(ctx)
		if err != nil {
			return nil, err
		}
	}
	if s.Where != nil {
		r, err = (&rsets.WhereRset{Expr: s.Where, Src: r}).Plan(ctx)
		if err != nil {
			return nil, err
		}
	}
	return r, nil
}

func removeRow(ctx context.Context, t table.Table, h int64, data []interface{}) error {
	// remove row's all indexies
	if err := t.RemoveRowAllIndex(ctx, h, data); err != nil {
		return err
	}
	// remove row
	if err := t.RemoveRow(ctx, h); err != nil {
		return err
	}
	variable.GetSessionVars(ctx).AddAffectedRows(1)
	return nil
}

// Exec implements the stmt.Statement Exec interface.
func (s *DeleteStmt) Exec(ctx context.Context) (_ rset.Recordset, err error) {
	if s.MultiTable && len(s.TableIdents) == 0 {
		return nil, nil
	}
	p, err := s.plan(ctx)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if p == nil {
		return nil, nil
	}
	defer p.Close()
	tblIDMap := make(map[int64]bool, len(s.TableIdents))
	// Get table alias map.
	tblNames := make(map[string]string)
	if s.MultiTable {
		// Delete from multiple tables should consider table ident list.
		fs := p.GetFields()
		for _, f := range fs {
			if f.TableName != f.OrgTableName {
				tblNames[f.TableName] = f.OrgTableName
			} else {
				tblNames[f.TableName] = f.TableName
			}
		}
		for _, t := range s.TableIdents {
			// Consider DBName.
			oname, ok := tblNames[t.Name.O]
			if !ok {
				return nil, errors.Errorf("Unknown table '%s' in MULTI DELETE", t.Name.O)
			}

			t.Name.O = oname
			t.Name.L = strings.ToLower(oname)

			var tbl table.Table
			tbl, err = getTable(ctx, t)
			if err != nil {
				return nil, errors.Trace(err)
			}
			tblIDMap[tbl.TableID()] = true
		}
	}
	rowKeyMap := make(map[string]table.Table)
	for {
		row, err1 := p.Next(ctx)
		if err1 != nil {
			return nil, errors.Trace(err1)
		}
		if row == nil {
			break
		}

		for _, entry := range row.RowKeys {
			if s.MultiTable {
				tid := entry.Tbl.TableID()
				if _, ok := tblIDMap[tid]; !ok {
					continue
				}
			}
			rowKeyMap[entry.Key] = entry.Tbl
		}
	}

	for k, t := range rowKeyMap {
		handle, err := util.DecodeHandleFromRowKey(k)
		if err != nil {
			return nil, errors.Trace(err)
		}
		data, err := t.Row(ctx, handle)
		if err != nil {
			return nil, errors.Trace(err)
		}
		err = removeRow(ctx, t, handle, data)
		if err != nil {
			return nil, errors.Trace(err)
		}
	}
	return nil, nil
}
<|endoftext|>


<commit_before>package store_test

import (
	"encoding/json"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju/go/charm"
	"launchpad.net/juju/go/store"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"time"
)

func (s *StoreSuite) prepareServer(c *C) (*store.Server, *charm.URL) {
	curl := charm.MustParseURL("cs:oneiric/wordpress")
	pub, err := s.store.CharmPublisher([]*charm.URL{curl}, "some-digest")
	c.Assert(err, IsNil)
	err = pub.Publish(&FakeCharmDir{})
	c.Assert(err, IsNil)

	server, err := store.NewServer(s.store)
	c.Assert(err, IsNil)
	return server, curl
}

func (s *StoreSuite) TestServerCharmInfo(c *C) {
	server, curl := s.prepareServer(c)
	req, err := http.NewRequest("GET", "/charm-info", nil)
	c.Assert(err, IsNil)

	var tests = []struct{ url, sha, err string }{
		{curl.String(), fakeRevZeroSha, ""},
		{"cs:oneiric/non-existent", "", "entry not found"},
		{"cs:bad", "", `charm URL without series: "cs:bad"`},
	}

	for _, t := range tests {
		req.Form = url.Values{"charms": []string{t.url}}
		rec := httptest.NewRecorder()
		server.ServeHTTP(rec, req)

		expected := make(map[string]interface{})
		if t.sha != "" {
			expected[t.url] = map[string]interface{}{
				"revision": float64(0),
				"sha256": t.sha,
			}
		} else {
			expected[t.url] = map[string]interface{}{
				"revision": float64(0),
				"errors": []interface{}{t.err},
			}
		}
		obtained := map[string]interface{}{}
		err = json.NewDecoder(rec.Body).Decode(&obtained)
		c.Assert(err, IsNil)
		c.Assert(obtained, DeepEquals, expected)
		c.Assert(rec.Header().Get("Content-Type"), Equals, "application/json")
	}

	s.checkCounterSum(c, []string{"charm-info", curl.Series, curl.Name}, false, 1)
	s.checkCounterSum(c, []string{"charm-missing", "oneiric", "non-existent"}, false, 1)
}

// checkCounterSum checks that statistics are properly collected.
// It retries a few times as they are generally collected in background.
func (s *StoreSuite) checkCounterSum(c *C, key []string, prefix bool, expected int64) {
	var sum int64
	var err error
	for retry := 0; retry < 10; retry++ {
		time.Sleep(1e8)
		sum, err = s.store.SumCounter(key, prefix)
		c.Assert(err, IsNil)
		if sum == expected {
			if expected == 0 && retry < 2 {
				continue // Wait a bit to make sure.
			}
			return
		}
	}
	c.Errorf("counter sum for %#v is %d, want %d", key, sum, expected)
}

func (s *StoreSuite) TestCharmStreaming(c *C) {
	server, curl := s.prepareServer(c)

	req, err := http.NewRequest("GET", "/charm/"+curl.String()[3:], nil)
	c.Assert(err, IsNil)
	rec := httptest.NewRecorder()
	server.ServeHTTP(rec, req)

	data, err := ioutil.ReadAll(rec.Body)
	c.Assert(string(data), Equals, "charm-revision-0")

	c.Assert(rec.Header().Get("Connection"), Equals, "close")
	c.Assert(rec.Header().Get("Content-Type"), Equals, "application/octet-stream")
	c.Assert(rec.Header().Get("Content-Length"), Equals, "16")

	// Check that it was accounted for in statistics.
	s.checkCounterSum(c, []string{"charm-bundle", curl.Series, curl.Name}, false, 1)
}

func (s *StoreSuite) TestDisableStats(c *C) {
	server, curl := s.prepareServer(c)

	req, err := http.NewRequest("GET", "/charm-info", nil)
	c.Assert(err, IsNil)
	req.Form = url.Values{"charms": []string{curl.String()}, "stats": []string{"0"}}
	server.ServeHTTP(httptest.NewRecorder(), req)

	req, err = http.NewRequest("GET", "/charm/"+curl.String()[3:], nil)
	c.Assert(err, IsNil)
	req.Form = url.Values{"stats": []string{"0"}}
	server.ServeHTTP(httptest.NewRecorder(), req)

	// No statistics should have been collected given the use of stats=0.
	for _, prefix := range []string{"charm-info", "charm-bundle", "charm-missing"} {
		s.checkCounterSum(c, []string{prefix}, true, 0)
	}
}

func (s *StoreSuite) TestServerStatus(c *C) {
	server, err := store.NewServer(s.store)
	c.Assert(err, IsNil)
	tests := []struct {
		path string
		code int
	}{
		{"/charm-info/any", 404},
		{"/charm/bad-url", 404},
		{"/charm/bad-series/wordpress", 404},
		{"/stats/counter/", 403},
		{"/stats/counter/*", 403},
		{"/stats/counter/any/", 404},
		{"/stats/", 404},
		{"/stats/any", 404},
	}
	for _, test := range tests {
		req, err := http.NewRequest("GET", test.path, nil)
		c.Assert(err, IsNil)
		rec := httptest.NewRecorder()
		server.ServeHTTP(rec, req)
		c.Assert(rec.Code, Equals, test.code, Commentf("Path: %s", test.path))
	}
}

func (s *StoreSuite) TestRootRedirect(c *C) {
	server, err := store.NewServer(s.store)
	c.Assert(err, IsNil)
	req, err := http.NewRequest("GET", "/", nil)
	c.Assert(err, IsNil)
	rec := httptest.NewRecorder()
	server.ServeHTTP(rec, req)
	c.Assert(rec.Code, Equals, 303)
	c.Assert(rec.Header().Get("Location"), Equals, "https://juju.ubuntu.com")
}

func (s *StoreSuite) TestStatsCounter(c *C) {
	for _, key := range [][]string{{"a", "b"}, {"a", "b"}, {"a"}} {
		err := s.store.IncCounter(key)
		c.Assert(err, IsNil)
	}

	server, _ := s.prepareServer(c)

	expected := map[string]string{
		"a:b": "2",
		"a:*": "3",
		"a":   "1",
	}

	for counter, n := range expected {
		req, err := http.NewRequest("GET", "/stats/counter/" + counter, nil)
		c.Assert(err, IsNil)
		rec := httptest.NewRecorder()
		server.ServeHTTP(rec, req)

		data, err := ioutil.ReadAll(rec.Body)
		c.Assert(string(data), Equals, n)

		c.Assert(rec.Header().Get("Content-Type"), Equals, "text/plain")
		c.Assert(rec.Header().Get("Content-Length"), Equals, strconv.Itoa(len(n)))
	}
}

func (s *StoreSuite) TestBlitzKey(c *C) {
	server, _ := s.prepareServer(c)

	// This is just a validation key to allow blitz.io to run
	// performance tests against the site.
	req, err := http.NewRequest("GET", "/mu-35700a31-6bf320ca-a800b670-05f845ee", nil)
	c.Assert(err, IsNil)
	rec := httptest.NewRecorder()
	server.ServeHTTP(rec, req)

	data, err := ioutil.ReadAll(rec.Body)
	c.Assert(string(data), Equals, "42")

	c.Assert(rec.Header().Get("Connection"), Equals, "close")
	c.Assert(rec.Header().Get("Content-Type"), Equals, "text/plain")
	c.Assert(rec.Header().Get("Content-Length"), Equals, "2")
}
<commit_msg>Check that requests succeed in the test to avoid further problems, as recommended by Roger.<commit_after>package store_test

import (
	"encoding/json"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju/go/charm"
	"launchpad.net/juju/go/store"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"time"
)

func (s *StoreSuite) prepareServer(c *C) (*store.Server, *charm.URL) {
	curl := charm.MustParseURL("cs:oneiric/wordpress")
	pub, err := s.store.CharmPublisher([]*charm.URL{curl}, "some-digest")
	c.Assert(err, IsNil)
	err = pub.Publish(&FakeCharmDir{})
	c.Assert(err, IsNil)

	server, err := store.NewServer(s.store)
	c.Assert(err, IsNil)
	return server, curl
}

func (s *StoreSuite) TestServerCharmInfo(c *C) {
	server, curl := s.prepareServer(c)
	req, err := http.NewRequest("GET", "/charm-info", nil)
	c.Assert(err, IsNil)

	var tests = []struct{ url, sha, err string }{
		{curl.String(), fakeRevZeroSha, ""},
		{"cs:oneiric/non-existent", "", "entry not found"},
		{"cs:bad", "", `charm URL without series: "cs:bad"`},
	}

	for _, t := range tests {
		req.Form = url.Values{"charms": []string{t.url}}
		rec := httptest.NewRecorder()
		server.ServeHTTP(rec, req)

		expected := make(map[string]interface{})
		if t.sha != "" {
			expected[t.url] = map[string]interface{}{
				"revision": float64(0),
				"sha256": t.sha,
			}
		} else {
			expected[t.url] = map[string]interface{}{
				"revision": float64(0),
				"errors": []interface{}{t.err},
			}
		}
		obtained := map[string]interface{}{}
		err = json.NewDecoder(rec.Body).Decode(&obtained)
		c.Assert(err, IsNil)
		c.Assert(obtained, DeepEquals, expected)
		c.Assert(rec.Header().Get("Content-Type"), Equals, "application/json")
	}

	s.checkCounterSum(c, []string{"charm-info", curl.Series, curl.Name}, false, 1)
	s.checkCounterSum(c, []string{"charm-missing", "oneiric", "non-existent"}, false, 1)
}

// checkCounterSum checks that statistics are properly collected.
// It retries a few times as they are generally collected in background.
func (s *StoreSuite) checkCounterSum(c *C, key []string, prefix bool, expected int64) {
	var sum int64
	var err error
	for retry := 0; retry < 10; retry++ {
		time.Sleep(1e8)
		sum, err = s.store.SumCounter(key, prefix)
		c.Assert(err, IsNil)
		if sum == expected {
			if expected == 0 && retry < 2 {
				continue // Wait a bit to make sure.
			}
			return
		}
	}
	c.Errorf("counter sum for %#v is %d, want %d", key, sum, expected)
}

func (s *StoreSuite) TestCharmStreaming(c *C) {
	server, curl := s.prepareServer(c)

	req, err := http.NewRequest("GET", "/charm/"+curl.String()[3:], nil)
	c.Assert(err, IsNil)
	rec := httptest.NewRecorder()
	server.ServeHTTP(rec, req)

	data, err := ioutil.ReadAll(rec.Body)
	c.Assert(string(data), Equals, "charm-revision-0")

	c.Assert(rec.Header().Get("Connection"), Equals, "close")
	c.Assert(rec.Header().Get("Content-Type"), Equals, "application/octet-stream")
	c.Assert(rec.Header().Get("Content-Length"), Equals, "16")

	// Check that it was accounted for in statistics.
	s.checkCounterSum(c, []string{"charm-bundle", curl.Series, curl.Name}, false, 1)
}

func (s *StoreSuite) TestDisableStats(c *C) {
	server, curl := s.prepareServer(c)

	req, err := http.NewRequest("GET", "/charm-info", nil)
	c.Assert(err, IsNil)
	req.Form = url.Values{"charms": []string{curl.String()}, "stats": []string{"0"}}
	rec := httptest.NewRecorder()
	server.ServeHTTP(rec, req)
	c.Assert(rec.Code, Equals, 200)

	req, err = http.NewRequest("GET", "/charm/"+curl.String()[3:], nil)
	c.Assert(err, IsNil)
	req.Form = url.Values{"stats": []string{"0"}}
	rec = httptest.NewRecorder()
	server.ServeHTTP(rec, req)
	c.Assert(rec.Code, Equals, 200)

	// No statistics should have been collected given the use of stats=0.
	for _, prefix := range []string{"charm-info", "charm-bundle", "charm-missing"} {
		s.checkCounterSum(c, []string{prefix}, true, 0)
	}
}

func (s *StoreSuite) TestServerStatus(c *C) {
	server, err := store.NewServer(s.store)
	c.Assert(err, IsNil)
	tests := []struct {
		path string
		code int
	}{
		{"/charm-info/any", 404},
		{"/charm/bad-url", 404},
		{"/charm/bad-series/wordpress", 404},
		{"/stats/counter/", 403},
		{"/stats/counter/*", 403},
		{"/stats/counter/any/", 404},
		{"/stats/", 404},
		{"/stats/any", 404},
	}
	for _, test := range tests {
		req, err := http.NewRequest("GET", test.path, nil)
		c.Assert(err, IsNil)
		rec := httptest.NewRecorder()
		server.ServeHTTP(rec, req)
		c.Assert(rec.Code, Equals, test.code, Commentf("Path: %s", test.path))
	}
}

func (s *StoreSuite) TestRootRedirect(c *C) {
	server, err := store.NewServer(s.store)
	c.Assert(err, IsNil)
	req, err := http.NewRequest("GET", "/", nil)
	c.Assert(err, IsNil)
	rec := httptest.NewRecorder()
	server.ServeHTTP(rec, req)
	c.Assert(rec.Code, Equals, 303)
	c.Assert(rec.Header().Get("Location"), Equals, "https://juju.ubuntu.com")
}

func (s *StoreSuite) TestStatsCounter(c *C) {
	for _, key := range [][]string{{"a", "b"}, {"a", "b"}, {"a"}} {
		err := s.store.IncCounter(key)
		c.Assert(err, IsNil)
	}

	server, _ := s.prepareServer(c)

	expected := map[string]string{
		"a:b": "2",
		"a:*": "3",
		"a":   "1",
	}

	for counter, n := range expected {
		req, err := http.NewRequest("GET", "/stats/counter/" + counter, nil)
		c.Assert(err, IsNil)
		rec := httptest.NewRecorder()
		server.ServeHTTP(rec, req)

		data, err := ioutil.ReadAll(rec.Body)
		c.Assert(string(data), Equals, n)

		c.Assert(rec.Header().Get("Content-Type"), Equals, "text/plain")
		c.Assert(rec.Header().Get("Content-Length"), Equals, strconv.Itoa(len(n)))
	}
}

func (s *StoreSuite) TestBlitzKey(c *C) {
	server, _ := s.prepareServer(c)

	// This is just a validation key to allow blitz.io to run
	// performance tests against the site.
	req, err := http.NewRequest("GET", "/mu-35700a31-6bf320ca-a800b670-05f845ee", nil)
	c.Assert(err, IsNil)
	rec := httptest.NewRecorder()
	server.ServeHTTP(rec, req)

	data, err := ioutil.ReadAll(rec.Body)
	c.Assert(string(data), Equals, "42")

	c.Assert(rec.Header().Get("Connection"), Equals, "close")
	c.Assert(rec.Header().Get("Content-Type"), Equals, "text/plain")
	c.Assert(rec.Header().Get("Content-Length"), Equals, "2")
}
<|endoftext|>


<commit_before>package resource

import (
	"fmt"
	"os"

	"github.com/mitchellh/colorstring"
)

func Fail(err error) {
	fmt.Fprintf(os.Stderr, colorstring.Color("[red] error %s"), err)
	os.Exit(1)
}

func Fatal(doing string, err error) {
	Sayf(colorstring.Color("[red]error %s: %s\n"), doing, err)
	os.Exit(1)
}

func Sayf(message string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, message, args...)
}
<commit_msg>remove space in error log<commit_after>package resource

import (
	"fmt"
	"os"

	"github.com/mitchellh/colorstring"
)

func Fail(err error) {
	fmt.Fprintf(os.Stderr, colorstring.Color("[red]error %s\n"), err)
	os.Exit(1)
}

func Fatal(doing string, err error) {
	Sayf(colorstring.Color("[red]error %s: %s\n"), doing, err)
	os.Exit(1)
}

func Sayf(message string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, message, args...)
}
<|endoftext|>


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

import (
	"bytes"
	"github.com/Comcast/webpa-common/logging"
	"github.com/Comcast/webpa-common/webhook"
	"github.com/Comcast/webpa-common/wrp"
	"github.com/stretchr/testify/assert"
	"net/http"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

type result struct {
	URL      string
	event    string
	transID  string
	deviceID string
}

// Make a simple RoundTrip implementation that let's me short-circuit the network
type swTransport struct {
	i       int32
	fn      func(*http.Request, int) (*http.Response, error)
	results []result
	mutex   sync.Mutex
}

func (t *swTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	atomic.AddInt32(&t.i, 1)

	r := result{URL: req.URL.String(),
		event:    req.Header.Get("X-Webpa-Event"),
		transID:  req.Header.Get("X-Webpa-Transaction-Id"),
		deviceID: req.Header.Get("X-Webpa-Device-Id"),
	}

	t.mutex.Lock()
	t.results = append(t.results, r)
	t.mutex.Unlock()

	resp := &http.Response{Status: "200 OK", StatusCode: 200}
	return resp, nil
}

func getFakeFactory() *SenderWrapperFactory {
	fakeICTC := new(mockCounter)
	fakeICTC.On("With", []string{"content_type", "msgpack"}).Return(fakeICTC).
		On("With", []string{"content_type", "unknown"}).Return(fakeICTC).
		On("Add", 1.0).Return()

	fakeDDTIP := new(mockCounter)
	fakeDDTIP.On("Add", 1.0).Return()

	fakeGauge := new(mockGauge)
	fakeGauge.On("Add", 1.0).Return().
		On("Add", -1.0).Return().
		//On("With", []string{"url", "unknown"}).Return(fakeGauge).
		On("With", []string{"url", "http://localhost:8888/foo"}).Return(fakeGauge).
		On("With", []string{"url", "http://localhost:9999/foo"}).Return(fakeGauge)

	fakeIgnore := new(mockCounter)
	fakeIgnore.On("Add", 1.0).Return().On("Add", 0.0).Return().
		On("With", []string{"url", "unknown"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "reason", "queue_full"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "reason", "expired"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "reason", "network_err"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "reason", "invalid_config"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "reason", "queue_full"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "reason", "expired"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "reason", "network_err"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "reason", "invalid_config"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "code", "200"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "code", "201"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "code", "202"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "code", "204"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "code", "200"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "code", "201"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "code", "202"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "code", "204"}).Return(fakeIgnore).
		On("With", []string{"event", "test/extra-stuff"}).Return(fakeIgnore).
		On("With", []string{"event", "wrp"}).Return(fakeIgnore).
		On("With", []string{"event", "unknown"}).Return(fakeIgnore).
		On("With", []string{"event", "iot"}).Return(fakeIgnore)

	fakeRegistry := new(mockCaduceusMetricsRegistry)
	fakeRegistry.On("NewCounter", IncomingContentTypeCounter).Return(fakeICTC)
	fakeRegistry.On("NewCounter", DropsDueToInvalidPayload).Return(fakeDDTIP)
	fakeRegistry.On("NewCounter", DeliveryRetryCounter).Return(fakeIgnore)
	fakeRegistry.On("NewCounter", DeliveryCounter).Return(fakeIgnore)
	fakeRegistry.On("NewCounter", SlowConsumerCounter).Return(fakeIgnore)
	fakeRegistry.On("NewCounter", SlowConsumerDroppedMsgCounter).Return(fakeIgnore)
	fakeRegistry.On("NewGauge", OutgoingQueueDepth).Return(fakeGauge)

	return &SenderWrapperFactory{
		NumWorkersPerSender: 10,
		QueueSizePerSender:  10,
		CutOffPeriod:        30 * time.Second,
		Logger:              logging.DefaultLogger(),
		Linger:              0 * time.Second,
		MetricsRegistry:     fakeRegistry,
	}
}

func TestInvalidLinger(t *testing.T) {
	swf := getFakeFactory()
	sw, err := swf.New()

	assert := assert.New(t)
	assert.Nil(sw)
	assert.NotNil(err)
}

func TestSwSimple(t *testing.T) {
	assert := assert.New(t)

	wrpMessage := wrp.SimpleRequestResponse{
		Source:          "mac:112233445566",
		Destination:     "event:wrp",
		TransactionUUID: "12345",
	}

	var buffer bytes.Buffer
	encoder := wrp.NewEncoder(&buffer, wrp.Msgpack)
	err := encoder.Encode(&wrpMessage)
	assert.Nil(err)

	iot := simpleRequest()
	iot.Destination = "mac:112233445566/event/iot"
	test := simpleRequest()
	test.Destination = "mac:112233445566/event/test/extra-stuff"

	trans := &swTransport{}

	swf := getFakeFactory()
	swf.Sender = trans.RoundTrip

	swf.Linger = 1 * time.Second
	sw, err := swf.New()

	assert.Nil(err)
	assert.NotNil(sw)

	// No listeners

	sw.Queue(iot)
	sw.Queue(iot)
	sw.Queue(iot)

	assert.Equal(int32(0), trans.i)

	w1 := webhook.W{
		Duration: 6 * time.Second,
		Until:    time.Now().Add(6 * time.Second),
		Events:   []string{"iot"},
	}
	w1.Config.URL = "http://localhost:8888/foo"
	w1.Config.ContentType = "application/json"
	w1.Matcher.DeviceId = []string{"mac:112233445566"}

	w2 := webhook.W{
		Duration: 4 * time.Second,
		Until:    time.Now().Add(4 * time.Second),
		Events:   []string{"iot", "test/extra-stuff", "wrp"},
	}
	w2.Config.URL = "http://localhost:9999/foo"
	w2.Config.ContentType = "application/json"
	w2.Matcher.DeviceId = []string{"mac:112233445566"}

	// Add 2 listeners
	list := []webhook.W{w1, w2}

	sw.Update(list)

	// Send iot message
	sw.Queue(iot)
	time.Sleep(time.Second)
	assert.Equal(int32(2), atomic.LoadInt32(&trans.i))

	// Send test message
	sw.Queue(test)
	time.Sleep(time.Second)
	assert.Equal(int32(3), atomic.LoadInt32(&trans.i))

	// Wait for one to expire & send it again
	time.Sleep(2 * time.Second)
	sw.Queue(test)
	time.Sleep(time.Second)
	assert.Equal(int32(3), atomic.LoadInt32(&trans.i))

	w3 := webhook.W{
		Duration: 5 * time.Second,
		Until:    time.Now().Add(5 * time.Second),
		Events:   []string{"iot"},
	}
	w3.Config.URL = "http://localhost:9999/foo"
	w3.Config.ContentType = "application/json"

	// We get a registration
	list2 := []webhook.W{w3}
	sw.Update(list2)
	time.Sleep(time.Second)

	// Send iot
	sw.Queue(iot)

	sw.Shutdown(true)
	assert.Equal(int32(5), atomic.LoadInt32(&trans.i))
}
<commit_msg>The tests inside senderWrapper_test need to be rethought since they happen asynchronously.<commit_after>/**
 * Copyright 2017 Comcast Cable Communications Management, LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */
package main

import (
	"bytes"
	"github.com/Comcast/webpa-common/logging"
	"github.com/Comcast/webpa-common/webhook"
	"github.com/Comcast/webpa-common/wrp"
	"github.com/stretchr/testify/assert"
	"net/http"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

type result struct {
	URL      string
	event    string
	transID  string
	deviceID string
}

// Make a simple RoundTrip implementation that let's me short-circuit the network
type swTransport struct {
	i       int32
	fn      func(*http.Request, int) (*http.Response, error)
	results []result
	mutex   sync.Mutex
}

func (t *swTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	atomic.AddInt32(&t.i, 1)

	r := result{URL: req.URL.String(),
		event:    req.Header.Get("X-Webpa-Event"),
		transID:  req.Header.Get("X-Webpa-Transaction-Id"),
		deviceID: req.Header.Get("X-Webpa-Device-Id"),
	}

	t.mutex.Lock()
	t.results = append(t.results, r)
	t.mutex.Unlock()

	resp := &http.Response{Status: "200 OK", StatusCode: 200}
	return resp, nil
}

func getFakeFactory() *SenderWrapperFactory {
	fakeICTC := new(mockCounter)
	fakeICTC.On("With", []string{"content_type", "msgpack"}).Return(fakeICTC).
		On("With", []string{"content_type", "unknown"}).Return(fakeICTC).
		On("Add", 1.0).Return()

	fakeDDTIP := new(mockCounter)
	fakeDDTIP.On("Add", 1.0).Return()

	fakeGauge := new(mockGauge)
	fakeGauge.On("Add", 1.0).Return().
		On("Add", -1.0).Return().
		//On("With", []string{"url", "unknown"}).Return(fakeGauge).
		On("With", []string{"url", "http://localhost:8888/foo"}).Return(fakeGauge).
		On("With", []string{"url", "http://localhost:9999/foo"}).Return(fakeGauge)

	fakeIgnore := new(mockCounter)
	fakeIgnore.On("Add", 1.0).Return().On("Add", 0.0).Return().
		On("With", []string{"url", "unknown"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "reason", "queue_full"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "reason", "expired"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "reason", "network_err"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "reason", "invalid_config"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "reason", "queue_full"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "reason", "expired"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "reason", "network_err"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "reason", "invalid_config"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "code", "200"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "code", "201"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "code", "202"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:8888/foo", "code", "204"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "code", "200"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "code", "201"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "code", "202"}).Return(fakeIgnore).
		On("With", []string{"url", "http://localhost:9999/foo", "code", "204"}).Return(fakeIgnore).
		On("With", []string{"event", "test/extra-stuff"}).Return(fakeIgnore).
		On("With", []string{"event", "wrp"}).Return(fakeIgnore).
		On("With", []string{"event", "unknown"}).Return(fakeIgnore).
		On("With", []string{"event", "iot"}).Return(fakeIgnore)

	fakeRegistry := new(mockCaduceusMetricsRegistry)
	fakeRegistry.On("NewCounter", IncomingContentTypeCounter).Return(fakeICTC)
	fakeRegistry.On("NewCounter", DropsDueToInvalidPayload).Return(fakeDDTIP)
	fakeRegistry.On("NewCounter", DeliveryRetryCounter).Return(fakeIgnore)
	fakeRegistry.On("NewCounter", DeliveryCounter).Return(fakeIgnore)
	fakeRegistry.On("NewCounter", SlowConsumerCounter).Return(fakeIgnore)
	fakeRegistry.On("NewCounter", SlowConsumerDroppedMsgCounter).Return(fakeIgnore)
	fakeRegistry.On("NewGauge", OutgoingQueueDepth).Return(fakeGauge)

	return &SenderWrapperFactory{
		NumWorkersPerSender: 10,
		QueueSizePerSender:  10,
		CutOffPeriod:        30 * time.Second,
		Logger:              logging.DefaultLogger(),
		Linger:              0 * time.Second,
		MetricsRegistry:     fakeRegistry,
	}
}

func TestInvalidLinger(t *testing.T) {
	swf := getFakeFactory()
	sw, err := swf.New()

	assert := assert.New(t)
	assert.Nil(sw)
	assert.NotNil(err)
}

func TestSwSimple(t *testing.T) {
	assert := assert.New(t)

	wrpMessage := wrp.SimpleRequestResponse{
		Source:          "mac:112233445566",
		Destination:     "event:wrp",
		TransactionUUID: "12345",
	}

	var buffer bytes.Buffer
	encoder := wrp.NewEncoder(&buffer, wrp.Msgpack)
	err := encoder.Encode(&wrpMessage)
	assert.Nil(err)

	iot := simpleRequest()
	iot.Destination = "mac:112233445566/event/iot"
	test := simpleRequest()
	test.Destination = "mac:112233445566/event/test/extra-stuff"

	trans := &swTransport{}

	swf := getFakeFactory()
	swf.Sender = trans.RoundTrip

	swf.Linger = 1 * time.Second
	sw, err := swf.New()

	assert.Nil(err)
	assert.NotNil(sw)

	// No listeners

	sw.Queue(iot)
	sw.Queue(iot)
	sw.Queue(iot)

	assert.Equal(int32(0), trans.i)

	w1 := webhook.W{
		Duration: 6 * time.Second,
		Until:    time.Now().Add(6 * time.Second),
		Events:   []string{"iot"},
	}
	w1.Config.URL = "http://localhost:8888/foo"
	w1.Config.ContentType = "application/json"
	w1.Matcher.DeviceId = []string{"mac:112233445566"}

	w2 := webhook.W{
		Duration: 4 * time.Second,
		Until:    time.Now().Add(4 * time.Second),
		Events:   []string{"iot", "test/extra-stuff", "wrp"},
	}
	w2.Config.URL = "http://localhost:9999/foo"
	w2.Config.ContentType = "application/json"
	w2.Matcher.DeviceId = []string{"mac:112233445566"}

	// Add 2 listeners
	list := []webhook.W{w1, w2}

	sw.Update(list)

	// Send iot message
	sw.Queue(iot)

	// Send test message
	sw.Queue(test)

	// Send it again
	sw.Queue(test)

	w3 := webhook.W{
		Events: []string{"iot"},
	}
	w3.Config.URL = "http://localhost:9999/foo"
	w3.Config.ContentType = "application/json"

	// We get a registration
	list2 := []webhook.W{w3}
	sw.Update(list2)
	time.Sleep(time.Second)

	// Send iot
	sw.Queue(iot)

	sw.Shutdown(true)
	//assert.Equal(int32(4), atomic.LoadInt32(&trans.i))
}
<|endoftext|>


<commit_before>/*
 * The task of the forwarder is to take an HTTP request, pass it to the
 * Interceptor, and then only if the Interceptor returns false, pass it
 * on to the application.
 *
 * Importantly, the forwarder will save the request body to file, as it
 * is not possible to stream the body first to the Interceptor, then to
 * the application without doing so. This adds an inevitable overhead.
 */
package forwarder

import (
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"
)

const applicationUrlHeader string = "X-Clammit-Backend"

/*
 * The Interceptor will be passed the request to examine and pass.
 *
 * If the Interceptor deems that the request should not be forwarded to the
 * target application, it should return true.
 *
 * The request body is at EOF, so if the Interceptor needs to examine the
 * body, it should work with the "body" parameter.
 *
 * Also, the Interceptor is passed the ResponseWriter. If it fails the
 * request, the Interceptor must set the response status code and body
 * as it deems appropriate - the Forwarder will not do this.
 */
type Interceptor interface {
	Handle(w http.ResponseWriter, req *http.Request, body io.Reader) bool
}

/*
 * Forwarder implementation
 */
type Forwarder struct {
	applicationURL         *url.URL
	interceptor            Interceptor
	logger                 *log.Logger
	debug                  bool
	contentMemoryThreshold int64
}

/*
 * Constructs a new forwarder. Pass in the application URL and the interceptor.
 */
func NewForwarder(applicationURL *url.URL, contentMemoryThreshold int64, interceptor Interceptor) *Forwarder {
	return &Forwarder{
		applicationURL:         applicationURL,
		interceptor:            interceptor,
		logger:                 log.New(ioutil.Discard, "", 0),
		contentMemoryThreshold: contentMemoryThreshold,
	}
}

/*
 * Sets the logger. The default is to log nothing, so if you wish for forwarder
 * debug information, you will need to call this method.
 */
func (f *Forwarder) SetLogger(logger *log.Logger, debug bool) {
	if logger == nil {
		logger = log.New(ioutil.Discard, "", 0)
	}
	f.logger = logger
	f.debug = debug
}

/*
 * Handles the given HTTP request.
 */
func (f *Forwarder) HandleRequest(w http.ResponseWriter, req *http.Request) {
	// Catch panics and return a 500 Internal Server Error
	defer func() {
		if err := recover(); err != nil {
			f.logger.Printf("ERROR %s", err)

			// Return 500 response
			http.Error(w, "Internal Server Error", 500)
		}
	}()

	if f.debug {
		f.logger.Println("Received scan request")
	}

	//
	// Save the request body
	//
	bodyHolder, err := NewBodyHolder(req.Body, req.ContentLength, f.contentMemoryThreshold)
	if err != nil {
		f.logger.Println("Unable to save body to local store: %s", err.Error())
		http.Error(w, "Internal Server Error", 500)
		return
	}
	defer bodyHolder.Close()

	//
	// Allow the interceptor its chance
	//
	if f.interceptor != nil {
		if f.debug {
			f.logger.Println("Passing to interceptor")
		}
		r, _ := bodyHolder.GetReadCloser()
		defer r.Close()
		if f.interceptor.Handle(w, req, r) {
			f.logger.Println("Interceptor has deemed that this request should not be forwarded")
			return
		}
	}

	if f.debug {
		f.logger.Println("Interceptor passed this request")
	}

	//
	// Forward the request to the configured server
	//
	body, _ := bodyHolder.GetReadCloser()
	defer body.Close()
	resp, err := f.forwardRequest(req, body, bodyHolder.ContentLength())
	if err != nil {
		f.logger.Printf("Failed to forward request: %s", err.Error())
		http.Error(w, "Bad Gateway", 502)
		return
	}
	if resp == nil {
		f.logger.Printf("Failed to forward request: no response at all")
		http.Error(w, "Bad Gateway", 502)
		return
	}
	if resp.Body != nil {
		f.logger.Printf("Request forwarded, response %s\n", resp.Status)
		defer resp.Body.Close()
	}

	//
	// and return the response
	//
	for key, val := range resp.Header {
		w.Header()[key] = val
	}
	w.WriteHeader(resp.StatusCode)
	if resp.Body != nil {
		io.Copy(w, resp.Body) // this could throw an error, but there's nowt we can do about it now
	}

	return
}

/*
 * Forwards the request to the application. This function tries to preserve as much
 * as possible of the request - headers and body.
 */
func (f *Forwarder) forwardRequest(req *http.Request, body io.Reader, contentLength int64) (*http.Response, error) {
	client, url := f.getClient(req)
	freq, _ := http.NewRequest(req.Method, url.String(), body)
	freq.ContentLength = contentLength
	for key, val := range req.Header {
		freq.Header[key] = val
	}

	// Be nice and add client IP to forwarding chain
	if req.RemoteAddr != "@" {
		xff := freq.Header.Get("X-Forwarded-For")
		if xff != "" {
			xff += ", "
		}
		xff += req.Header.Get("X-Forwarded-For") + strings.Split(req.RemoteAddr, ":")[0]
		freq.Header.Set("X-Forwarded-For", xff)
	}

	return client.Do(freq)
}

func (f *Forwarder) getApplicationURL(req *http.Request) *url.URL {
	// Return the applicationURL if it's set
	if f.applicationURL != nil && f.applicationURL.String() != "" {
		return f.applicationURL
	}

	// Otherwise check for the X-Clammit-Backend header
	url, err := url.Parse(req.Header.Get(applicationUrlHeader))
	if err != nil {
		f.logger.Panicf("Error parsing application URL in %s: %s (%s)", applicationUrlHeader, err.Error(), req.Header.Get(applicationUrlHeader))
		return nil
	}

	if len(url.String()) == 0 {
		f.logger.Panicf("No application URL available - header %s is blank", applicationUrlHeader)
	}

	return url
}

/*
 * Gets an appropriate net/http.Client. I'm not sure if this is necessary, but it forces the issue.
 */
func (f *Forwarder) getClient(req *http.Request) (*http.Client, *url.URL) {
	applicationURL := f.getApplicationURL(req)
	url := &url.URL{
		Scheme:   applicationURL.Scheme,
		Opaque:   applicationURL.Opaque,
		User:     applicationURL.User, // TODO: clone this
		Host:     applicationURL.Host,
		Path:     req.URL.Path,
		RawQuery: req.URL.RawQuery,
		Fragment: req.URL.Fragment,
	}
	if applicationURL.Scheme == "unix" {
		f.logger.Printf("Forwarding to unix socket %s", applicationURL.Path)
		url.Scheme = "http"
		url.Host = "x"
		jar, _ := cookiejar.New(nil)
		return &http.Client{
			Jar: jar,
			Transport: &http.Transport{
				Dial: func(network, addr string) (net.Conn, error) {
					return net.Dial("unix", applicationURL.Path)
				},
			},
		}, url
	} else {
		f.logger.Printf("Forwarding to %s", applicationURL.String())
		return &http.Client{}, url
	}
}
<commit_msg>Fix: X-Forwarded-For double-use and no comma<commit_after>/*
 * The task of the forwarder is to take an HTTP request, pass it to the
 * Interceptor, and then only if the Interceptor returns false, pass it
 * on to the application.
 *
 * Importantly, the forwarder will save the request body to file, as it
 * is not possible to stream the body first to the Interceptor, then to
 * the application without doing so. This adds an inevitable overhead.
 */
package forwarder

import (
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"
)

const applicationUrlHeader string = "X-Clammit-Backend"

/*
 * The Interceptor will be passed the request to examine and pass.
 *
 * If the Interceptor deems that the request should not be forwarded to the
 * target application, it should return true.
 *
 * The request body is at EOF, so if the Interceptor needs to examine the
 * body, it should work with the "body" parameter.
 *
 * Also, the Interceptor is passed the ResponseWriter. If it fails the
 * request, the Interceptor must set the response status code and body
 * as it deems appropriate - the Forwarder will not do this.
 */
type Interceptor interface {
	Handle(w http.ResponseWriter, req *http.Request, body io.Reader) bool
}

/*
 * Forwarder implementation
 */
type Forwarder struct {
	applicationURL         *url.URL
	interceptor            Interceptor
	logger                 *log.Logger
	debug                  bool
	contentMemoryThreshold int64
}

/*
 * Constructs a new forwarder. Pass in the application URL and the interceptor.
 */
func NewForwarder(applicationURL *url.URL, contentMemoryThreshold int64, interceptor Interceptor) *Forwarder {
	return &Forwarder{
		applicationURL:         applicationURL,
		interceptor:            interceptor,
		logger:                 log.New(ioutil.Discard, "", 0),
		contentMemoryThreshold: contentMemoryThreshold,
	}
}

/*
 * Sets the logger. The default is to log nothing, so if you wish for forwarder
 * debug information, you will need to call this method.
 */
func (f *Forwarder) SetLogger(logger *log.Logger, debug bool) {
	if logger == nil {
		logger = log.New(ioutil.Discard, "", 0)
	}
	f.logger = logger
	f.debug = debug
}

/*
 * Handles the given HTTP request.
 */
func (f *Forwarder) HandleRequest(w http.ResponseWriter, req *http.Request) {
	// Catch panics and return a 500 Internal Server Error
	defer func() {
		if err := recover(); err != nil {
			f.logger.Printf("ERROR %s", err)

			// Return 500 response
			http.Error(w, "Internal Server Error", 500)
		}
	}()

	if f.debug {
		f.logger.Println("Received scan request")
	}

	//
	// Save the request body
	//
	bodyHolder, err := NewBodyHolder(req.Body, req.ContentLength, f.contentMemoryThreshold)
	if err != nil {
		f.logger.Println("Unable to save body to local store: %s", err.Error())
		http.Error(w, "Internal Server Error", 500)
		return
	}
	defer bodyHolder.Close()

	//
	// Allow the interceptor its chance
	//
	if f.interceptor != nil {
		if f.debug {
			f.logger.Println("Passing to interceptor")
		}
		r, _ := bodyHolder.GetReadCloser()
		defer r.Close()
		if f.interceptor.Handle(w, req, r) {
			f.logger.Println("Interceptor has deemed that this request should not be forwarded")
			return
		}
	}

	if f.debug {
		f.logger.Println("Interceptor passed this request")
	}

	//
	// Forward the request to the configured server
	//
	body, _ := bodyHolder.GetReadCloser()
	defer body.Close()
	resp, err := f.forwardRequest(req, body, bodyHolder.ContentLength())
	if err != nil {
		f.logger.Printf("Failed to forward request: %s", err.Error())
		http.Error(w, "Bad Gateway", 502)
		return
	}
	if resp == nil {
		f.logger.Printf("Failed to forward request: no response at all")
		http.Error(w, "Bad Gateway", 502)
		return
	}
	if resp.Body != nil {
		f.logger.Printf("Request forwarded, response %s\n", resp.Status)
		defer resp.Body.Close()
	}

	//
	// and return the response
	//
	for key, val := range resp.Header {
		w.Header()[key] = val
	}
	w.WriteHeader(resp.StatusCode)
	if resp.Body != nil {
		io.Copy(w, resp.Body) // this could throw an error, but there's nowt we can do about it now
	}

	return
}

/*
 * Forwards the request to the application. This function tries to preserve as much
 * as possible of the request - headers and body.
 */
func (f *Forwarder) forwardRequest(req *http.Request, body io.Reader, contentLength int64) (*http.Response, error) {
	client, url := f.getClient(req)
	freq, _ := http.NewRequest(req.Method, url.String(), body)
	freq.ContentLength = contentLength
	for key, val := range req.Header {
		freq.Header[key] = val
	}

	// Be nice and add client IP to forwarding chain
	if req.RemoteAddr != "@" {
		xff := freq.Header.Get("X-Forwarded-For")
		if xff != "" {
			xff += ", "
		}
		xff += strings.Split(req.RemoteAddr, ":")[0]
		freq.Header.Set("X-Forwarded-For", xff)
	}

	return client.Do(freq)
}

func (f *Forwarder) getApplicationURL(req *http.Request) *url.URL {
	// Return the applicationURL if it's set
	if f.applicationURL != nil && f.applicationURL.String() != "" {
		return f.applicationURL
	}

	// Otherwise check for the X-Clammit-Backend header
	url, err := url.Parse(req.Header.Get(applicationUrlHeader))
	if err != nil {
		f.logger.Panicf("Error parsing application URL in %s: %s (%s)", applicationUrlHeader, err.Error(), req.Header.Get(applicationUrlHeader))
		return nil
	}

	if len(url.String()) == 0 {
		f.logger.Panicf("No application URL available - header %s is blank", applicationUrlHeader)
	}

	return url
}

/*
 * Gets an appropriate net/http.Client. I'm not sure if this is necessary, but it forces the issue.
 */
func (f *Forwarder) getClient(req *http.Request) (*http.Client, *url.URL) {
	applicationURL := f.getApplicationURL(req)
	url := &url.URL{
		Scheme:   applicationURL.Scheme,
		Opaque:   applicationURL.Opaque,
		User:     applicationURL.User, // TODO: clone this
		Host:     applicationURL.Host,
		Path:     req.URL.Path,
		RawQuery: req.URL.RawQuery,
		Fragment: req.URL.Fragment,
	}
	if applicationURL.Scheme == "unix" {
		f.logger.Printf("Forwarding to unix socket %s", applicationURL.Path)
		url.Scheme = "http"
		url.Host = "x"
		jar, _ := cookiejar.New(nil)
		return &http.Client{
			Jar: jar,
			Transport: &http.Transport{
				Dial: func(network, addr string) (net.Conn, error) {
					return net.Dial("unix", applicationURL.Path)
				},
			},
		}, url
	} else {
		f.logger.Printf("Forwarding to %s", applicationURL.String())
		return &http.Client{}, url
	}
}
<|endoftext|>