<commit_before>package ndp

import (
	"errors"
	"fmt"
	"net"
	"runtime"
	"time"

	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv6"
)

// HopLimit is the expected IPv6 hop limit for all NDP messages.
const HopLimit = 255

// A Conn is a Neighbor Discovery Protocol connection.
type Conn struct {
	pc *ipv6.PacketConn
	cm *ipv6.ControlMessage

	ifi  *net.Interface
	addr *net.IPAddr

	// Used only in tests:
	//
	// icmpTest disables the self-filtering mechanism in ReadFrom, and
	// udpTestPort enables the Conn to run over UDP for easier unprivileged
	// tests.
	icmpTest    bool
	udpTestPort int
}

// Dial dials a NDP connection using the specified interface and address type.
//
// As a special case, literal IPv6 addresses may be specified to bind to a
// specific address for an interface.  If the IPv6 address does not exist on
// the interface, an error will be returned.
//
// Dial returns a Conn and the chosen IPv6 address of the interface.
func Dial(ifi *net.Interface, addr Addr) (*Conn, net.IP, error) {
	addrs, err := ifi.Addrs()
	if err != nil {
		return nil, nil, err
	}

	ipAddr, err := chooseAddr(addrs, ifi.Name, addr)
	if err != nil {
		return nil, nil, err
	}

	ic, err := icmp.ListenPacket("ip6:ipv6-icmp", ipAddr.String())
	if err != nil {
		return nil, nil, err
	}

	pc := ic.IPv6PacketConn()

	// Hop limit is always 255, per RFC 4861.
	if err := pc.SetHopLimit(HopLimit); err != nil {
		return nil, nil, err
	}
	if err := pc.SetMulticastHopLimit(HopLimit); err != nil {
		return nil, nil, err
	}

	if runtime.GOOS != "windows" {
		// Calculate and place ICMPv6 checksum at correct offset in all
		// messages (not implemented by golang.org/x/net/ipv6 on Windows).
		const chkOff = 2
		if err := pc.SetChecksum(true, chkOff); err != nil {
			return nil, nil, fmt.Errorf("set checksum: %w", err)
		}
	}

	return newConn(pc, ipAddr, ifi)
}

// newConn is an internal test constructor used for creating a Conn from an
// arbitrary ipv6.PacketConn.
func newConn(pc *ipv6.PacketConn, src *net.IPAddr, ifi *net.Interface) (*Conn, net.IP, error) {
	c := &Conn{
		pc: pc,

		// The default control message used when none is specified.
		cm: &ipv6.ControlMessage{
			HopLimit: HopLimit,
			Src:      src.IP,
			IfIndex:  ifi.Index,
		},

		ifi:  ifi,
		addr: src,
	}

	return c, src.IP, nil
}

// Close closes the Conn's underlying connection.
func (c *Conn) Close() error {
	return c.pc.Close()
}

// SetReadDeadline sets a deadline for the next NDP message to arrive.
func (c *Conn) SetReadDeadline(t time.Time) error {
	return c.pc.SetReadDeadline(t)
}

// JoinGroup joins the specified multicast group.
func (c *Conn) JoinGroup(group net.IP) error {
	return c.pc.JoinGroup(c.ifi, &net.IPAddr{
		IP:   group,
		Zone: c.ifi.Name,
	})
}

// LeaveGroup leaves the specified multicast group.
func (c *Conn) LeaveGroup(group net.IP) error {
	return c.pc.LeaveGroup(c.ifi, &net.IPAddr{
		IP:   group,
		Zone: c.ifi.Name,
	})
}

// SetICMPFilter applies the specified ICMP filter. This option can be used
// to ensure a Conn only accepts certain kinds of NDP messages.
func (c *Conn) SetICMPFilter(f *ipv6.ICMPFilter) error {
	return c.pc.SetICMPFilter(f)
}

// SetControlMessage enables the reception of *ipv6.ControlMessages based on
// the specified flags.
func (c *Conn) SetControlMessage(cf ipv6.ControlFlags, on bool) error {
	return c.pc.SetControlMessage(cf, on)
}

// ReadFrom reads a Message from the Conn and returns its control message and
// source network address. Messages sourced from this machine and malformed or
// unrecognized ICMPv6 messages are filtered.
//
// If more control and/or a more efficient low-level API are required, see
// ReadRaw.
func (c *Conn) ReadFrom() (Message, *ipv6.ControlMessage, net.IP, error) {
	b := make([]byte, c.ifi.MTU)
	for {
		n, cm, ip, err := c.ReadRaw(b)
		if err != nil {
			return nil, nil, nil, err
		}

		// Filter message if:
		//   - not testing the Conn implementation.
		//   - this address sent this message.
		if !c.test() && ip.Equal(c.addr.IP) {
			continue
		}

		m, err := ParseMessage(b[:n])
		if err != nil {
			// Filter parsing errors on the caller's behalf.
			if errors.Is(err, errParseMessage) {
				continue
			}

			return nil, nil, nil, err
		}

		return m, cm, ip, nil
	}
}

// ReadRaw reads ICMPv6 message bytes into b from the Conn and returns the
// number of bytes read, the control message, and the source network address.
//
// Most callers should use ReadFrom instead, which parses bytes into Messages
// and also handles malformed and unrecognized ICMPv6 messages.
func (c *Conn) ReadRaw(b []byte) (int, *ipv6.ControlMessage, net.IP, error) {
	n, cm, src, err := c.pc.ReadFrom(b)
	if err != nil {
		return n, nil, nil, err
	}

	return n, cm, srcIP(src), nil
}

// WriteTo writes a Message to the Conn, with an optional control message and
// destination network address.
//
// If cm is nil, a default control message will be sent.
func (c *Conn) WriteTo(m Message, cm *ipv6.ControlMessage, dst net.IP) error {
	b, err := MarshalMessage(m)
	if err != nil {
		return err
	}

	return c.writeRaw(b, cm, dst)
}

// writeRaw allows writing raw bytes with a Conn.
func (c *Conn) writeRaw(b []byte, cm *ipv6.ControlMessage, dst net.IP) error {
	// Set reasonable defaults if control message is nil.
	if cm == nil {
		cm = c.cm
	}

	_, err := c.pc.WriteTo(b, cm, c.dstAddr(dst, c.ifi.Name))
	return err
}

// dstAddr returns a different net.Addr type depending on if the Conn is
// configured for testing.
func (c *Conn) dstAddr(ip net.IP, zone string) net.Addr {
	if !c.test() || c.udpTestPort == 0 {
		return &net.IPAddr{
			IP:   ip,
			Zone: zone,
		}
	}

	return &net.UDPAddr{
		IP:   ip,
		Port: c.udpTestPort,
		Zone: c.ifi.Name,
	}
}

// test determines if Conn is configured for testing.
func (c *Conn) test() bool {
	return c.icmpTest || c.udpTestPort != 0
}

// srcIP retrieves the net.IP from possible net.Addr types used in a Conn.
func srcIP(addr net.Addr) net.IP {
	switch a := addr.(type) {
	case *net.IPAddr:
		return a.IP
	case *net.UDPAddr:
		return a.IP
	default:
		panic(fmt.Sprintf("ndp: unhandled source net.Addr: %#v", addr))
	}
}

// SolicitedNodeMulticast returns the solicited-node multicast address for
// an IPv6 address.
func SolicitedNodeMulticast(ip net.IP) (net.IP, error) {
	if err := checkIPv6(ip); err != nil {
		return nil, err
	}

	// Fixed prefix, and low 24 bits taken from input address.
	snm := net.ParseIP("ff02::1:ff00:0")
	for i := 13; i < 16; i++ {
		snm[i] = ip[i]
	}

	return snm, nil
}

// TestConns sets up a pair of testing NDP peer Conns over UDP using the
// specified interface, and returns the address which can be used to send
// messages between them.
//
// TestConns is useful for environments and tests which do not allow direct
// ICMPv6 communications.
func TestConns(ifi *net.Interface) (*Conn, *Conn, net.IP, error) {
	addrs, err := ifi.Addrs()
	if err != nil {
		return nil, nil, nil, fmt.Errorf("ndp: failed to get interface %q addresses: %v", ifi.Name, err)
	}

	addr, err := chooseAddr(addrs, ifi.Name, LinkLocal)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("ndp: failed to find link-local address for %q: %v", ifi.Name, err)
	}

	// Create two UDPv6 connections and instruct them to communicate
	// with each other for Conn tests.
	c1, p1, err := udpConn(addr, ifi)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("ndp: failed to set up first test connection: %v", err)
	}

	c2, p2, err := udpConn(addr, ifi)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("ndp: failed to set up second test connection: %v", err)
	}

	c1.udpTestPort = p2
	c2.udpTestPort = p1

	return c1, c2, addr.IP, nil
}

// udpConn creates a single test Conn over UDP, and returns the port used to
// send messages to it.
func udpConn(addr *net.IPAddr, ifi *net.Interface) (*Conn, int, error) {
	laddr := &net.UDPAddr{
		IP: addr.IP,
		// Port omitted so it will be assigned automatically.
		Zone: addr.Zone,
	}

	uc, err := net.ListenUDP("udp6", laddr)
	if err != nil {
		return nil, 0, fmt.Errorf("ndp: failed to listen UDPv6: %v", err)
	}

	pc := ipv6.NewPacketConn(uc)

	c, _, err := newConn(pc, addr, ifi)
	if err != nil {
		return nil, 0, fmt.Errorf("ndp: failed to create test NDP conn: %v", err)
	}

	return c, uc.LocalAddr().(*net.UDPAddr).Port, nil
}
<commit_msg>ndp: don't wrap Dial SetChecksum error<commit_after>package ndp

import (
	"errors"
	"fmt"
	"net"
	"runtime"
	"time"

	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv6"
)

// HopLimit is the expected IPv6 hop limit for all NDP messages.
const HopLimit = 255

// A Conn is a Neighbor Discovery Protocol connection.
type Conn struct {
	pc *ipv6.PacketConn
	cm *ipv6.ControlMessage

	ifi  *net.Interface
	addr *net.IPAddr

	// Used only in tests:
	//
	// icmpTest disables the self-filtering mechanism in ReadFrom, and
	// udpTestPort enables the Conn to run over UDP for easier unprivileged
	// tests.
	icmpTest    bool
	udpTestPort int
}

// Dial dials a NDP connection using the specified interface and address type.
//
// As a special case, literal IPv6 addresses may be specified to bind to a
// specific address for an interface.  If the IPv6 address does not exist on
// the interface, an error will be returned.
//
// Dial returns a Conn and the chosen IPv6 address of the interface.
func Dial(ifi *net.Interface, addr Addr) (*Conn, net.IP, error) {
	addrs, err := ifi.Addrs()
	if err != nil {
		return nil, nil, err
	}

	ipAddr, err := chooseAddr(addrs, ifi.Name, addr)
	if err != nil {
		return nil, nil, err
	}

	ic, err := icmp.ListenPacket("ip6:ipv6-icmp", ipAddr.String())
	if err != nil {
		return nil, nil, err
	}

	pc := ic.IPv6PacketConn()

	// Hop limit is always 255, per RFC 4861.
	if err := pc.SetHopLimit(HopLimit); err != nil {
		return nil, nil, err
	}
	if err := pc.SetMulticastHopLimit(HopLimit); err != nil {
		return nil, nil, err
	}

	if runtime.GOOS != "windows" {
		// Calculate and place ICMPv6 checksum at correct offset in all
		// messages (not implemented by golang.org/x/net/ipv6 on Windows).
		const chkOff = 2
		if err := pc.SetChecksum(true, chkOff); err != nil {
			return nil, nil, err
		}
	}

	return newConn(pc, ipAddr, ifi)
}

// newConn is an internal test constructor used for creating a Conn from an
// arbitrary ipv6.PacketConn.
func newConn(pc *ipv6.PacketConn, src *net.IPAddr, ifi *net.Interface) (*Conn, net.IP, error) {
	c := &Conn{
		pc: pc,

		// The default control message used when none is specified.
		cm: &ipv6.ControlMessage{
			HopLimit: HopLimit,
			Src:      src.IP,
			IfIndex:  ifi.Index,
		},

		ifi:  ifi,
		addr: src,
	}

	return c, src.IP, nil
}

// Close closes the Conn's underlying connection.
func (c *Conn) Close() error {
	return c.pc.Close()
}

// SetReadDeadline sets a deadline for the next NDP message to arrive.
func (c *Conn) SetReadDeadline(t time.Time) error {
	return c.pc.SetReadDeadline(t)
}

// JoinGroup joins the specified multicast group.
func (c *Conn) JoinGroup(group net.IP) error {
	return c.pc.JoinGroup(c.ifi, &net.IPAddr{
		IP:   group,
		Zone: c.ifi.Name,
	})
}

// LeaveGroup leaves the specified multicast group.
func (c *Conn) LeaveGroup(group net.IP) error {
	return c.pc.LeaveGroup(c.ifi, &net.IPAddr{
		IP:   group,
		Zone: c.ifi.Name,
	})
}

// SetICMPFilter applies the specified ICMP filter. This option can be used
// to ensure a Conn only accepts certain kinds of NDP messages.
func (c *Conn) SetICMPFilter(f *ipv6.ICMPFilter) error {
	return c.pc.SetICMPFilter(f)
}

// SetControlMessage enables the reception of *ipv6.ControlMessages based on
// the specified flags.
func (c *Conn) SetControlMessage(cf ipv6.ControlFlags, on bool) error {
	return c.pc.SetControlMessage(cf, on)
}

// ReadFrom reads a Message from the Conn and returns its control message and
// source network address. Messages sourced from this machine and malformed or
// unrecognized ICMPv6 messages are filtered.
//
// If more control and/or a more efficient low-level API are required, see
// ReadRaw.
func (c *Conn) ReadFrom() (Message, *ipv6.ControlMessage, net.IP, error) {
	b := make([]byte, c.ifi.MTU)
	for {
		n, cm, ip, err := c.ReadRaw(b)
		if err != nil {
			return nil, nil, nil, err
		}

		// Filter message if:
		//   - not testing the Conn implementation.
		//   - this address sent this message.
		if !c.test() && ip.Equal(c.addr.IP) {
			continue
		}

		m, err := ParseMessage(b[:n])
		if err != nil {
			// Filter parsing errors on the caller's behalf.
			if errors.Is(err, errParseMessage) {
				continue
			}

			return nil, nil, nil, err
		}

		return m, cm, ip, nil
	}
}

// ReadRaw reads ICMPv6 message bytes into b from the Conn and returns the
// number of bytes read, the control message, and the source network address.
//
// Most callers should use ReadFrom instead, which parses bytes into Messages
// and also handles malformed and unrecognized ICMPv6 messages.
func (c *Conn) ReadRaw(b []byte) (int, *ipv6.ControlMessage, net.IP, error) {
	n, cm, src, err := c.pc.ReadFrom(b)
	if err != nil {
		return n, nil, nil, err
	}

	return n, cm, srcIP(src), nil
}

// WriteTo writes a Message to the Conn, with an optional control message and
// destination network address.
//
// If cm is nil, a default control message will be sent.
func (c *Conn) WriteTo(m Message, cm *ipv6.ControlMessage, dst net.IP) error {
	b, err := MarshalMessage(m)
	if err != nil {
		return err
	}

	return c.writeRaw(b, cm, dst)
}

// writeRaw allows writing raw bytes with a Conn.
func (c *Conn) writeRaw(b []byte, cm *ipv6.ControlMessage, dst net.IP) error {
	// Set reasonable defaults if control message is nil.
	if cm == nil {
		cm = c.cm
	}

	_, err := c.pc.WriteTo(b, cm, c.dstAddr(dst, c.ifi.Name))
	return err
}

// dstAddr returns a different net.Addr type depending on if the Conn is
// configured for testing.
func (c *Conn) dstAddr(ip net.IP, zone string) net.Addr {
	if !c.test() || c.udpTestPort == 0 {
		return &net.IPAddr{
			IP:   ip,
			Zone: zone,
		}
	}

	return &net.UDPAddr{
		IP:   ip,
		Port: c.udpTestPort,
		Zone: c.ifi.Name,
	}
}

// test determines if Conn is configured for testing.
func (c *Conn) test() bool {
	return c.icmpTest || c.udpTestPort != 0
}

// srcIP retrieves the net.IP from possible net.Addr types used in a Conn.
func srcIP(addr net.Addr) net.IP {
	switch a := addr.(type) {
	case *net.IPAddr:
		return a.IP
	case *net.UDPAddr:
		return a.IP
	default:
		panic(fmt.Sprintf("ndp: unhandled source net.Addr: %#v", addr))
	}
}

// SolicitedNodeMulticast returns the solicited-node multicast address for
// an IPv6 address.
func SolicitedNodeMulticast(ip net.IP) (net.IP, error) {
	if err := checkIPv6(ip); err != nil {
		return nil, err
	}

	// Fixed prefix, and low 24 bits taken from input address.
	snm := net.ParseIP("ff02::1:ff00:0")
	for i := 13; i < 16; i++ {
		snm[i] = ip[i]
	}

	return snm, nil
}

// TestConns sets up a pair of testing NDP peer Conns over UDP using the
// specified interface, and returns the address which can be used to send
// messages between them.
//
// TestConns is useful for environments and tests which do not allow direct
// ICMPv6 communications.
func TestConns(ifi *net.Interface) (*Conn, *Conn, net.IP, error) {
	addrs, err := ifi.Addrs()
	if err != nil {
		return nil, nil, nil, fmt.Errorf("ndp: failed to get interface %q addresses: %v", ifi.Name, err)
	}

	addr, err := chooseAddr(addrs, ifi.Name, LinkLocal)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("ndp: failed to find link-local address for %q: %v", ifi.Name, err)
	}

	// Create two UDPv6 connections and instruct them to communicate
	// with each other for Conn tests.
	c1, p1, err := udpConn(addr, ifi)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("ndp: failed to set up first test connection: %v", err)
	}

	c2, p2, err := udpConn(addr, ifi)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("ndp: failed to set up second test connection: %v", err)
	}

	c1.udpTestPort = p2
	c2.udpTestPort = p1

	return c1, c2, addr.IP, nil
}

// udpConn creates a single test Conn over UDP, and returns the port used to
// send messages to it.
func udpConn(addr *net.IPAddr, ifi *net.Interface) (*Conn, int, error) {
	laddr := &net.UDPAddr{
		IP: addr.IP,
		// Port omitted so it will be assigned automatically.
		Zone: addr.Zone,
	}

	uc, err := net.ListenUDP("udp6", laddr)
	if err != nil {
		return nil, 0, fmt.Errorf("ndp: failed to listen UDPv6: %v", err)
	}

	pc := ipv6.NewPacketConn(uc)

	c, _, err := newConn(pc, addr, ifi)
	if err != nil {
		return nil, 0, fmt.Errorf("ndp: failed to create test NDP conn: %v", err)
	}

	return c, uc.LocalAddr().(*net.UDPAddr).Port, nil
}
<|endoftext|>


<commit_before>package wemdigo

import (
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// Link instances wrap a Gorilla websocket and provkeye a link to
// a specific Mkeydle instance.  The Link is responsible for keeping
// the underlying websocket connection alive as well as reading / writing
// messages.
type Conn struct {
	ws      *websocket.Conn // hkeye
	wsConf  *WSConfig
	wgSend  sync.WaitGroup
	wgRec   sync.WaitGroup
	key     string
	deps    []string
	hubs    []string
	targets []string
	// Conn owned channels.
	read chan *Message
	send chan *Message // Messages to write to the websocket.
	done chan struct{} // Terminates the main loop.
	// Middle instance channels shared by all connections.
	mid   chan *Message // Channel instance uses to speak to the Middle.
	unreg chan *Conn
}

func (c *Conn) Init(cc ConnConfig, mid chan *Message, unreg chan *Conn) {
	c.ws = cc.Conn
	if cc.WSConf == nil {
		cc.WSConf = new(WSConfig)
		cc.WSConf.init()
	}
	c.key = cc.Key
	c.hubs = cc.Hubs
	c.targets = cc.Targets
	c.deps = cc.Deps
	c.wsConf = cc.WSConf
	c.mid = mid
	c.unreg = unreg
}

// NewConn to be initialized further by a Middle instance.
func NewConn() *Conn {
	return &Conn{
		read: make(chan *Message),
		send: make(chan *Message),
		done: make(chan struct{}),
	}
}

// Key returns the connection's key keyentifier and a bool indicating
// whether this value is set.  If the boolean is false, then this
// connection is not indirectly targetable.
func (c Conn) Key() (string, bool) {
	return c.key, c.key != ""
}

// HasKey reports whether the underlying connection is targetable.
func (c Conn) HasKey() bool {
	_, ok := c.Key()
	return ok
}

// Targets the connection subscribes to.
func (c *Conn) Targets() []string { return c.targets }

// Hugs the target subscribes to.
func (c *Conn) Hubs() []string { return c.hubs }

// Deps are target connections on which the instance depends.
func (c *Conn) Deps() []string { return c.deps }

// UnderlyingConn returns the underlying Gorilla websocket connection.
func (c Conn) UnderlyingConn() *websocket.Conn { return c.ws }

func (c *Conn) setReadDeadline() {
	var t time.Time
	if c.wsConf.PongWait != 0 {
		t = time.Now().Add(c.wsConf.PongWait)
	}
	c.ws.SetReadDeadline(t)
}

func (c *Conn) write(messageType int, data []byte) error {
	c.ws.SetWriteDeadline(time.Now().Add(c.wsConf.WriteWait))
	return c.ws.WriteMessage(messageType, data)
}

// readLoop pumps messages from the websocket Link to the Mkeydle.
func (c *Conn) readLoop() {
	defer func() {
		dlog("Conn %s sending to unregister chan.", c.key)
		c.ws.Close()
		close(c.read)

		// Wait until all sent messages have been evaluated by the handler
		// before unregistering with the Middle, as we do not want the
		// connection's dependencies to close before having a chance to write
		// all relevant messages.
		c.wgSend.Wait()
		c.unreg <- c
	}()

	ponghandler := func(appData string) error {
		c.setReadDeadline()
		dlog("Conn %s saw a pong.", c.key)
		return nil
	}

	c.ws.SetPongHandler(ponghandler)

	for {
		mt, data, err := c.ws.ReadMessage()
		if err != nil {
			return
		}
		dlog("Conn %s read a message", c.key)
		msg := &Message{Type: mt, Data: data, origin: c}
		c.read <- msg
	}
}

// writeLoop is responsible for writing to the peer.  This is the only
// location where writing occurs.
func (c *Conn) writeLoop() {

	defer func() {
		c.write(websocket.CloseMessage, nil)
		c.ws.Close()
		dlog("Conn %s no longer writing messages.", c.key)
	}()

	for msg := range c.send {
		if err := c.write(msg.Type, msg.Data); err != nil {
			return
		}
	}
}

// main loop handles all communication with the Mkeydle layer.  Messages
// read by the underlying websocket and those to be written by it
// are first funneled through here.
func (c *Conn) mainLoop() {
	pinger := time.NewTicker(c.wsConf.PingPeriod)
	defer func() {
		c.ws.Close()
		pinger.Stop()
		close(c.send)
	}()

	for {
		select {
		case msg, ok := <-c.read:
			if ok {
				c.mid <- msg
			}
		case <-pinger.C:
			msg := &Message{Type: websocket.PingMessage, Data: nil}
			c.send <- msg
		case <-c.done:
			dlog("Conn %s received on done channel.", c.key)
			return
		}
	}
}

func (c *Conn) run() {
	go c.writeLoop()
	go c.readLoop()
	go c.mainLoop()
}
<commit_msg>Fix double unreg issue.<commit_after>package wemdigo

import (
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// Link instances wrap a Gorilla websocket and provkeye a link to
// a specific Mkeydle instance.  The Link is responsible for keeping
// the underlying websocket connection alive as well as reading / writing
// messages.
type Conn struct {
	ws        *websocket.Conn // hkeye
	wsConf    *WSConfig
	wgSend    sync.WaitGroup
	wgRec     sync.WaitGroup
	key       string
	deps      []string
	hubs      []string
	targets   []string
	unregSent bool
	// Conn owned channels.
	read chan *Message
	send chan *Message // Messages to write to the websocket.
	done chan struct{} // Terminates the main loop.
	err  chan struct{} // Websocket error channel.
	// Middle instance channels shared by all connections.
	mid   chan *Message // Channel instance uses to speak to the Middle.
	unreg chan *Conn    // Middle's unregistration channel.
}

func (c *Conn) Init(cc ConnConfig, mid chan *Message, unreg chan *Conn) {
	c.ws = cc.Conn
	if cc.WSConf == nil {
		cc.WSConf = new(WSConfig)
		cc.WSConf.init()
	}
	c.key = cc.Key
	c.hubs = cc.Hubs
	c.targets = cc.Targets
	c.deps = cc.Deps
	c.wsConf = cc.WSConf
	c.mid = mid
	c.unreg = unreg
}

// NewConn to be initialized further by a Middle instance.
func NewConn() *Conn {
	return &Conn{
		err:  make(chan struct{}),
		read: make(chan *Message),
		send: make(chan *Message),
		done: make(chan struct{}),
	}
}

// Key returns the connection's key keyentifier and a bool indicating
// whether this value is set.  If the boolean is false, then this
// connection is not indirectly targetable.
func (c Conn) Key() (string, bool) {
	return c.key, c.key != ""
}

// HasKey reports whether the underlying connection is targetable.
func (c Conn) HasKey() bool {
	_, ok := c.Key()
	return ok
}

// Targets the connection subscribes to.
func (c *Conn) Targets() []string { return c.targets }

// Hugs the target subscribes to.
func (c *Conn) Hubs() []string { return c.hubs }

// Deps are target connections on which the instance depends.
func (c *Conn) Deps() []string { return c.deps }

// UnderlyingConn returns the underlying Gorilla websocket connection.
func (c Conn) UnderlyingConn() *websocket.Conn { return c.ws }

func (c *Conn) setReadDeadline() {
	var t time.Time
	if c.wsConf.PongWait != 0 {
		t = time.Now().Add(c.wsConf.PongWait)
	}
	c.ws.SetReadDeadline(t)
}

func (c *Conn) write(messageType int, data []byte) error {
	c.ws.SetWriteDeadline(time.Now().Add(c.wsConf.WriteWait))
	return c.ws.WriteMessage(messageType, data)
}

// readLoop pumps messages from the websocket Link to the Mkeydle.
func (c *Conn) readLoop() {
	defer func() {
		dlog("Conn %s sending to unregister chan.", c.key)
		c.ws.Close()
		close(c.read)

		// Wait until all sent messages have been evaluated by the handler
		// before unregistering with the Middle, as we do not want the
		// connection's dependencies to close before having a chance to write
		// all relevant messages.
		c.wgSend.Wait()
		c.err <- struct{}{}
		close(c.err)
	}()

	ponghandler := func(appData string) error {
		c.setReadDeadline()
		dlog("Conn %s saw a pong.", c.key)
		return nil
	}

	c.ws.SetPongHandler(ponghandler)

	for {
		mt, data, err := c.ws.ReadMessage()
		if err != nil {
			return
		}
		dlog("Conn %s read a message", c.key)
		msg := &Message{Type: mt, Data: data, origin: c}
		c.read <- msg
	}
}

// writeLoop is responsible for writing to the peer.  This is the only
// location where writing occurs.
func (c *Conn) writeLoop() {

	defer func() {
		c.write(websocket.CloseMessage, nil)
		c.ws.Close()
		dlog("Conn %s no longer writing messages.", c.key)
	}()

	for msg := range c.send {
		if err := c.write(msg.Type, msg.Data); err != nil {
			return
		}
	}
}

// main loop handles all communication with the Mkeydle layer.  Messages
// read by the underlying websocket and those to be written by it
// are first funneled through here.
func (c *Conn) mainLoop() {
	pinger := time.NewTicker(c.wsConf.PingPeriod)
	defer func() {
		c.ws.Close()
		pinger.Stop()
		close(c.send)
	}()

	for {
		select {
		case msg, ok := <-c.read:
			if ok {
				c.mid <- msg
			}
		case <-pinger.C:
			msg := &Message{Type: websocket.PingMessage, Data: nil}
			c.send <- msg
		case <-c.err:
			// Websocket can no longer read.
			if !c.unregSent {
				c.unregSent = true
				c.unreg <- c
			}
		case <-c.done:
			dlog("Conn %s received on done channel.", c.key)
			return
		}
	}
}

func (c *Conn) run() {
	go c.writeLoop()
	go c.readLoop()
	go c.mainLoop()
}
<|endoftext|>


<commit_before>package noodle

import (
	"golang.org/x/net/context"
	"net/http"
)

// Handler provides context-aware http.Handler with error return value for
// enhanced chaining
type Handler func(context.Context, http.ResponseWriter, *http.Request) error

// Middleware behaves like standard closure middleware pattern, only with
// context-aware handler type
type Middleware func(Handler) Handler

// Chain composes middlewares into a single context-aware handler
type Chain []Middleware

// New creates new middleware Chain and initalizes it with its parameters
func New(mws ...Middleware) Chain {
	return mws
}

// Use appends its parameters to middleware chain. Returns new separate
// middleware chain
func (c Chain) Use(mws ...Middleware) (res Chain) {
	res = make([]Middleware, len(c)+len(mws))
	copy(res[:len(c)], c)
	copy(res[len(c):], mws)
	return
}

// Then finalizes middleware Chain converting it to context-aware Handler
func (c Chain) Then(final Handler) Handler {
	for i := len(c) - 1; i >= 0; i-- {
		final = c[i](final)
	}
	return final
}

// origin is the root context for all requests. By default it contains
// reference to global thread-safe Store. Origin can be extended or overwritten to
// provide common application-wide initial context.
var origin = context.TODO()

// ServeHTTP creates empty context and applies Handler to it, satisfying
// http.Handler interface

func (h Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	_ = h(origin, w, r)
}
<commit_msg>documentation fix<commit_after>package noodle

import (
	"golang.org/x/net/context"
	"net/http"
)

// Handler provides context-aware http.Handler with error return value for
// enhanced chaining
type Handler func(context.Context, http.ResponseWriter, *http.Request) error

// Middleware behaves like standard closure middleware pattern, only with
// context-aware handler type
type Middleware func(Handler) Handler

// Chain composes middlewares into a single context-aware handler
type Chain []Middleware

// New creates new middleware Chain and initalizes it with its parameters
func New(mws ...Middleware) Chain {
	return mws
}

// Use appends its parameters to middleware chain. Returns new separate
// middleware chain
func (c Chain) Use(mws ...Middleware) (res Chain) {
	res = make([]Middleware, len(c)+len(mws))
	copy(res[:len(c)], c)
	copy(res[len(c):], mws)
	return
}

// Then finalizes middleware Chain converting it to context-aware Handler
func (c Chain) Then(final Handler) Handler {
	for i := len(c) - 1; i >= 0; i-- {
		final = c[i](final)
	}
	return final
}

// origin is the root context for all requests
var origin = context.TODO()

// ServeHTTP creates empty context and applies Handler to it, satisfying
// http.Handler interface
func (h Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	_ = h(origin, w, r)
}
<|endoftext|>


<commit_before>package cors

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

type CorsHandler struct {
	ALLOWED_METHODS   []string
	ALLOWED_ORIGINS   []string
	ALLOWED_HEADERS   []string
	ALLOW_CREDENTIALS string
	MAX_AGE           float64
	handler           http.Handler
}

func New(handler http.Handler) *CorsHandler {
	return &CorsHandler{
		ALLOWED_METHODS:   []string{"GET", "POST", "PUT", "DELETE", "HEAD", "OPTIONS"},
		ALLOWED_ORIGINS:   []string{"*"},
		ALLOWED_HEADERS:   []string{"Content-Type"},
		ALLOW_CREDENTIALS: "true",
		MAX_AGE:           0,
		handler:           handler,
	}
}
func (cors *CorsHandler) AllowOrigin(origin string) {
	if origin == "*" {
		cors.ALLOWED_ORIGINS = []string{"*"}
	}
	for i := 0; i < len(cors.ALLOWED_ORIGINS); i++ {
		if origin == cors.ALLOWED_ORIGINS[i] {
			return
		}
	}
	cors.ALLOWED_ORIGINS = append(cors.ALLOWED_ORIGINS, origin)
}
func (cors *CorsHandler) AllowMethod(method string) {
	method = strings.ToUpper(method)
	for i := 0; i < len(cors.ALLOWED_METHODS); i++ {
		if method == cors.ALLOWED_METHODS[i] {
			return
		}
	}
	cors.ALLOWED_METHODS = append(cors.ALLOWED_METHODS, method)
}
func (cors *CorsHandler) AllowHeader(header string) {
	for i := 0; i < len(cors.ALLOWED_HEADERS); i++ {
		if header == cors.ALLOWED_HEADERS[i] {
			return
		}
	}
	cors.ALLOWED_HEADERS = append(cors.ALLOWED_HEADERS, header)

}
func (cors *CorsHandler) AllowCredentials(creds bool) {
	if creds {
		cors.ALLOW_CREDENTIALS = "true"
	} else {
		cors.ALLOW_CREDENTIALS = "false"
	}
}
func (cors *CorsHandler) RemoveOrigin(origin string) {
	for i := 0; i < len(cors.ALLOWED_ORIGINS); i++ {
		if origin == cors.ALLOWED_ORIGINS[i] {
			cors.ALLOWED_ORIGINS = cors.ALLOWED_ORIGINS[:i+copy(cors.ALLOWED_ORIGINS[i:], cors.ALLOWED_ORIGINS[i+1:])]
		}
	}

}
func (cors *CorsHandler) RemoveMethod(method string) {
	method = strings.ToUpper(method)
	for i := 0; i < len(cors.ALLOWED_METHODS); i++ {
		if method == cors.ALLOWED_METHODS[i] {
			cors.ALLOWED_METHODS = cors.ALLOWED_METHODS[:i+copy(cors.ALLOWED_METHODS[i:], cors.ALLOWED_METHODS[i+1:])]
		}
	}
}
func (cors *CorsHandler) RemoveHeader(header string) {
	for i := 0; i < len(cors.ALLOWED_HEADERS); i++ {
		if header == cors.ALLOWED_HEADERS[i] {
			cors.ALLOWED_HEADERS = cors.ALLOWED_HEADERS[:i+copy(cors.ALLOWED_HEADERS[i:], cors.ALLOWED_HEADERS[i+1:])]
		}
	}
}
func (cors *CorsHandler) IsOriginAllowed(origin string) bool {
	for i := 0; i < len(cors.ALLOWED_ORIGINS); i++ {
		if "*" == cors.ALLOWED_ORIGINS[i] {
			return true
		} else if origin == cors.ALLOWED_ORIGINS[i] {
			return true
		}
	}
	return false
}
func (cors *CorsHandler) IsMethodAllowed(method string) bool {
	method = strings.ToUpper(method)
	for i := 0; i < len(cors.ALLOWED_METHODS); i++ {
		if method == cors.ALLOWED_METHODS[i] {
			return true
		}
	}
	return false
}
func (cors *CorsHandler) IsHeaderAllowed(header string) bool {
	for i := 0; i < len(cors.ALLOWED_HEADERS); i++ {
		if header == cors.ALLOWED_HEADERS[i] {
			return true
		}
	}
	return false
}
func (cors *CorsHandler) AllowedMethods() string {
	methods := ""
	for i := 0; i < len(cors.ALLOWED_METHODS); i++ {
		if methods == "" {
			methods = cors.ALLOWED_METHODS[i]
		} else {
			methods = fmt.Sprintf("%s, %s", methods, cors.ALLOWED_METHODS[i])
		}
	}
	return methods
}
func (cors *CorsHandler) AllowedHeaders() string {
	headers := ""
	for i := 0; i < len(cors.ALLOWED_HEADERS); i++ {
		if headers == "" {
			headers = cors.ALLOWED_HEADERS[i]
		} else {
			headers = fmt.Sprintf("%s, %s", headers, cors.ALLOWED_HEADERS[i])
		}
	}
	return headers
}
func (cors *CorsHandler) SetMaxAge(age float64) {
	cors.MAX_AGE = age
}
func (cors *CorsHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	origin := req.Header.Get("Origin")
	if origin != "" && cors.IsOriginAllowed(origin) {
		w.Header().Add("Access-Control-Allow-Origin", origin)
		w.Header().Add("Access-Control-Allow-Methods", cors.AllowedMethods())
		w.Header().Add("Access-Control-Allow-Headers", cors.AllowedHeaders())
		w.Header().Add("Access-Control-Allow-Credentials", cors.ALLOW_CREDENTIALS)
		if cors.MAX_AGE > 0 {
			w.Header().Add("Access-Control-Max-Age", fmt.Sprintf("%s", cors.MAX_AGE))
		}
	}
	cors.handler.ServeHTTP(w, req)
}
<commit_msg>Fixing max age output<commit_after>package cors

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

type CorsHandler struct {
	ALLOWED_METHODS   []string
	ALLOWED_ORIGINS   []string
	ALLOWED_HEADERS   []string
	ALLOW_CREDENTIALS string
	MAX_AGE           float64
	handler           http.Handler
}

func New(handler http.Handler) *CorsHandler {
	return &CorsHandler{
		ALLOWED_METHODS:   []string{"GET", "POST", "PUT", "DELETE", "HEAD", "OPTIONS"},
		ALLOWED_ORIGINS:   []string{"*"},
		ALLOWED_HEADERS:   []string{"Content-Type"},
		ALLOW_CREDENTIALS: "true",
		MAX_AGE:           0,
		handler:           handler,
	}
}
func (cors *CorsHandler) AllowOrigin(origin string) {
	if origin == "*" {
		cors.ALLOWED_ORIGINS = []string{"*"}
	}
	for i := 0; i < len(cors.ALLOWED_ORIGINS); i++ {
		if origin == cors.ALLOWED_ORIGINS[i] {
			return
		}
	}
	cors.ALLOWED_ORIGINS = append(cors.ALLOWED_ORIGINS, origin)
}
func (cors *CorsHandler) AllowMethod(method string) {
	method = strings.ToUpper(method)
	for i := 0; i < len(cors.ALLOWED_METHODS); i++ {
		if method == cors.ALLOWED_METHODS[i] {
			return
		}
	}
	cors.ALLOWED_METHODS = append(cors.ALLOWED_METHODS, method)
}
func (cors *CorsHandler) AllowHeader(header string) {
	for i := 0; i < len(cors.ALLOWED_HEADERS); i++ {
		if header == cors.ALLOWED_HEADERS[i] {
			return
		}
	}
	cors.ALLOWED_HEADERS = append(cors.ALLOWED_HEADERS, header)

}
func (cors *CorsHandler) AllowCredentials(creds bool) {
	if creds {
		cors.ALLOW_CREDENTIALS = "true"
	} else {
		cors.ALLOW_CREDENTIALS = "false"
	}
}
func (cors *CorsHandler) RemoveOrigin(origin string) {
	for i := 0; i < len(cors.ALLOWED_ORIGINS); i++ {
		if origin == cors.ALLOWED_ORIGINS[i] {
			cors.ALLOWED_ORIGINS = cors.ALLOWED_ORIGINS[:i+copy(cors.ALLOWED_ORIGINS[i:], cors.ALLOWED_ORIGINS[i+1:])]
		}
	}

}
func (cors *CorsHandler) RemoveMethod(method string) {
	method = strings.ToUpper(method)
	for i := 0; i < len(cors.ALLOWED_METHODS); i++ {
		if method == cors.ALLOWED_METHODS[i] {
			cors.ALLOWED_METHODS = cors.ALLOWED_METHODS[:i+copy(cors.ALLOWED_METHODS[i:], cors.ALLOWED_METHODS[i+1:])]
		}
	}
}
func (cors *CorsHandler) RemoveHeader(header string) {
	for i := 0; i < len(cors.ALLOWED_HEADERS); i++ {
		if header == cors.ALLOWED_HEADERS[i] {
			cors.ALLOWED_HEADERS = cors.ALLOWED_HEADERS[:i+copy(cors.ALLOWED_HEADERS[i:], cors.ALLOWED_HEADERS[i+1:])]
		}
	}
}
func (cors *CorsHandler) IsOriginAllowed(origin string) bool {
	for i := 0; i < len(cors.ALLOWED_ORIGINS); i++ {
		if "*" == cors.ALLOWED_ORIGINS[i] {
			return true
		} else if origin == cors.ALLOWED_ORIGINS[i] {
			return true
		}
	}
	return false
}
func (cors *CorsHandler) IsMethodAllowed(method string) bool {
	method = strings.ToUpper(method)
	for i := 0; i < len(cors.ALLOWED_METHODS); i++ {
		if method == cors.ALLOWED_METHODS[i] {
			return true
		}
	}
	return false
}
func (cors *CorsHandler) IsHeaderAllowed(header string) bool {
	for i := 0; i < len(cors.ALLOWED_HEADERS); i++ {
		if header == cors.ALLOWED_HEADERS[i] {
			return true
		}
	}
	return false
}
func (cors *CorsHandler) AllowedMethods() string {
	methods := ""
	for i := 0; i < len(cors.ALLOWED_METHODS); i++ {
		if methods == "" {
			methods = cors.ALLOWED_METHODS[i]
		} else {
			methods = fmt.Sprintf("%s, %s", methods, cors.ALLOWED_METHODS[i])
		}
	}
	return methods
}
func (cors *CorsHandler) AllowedHeaders() string {
	headers := ""
	for i := 0; i < len(cors.ALLOWED_HEADERS); i++ {
		if headers == "" {
			headers = cors.ALLOWED_HEADERS[i]
		} else {
			headers = fmt.Sprintf("%s, %s", headers, cors.ALLOWED_HEADERS[i])
		}
	}
	return headers
}
func (cors *CorsHandler) SetMaxAge(age float64) {
	cors.MAX_AGE = age
}
func (cors *CorsHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	origin := req.Header.Get("Origin")
	if origin != "" && cors.IsOriginAllowed(origin) {
		w.Header().Add("Access-Control-Allow-Origin", origin)
		w.Header().Add("Access-Control-Allow-Methods", cors.AllowedMethods())
		w.Header().Add("Access-Control-Allow-Headers", cors.AllowedHeaders())
		w.Header().Add("Access-Control-Allow-Credentials", cors.ALLOW_CREDENTIALS)
		if cors.MAX_AGE > 0 {
			w.Header().Add("Access-Control-Max-Age", fmt.Sprintf("%9.f", cors.MAX_AGE))
		}
	}
	cors.handler.ServeHTTP(w, req)
}
<|endoftext|>


<commit_before>package melware

import (
	"time"
	"strings"
	"net/http"
	"fmt"
	"github.com/ridewindx/mel"
)

type Cors struct {
	// AllowedOrigins is a slice of origins that a cors request can be executed from.
	// An origin may contain a wildcard (*) to replace 0 or more characters
	// (e.g., http://*.domain.com). Only one wildcard can be used per origin.
	// Default value is ["*"], i.e., all origins are allowed.
	AllowOrigins []string

	// AllowOriginFunc is a custom function to validate the origin. It take the origin
	// as argument and returns true if allowed or false otherwise.
	// It has lower precedence than AllowOrigins.
	AllowOriginFunc func(origin string) bool

	// AllowedMethods is a slice of methods the client is allowed to use with
	// cross-domain requests.
	// Default to {"GET", "POST", "HEAD"}.
	AllowMethods []string

	// AllowedHeaders is slice of non simple headers the client is allowed to use with
	// cross-domain requests.
	// Default to {"Origin", "Accept", "Content-Type"}.
	AllowHeaders []string

	// AllowCredentials indicates whether the request can include user credentials like
	// cookies, HTTP authentication or client side SSL certificates.
	AllowCredentials bool

	// ExposedHeaders indicates which headers are safe to expose to the API of a CORS
	// API specification
	ExposeHeaders []string

	// MaxAge indicates how long (in seconds) the results of a preflight request
	// can be cached
	MaxAge time.Duration

	allowAllOrigins bool
	normalHeaders    http.Header
	preflightHeaders http.Header
}

func NewCors() *Cors {
	return &Cors{
		AllowOrigins: []string{"*"},
		AllowMethods: []string{"GET", "POST", "HEAD"},
		AllowHeaders: []string{"Origin", "Accept", "Content-Type"},
		AllowCredentials: false,
		MaxAge: 12 * time.Hour,
	}
}

func (c *Cors) Middleware() mel.Handler {
    c.validateAllowOrigins()

	c.normalHeaders = c.generateNormalHeaders()
	c.preflightHeaders = c.generatePreflightHeaders()

	return func(ctx *mel.Context) {
		origin := ctx.Request.Header.Get("Origin")
		if len(origin) == 0 { // request is not a CORS request
			return
		}
		if !c.validateOrigin(origin) {
			ctx.AbortWithStatus(http.StatusForbidden)
			return
		}

		if ctx.Request.Method == "OPTIONS" {
			for key, value := range c.preflightHeaders {
				ctx.Header(key, value[0])
			}
			defer ctx.AbortWithStatus(http.StatusOK)
		} else {
			for key, value := range c.normalHeaders {
				ctx.Header(key, value[0])
			}
		}

		if !c.allowAllOrigins && !c.AllowCredentials {
			ctx.Header("Access-Control-Allow-Origin", origin)
		}

		ctx.Next()
	}
}

func (c *Cors) validateAllowOrigins() {
	c.AllowOrigins = c.normalizeStrs(c.AllowOrigins)
	if len(c.AllowOrigins) == 1 && c.AllowOrigins[0] == "*" {
		c.allowAllOrigins = true
		if c.AllowOriginFunc != nil {
			panic("All origins are allowed, no predicate function needed")
		}
	} else if len(c.AllowOrigins) > 0 {
		for _, origin := range c.AllowOrigins {
			if origin == "*" {
				panic("All origins for cors are allowed, no individual origins needed")
			} else if !strings.HasPrefix(origin, "http://") && !strings.HasPrefix(origin, "https://") {
				panic("Origin must have prefix 'http://' or 'https://'")
			}
		}
	} else if c.AllowOriginFunc == nil {
		panic("No origin is allowed")
	}
}

func (c *Cors) validateOrigin(origin string) bool {
	if c.allowAllOrigins {
		return true
	}
	for _, value := range c.AllowOrigins {
		if value == origin {
			return true
		}
	}
	if c.AllowOriginFunc != nil {
		return c.AllowOriginFunc(origin)
	}
	return false
}

func (c *Cors) normalizeStrs(strs []string) []string {
	if strs == nil {
		return nil
	}
	set := make(map[string]bool)
	var normalized []string
	for _, str := range strs {
		str = strings.TrimSpace(str)
		str = strings.ToLower(str)
		if _, seen := set[str]; !seen {
			normalized = append(normalized, str)
			set[str] = true
		}
	}
	return normalized
}

func (c *Cors) generateNormalHeaders() http.Header {
	headers := make(http.Header)
	if c.AllowCredentials {
		headers.Set("Access-Control-Allow-Credentials", "true")
	}
	if len(c.ExposeHeaders) > 0 {
		exposeHeaders := c.convert(c.normalizeStrs(c.ExposeHeaders), http.CanonicalHeaderKey)
		headers.Set("Access-Control-Expose-Headers", strings.Join(exposeHeaders, ","))
	}
	if c.allowAllOrigins {
		headers.Set("Access-Control-Allow-Origin", "*")
	} else {
		headers.Set("Vary", "Origin")
	}
	return headers
}

func (c *Cors) generatePreflightHeaders() http.Header {
	headers := make(http.Header)
	if c.AllowCredentials {
		headers.Set("Access-Control-Allow-Credentials", "true")
	}
	if len(c.AllowMethods) > 0 {
		allowMethods := c.convert(c.normalizeStrs(c.AllowMethods), strings.ToUpper)
		headers.Set("Access-Control-Allow-Methods", strings.Join(allowMethods, ","))
	}
	if len(c.AllowHeaders) > 0 {
		allowHeaders := c.convert(c.normalizeStrs(c.AllowHeaders), http.CanonicalHeaderKey)
		headers.Set("Access-Control-Allow-Headers", strings.Join(allowHeaders, ","))
	}
	if c.MaxAge > time.Duration(0) {
		headers.Set("Access-Control-Max-Age", fmt.Sprintf("%d", c.MaxAge/time.Second))
	}
	if c.allowAllOrigins {
		headers.Set("Access-Control-Allow-Origin", "*")
	} else {
		// If the server specifies an origin host rather than "*",
		// then it could also include Origin in the Vary response header
		// to indicate to clients that server responses will differ based
		// on the value of the Origin request header.
		headers.Add("Vary", "Origin")
		headers.Add("Vary", "Access-Control-Request-Method")
		headers.Add("Vary", "Access-Control-Request-Headers")
	}
	return headers
}

func (c *Cors) convert(strs []string, f func(string) string) []string {
	var result []string
	for _, str := range strs {
		result = append(result, f(str))
	}
	return result
}
<commit_msg>Update files<commit_after>package melware

import (
	"net/http"
	"github.com/rs/cors"
	"github.com/ridewindx/mel"
)

type Cors struct {
	// AllowedOrigins is a list of origins a cross-domain request can be executed from.
	// If the special "*" value is present in the list, all origins will be allowed.
	// An origin may contain a wildcard (*) to replace 0 or more characters
	// (i.e.: http://*.domain.com). Usage of wildcards implies a small performance penalty.
	// Only one wildcard can be used per origin.
	// Default value is ["*"]
	AllowedOrigins []string
	// AllowOriginFunc is a custom function to validate the origin. It take the origin
	// as argument and returns true if allowed or false otherwise. If this option is
	// set, the content of AllowedOrigins is ignored.
	AllowOriginFunc func(origin string) bool
	// AllowedMethods is a list of methods the client is allowed to use with
	// cross-domain requests. Default value is simple methods (HEAD, GET and POST).
	AllowedMethods []string
	// AllowedHeaders is list of non simple headers the client is allowed to use with
	// cross-domain requests.
	// If the special "*" value is present in the list, all headers will be allowed.
	// Default value is [] but "Origin" is always appended to the list.
	AllowedHeaders []string
	// ExposedHeaders indicates which headers are safe to expose to the API of a CORS
	// API specification
	ExposedHeaders []string
	// AllowCredentials indicates whether the request can include user credentials like
	// cookies, HTTP authentication or client side SSL certificates.
	AllowCredentials bool
	// MaxAge indicates how long (in seconds) the results of a preflight request
	// can be cached
	MaxAge int
	// OptionsPassthrough instructs preflight to let other potential next handlers to
	// process the OPTIONS method. Turn this on if your application handles OPTIONS.
	OptionsPassthrough bool
	// Debugging flag adds additional output to debug server side CORS issues
	Debug bool

	*cors.Cors
}

func CorsAllowAll() *Cors {
	return &Cors{
		AllowedOrigins:   []string{"*"},
		AllowedMethods:   []string{"HEAD", "GET", "POST", "PUT", "PATCH", "DELETE"},
		AllowedHeaders:   []string{"*"},
		AllowCredentials: true,
	}
}

func (c *Cors) Middleware() mel.Handler {
	if c.Cors == nil {
		c.Cors = cors.New(cors.Options{
			AllowedOrigins: c.AllowedOrigins,
			AllowOriginFunc: c.AllowOriginFunc,
			AllowedMethods: c.AllowedMethods,
			AllowedHeaders: c.AllowedHeaders,
			ExposedHeaders: c.ExposedHeaders,
			AllowCredentials: c.AllowCredentials,
			MaxAge: c.MaxAge,
			OptionsPassthrough: c.OptionsPassthrough,
			Debug: c.Debug,
		})
	}

	return func(ctx *mel.Context) {
		c.ServeHTTP(ctx.Writer, ctx.Request, func(writer http.ResponseWriter, request *http.Request) {
			ctx.Next()
		})
	}
}
<|endoftext|>


<commit_before>package csrf

import (
	"crypto/sha1"
	"encoding/base64"
	"errors"
	"io"

	"github.com/dchest/uniuri"
	"github.com/gin-contrib/sessions"
	"github.com/gin-gonic/gin"
)

const (
	csrfSecret = "csrfSecret"
	csrfSalt   = "csrfSalt"
	csrfToken  = "csrfToken"
)

var defaultIgnoreMethods = []string{"GET", "HEAD", "OPTIONS"}

var defaultErrorFunc = func(c *gin.Context) {
	panic(errors.New("CSRF token mismatch"))
}

var defaultTokenGetter = func(c *gin.Context) string {
	r := c.Request

	if t := r.FormValue("_csrf"); len(t) > 0 {
		return t
	} else if t := r.URL.Query().Get("_csrf"); len(t) > 0 {
		return t
	} else if t := r.Header.Get("X-CSRF-TOKEN"); len(t) > 0 {
		return t
	} else if t := r.Header.Get("X-XSRF-TOKEN"); len(t) > 0 {
		return t
	}

	return ""
}

// Options stores configurations for a CSRF middleware.
type Options struct {
	Secret        string
	IgnoreMethods []string
	ErrorFunc     gin.HandlerFunc
	TokenGetter   func(c *gin.Context) string
}

func tokenize(secret, salt string) string {
	h := sha1.New()
	io.WriteString(h, salt+"-"+secret)
	hash := base64.URLEncoding.EncodeToString(h.Sum(nil))

	return hash
}

func inArray(arr []string, value string) bool {
	inarr := false

	for _, v := range arr {
		if v == value {
			inarr = true
			break
		}
	}

	return inarr
}

// Middleware validates CSRF token.
func Middleware(options Options) gin.HandlerFunc {
	ignoreMethods := options.IgnoreMethods
	errorFunc := options.ErrorFunc
	tokenGetter := options.TokenGetter

	if ignoreMethods == nil {
		ignoreMethods = defaultIgnoreMethods
	}

	if errorFunc == nil {
		errorFunc = defaultErrorFunc
	}

	if tokenGetter == nil {
		tokenGetter = defaultTokenGetter
	}

	return func(c *gin.Context) {
		session := sessions.Default(c)
		c.Set(csrfSecret, options.Secret)

		if inArray(ignoreMethods, c.Request.Method) {
			c.Next()
			return
		}

		var salt string

		if s, ok := session.Get(csrfSalt).(string); !ok || len(s) == 0 {
			c.Next()
			return
		} else {
			salt = s
		}

		session.Delete(csrfSalt)

		token := tokenGetter(c)

		if tokenize(options.Secret, salt) != token {
			errorFunc(c)
			return
		}

		c.Next()
	}
}

// GetToken returns a CSRF token.
func GetToken(c *gin.Context) string {
	session := sessions.Default(c)
	secret := c.MustGet(csrfSecret).(string)

	if t, ok := c.Get(csrfToken); ok {
		return t.(string)
	}

	salt := uniuri.New()
	token := tokenize(secret, salt)
	session.Set(csrfSalt, salt)
	session.Save()
	c.Set(csrfToken, token)

	return token
}
<commit_msg>Fixed golint warning<commit_after>package csrf

import (
	"crypto/sha1"
	"encoding/base64"
	"errors"
	"io"

	"github.com/dchest/uniuri"
	"github.com/gin-contrib/sessions"
	"github.com/gin-gonic/gin"
)

const (
	csrfSecret = "csrfSecret"
	csrfSalt   = "csrfSalt"
	csrfToken  = "csrfToken"
)

var defaultIgnoreMethods = []string{"GET", "HEAD", "OPTIONS"}

var defaultErrorFunc = func(c *gin.Context) {
	panic(errors.New("CSRF token mismatch"))
}

var defaultTokenGetter = func(c *gin.Context) string {
	r := c.Request

	if t := r.FormValue("_csrf"); len(t) > 0 {
		return t
	} else if t := r.URL.Query().Get("_csrf"); len(t) > 0 {
		return t
	} else if t := r.Header.Get("X-CSRF-TOKEN"); len(t) > 0 {
		return t
	} else if t := r.Header.Get("X-XSRF-TOKEN"); len(t) > 0 {
		return t
	}

	return ""
}

// Options stores configurations for a CSRF middleware.
type Options struct {
	Secret        string
	IgnoreMethods []string
	ErrorFunc     gin.HandlerFunc
	TokenGetter   func(c *gin.Context) string
}

func tokenize(secret, salt string) string {
	h := sha1.New()
	io.WriteString(h, salt+"-"+secret)
	hash := base64.URLEncoding.EncodeToString(h.Sum(nil))

	return hash
}

func inArray(arr []string, value string) bool {
	inarr := false

	for _, v := range arr {
		if v == value {
			inarr = true
			break
		}
	}

	return inarr
}

// Middleware validates CSRF token.
func Middleware(options Options) gin.HandlerFunc {
	ignoreMethods := options.IgnoreMethods
	errorFunc := options.ErrorFunc
	tokenGetter := options.TokenGetter

	if ignoreMethods == nil {
		ignoreMethods = defaultIgnoreMethods
	}

	if errorFunc == nil {
		errorFunc = defaultErrorFunc
	}

	if tokenGetter == nil {
		tokenGetter = defaultTokenGetter
	}

	return func(c *gin.Context) {
		session := sessions.Default(c)
		c.Set(csrfSecret, options.Secret)

		if inArray(ignoreMethods, c.Request.Method) {
			c.Next()
			return
		}

		var salt string

		s, ok := session.Get(csrfSalt).(string)

		if !ok || len(s) == 0 {
			c.Next()
			return
		}

		salt = s

		session.Delete(csrfSalt)

		token := tokenGetter(c)

		if tokenize(options.Secret, salt) != token {
			errorFunc(c)
			return
		}

		c.Next()
	}
}

// GetToken returns a CSRF token.
func GetToken(c *gin.Context) string {
	session := sessions.Default(c)
	secret := c.MustGet(csrfSecret).(string)

	if t, ok := c.Get(csrfToken); ok {
		return t.(string)
	}

	salt := uniuri.New()
	token := tokenize(secret, salt)
	session.Set(csrfSalt, salt)
	session.Save()
	c.Set(csrfToken, token)

	return token
}
<|endoftext|>


<commit_before>package cuid

import (
	"math"
	"math/rand"
	"os"
	"strconv"
	"strings"
	"time"
)

type Counter interface {
	Next() int64
}

const (
	BLOCK_SIZE = 4
	BASE       = 36
)

var (
	defaultCounter Counter = nil
	defaultRandom          = rand.New(rand.NewSource(time.Now().Unix()))
	discreteValues         = int64(math.Pow(BASE, BLOCK_SIZE))
	padding                = strings.Repeat("0", BLOCK_SIZE)
	fingerprint            = ""
)

func init() {
	hostname, err := os.Hostname()
	if err != nil {
		hostname = "dummy-host"
	}
	acc := int64(len(hostname) + BASE)
	for i := range hostname {
		acc = acc + int64(hostname[i])
	}

	hostID := pad(strconv.FormatInt(int64(os.Getpid()), BASE), 2)
	host := pad(strconv.FormatInt(acc, 10), 2)
	fingerprint = hostID + host
}

func New() string {
	if defaultCounter == nil {
		defaultCounter = NewDefaultCounter()
	}

	timestampBlock := strconv.FormatInt(time.Now().Unix()*1000, BASE)
	counterBlock := pad(strconv.FormatInt(defaultCounter.Next(), BASE), BLOCK_SIZE)
	randomBlock1 := pad(strconv.FormatInt(defaultRandom.Int63n(discreteValues), BASE), BLOCK_SIZE)
	randomBlock2 := pad(strconv.FormatInt(defaultRandom.Int63n(discreteValues), BASE), BLOCK_SIZE)

	return "c" + timestampBlock + counterBlock + fingerprint + randomBlock1 + randomBlock2
}

func pad(str string, size int) string {
	if len(str) == size {
		return str
	}

	if len(str) < size {
		str = padding + str
	}

	i := len(str) - size

	return str[i:]
}

// Default counter implementation
// The default counter is a simply generator running in a goroutine
// and providing values through a channel.

type DefaultCounter struct {
	counterChan chan int64
}

func NewDefaultCounter() *DefaultCounter {
	counter := &DefaultCounter{make(chan int64)}
	go counter.loop()
	<-counter.counterChan

	return counter
}

func (c *DefaultCounter) Next() int64 {
	return <-c.counterChan
}

func (c *DefaultCounter) loop() {
	var count int64 = -1
	for {
		c.counterChan <- count
		count = count + 1
		if count >= discreteValues {
			count = 0
		}
	}
}
<commit_msg>Counter optimisation.<commit_after>package cuid

import (
	"math"
	"math/rand"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	BLOCK_SIZE = 4
	BASE       = 36
)

var (
	defaultCounter Counter = nil
	defaultRandom          = rand.New(rand.NewSource(time.Now().Unix()))
	discreteValues         = int64(math.Pow(BASE, BLOCK_SIZE))
	padding                = strings.Repeat("0", BLOCK_SIZE)
	fingerprint            = ""
)

func init() {
	hostname, err := os.Hostname()
	if err != nil {
		hostname = "dummy-host"
	}
	acc := int64(len(hostname) + BASE)
	for i := range hostname {
		acc = acc + int64(hostname[i])
	}

	hostID := pad(strconv.FormatInt(int64(os.Getpid()), BASE), 2)
	host := pad(strconv.FormatInt(acc, 10), 2)
	fingerprint = hostID + host
}

func New() string {
	if defaultCounter == nil {
		defaultCounter = &DefaultCounter{}
	}

	timestampBlock := strconv.FormatInt(time.Now().Unix()*1000, BASE)
	counterBlock := pad(strconv.FormatInt(defaultCounter.Next(), BASE), BLOCK_SIZE)
	randomBlock1 := pad(strconv.FormatInt(defaultRandom.Int63n(discreteValues), BASE), BLOCK_SIZE)
	randomBlock2 := pad(strconv.FormatInt(defaultRandom.Int63n(discreteValues), BASE), BLOCK_SIZE)

	return "c" + timestampBlock + counterBlock + fingerprint + randomBlock1 + randomBlock2
}

func pad(str string, size int) string {
	if len(str) == size {
		return str
	}

	if len(str) < size {
		str = padding + str
	}

	i := len(str) - size

	return str[i:]
}

// Default counter implementation

type Counter interface {
	Next() int64
}

type DefaultCounter struct {
	count int64
	mutex sync.Mutex
}

func (c *DefaultCounter) Next() int64 {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	counterValue := c.count

	c.count = c.count + 1
	if c.count >= discreteValues {
		c.count = 0
	}

	return counterValue
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"time"
)

var (
	addObservation *sql.Stmt
	addSite        *sql.Stmt
)

// initData should be called after the db is available.
func initData() (err error) {
	// siteID, typeID, methodID, sampleID, systemID, time, value, error
	addObservation, err = db.Prepare("SELECT fits.add_observation($1, $2, $3, $4, $5, $6, $7, $8)")

	// siteID, name, longitude, latitude, height, ground_relationship
	addSite, err = db.Prepare("SELECT fits.add_site($1, $2, $3, $4, $5, $6)")
	if err != nil {
		return err
	}

	return
}

type data struct {
	sourceFile, observationFile string
	source
	observation
}

func (d *data) parseAndValidate() (err error) {

	b, err := ioutil.ReadFile(d.sourceFile)
	if err != nil {
		return err
	}

	if err = d.unmarshall(b); err != nil {
		return err
	}

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

	if err = d.read(f); err != nil {
		return err
	}
	f.Close()

	if !locValid {
		if err = d.valid(); err != nil {
			return err
		}
	}

	return err
}

// updateOrAdd saves data to by d to the FITS DB.  If
// an observation already exists for the source timestamp then the value and error are updated
// otherwise the data is inserted.
func (d *data) updateOrAdd() (err error) {
	for _, o := range d.obs {
		_, err = addObservation.Exec(
			d.Properties.SiteID,
			d.Properties.TypeID,
			d.Properties.MethodID,
			d.Properties.SampleID,
			d.Properties.SystemID,
			o.t,
			o.v,
			o.e)
		if err != nil {
			return err
		}
	}

	return err
}

// deleteThenSave saves data to the FITS db.  Observations for the source are first deleted and then
// values in *obs added.  This is done in a transaction.
func (d *data) deleteThenSave() (err error) {

	tx, err := db.Begin()

	var sitePK int
	err = tx.QueryRow(`SELECT DISTINCT ON (sitepk) sitepk
				FROM fits.site WHERE siteid = $1`, d.Properties.SiteID).Scan(&sitePK)
	if err == sql.ErrNoRows {
		return fmt.Errorf("couldn't get sitePK for %s", d.Properties.SiteID)
	}
	if err != nil {
		return err
	}

	var samplePK int
	err = tx.QueryRow(`SELECT DISTINCT ON (samplePK) samplePK
				FROM fits.sample join fits.system using (systempk)
				WHERE sampleID = $1
				AND
				systemID = $2`, d.Properties.SampleID, d.Properties.SystemID).Scan(&samplePK)
	if err == sql.ErrNoRows {
		return fmt.Errorf("couldn't get samplePK for %s.%s", d.Properties.SampleID, d.Properties.SystemID)
	}
	if err != nil {
		return err
	}

	var methodPK int
	err = tx.QueryRow(`SELECT methodPK FROM fits.method WHERE methodID = $1`, d.Properties.MethodID).Scan(&methodPK)
	if err == sql.ErrNoRows {
		return fmt.Errorf("couldn't get methodPK for %s", d.Properties.MethodID)
	}
	if err != nil {
		return err
	}

	// also checks that the type is valid for this method.
	var typePK int
	err = tx.QueryRow(`SELECT DISTINCT ON (typePK) typePK
				FROM fits.type
				JOIN fits.type_method USING (typepk)
				JOIN fits.method USING (methodpk)
				WHERE
				typeid = $1
				AND
				 methodid = $2`, d.Properties.TypeID, d.Properties.MethodID).Scan(&typePK)
	if err == sql.ErrNoRows {
		return fmt.Errorf("couldn't get typePK for %s.%s", d.Properties.TypeID, d.Properties.MethodID)
	}
	if err != nil {
		return err
	}

	obsDelete, err := tx.Prepare(`DELETE FROM fits.observation
					WHERE
					sitepk = (SELECT DISTINCT ON (sitepk) sitepk FROM fits.site WHERE siteid = $1)
					AND
					typePK = (SELECT DISTINCT ON (typepk)  typepk FROM fits.type WHERE typeid = $2)
					`)
	if err != nil {
		return err
	}
	defer obsDelete.Close()

	insert := `INSERT INTO fits.observation(sitePK, typePK, methodPK, samplePK, time, value, error) VALUES `

	var row string
	for _, v := range d.obs {
		row = fmt.Sprintf("(%d, %d, %d, %d, '%s'::timestamptz, %f, %f),", sitePK, typePK, methodPK, samplePK, v.t.Format(time.RFC3339), v.v, v.e)
		insert += row
	}

	insert = strings.TrimSuffix(insert, ",")

	obsInsert, err := tx.Prepare(insert)
	if err != nil {
		return err
	}
	defer obsInsert.Close()

	_, err = obsDelete.Exec(d.Properties.SiteID, d.Properties.TypeID)
	if err != nil {
		tx.Rollback()
		return err
	}

	_, err = obsInsert.Exec()
	if err != nil {
		tx.Rollback()
		return err
	}

	return tx.Commit()
}
<commit_msg>Add some error handling to fix lint warnings<commit_after>package main

import (
	"database/sql"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"time"
)

var (
	addObservation *sql.Stmt
	addSite        *sql.Stmt
)

// initData should be called after the db is available.
func initData() (err error) {
	// siteID, typeID, methodID, sampleID, systemID, time, value, error
	addObservation, err = db.Prepare("SELECT fits.add_observation($1, $2, $3, $4, $5, $6, $7, $8)")
	if err != nil {
		return err
	}

	// siteID, name, longitude, latitude, height, ground_relationship
	addSite, err = db.Prepare("SELECT fits.add_site($1, $2, $3, $4, $5, $6)")
	if err != nil {
		return err
	}

	return
}

type data struct {
	sourceFile, observationFile string
	source
	observation
}

func (d *data) parseAndValidate() (err error) {

	b, err := ioutil.ReadFile(d.sourceFile)
	if err != nil {
		return err
	}

	if err = d.unmarshall(b); err != nil {
		return err
	}

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

	if err = d.read(f); err != nil {
		return err
	}
	f.Close()

	if !locValid {
		if err = d.valid(); err != nil {
			return err
		}
	}

	return err
}

// updateOrAdd saves data to by d to the FITS DB.  If
// an observation already exists for the source timestamp then the value and error are updated
// otherwise the data is inserted.
func (d *data) updateOrAdd() (err error) {
	for _, o := range d.obs {
		_, err = addObservation.Exec(
			d.Properties.SiteID,
			d.Properties.TypeID,
			d.Properties.MethodID,
			d.Properties.SampleID,
			d.Properties.SystemID,
			o.t,
			o.v,
			o.e)
		if err != nil {
			return err
		}
	}

	return err
}

// deleteThenSave saves data to the FITS db.  Observations for the source are first deleted and then
// values in *obs added.  This is done in a transaction.
func (d *data) deleteThenSave() (err error) {

	tx, err := db.Begin()

	var sitePK int
	err = tx.QueryRow(`SELECT DISTINCT ON (sitepk) sitepk
				FROM fits.site WHERE siteid = $1`, d.Properties.SiteID).Scan(&sitePK)
	if err == sql.ErrNoRows {
		return fmt.Errorf("couldn't get sitePK for %s", d.Properties.SiteID)
	}
	if err != nil {
		return err
	}

	var samplePK int
	err = tx.QueryRow(`SELECT DISTINCT ON (samplePK) samplePK
				FROM fits.sample join fits.system using (systempk)
				WHERE sampleID = $1
				AND
				systemID = $2`, d.Properties.SampleID, d.Properties.SystemID).Scan(&samplePK)
	if err == sql.ErrNoRows {
		return fmt.Errorf("couldn't get samplePK for %s.%s", d.Properties.SampleID, d.Properties.SystemID)
	}
	if err != nil {
		return err
	}

	var methodPK int
	err = tx.QueryRow(`SELECT methodPK FROM fits.method WHERE methodID = $1`, d.Properties.MethodID).Scan(&methodPK)
	if err == sql.ErrNoRows {
		return fmt.Errorf("couldn't get methodPK for %s", d.Properties.MethodID)
	}
	if err != nil {
		return err
	}

	// also checks that the type is valid for this method.
	var typePK int
	err = tx.QueryRow(`SELECT DISTINCT ON (typePK) typePK
				FROM fits.type
				JOIN fits.type_method USING (typepk)
				JOIN fits.method USING (methodpk)
				WHERE
				typeid = $1
				AND
				 methodid = $2`, d.Properties.TypeID, d.Properties.MethodID).Scan(&typePK)
	if err == sql.ErrNoRows {
		return fmt.Errorf("couldn't get typePK for %s.%s", d.Properties.TypeID, d.Properties.MethodID)
	}
	if err != nil {
		return err
	}

	obsDelete, err := tx.Prepare(`DELETE FROM fits.observation
					WHERE
					sitepk = (SELECT DISTINCT ON (sitepk) sitepk FROM fits.site WHERE siteid = $1)
					AND
					typePK = (SELECT DISTINCT ON (typepk)  typepk FROM fits.type WHERE typeid = $2)
					`)
	if err != nil {
		return err
	}
	defer obsDelete.Close()

	insert := `INSERT INTO fits.observation(sitePK, typePK, methodPK, samplePK, time, value, error) VALUES `

	var row string
	for _, v := range d.obs {
		row = fmt.Sprintf("(%d, %d, %d, %d, '%s'::timestamptz, %f, %f),", sitePK, typePK, methodPK, samplePK, v.t.Format(time.RFC3339), v.v, v.e)
		insert += row
	}

	insert = strings.TrimSuffix(insert, ",")

	obsInsert, err := tx.Prepare(insert)
	if err != nil {
		return err
	}
	defer obsInsert.Close()

	_, err = obsDelete.Exec(d.Properties.SiteID, d.Properties.TypeID)
	if err != nil {
		rollbackErr := tx.Rollback()
		if rollbackErr != nil {
			fmt.Printf("error in rollback of DB delete transaction: %v\n", rollbackErr)
		}
		return err
	}

	_, err = obsInsert.Exec()
	if err != nil {
		rollbackErr := tx.Rollback()
		if rollbackErr != nil {
			fmt.Printf("error in rollback of DB insert transaction: %v\n", rollbackErr)
		}
		return err
	}

	return tx.Commit()
}
<|endoftext|>


<commit_before>// +build linux

package main

import (
	"os"
	"runtime"

	"github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/coreos/go-systemd/activation"
	"github.com/opencontainers/runc/libcontainer"
	"github.com/opencontainers/specs"
)

// default action is to start a container
var startCommand = cli.Command{
	Name:  "start",
	Usage: "create and run a container",
	ArgsUsage: `<container-id>

Where "<container-id>" is your name for the instance of the container that you
are starting. The name you provide for the container instance must be unique on
your host.`,
	Description: `The start command creates an instance of a container for a bundle. The bundle
is a directory with a specification file and a root filesystem.`,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "bundle, b",
			Value: "",
			Usage: `path to the root of the bundle directory, defaults to the current directory`,
		},
		cli.StringFlag{
			Name:  "console",
			Value: "",
			Usage: "specify the pty slave path for use with the container",
		},
		cli.BoolFlag{
			Name:  "detach,d",
			Usage: "detach from the container's process",
		},
		cli.StringFlag{
			Name:  "pid-file",
			Value: "",
			Usage: "specify the file to write the process id to",
		},
	},
	Action: func(context *cli.Context) {
		bundle := context.String("bundle")
		if bundle != "" {
			if err := os.Chdir(bundle); err != nil {
				fatal(err)
			}
		}
		spec, err := loadSpec(specConfig)
		if err != nil {
			fatal(err)
		}

		notifySocket := os.Getenv("NOTIFY_SOCKET")
		if notifySocket != "" {
			setupSdNotify(spec, notifySocket)
		}

		if os.Geteuid() != 0 {
			logrus.Fatal("runc should be run as root")
		}

		status, err := startContainer(context, spec)
		if err != nil {
			logrus.Fatalf("Container start failed: %v", err)
		}
		// exit with the container's exit status so any external supervisor is
		// notified of the exit with the correct exit status.
		os.Exit(status)
	},
}

func init() {
	if len(os.Args) > 1 && os.Args[1] == "init" {
		runtime.GOMAXPROCS(1)
		runtime.LockOSThread()
		factory, _ := libcontainer.New("")
		if err := factory.StartInitialization(); err != nil {
			fatal(err)
		}
		panic("--this line should have never been executed, congratulations--")
	}
}

func startContainer(context *cli.Context, spec *specs.LinuxSpec) (int, error) {
	id := context.Args().First()
	if id == "" {
		return -1, errEmptyID
	}
	container, err := createContainer(context, id, spec)
	if err != nil {
		return -1, err
	}

	// ensure that the container is always removed if we were the process
	// that created it.
	detach := context.Bool("detach")
	if !detach {
		defer destroy(container)
	}

	// Support on-demand socket activation by passing file descriptors into the container init process.
	listenFDs := []*os.File{}
	if os.Getenv("LISTEN_FDS") != "" {
		listenFDs = activation.Files(false)
	}

	return runProcess(container, &spec.Process, listenFDs, context.String("console"), context.String("pid-file"), detach)
}
<commit_msg>Make sure container is destroyed on error<commit_after>// +build linux

package main

import (
	"os"
	"runtime"

	"github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/coreos/go-systemd/activation"
	"github.com/opencontainers/runc/libcontainer"
	"github.com/opencontainers/specs"
)

// default action is to start a container
var startCommand = cli.Command{
	Name:  "start",
	Usage: "create and run a container",
	ArgsUsage: `<container-id>

Where "<container-id>" is your name for the instance of the container that you
are starting. The name you provide for the container instance must be unique on
your host.`,
	Description: `The start command creates an instance of a container for a bundle. The bundle
is a directory with a specification file and a root filesystem.`,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "bundle, b",
			Value: "",
			Usage: `path to the root of the bundle directory, defaults to the current directory`,
		},
		cli.StringFlag{
			Name:  "console",
			Value: "",
			Usage: "specify the pty slave path for use with the container",
		},
		cli.BoolFlag{
			Name:  "detach,d",
			Usage: "detach from the container's process",
		},
		cli.StringFlag{
			Name:  "pid-file",
			Value: "",
			Usage: "specify the file to write the process id to",
		},
	},
	Action: func(context *cli.Context) {
		bundle := context.String("bundle")
		if bundle != "" {
			if err := os.Chdir(bundle); err != nil {
				fatal(err)
			}
		}
		spec, err := loadSpec(specConfig)
		if err != nil {
			fatal(err)
		}

		notifySocket := os.Getenv("NOTIFY_SOCKET")
		if notifySocket != "" {
			setupSdNotify(spec, notifySocket)
		}

		if os.Geteuid() != 0 {
			logrus.Fatal("runc should be run as root")
		}

		status, err := startContainer(context, spec)
		if err != nil {
			logrus.Fatalf("Container start failed: %v", err)
		}
		// exit with the container's exit status so any external supervisor is
		// notified of the exit with the correct exit status.
		os.Exit(status)
	},
}

func init() {
	if len(os.Args) > 1 && os.Args[1] == "init" {
		runtime.GOMAXPROCS(1)
		runtime.LockOSThread()
		factory, _ := libcontainer.New("")
		if err := factory.StartInitialization(); err != nil {
			fatal(err)
		}
		panic("--this line should have never been executed, congratulations--")
	}
}

func startContainer(context *cli.Context, spec *specs.LinuxSpec) (int, error) {
	id := context.Args().First()
	if id == "" {
		return -1, errEmptyID
	}
	container, err := createContainer(context, id, spec)
	if err != nil {
		return -1, err
	}

	// ensure that the container is always removed if we were the process
	// that created it.
	detach := context.Bool("detach")
	if !detach {
		defer destroy(container)
	}

	// Support on-demand socket activation by passing file descriptors into the container init process.
	listenFDs := []*os.File{}
	if os.Getenv("LISTEN_FDS") != "" {
		listenFDs = activation.Files(false)
	}

	status, err := runProcess(container, &spec.Process, listenFDs, context.String("console"), context.String("pid-file"), detach)
	if err != nil {
		destroy(container)
		return -1, err
	}
	return status, nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

const shutdownGraceTime = 3 * time.Second

var flagPort int
var flagConcurrency string
var flagRestart bool

var cmdStart = &Command{
	Run:   runStart,
	Usage: "start [process name] [-f procfile] [-e env] [-c concurrency] [-p port] [-r]",
	Short: "Start the application",
	Long: `
Start the application specified by a Procfile (defaults to ./Procfile)

Examples:

  forego start
  forego start web
  forego start -f Procfile.test -e .env.test
`,
}

func init() {
	cmdStart.Flag.StringVar(&flagProcfile, "f", "Procfile", "procfile")
	cmdStart.Flag.StringVar(&flagEnv, "e", "", "env")
	cmdStart.Flag.IntVar(&flagPort, "p", 5000, "port")
	cmdStart.Flag.StringVar(&flagConcurrency, "c", "", "concurrency")
	cmdStart.Flag.BoolVar(&flagRestart, "r", false, "restart")
}

func parseConcurrency(value string) (map[string]int, error) {
	concurrency := map[string]int{}
	if strings.TrimSpace(value) == "" {
		return concurrency, nil
	}

	parts := strings.Split(value, ",")
	for _, part := range parts {
		if !strings.Contains(part, "=") {
			return concurrency, errors.New("Parsing concurency")
		}

		nameValue := strings.Split(part, "=")
		n, v := strings.TrimSpace(nameValue[0]), strings.TrimSpace(nameValue[1])
		if n == "" || v == "" {
			return concurrency, errors.New("Parsing concurency")
		}

		numProcs, err := strconv.ParseInt(v, 10, 16)
		if err != nil {
			return concurrency, err
		}

		concurrency[n] = int(numProcs)
	}
	return concurrency, nil
}

type Forego struct {
	shutdown sync.Once     // Closes teardown exactly once
	teardown chan struct{} // barrier: closed when shutting down

	wg sync.WaitGroup
}

func (f *Forego) SignalShutdown() {
	f.shutdown.Do(func() { close(f.teardown) })
}

func (f *Forego) monitorInterrupt() {
	handler := make(chan os.Signal, 1)
	signal.Notify(handler, os.Interrupt)

	first := true

	for sig := range handler {
		switch sig {
		case os.Interrupt:
			fmt.Println("      | ctrl-c detected")

			if !first {

			}
			f.SignalShutdown()
		}
	}
}

func (f *Forego) startProcess(idx, procNum int, proc ProcfileEntry, env Env, of *OutletFactory) {
	port := flagPort + (idx * 100)

	ps := NewProcess(proc.Command, env)
	procName := fmt.Sprint(proc.Name, ".", procNum+1)
	ps.Env["PORT"] = strconv.Itoa(port)
	ps.Root = filepath.Dir(flagProcfile)
	ps.Stdin = nil
	ps.Stdout = of.CreateOutlet(procName, idx, false)
	ps.Stderr = of.CreateOutlet(procName, idx, true)

	of.SystemOutput(fmt.Sprintf("starting %s on port %d", procName, port))

	finished := make(chan struct{}) // closed on process exit

	ps.Start()
	go func() {
		defer close(finished)
		ps.Wait()
	}()

	f.wg.Add(1)
	go func() {
		defer f.wg.Done()

		// Prevent goroutine from exiting before process has finished.
		defer func() { <-finished }()

		select {
		case <-finished:
			if flagRestart {
				f.startProcess(idx, procNum, proc, env, of)
				return
			} else {
				f.SignalShutdown()
			}

		case <-f.teardown:
			// Forego tearing down

			if !osHaveSigTerm {
				of.SystemOutput(fmt.Sprintf("Killing %s", procName))
				ps.cmd.Process.Kill()
				return
			}

			of.SystemOutput(fmt.Sprintf("sending SIGTERM to %s", procName))
			ps.SendSigTerm()

			// Give the process a chance to exit, otherwise kill it.
			select {
			case <-time.After(shutdownGraceTime):
				of.SystemOutput(fmt.Sprintf("Killing %s", procName))
				ps.SendSigKill()
			case <-finished:
			}
		}
	}()
}

func runStart(cmd *Command, args []string) {
	root := filepath.Dir(flagProcfile)

	if flagEnv == "" {
		flagEnv = filepath.Join(root, ".env")
	}

	pf, err := ReadProcfile(flagProcfile)
	handleError(err)

	env, err := ReadEnv(flagEnv)
	handleError(err)

	concurrency, err := parseConcurrency(flagConcurrency)
	handleError(err)

	of := NewOutletFactory()
	of.Padding = pf.LongestProcessName()

	f := &Forego{
		teardown: make(chan struct{}),
	}

	go f.monitorInterrupt()

	var singleton string = ""
	if len(args) > 0 {
		singleton = args[0]
		if !pf.HasProcess(singleton) {
			of.ErrorOutput(fmt.Sprintf("no such process: %s", singleton))
		}
	}

	for idx, proc := range pf.Entries {
		numProcs := 1
		if value, ok := concurrency[proc.Name]; ok {
			numProcs = value
		}
		for i := 0; i < numProcs; i++ {
			if (singleton == "") || (singleton == proc.Name) {
				f.startProcess(idx, i, proc, env, of)
			}
		}
	}

	<-f.teardown
	of.SystemOutput("shutting down")

	f.wg.Wait()
}
<commit_msg>Implement teardownNow with Kill on 2nd CTRL-C<commit_after>package main

import (
	"errors"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

const shutdownGraceTime = 3 * time.Second

var flagPort int
var flagConcurrency string
var flagRestart bool

var cmdStart = &Command{
	Run:   runStart,
	Usage: "start [process name] [-f procfile] [-e env] [-c concurrency] [-p port] [-r]",
	Short: "Start the application",
	Long: `
Start the application specified by a Procfile (defaults to ./Procfile)

Examples:

  forego start
  forego start web
  forego start -f Procfile.test -e .env.test
`,
}

func init() {
	cmdStart.Flag.StringVar(&flagProcfile, "f", "Procfile", "procfile")
	cmdStart.Flag.StringVar(&flagEnv, "e", "", "env")
	cmdStart.Flag.IntVar(&flagPort, "p", 5000, "port")
	cmdStart.Flag.StringVar(&flagConcurrency, "c", "", "concurrency")
	cmdStart.Flag.BoolVar(&flagRestart, "r", false, "restart")
}

func parseConcurrency(value string) (map[string]int, error) {
	concurrency := map[string]int{}
	if strings.TrimSpace(value) == "" {
		return concurrency, nil
	}

	parts := strings.Split(value, ",")
	for _, part := range parts {
		if !strings.Contains(part, "=") {
			return concurrency, errors.New("Parsing concurency")
		}

		nameValue := strings.Split(part, "=")
		n, v := strings.TrimSpace(nameValue[0]), strings.TrimSpace(nameValue[1])
		if n == "" || v == "" {
			return concurrency, errors.New("Parsing concurency")
		}

		numProcs, err := strconv.ParseInt(v, 10, 16)
		if err != nil {
			return concurrency, err
		}

		concurrency[n] = int(numProcs)
	}
	return concurrency, nil
}

type Forego struct {
	shutdown    sync.Once     // Closes teardown exactly once
	teardown    chan struct{} // barrier: closed when shutting down
	teardownNow chan struct{} // barrier: second CTRL-C. More urgent.

	wg sync.WaitGroup
}

func (f *Forego) SignalShutdown() {
	f.shutdown.Do(func() { close(f.teardown) })
}

func (f *Forego) monitorInterrupt() {
	handler := make(chan os.Signal, 1)
	signal.Notify(handler, os.Interrupt)

	first := true
	var once sync.Once

	for sig := range handler {
		switch sig {
		case os.Interrupt:
			fmt.Println("      | ctrl-c detected")

			if !first {
				once.Do(func() { close(f.teardownNow) })
			}
			f.SignalShutdown()
			first = false
		}
	}
}

func (f *Forego) startProcess(idx, procNum int, proc ProcfileEntry, env Env, of *OutletFactory) {
	port := flagPort + (idx * 100)

	ps := NewProcess(proc.Command, env)
	procName := fmt.Sprint(proc.Name, ".", procNum+1)
	ps.Env["PORT"] = strconv.Itoa(port)
	ps.Root = filepath.Dir(flagProcfile)
	ps.Stdin = nil
	ps.Stdout = of.CreateOutlet(procName, idx, false)
	ps.Stderr = of.CreateOutlet(procName, idx, true)

	of.SystemOutput(fmt.Sprintf("starting %s on port %d", procName, port))

	finished := make(chan struct{}) // closed on process exit

	ps.Start()
	go func() {
		defer close(finished)
		ps.Wait()
	}()

	f.wg.Add(1)
	go func() {
		defer f.wg.Done()

		// Prevent goroutine from exiting before process has finished.
		defer func() { <-finished }()

		select {
		case <-finished:
			if flagRestart {
				f.startProcess(idx, procNum, proc, env, of)
				return
			} else {
				f.SignalShutdown()
			}

		case <-f.teardown:
			// Forego tearing down

			if !osHaveSigTerm {
				of.SystemOutput(fmt.Sprintf("Killing %s", procName))
				ps.cmd.Process.Kill()
				return
			}

			of.SystemOutput(fmt.Sprintf("sending SIGTERM to %s", procName))
			ps.SendSigTerm()

			// Give the process a chance to exit, otherwise kill it.
			select {
			case <-time.After(shutdownGraceTime):
				of.SystemOutput(fmt.Sprintf("Killing %s", procName))
				ps.SendSigKill()
			case <-f.teardownNow:
				of.SystemOutput(fmt.Sprintf("Killing %s", procName))
				ps.SendSigKill()
			case <-finished:
			}
		}
	}()
}

func runStart(cmd *Command, args []string) {
	root := filepath.Dir(flagProcfile)

	if flagEnv == "" {
		flagEnv = filepath.Join(root, ".env")
	}

	pf, err := ReadProcfile(flagProcfile)
	handleError(err)

	env, err := ReadEnv(flagEnv)
	handleError(err)

	concurrency, err := parseConcurrency(flagConcurrency)
	handleError(err)

	of := NewOutletFactory()
	of.Padding = pf.LongestProcessName()

	f := &Forego{
		teardown:    make(chan struct{}),
		teardownNow: make(chan struct{}),
	}

	go f.monitorInterrupt()

	var singleton string = ""
	if len(args) > 0 {
		singleton = args[0]
		if !pf.HasProcess(singleton) {
			of.ErrorOutput(fmt.Sprintf("no such process: %s", singleton))
		}
	}

	for idx, proc := range pf.Entries {
		numProcs := 1
		if value, ok := concurrency[proc.Name]; ok {
			numProcs = value
		}
		for i := 0; i < numProcs; i++ {
			if (singleton == "") || (singleton == proc.Name) {
				f.startProcess(idx, i, proc, env, of)
			}
		}
	}

	<-f.teardown
	of.SystemOutput("shutting down")

	f.wg.Wait()
}
<|endoftext|>


<commit_before>package templar

import (
	"fmt"
	"github.com/amir/raidman"
	"net/http"
	"strings"
	"time"
)

type DebugStats struct{}

func (d *DebugStats) StartRequest(req *http.Request) {
	fmt.Printf("[%s] S %s %s\n", time.Now().Format(time.RFC3339Nano), req.Method, req.URL)
}

func (d *DebugStats) Emit(req *http.Request, dur time.Duration) {
	fmt.Printf("[%s] E %s %s (%s)\n", time.Now().Format(time.RFC3339Nano), req.Method, req.URL, dur)
}

func (d *DebugStats) RequestTimeout(req *http.Request, timeout time.Duration) {
	fmt.Printf("[%s] T %s %s (%s)\n", time.Now().Format(time.RFC3339Nano), req.Method, req.URL, timeout)
}

var _ = Stats(&DebugStats{})

type StatsdOutput struct {
	client StatsdClient
}

var _ = Stats(&StatsdOutput{})

func NewStatsdOutput(client StatsdClient) *StatsdOutput {
	return &StatsdOutput{client}
}

func (s *StatsdOutput) url(req *http.Request) string {
	return req.Host + strings.Replace(req.URL.Path, "/", "-", -1)
}

func (s *StatsdOutput) StartRequest(req *http.Request) {
	s.client.Incr("templar.request.method."+req.Method, 1)
	s.client.Incr("templar.request.host."+req.Host, 1)
	s.client.Incr("templar.request.url."+s.url(req), 1)
	s.client.GaugeDelta("templar.requests.active", 1)
}

func (s *StatsdOutput) Emit(req *http.Request, delta time.Duration) {
	s.client.GaugeDelta("templar.requests.active", -1)
	s.client.PrecisionTiming("templar.request.url."+s.url(req), delta)
}

func (s *StatsdOutput) RequestTimeout(req *http.Request, timeout time.Duration) {
	s.client.Incr("templar.timeout.host."+req.Host, 1)
	s.client.Incr("templar.timeout.url."+s.url(req), 1)
}

type RiemannOutput struct {
	client RiemannClient
}

func NewRiemannOutput(client RiemannClient) *RiemannOutput {
	return &RiemannOutput{client}
}

func (r *RiemannOutput) StartRequest(req *http.Request) {
	attributes := make(map[string]string)
	attributes["method"] = req.Method
	attributes["http-host"] = req.Host
	attributes["path"] = req.URL.Path
	var event = &raidman.Event{
		State:      "ok",
		Service:    "templar request",
		Metric:     1,
		Attributes: attributes,
	}
	r.client.Send(event)
}

func (r *RiemannOutput) Emit(req *http.Request, delta time.Duration) {
	attributes := make(map[string]string)
	attributes["method"] = req.Method
	attributes["http-host"] = req.Host
	attributes["path"] = req.URL.Path
	var event = &raidman.Event{
		State:      "ok",
		Service:    "templar response",
		Metric:     1000.0 * delta.Seconds(),
		Attributes: attributes,
	}
	r.client.Send(event)
}

func (r *RiemannOutput) RequestTimeout(req *http.Request, timeout time.Duration) {
	attributes := make(map[string]string)
	attributes["method"] = req.Method
	attributes["host"] = req.Host
	attributes["path"] = req.URL.Path
	var event = &raidman.Event{
		State:      "warning",
		Service:    "templar timeout",
		Metric:     timeout.Seconds() * 1000.0,
		Attributes: attributes,
	}
	r.client.Send(event)
}

type MultiStats []Stats

var _ = Stats(MultiStats{})

func (m MultiStats) StartRequest(req *http.Request) {
	for _, s := range m {
		s.StartRequest(req)
	}
}

func (m MultiStats) Emit(req *http.Request, t time.Duration) {
	for _, s := range m {
		s.Emit(req, t)
	}
}

func (m MultiStats) RequestTimeout(req *http.Request, timeout time.Duration) {
	for _, s := range m {
		s.RequestTimeout(req, timeout)
	}
}
<commit_msg>prefix all riemann http custom fields with http-<commit_after>package templar

import (
	"fmt"
	"github.com/amir/raidman"
	"net/http"
	"strings"
	"time"
)

type DebugStats struct{}

func (d *DebugStats) StartRequest(req *http.Request) {
	fmt.Printf("[%s] S %s %s\n", time.Now().Format(time.RFC3339Nano), req.Method, req.URL)
}

func (d *DebugStats) Emit(req *http.Request, dur time.Duration) {
	fmt.Printf("[%s] E %s %s (%s)\n", time.Now().Format(time.RFC3339Nano), req.Method, req.URL, dur)
}

func (d *DebugStats) RequestTimeout(req *http.Request, timeout time.Duration) {
	fmt.Printf("[%s] T %s %s (%s)\n", time.Now().Format(time.RFC3339Nano), req.Method, req.URL, timeout)
}

var _ = Stats(&DebugStats{})

type StatsdOutput struct {
	client StatsdClient
}

var _ = Stats(&StatsdOutput{})

func NewStatsdOutput(client StatsdClient) *StatsdOutput {
	return &StatsdOutput{client}
}

func (s *StatsdOutput) url(req *http.Request) string {
	return req.Host + strings.Replace(req.URL.Path, "/", "-", -1)
}

func (s *StatsdOutput) StartRequest(req *http.Request) {
	s.client.Incr("templar.request.method."+req.Method, 1)
	s.client.Incr("templar.request.host."+req.Host, 1)
	s.client.Incr("templar.request.url."+s.url(req), 1)
	s.client.GaugeDelta("templar.requests.active", 1)
}

func (s *StatsdOutput) Emit(req *http.Request, delta time.Duration) {
	s.client.GaugeDelta("templar.requests.active", -1)
	s.client.PrecisionTiming("templar.request.url."+s.url(req), delta)
}

func (s *StatsdOutput) RequestTimeout(req *http.Request, timeout time.Duration) {
	s.client.Incr("templar.timeout.host."+req.Host, 1)
	s.client.Incr("templar.timeout.url."+s.url(req), 1)
}

type RiemannOutput struct {
	client RiemannClient
}

func NewRiemannOutput(client RiemannClient) *RiemannOutput {
	return &RiemannOutput{client}
}

func (r *RiemannOutput) StartRequest(req *http.Request) {
	attributes := make(map[string]string)
	attributes["http-method"] = req.Method
	attributes["http-host"] = req.Host
	attributes["http-path"] = req.URL.Path
	var event = &raidman.Event{
		State:      "ok",
		Service:    "templar request",
		Metric:     1,
		Attributes: attributes,
	}
	r.client.Send(event)
}

func (r *RiemannOutput) Emit(req *http.Request, delta time.Duration) {
	attributes := make(map[string]string)
	attributes["http-method"] = req.Method
	attributes["http-host"] = req.Host
	attributes["http-path"] = req.URL.Path
	var event = &raidman.Event{
		State:      "ok",
		Service:    "templar response",
		Metric:     1000.0 * delta.Seconds(),
		Attributes: attributes,
	}
	r.client.Send(event)
}

func (r *RiemannOutput) RequestTimeout(req *http.Request, timeout time.Duration) {
	attributes := make(map[string]string)
	attributes["method"] = req.Method
	attributes["host"] = req.Host
	attributes["path"] = req.URL.Path
	var event = &raidman.Event{
		State:      "warning",
		Service:    "templar timeout",
		Metric:     timeout.Seconds() * 1000.0,
		Attributes: attributes,
	}
	r.client.Send(event)
}

type MultiStats []Stats

var _ = Stats(MultiStats{})

func (m MultiStats) StartRequest(req *http.Request) {
	for _, s := range m {
		s.StartRequest(req)
	}
}

func (m MultiStats) Emit(req *http.Request, t time.Duration) {
	for _, s := range m {
		s.Emit(req, t)
	}
}

func (m MultiStats) RequestTimeout(req *http.Request, timeout time.Duration) {
	for _, s := range m {
		s.RequestTimeout(req, timeout)
	}
}
<|endoftext|>


<commit_before>package memkv

import (
	"path"
	"sort"
	"strings"
	"sync"

	"github.com/derekparker/trie"
)

// A Store represents an in-memory key-value store safe for
// concurrent access.
type Store struct {
	FuncMap map[string]interface{}
	*sync.RWMutex
	t *trie.Trie
}

func New() Store {
	s := Store{
		t:       trie.New(),
		RWMutex: &sync.RWMutex{},
	}
	s.FuncMap = map[string]interface{}{
		"exists":    s.Exists,
		"ls":        s.List,
		"lsdir":     s.ListDir,
		"get":       s.Get,
		"gets":      s.GetAll,
		"getallkvs": s.GetAllKVs,
		"getv":      s.GetValue,
		"getvs":     s.GetAllValues,
	}
	return s
}

// Delete deletes the KVPair associated with key.
func (s Store) Del(key string) {
	s.Lock()
	s.t.Remove(key)
	s.Unlock()
}

// Exists checks for the existence of key in the store.
func (s Store) Exists(key string) bool {
	s.RLock()
	_, ok := s.t.Find(key)
	s.RUnlock()
	return ok
}

// Get gets the KVPair associated with key. If there is no KVPair
// associated with key, Get returns KVPair{}.
func (s Store) Get(key string) KVPair {
	s.RLock()
	defer s.RUnlock()
	node, ok := s.t.Find(key)
	if !ok {
		return KVPair{}
	}
	return node.Meta().(KVPair)
}

// GetAll returns a KVPair for all nodes with keys matching pattern.
// The syntax of patterns is the same as in path.Match.
func (s Store) GetAll(pattern string) KVPairs {
	ks := make(KVPairs, 0)
	s.RLock()
	defer s.RUnlock()
	for _, k := range s.t.Keys() {
		m, err := path.Match(pattern, k)
		if err != nil {
			return nil
		}
		kv := s.Get(k)
		if m {
			ks = append(ks, kv)
		}
	}
	if len(ks) == 0 {
		return nil
	}
	sort.Sort(ks)
	return ks
}

func (s Store) GetAllValues(pattern string) []string {
	vs := make([]string, 0)
	for _, kv := range s.GetAll(pattern) {
		vs = append(vs, kv.Value)
	}
	sort.Strings(vs)
	return vs
}

// GetAllKVs returns all KV-Pairs
func (s Store) GetAllKVs() KVPairs {
	ks := make(KVPairs, 0)
	s.RLock()
	defer s.RUnlock()
	for _, k := range s.t.Keys() {
		ks = append(ks, s.Get(k))
	}
	sort.Sort(ks)
	return ks
}

// GetValue gets the value associated with key. If there are no values
// associated with key, GetValue returns "".
func (s Store) GetValue(key string, v ...string) string {
	defaultValue := ""
	if len(v) > 0 {
		defaultValue = v[0]
	}
	kv := s.Get(key)
	if kv.Key == "" {
		return defaultValue
	}
	return kv.Value
}

func (s Store) list(filePath string, dir bool) []string {
	vs := make([]string, 0)
	m := make(map[string]bool)
	// The prefix search should only return dirs
	filePath = path.Clean(filePath) + "/"
	s.RLock()
	defer s.RUnlock()
	for _, k := range s.t.PrefixSearch(filePath) {
		items := strings.Split(stripKey(k, filePath), "/")
		if dir {
			if len(items) < 2 {
				continue
			}
		}
		m[items[0]] = true
	}
	for k := range m {
		vs = append(vs, k)
	}
	sort.Strings(vs)
	return vs
}

func (s Store) List(filePath string) []string {
	return s.list(filePath, false)
}

func (s Store) ListDir(filePath string) []string {
	return s.list(filePath, true)
}

// Set sets the KVPair entry associated with key to value.
func (s Store) Set(key string, value string) {
	s.Lock()
	s.t.Add(key, KVPair{key, value})
	s.Unlock()
}

func (s Store) Purge() {
	s.Lock()
	for _, k := range s.t.Keys() {
		s.t.Remove(k)
	}
	s.Unlock()
}

func stripKey(key, prefix string) string {
	return strings.TrimPrefix(strings.TrimPrefix(key, prefix), "/")
}
<commit_msg>use my fork of trie for now so that we can delete the root node<commit_after>package memkv

import (
	"path"
	"sort"
	"strings"
	"sync"

	"github.com/HeavyHorst/trie"
)

// A Store represents an in-memory key-value store safe for
// concurrent access.
type Store struct {
	FuncMap map[string]interface{}
	*sync.RWMutex
	t *trie.Trie
}

func New() Store {
	s := Store{
		t:       trie.New(),
		RWMutex: &sync.RWMutex{},
	}
	s.FuncMap = map[string]interface{}{
		"exists":    s.Exists,
		"ls":        s.List,
		"lsdir":     s.ListDir,
		"get":       s.Get,
		"gets":      s.GetAll,
		"getallkvs": s.GetAllKVs,
		"getv":      s.GetValue,
		"getvs":     s.GetAllValues,
	}
	return s
}

// Delete deletes the KVPair associated with key.
func (s Store) Del(key string) {
	s.Lock()
	s.t.Remove(key)
	s.Unlock()
}

// Exists checks for the existence of key in the store.
func (s Store) Exists(key string) bool {
	s.RLock()
	_, ok := s.t.Find(key)
	s.RUnlock()
	return ok
}

// Get gets the KVPair associated with key. If there is no KVPair
// associated with key, Get returns KVPair{}.
func (s Store) Get(key string) KVPair {
	s.RLock()
	defer s.RUnlock()
	node, ok := s.t.Find(key)
	if !ok {
		return KVPair{}
	}
	return node.Meta().(KVPair)
}

// GetAll returns a KVPair for all nodes with keys matching pattern.
// The syntax of patterns is the same as in path.Match.
func (s Store) GetAll(pattern string) KVPairs {
	ks := make(KVPairs, 0)
	s.RLock()
	defer s.RUnlock()
	for _, k := range s.t.Keys() {
		m, err := path.Match(pattern, k)
		if err != nil {
			return nil
		}
		kv := s.Get(k)
		if m {
			ks = append(ks, kv)
		}
	}
	if len(ks) == 0 {
		return nil
	}
	sort.Sort(ks)
	return ks
}

func (s Store) GetAllValues(pattern string) []string {
	vs := make([]string, 0)
	for _, kv := range s.GetAll(pattern) {
		vs = append(vs, kv.Value)
	}
	sort.Strings(vs)
	return vs
}

// GetAllKVs returns all KV-Pairs
func (s Store) GetAllKVs() KVPairs {
	ks := make(KVPairs, 0)
	s.RLock()
	defer s.RUnlock()
	for _, k := range s.t.Keys() {
		ks = append(ks, s.Get(k))
	}
	sort.Sort(ks)
	return ks
}

// GetValue gets the value associated with key. If there are no values
// associated with key, GetValue returns "".
func (s Store) GetValue(key string, v ...string) string {
	defaultValue := ""
	if len(v) > 0 {
		defaultValue = v[0]
	}
	kv := s.Get(key)
	if kv.Key == "" {
		return defaultValue
	}
	return kv.Value
}

func (s Store) list(filePath string, dir bool) []string {
	vs := make([]string, 0)
	m := make(map[string]bool)
	// The prefix search should only return dirs
	filePath = path.Clean(filePath) + "/"
	s.RLock()
	defer s.RUnlock()
	for _, k := range s.t.PrefixSearch(filePath) {
		items := strings.Split(stripKey(k, filePath), "/")
		if dir {
			if len(items) < 2 {
				continue
			}
		}
		m[items[0]] = true
	}
	for k := range m {
		vs = append(vs, k)
	}
	sort.Strings(vs)
	return vs
}

func (s Store) List(filePath string) []string {
	return s.list(filePath, false)
}

func (s Store) ListDir(filePath string) []string {
	return s.list(filePath, true)
}

// Set sets the KVPair entry associated with key to value.
func (s Store) Set(key string, value string) {
	s.Lock()
	s.t.Add(key, KVPair{key, value})
	s.Unlock()
}

func (s Store) Purge() {
	s.Lock()
	for _, k := range s.t.Keys() {
		s.t.Remove(k)
	}
	s.Unlock()
}

func stripKey(key, prefix string) string {
	return strings.TrimPrefix(strings.TrimPrefix(key, prefix), "/")
}
<|endoftext|>


<commit_before>package arel

import (
	"bytes"
	"strconv"
)

type Table struct {
	Name    string
	Engine  Engine
	Aliases *[]TableAliasNode
}

func NewTable(name string, e Engine) Table {
	table := Table{Name: name, Engine: e}
	return table
}

func (t *Table) Project(a ...AstNode) *SelectManager {
	return t.SelectManager().Project(a...)
}

func (t *Table) Take(i int) *SelectManager {
	return t.SelectManager().Take(i)
}

func (t *Table) Where(n AstNode) *SelectManager {
	return t.SelectManager().Where(n)
}

func (t *Table) Skip(i int) *SelectManager {
	return t.SelectManager().Skip(i)
}

func (t *Table) Offset(i int) *SelectManager {
	return t.SelectManager().Offset(i)
}

func (t *Table) Having(a ...AstNode) *SelectManager {
	return t.SelectManager().Having(a...)
}

func (t *Table) Group(a ...AstNode) *SelectManager {
	return t.SelectManager().Group(a...)
}

func (t *Table) Order(exprs ...interface{}) *SelectManager {
	return t.SelectManager().Order(exprs...)
}

func (t *Table) CreateStringJoin(left string) StringJoinNode {
	return StringJoinNode{
		Left: left,
	}
}

func (t *Table) CreateInnerJoin(left *Table, right *Table) InnerJoinNode {
	return InnerJoinNode{
		Left:  left,
		Right: right,
	}
}

func (t *Table) CreateOuterJoin(left *Table, right *Table) OuterJoinNode {
	return OuterJoinNode{
		Left:  left,
		Right: right,
	}
}

func (t *Table) SelectManager() *SelectManager {
	manager := NewSelectManager(t)
	return &manager
}

func (t *Table) InsertManager() *InsertManager {
	manager := NewInsertManager(t)
	return &manager
}

func (t *Table) Alias() TableAliasNode {
	var buf bytes.Buffer
	if t.Aliases == nil {
		aliases := make([]TableAliasNode, 0)
		t.Aliases = &aliases
	}
	n := len(*t.Aliases)
	buf.WriteString(t.Name)
	buf.WriteString("_")
	buf.WriteString(strconv.Itoa(n + 2))
	alias := NewTableAliasNode(t, buf.String())
	*t.Aliases = append(*t.Aliases, alias)
	return alias
}

func (t *Table) Attr(name string) AttributeNode {
	return NewAttributeNode(name, t)
}
<commit_msg>Created Table#Select as an alias to Table#Project<commit_after>package arel

import (
	"bytes"
	"strconv"
)

type Table struct {
	Name    string
	Engine  Engine
	Aliases *[]TableAliasNode
}

func NewTable(name string, e Engine) Table {
	table := Table{Name: name, Engine: e}
	return table
}

func (t *Table) Project(a ...AstNode) *SelectManager {
	return t.SelectManager().Project(a...)
}

func (t *Table) Select(a ...AstNode) *SelectManager {
	return t.Project(a...)
}

func (t *Table) Take(i int) *SelectManager {
	return t.SelectManager().Take(i)
}

func (t *Table) Where(n AstNode) *SelectManager {
	return t.SelectManager().Where(n)
}

func (t *Table) Skip(i int) *SelectManager {
	return t.SelectManager().Skip(i)
}

func (t *Table) Offset(i int) *SelectManager {
	return t.SelectManager().Offset(i)
}

func (t *Table) Having(a ...AstNode) *SelectManager {
	return t.SelectManager().Having(a...)
}

func (t *Table) Group(a ...AstNode) *SelectManager {
	return t.SelectManager().Group(a...)
}

func (t *Table) Order(exprs ...interface{}) *SelectManager {
	return t.SelectManager().Order(exprs...)
}

func (t *Table) CreateStringJoin(left string) StringJoinNode {
	return StringJoinNode{
		Left: left,
	}
}

func (t *Table) CreateInnerJoin(left *Table, right *Table) InnerJoinNode {
	return InnerJoinNode{
		Left:  left,
		Right: right,
	}
}

func (t *Table) CreateOuterJoin(left *Table, right *Table) OuterJoinNode {
	return OuterJoinNode{
		Left:  left,
		Right: right,
	}
}

func (t *Table) SelectManager() *SelectManager {
	manager := NewSelectManager(t)
	return &manager
}

func (t *Table) InsertManager() *InsertManager {
	manager := NewInsertManager(t)
	return &manager
}

func (t *Table) Alias() TableAliasNode {
	var buf bytes.Buffer
	if t.Aliases == nil {
		aliases := make([]TableAliasNode, 0)
		t.Aliases = &aliases
	}
	n := len(*t.Aliases)
	buf.WriteString(t.Name)
	buf.WriteString("_")
	buf.WriteString(strconv.Itoa(n + 2))
	alias := NewTableAliasNode(t, buf.String())
	*t.Aliases = append(*t.Aliases, alias)
	return alias
}

func (t *Table) Attr(name string) AttributeNode {
	return NewAttributeNode(name, t)
}
<|endoftext|>


<commit_before>package sqlite3

import "C"
import "fmt"

type Table struct {
	Name		string
	ColumnSpec	string
}

func (t *Table) Create(db *Database) (e error) {
	sql := fmt.Sprintf("CREATE TABLE %v (%v);", t.Name, t.ColumnSpec)
	_, e = db.Execute(sql)
	return
}

func (t *Table) Drop(db *Database) (e error) {
	sql := fmt.Sprintf("DROP TABLE IF EXISTS %v;", t.Name, t.ColumnSpec)
	_, e = db.Execute(sql)
	return
}

func (t *Table) Rows(db *Database) (c int, e error) {
	sql := fmt.Sprintf("SELECT Count(*) FROM %v;", t.Name)
	_, e = db.Execute(sql, func(s *Statement, values ...interface{}) {
		c = int(values[0].(int64))
	})
	return
}<commit_msg>Table documentation added.<commit_after>package sqlite3

import "C"
import "fmt"

// Table implements a high level view of a SQL table.
type Table struct {
	Name		string
	ColumnSpec	string
}

// Create is used to create a SQL table.
func (t *Table) Create(db *Database) (e error) {
	sql := fmt.Sprintf("CREATE TABLE %v (%v);", t.Name, t.ColumnSpec)
	_, e = db.Execute(sql)
	return
}

// Drop is used to delete a SQL table.
func (t *Table) Drop(db *Database) (e error) {
	sql := fmt.Sprintf("DROP TABLE IF EXISTS %v;", t.Name, t.ColumnSpec)
	_, e = db.Execute(sql)
	return
}

// Rows returns the number of rows in the table.
func (t *Table) Rows(db *Database) (c int, e error) {
	sql := fmt.Sprintf("SELECT Count(*) FROM %v;", t.Name)
	_, e = db.Execute(sql, func(s *Statement, values ...interface{}) {
		c = int(values[0].(int64))
	})
	return
}
<|endoftext|>