<commit_before>package router

import (
	"bytes"
	"code.google.com/p/gopacket/layers"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"net"
	"syscall"
	"time"
)

const macMaxAge = 10 * time.Minute // [1]

// [1] should be greater than typical ARP cache expiries, i.e. > 3/2 *
// /proc/sys/net/ipv4_neigh/*/base_reachable_time_ms on Linux

type Router struct {
	Iface           *net.Interface
	Ourself         *LocalPeer
	Macs            *MacCache
	Peers           *Peers
	Routes          *Routes
	ConnectionMaker *ConnectionMaker
	GossipChannels  map[uint32]*GossipChannel
	TopologyGossip  Gossip
	UDPListener     *net.UDPConn
	Password        []byte
	ConnLimit       int
	BufSz           int
	LogFrame        func(string, []byte, *layers.Ethernet)
}

type PacketSource interface {
	ReadPacket() ([]byte, error)
}

type PacketSink interface {
	WritePacket([]byte) error
}

type PacketSourceSink interface {
	PacketSource
	PacketSink
}

func NewRouter(iface *net.Interface, name PeerName, nickName string, password []byte, connLimit int, bufSz int, logFrame func(string, []byte, *layers.Ethernet)) *Router {
	router := &Router{
		Iface:          iface,
		GossipChannels: make(map[uint32]*GossipChannel),
		Password:       password,
		ConnLimit:      connLimit,
		BufSz:          bufSz,
		LogFrame:       logFrame}
	onMacExpiry := func(mac net.HardwareAddr, peer *Peer) {
		log.Println("Expired MAC", mac, "at", peer.FullName())
	}
	onPeerGC := func(peer *Peer) {
		router.Macs.Delete(peer)
		log.Println("Removed unreachable peer", peer.FullName())
	}
	router.Ourself = NewLocalPeer(name, nickName, router)
	router.Macs = NewMacCache(macMaxAge, onMacExpiry)
	router.Peers = NewPeers(router.Ourself.Peer, onPeerGC)
	router.Peers.FetchWithDefault(router.Ourself.Peer)
	router.Routes = NewRoutes(router.Ourself.Peer, router.Peers)
	router.ConnectionMaker = NewConnectionMaker(router.Ourself, router.Peers)
	router.TopologyGossip = router.NewGossip("topology", router)
	return router
}

func (router *Router) Start() {
	// we need two pcap handles since they aren't thread-safe
	pio, err := NewPcapIO(router.Iface.Name, router.BufSz)
	checkFatal(err)
	po, err := NewPcapO(router.Iface.Name)
	checkFatal(err)
	router.Ourself.Start()
	router.Macs.Start()
	router.Routes.Start()
	router.ConnectionMaker.Start()
	router.UDPListener = router.listenUDP(Port, po)
	router.listenTCP(Port)
	router.sniff(pio)
}

func (router *Router) UsingPassword() bool {
	return router.Password != nil
}

func (router *Router) Status() string {
	var buf bytes.Buffer
	fmt.Fprintln(&buf, "Our name is", router.Ourself.FullName())
	fmt.Fprintln(&buf, "Sniffing traffic on", router.Iface)
	fmt.Fprintf(&buf, "MACs:\n%s", router.Macs)
	fmt.Fprintf(&buf, "Peers:\n%s", router.Peers)
	fmt.Fprintf(&buf, "Routes:\n%s", router.Routes)
	fmt.Fprintf(&buf, "Reconnects:\n%s", router.ConnectionMaker)
	return buf.String()
}

func (router *Router) sniff(pio PacketSourceSink) {
	log.Println("Sniffing traffic on", router.Iface)

	dec := NewEthernetDecoder()
	mac := router.Iface.HardwareAddr
	if router.Macs.Enter(mac, router.Ourself.Peer) {
		log.Println("Discovered our MAC", mac)
	}
	go func() {
		for {
			pkt, err := pio.ReadPacket()
			checkFatal(err)
			router.LogFrame("Sniffed", pkt, nil)
			router.handleCapturedPacket(pkt, dec, pio)
		}
	}()
}

func (router *Router) handleCapturedPacket(frameData []byte, dec *EthernetDecoder, po PacketSink) {
	dec.DecodeLayers(frameData)
	decodedLen := len(dec.decoded)
	if decodedLen == 0 {
		return
	}
	srcMac := dec.eth.SrcMAC
	srcPeer, found := router.Macs.Lookup(srcMac)
	// We need to filter out frames we injected ourselves. For such
	// frames, the srcMAC will have been recorded as associated with a
	// different peer.
	if found && srcPeer != router.Ourself.Peer {
		return
	}
	if router.Macs.Enter(srcMac, router.Ourself.Peer) {
		log.Println("Discovered local MAC", srcMac)
	}
	if dec.DropFrame() {
		return
	}
	dstMac := dec.eth.DstMAC
	dstPeer, found := router.Macs.Lookup(dstMac)
	if found && dstPeer == router.Ourself.Peer {
		return
	}
	df := decodedLen == 2 && (dec.ip.Flags&layers.IPv4DontFragment != 0)
	if df {
		router.LogFrame("Forwarding DF", frameData, &dec.eth)
	} else {
		router.LogFrame("Forwarding", frameData, &dec.eth)
	}
	// at this point we are handing over the frame to forwarders, so
	// we need to make a copy of it in order to prevent the next
	// capture from overwriting the data
	frameLen := len(frameData)
	frameCopy := make([]byte, frameLen, frameLen)
	copy(frameCopy, frameData)

	if !found {
		router.Ourself.Broadcast(df, frameCopy, dec)
		return
	}
	err := router.Ourself.Forward(dstPeer, df, frameCopy, dec)
	if ftbe, ok := err.(FrameTooBigError); ok {
		err = dec.sendICMPFragNeeded(ftbe.EPMTU, po.WritePacket)
	}
	checkWarn(err)
}

func (router *Router) listenTCP(localPort int) {
	localAddr, err := net.ResolveTCPAddr("tcp4", fmt.Sprint(":", localPort))
	checkFatal(err)
	ln, err := net.ListenTCP("tcp4", localAddr)
	checkFatal(err)
	go func() {
		defer ln.Close()
		for {
			tcpConn, err := ln.AcceptTCP()
			if err != nil {
				log.Println(err)
				continue
			}
			router.acceptTCP(tcpConn)
		}
	}()
}

func (router *Router) acceptTCP(tcpConn *net.TCPConn) {
	// someone else is dialing us, so our udp sender is the conn
	// on Port and we wait for them to send us something on UDP to
	// start.
	remoteAddrStr := tcpConn.RemoteAddr().String()
	log.Printf("->[%s] connection accepted\n", remoteAddrStr)
	connRemote := NewRemoteConnection(router.Ourself.Peer, nil, remoteAddrStr, false, false)
	connLocal := NewLocalConnection(connRemote, tcpConn, nil, router)
	connLocal.Start(true)
}

func (router *Router) listenUDP(localPort int, po PacketSink) *net.UDPConn {
	localAddr, err := net.ResolveUDPAddr("udp4", fmt.Sprint(":", localPort))
	checkFatal(err)
	conn, err := net.ListenUDP("udp4", localAddr)
	checkFatal(err)
	f, err := conn.File()
	defer f.Close()
	checkFatal(err)
	fd := int(f.Fd())
	// This one makes sure all packets we send out do not have DF set on them.
	err = syscall.SetsockoptInt(fd, syscall.IPPROTO_IP, syscall.IP_MTU_DISCOVER, syscall.IP_PMTUDISC_DONT)
	checkFatal(err)
	go router.udpReader(conn, po)
	return conn
}

func (router *Router) udpReader(conn *net.UDPConn, po PacketSink) {
	defer conn.Close()
	dec := NewEthernetDecoder()
	buf := make([]byte, MaxUDPPacketSize)
	for {
		n, sender, err := conn.ReadFromUDP(buf)
		if err == io.EOF {
			return
		} else if err != nil {
			log.Println("ignoring UDP read error", err)
			continue
		} else if n < NameSize {
			log.Println("ignoring too short UDP packet from", sender)
			continue
		}
		name := PeerNameFromBin(buf[:NameSize])
		packet := make([]byte, n-NameSize)
		copy(packet, buf[NameSize:n])
		peerConn, found := router.Ourself.ConnectionTo(name)
		if !found {
			continue
		}
		relayConn, ok := peerConn.(*LocalConnection)
		if !ok {
			continue
		}
		err = relayConn.Decryptor.IterateFrames(packet, router.handleUDPPacketFunc(dec, sender, po))
		if pde, ok := err.(PacketDecodingError); ok {
			if pde.Fatal {
				relayConn.Shutdown(pde)
			} else {
				relayConn.Log(pde.Error())
			}
		} else {
			checkWarn(err)
		}
	}
}

func (router *Router) handleUDPPacketFunc(dec *EthernetDecoder, sender *net.UDPAddr, po PacketSink) FrameConsumer {
	return func(relayConn *LocalConnection, srcNameByte, dstNameByte []byte, frame []byte) {
		srcPeer, found := router.Peers.Fetch(PeerNameFromBin(srcNameByte))
		if !found {
			return
		}
		dstPeer, found := router.Peers.Fetch(PeerNameFromBin(dstNameByte))
		if !found {
			return
		}

		dec.DecodeLayers(frame)
		decodedLen := len(dec.decoded)
		if decodedLen == 0 {
			return
		}
		// Handle special frames produced internally (rather than
		// captured/forwarded) by the remote router.
		//
		// We really shouldn't be decoding these above, since they are
		// not genuine Ethernet frames. However, it is actually more
		// efficient to do so, as we want to optimise for the common
		// (i.e. non-special) frames. These always need decoding, and
		// detecting special frames is cheaper post decoding than pre.
		if decodedLen == 1 && dec.IsSpecial() {
			if srcPeer == relayConn.Remote() && dstPeer == router.Ourself.Peer {
				handleSpecialFrame(relayConn, sender, frame)
			}
		}

		df := decodedLen == 2 && (dec.ip.Flags&layers.IPv4DontFragment != 0)

		if dstPeer != router.Ourself.Peer {
			// it's not for us, we're just relaying it
			if df {
				router.LogFrame("Relaying DF", frame, &dec.eth)
			} else {
				router.LogFrame("Relaying", frame, &dec.eth)
			}

			err := router.Ourself.Relay(srcPeer, dstPeer, df, frame, dec)
			if ftbe, ok := err.(FrameTooBigError); ok {
				err = dec.sendICMPFragNeeded(ftbe.EPMTU, func(icmpFrame []byte) error {
					return router.Ourself.Forward(srcPeer, false, icmpFrame, nil)
				})
			}

			checkWarn(err)
			return
		}

		srcMac := dec.eth.SrcMAC
		dstMac := dec.eth.DstMAC

		if router.Macs.Enter(srcMac, srcPeer) {
			log.Println("Discovered remote MAC", srcMac, "at", srcPeer.FullName())
		}
		router.LogFrame("Injecting", frame, &dec.eth)
		checkWarn(po.WritePacket(frame))

		dstPeer, found = router.Macs.Lookup(dstMac)
		if !found || dstPeer != router.Ourself.Peer {
			router.LogFrame("Relaying broadcast", frame, &dec.eth)
			router.Ourself.RelayBroadcast(srcPeer, df, frame, dec)
		}
	}
}

func handleSpecialFrame(relayConn *LocalConnection, sender *net.UDPAddr, frame []byte) {
	frameLen := len(frame)
	switch {
	case frameLen == EthernetOverhead+8:
		relayConn.ReceivedHeartbeat(sender, binary.BigEndian.Uint64(frame[EthernetOverhead:]))
	case frameLen == FragTestSize && bytes.Equal(frame, FragTest):
		relayConn.SendProtocolMsg(ProtocolMsg{ProtocolFragmentationReceived, nil})
	case frameLen == PMTUDiscoverySize && bytes.Equal(frame, PMTUDiscovery):
	default:
		frameLenBytes := []byte{0, 0}
		binary.BigEndian.PutUint16(frameLenBytes, uint16(frameLen-EthernetOverhead))
		relayConn.SendProtocolMsg(ProtocolMsg{ProtocolPMTUVerified, frameLenBytes})
	}
}

// Gossiper methods - the Router is the topology Gossiper

func (router *Router) OnGossipUnicast(sender PeerName, msg []byte) error {
	return fmt.Errorf("unexpected topology gossip unicast: %v", msg)
}

func (router *Router) OnGossipBroadcast(msg []byte) error {
	return fmt.Errorf("unexpected topology gossip broadcast: %v", msg)
}

type PeerNameSet map[PeerName]bool

type TopologyGossipData struct {
	peers  *Peers
	update PeerNameSet
}

func NewTopologyGossipData(peers *Peers, update ...*Peer) *TopologyGossipData {
	names := make(PeerNameSet)
	for _, p := range update {
		names[p.Name] = true
	}
	return &TopologyGossipData{peers: peers, update: names}
}

func (d *TopologyGossipData) Merge(other GossipData) {
	for name := range other.(*TopologyGossipData).update {
		d.update[name] = true
	}
}

func (d *TopologyGossipData) Encode() []byte {
	return d.peers.EncodePeers(d.update)
}

func (router *Router) Gossip() GossipData {
	return &TopologyGossipData{peers: router.Peers, update: router.Peers.Names()}
}

func (router *Router) OnGossip(buf []byte) (GossipData, error) {
	newUpdate, err := router.Peers.ApplyUpdate(buf)
	if _, ok := err.(UnknownPeerError); err != nil && ok {
		// That update contained a reference to a peer which wasn't
		// itself included in the update, and we didn't know about
		// already. We ignore this; eventually we should receive an
		// update containing a complete topology.
		log.Println("Topology gossip:", err)
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	if len(newUpdate) == 0 {
		return nil, nil
	}
	router.ConnectionMaker.Refresh()
	router.Routes.Recalculate()
	return &TopologyGossipData{peers: router.Peers, update: newUpdate}, nil
}
<commit_msg>oops I broke this in 1eca5e1097d7480d82248bf8cdb297d716f79506<commit_after>package router

import (
	"bytes"
	"code.google.com/p/gopacket/layers"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"net"
	"syscall"
	"time"
)

const macMaxAge = 10 * time.Minute // [1]

// [1] should be greater than typical ARP cache expiries, i.e. > 3/2 *
// /proc/sys/net/ipv4_neigh/*/base_reachable_time_ms on Linux

type Router struct {
	Iface           *net.Interface
	Ourself         *LocalPeer
	Macs            *MacCache
	Peers           *Peers
	Routes          *Routes
	ConnectionMaker *ConnectionMaker
	GossipChannels  map[uint32]*GossipChannel
	TopologyGossip  Gossip
	UDPListener     *net.UDPConn
	Password        []byte
	ConnLimit       int
	BufSz           int
	LogFrame        func(string, []byte, *layers.Ethernet)
}

type PacketSource interface {
	ReadPacket() ([]byte, error)
}

type PacketSink interface {
	WritePacket([]byte) error
}

type PacketSourceSink interface {
	PacketSource
	PacketSink
}

func NewRouter(iface *net.Interface, name PeerName, nickName string, password []byte, connLimit int, bufSz int, logFrame func(string, []byte, *layers.Ethernet)) *Router {
	router := &Router{
		Iface:          iface,
		GossipChannels: make(map[uint32]*GossipChannel),
		Password:       password,
		ConnLimit:      connLimit,
		BufSz:          bufSz,
		LogFrame:       logFrame}
	onMacExpiry := func(mac net.HardwareAddr, peer *Peer) {
		log.Println("Expired MAC", mac, "at", peer.FullName())
	}
	onPeerGC := func(peer *Peer) {
		router.Macs.Delete(peer)
		log.Println("Removed unreachable peer", peer.FullName())
	}
	router.Ourself = NewLocalPeer(name, nickName, router)
	router.Macs = NewMacCache(macMaxAge, onMacExpiry)
	router.Peers = NewPeers(router.Ourself.Peer, onPeerGC)
	router.Peers.FetchWithDefault(router.Ourself.Peer)
	router.Routes = NewRoutes(router.Ourself.Peer, router.Peers)
	router.ConnectionMaker = NewConnectionMaker(router.Ourself, router.Peers)
	router.TopologyGossip = router.NewGossip("topology", router)
	return router
}

func (router *Router) Start() {
	// we need two pcap handles since they aren't thread-safe
	pio, err := NewPcapIO(router.Iface.Name, router.BufSz)
	checkFatal(err)
	po, err := NewPcapO(router.Iface.Name)
	checkFatal(err)
	router.Ourself.Start()
	router.Macs.Start()
	router.Routes.Start()
	router.ConnectionMaker.Start()
	router.UDPListener = router.listenUDP(Port, po)
	router.listenTCP(Port)
	router.sniff(pio)
}

func (router *Router) UsingPassword() bool {
	return router.Password != nil
}

func (router *Router) Status() string {
	var buf bytes.Buffer
	fmt.Fprintln(&buf, "Our name is", router.Ourself.FullName())
	fmt.Fprintln(&buf, "Sniffing traffic on", router.Iface)
	fmt.Fprintf(&buf, "MACs:\n%s", router.Macs)
	fmt.Fprintf(&buf, "Peers:\n%s", router.Peers)
	fmt.Fprintf(&buf, "Routes:\n%s", router.Routes)
	fmt.Fprintf(&buf, "Reconnects:\n%s", router.ConnectionMaker)
	return buf.String()
}

func (router *Router) sniff(pio PacketSourceSink) {
	log.Println("Sniffing traffic on", router.Iface)

	dec := NewEthernetDecoder()
	mac := router.Iface.HardwareAddr
	if router.Macs.Enter(mac, router.Ourself.Peer) {
		log.Println("Discovered our MAC", mac)
	}
	go func() {
		for {
			pkt, err := pio.ReadPacket()
			checkFatal(err)
			router.LogFrame("Sniffed", pkt, nil)
			router.handleCapturedPacket(pkt, dec, pio)
		}
	}()
}

func (router *Router) handleCapturedPacket(frameData []byte, dec *EthernetDecoder, po PacketSink) {
	dec.DecodeLayers(frameData)
	decodedLen := len(dec.decoded)
	if decodedLen == 0 {
		return
	}
	srcMac := dec.eth.SrcMAC
	srcPeer, found := router.Macs.Lookup(srcMac)
	// We need to filter out frames we injected ourselves. For such
	// frames, the srcMAC will have been recorded as associated with a
	// different peer.
	if found && srcPeer != router.Ourself.Peer {
		return
	}
	if router.Macs.Enter(srcMac, router.Ourself.Peer) {
		log.Println("Discovered local MAC", srcMac)
	}
	if dec.DropFrame() {
		return
	}
	dstMac := dec.eth.DstMAC
	dstPeer, found := router.Macs.Lookup(dstMac)
	if found && dstPeer == router.Ourself.Peer {
		return
	}
	df := decodedLen == 2 && (dec.ip.Flags&layers.IPv4DontFragment != 0)
	if df {
		router.LogFrame("Forwarding DF", frameData, &dec.eth)
	} else {
		router.LogFrame("Forwarding", frameData, &dec.eth)
	}
	// at this point we are handing over the frame to forwarders, so
	// we need to make a copy of it in order to prevent the next
	// capture from overwriting the data
	frameLen := len(frameData)
	frameCopy := make([]byte, frameLen, frameLen)
	copy(frameCopy, frameData)

	if !found {
		router.Ourself.Broadcast(df, frameCopy, dec)
		return
	}
	err := router.Ourself.Forward(dstPeer, df, frameCopy, dec)
	if ftbe, ok := err.(FrameTooBigError); ok {
		err = dec.sendICMPFragNeeded(ftbe.EPMTU, po.WritePacket)
	}
	checkWarn(err)
}

func (router *Router) listenTCP(localPort int) {
	localAddr, err := net.ResolveTCPAddr("tcp4", fmt.Sprint(":", localPort))
	checkFatal(err)
	ln, err := net.ListenTCP("tcp4", localAddr)
	checkFatal(err)
	go func() {
		defer ln.Close()
		for {
			tcpConn, err := ln.AcceptTCP()
			if err != nil {
				log.Println(err)
				continue
			}
			router.acceptTCP(tcpConn)
		}
	}()
}

func (router *Router) acceptTCP(tcpConn *net.TCPConn) {
	// someone else is dialing us, so our udp sender is the conn
	// on Port and we wait for them to send us something on UDP to
	// start.
	remoteAddrStr := tcpConn.RemoteAddr().String()
	log.Printf("->[%s] connection accepted\n", remoteAddrStr)
	connRemote := NewRemoteConnection(router.Ourself.Peer, nil, remoteAddrStr, false, false)
	connLocal := NewLocalConnection(connRemote, tcpConn, nil, router)
	connLocal.Start(true)
}

func (router *Router) listenUDP(localPort int, po PacketSink) *net.UDPConn {
	localAddr, err := net.ResolveUDPAddr("udp4", fmt.Sprint(":", localPort))
	checkFatal(err)
	conn, err := net.ListenUDP("udp4", localAddr)
	checkFatal(err)
	f, err := conn.File()
	defer f.Close()
	checkFatal(err)
	fd := int(f.Fd())
	// This one makes sure all packets we send out do not have DF set on them.
	err = syscall.SetsockoptInt(fd, syscall.IPPROTO_IP, syscall.IP_MTU_DISCOVER, syscall.IP_PMTUDISC_DONT)
	checkFatal(err)
	go router.udpReader(conn, po)
	return conn
}

func (router *Router) udpReader(conn *net.UDPConn, po PacketSink) {
	defer conn.Close()
	dec := NewEthernetDecoder()
	buf := make([]byte, MaxUDPPacketSize)
	for {
		n, sender, err := conn.ReadFromUDP(buf)
		if err == io.EOF {
			return
		} else if err != nil {
			log.Println("ignoring UDP read error", err)
			continue
		} else if n < NameSize {
			log.Println("ignoring too short UDP packet from", sender)
			continue
		}
		name := PeerNameFromBin(buf[:NameSize])
		packet := make([]byte, n-NameSize)
		copy(packet, buf[NameSize:n])
		peerConn, found := router.Ourself.ConnectionTo(name)
		if !found {
			continue
		}
		relayConn, ok := peerConn.(*LocalConnection)
		if !ok {
			continue
		}
		err = relayConn.Decryptor.IterateFrames(packet, router.handleUDPPacketFunc(dec, sender, po))
		if pde, ok := err.(PacketDecodingError); ok {
			if pde.Fatal {
				relayConn.Shutdown(pde)
			} else {
				relayConn.Log(pde.Error())
			}
		} else {
			checkWarn(err)
		}
	}
}

func (router *Router) handleUDPPacketFunc(dec *EthernetDecoder, sender *net.UDPAddr, po PacketSink) FrameConsumer {
	return func(relayConn *LocalConnection, srcNameByte, dstNameByte []byte, frame []byte) {
		srcPeer, found := router.Peers.Fetch(PeerNameFromBin(srcNameByte))
		if !found {
			return
		}
		dstPeer, found := router.Peers.Fetch(PeerNameFromBin(dstNameByte))
		if !found {
			return
		}

		dec.DecodeLayers(frame)
		decodedLen := len(dec.decoded)
		if decodedLen == 0 {
			return
		}
		// Handle special frames produced internally (rather than
		// captured/forwarded) by the remote router.
		//
		// We really shouldn't be decoding these above, since they are
		// not genuine Ethernet frames. However, it is actually more
		// efficient to do so, as we want to optimise for the common
		// (i.e. non-special) frames. These always need decoding, and
		// detecting special frames is cheaper post decoding than pre.
		if decodedLen == 1 && dec.IsSpecial() {
			if srcPeer == relayConn.Remote() && dstPeer == router.Ourself.Peer {
				handleSpecialFrame(relayConn, sender, frame)
			}
			return
		}

		df := decodedLen == 2 && (dec.ip.Flags&layers.IPv4DontFragment != 0)

		if dstPeer != router.Ourself.Peer {
			// it's not for us, we're just relaying it
			if df {
				router.LogFrame("Relaying DF", frame, &dec.eth)
			} else {
				router.LogFrame("Relaying", frame, &dec.eth)
			}

			err := router.Ourself.Relay(srcPeer, dstPeer, df, frame, dec)
			if ftbe, ok := err.(FrameTooBigError); ok {
				err = dec.sendICMPFragNeeded(ftbe.EPMTU, func(icmpFrame []byte) error {
					return router.Ourself.Forward(srcPeer, false, icmpFrame, nil)
				})
			}

			checkWarn(err)
			return
		}

		srcMac := dec.eth.SrcMAC
		dstMac := dec.eth.DstMAC

		if router.Macs.Enter(srcMac, srcPeer) {
			log.Println("Discovered remote MAC", srcMac, "at", srcPeer.FullName())
		}
		router.LogFrame("Injecting", frame, &dec.eth)
		checkWarn(po.WritePacket(frame))

		dstPeer, found = router.Macs.Lookup(dstMac)
		if !found || dstPeer != router.Ourself.Peer {
			router.LogFrame("Relaying broadcast", frame, &dec.eth)
			router.Ourself.RelayBroadcast(srcPeer, df, frame, dec)
		}
	}
}

func handleSpecialFrame(relayConn *LocalConnection, sender *net.UDPAddr, frame []byte) {
	frameLen := len(frame)
	switch {
	case frameLen == EthernetOverhead+8:
		relayConn.ReceivedHeartbeat(sender, binary.BigEndian.Uint64(frame[EthernetOverhead:]))
	case frameLen == FragTestSize && bytes.Equal(frame, FragTest):
		relayConn.SendProtocolMsg(ProtocolMsg{ProtocolFragmentationReceived, nil})
	case frameLen == PMTUDiscoverySize && bytes.Equal(frame, PMTUDiscovery):
	default:
		frameLenBytes := []byte{0, 0}
		binary.BigEndian.PutUint16(frameLenBytes, uint16(frameLen-EthernetOverhead))
		relayConn.SendProtocolMsg(ProtocolMsg{ProtocolPMTUVerified, frameLenBytes})
	}
}

// Gossiper methods - the Router is the topology Gossiper

func (router *Router) OnGossipUnicast(sender PeerName, msg []byte) error {
	return fmt.Errorf("unexpected topology gossip unicast: %v", msg)
}

func (router *Router) OnGossipBroadcast(msg []byte) error {
	return fmt.Errorf("unexpected topology gossip broadcast: %v", msg)
}

type PeerNameSet map[PeerName]bool

type TopologyGossipData struct {
	peers  *Peers
	update PeerNameSet
}

func NewTopologyGossipData(peers *Peers, update ...*Peer) *TopologyGossipData {
	names := make(PeerNameSet)
	for _, p := range update {
		names[p.Name] = true
	}
	return &TopologyGossipData{peers: peers, update: names}
}

func (d *TopologyGossipData) Merge(other GossipData) {
	for name := range other.(*TopologyGossipData).update {
		d.update[name] = true
	}
}

func (d *TopologyGossipData) Encode() []byte {
	return d.peers.EncodePeers(d.update)
}

func (router *Router) Gossip() GossipData {
	return &TopologyGossipData{peers: router.Peers, update: router.Peers.Names()}
}

func (router *Router) OnGossip(buf []byte) (GossipData, error) {
	newUpdate, err := router.Peers.ApplyUpdate(buf)
	if _, ok := err.(UnknownPeerError); err != nil && ok {
		// That update contained a reference to a peer which wasn't
		// itself included in the update, and we didn't know about
		// already. We ignore this; eventually we should receive an
		// update containing a complete topology.
		log.Println("Topology gossip:", err)
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	if len(newUpdate) == 0 {
		return nil, nil
	}
	router.ConnectionMaker.Refresh()
	router.Routes.Recalculate()
	return &TopologyGossipData{peers: router.Peers, update: newUpdate}, nil
}
<|endoftext|>


<commit_before>// Defines the Route interface, and registers routes to a server
package router

import (
	"net/http"
	"time"

	"github.com/KyleBanks/go-kit/log"
)

// Interface for the provided server to comply with
type Server interface {
	HandleFunc(string, func(http.ResponseWriter, *http.Request))
}

// Defines an executable Route
type Route struct {
	Path   string // The URL path to listen for (i.e. "/api")
	Handle func(w http.ResponseWriter, r *http.Request)
}

// Register registers each Route with the Server provided.
//
// Each Route will be wrapped in a middleware function that adds trace logging.
func Register(s Server, routes []Route) {
	for _, route := range routes {
		log.Info("Registering route:", route.Path)
		s.HandleFunc(route.Path, handleWrapper(route))
	}
}

// handleWrapper returns a request handling function that wraps the provided route.
func handleWrapper(route Route) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		// Enable CORS
		w.Header().Set("Access-Control-Allow-Origin", "*")

		start := time.Now().Unix()

		log.Info("START:", r.URL.Path, r.URL.RawQuery, r.PostForm)
		route.Handle(w, r)
		log.Info("END:", r.URL.Path, r.URL.RawQuery, r.PostForm, time.Now().Unix()-start)
	}
}

// Param returns a POST/GET parameter from the request.
//
// If the parameter is found in the POST and the GET parameter set, the POST parameter
// will be given priority.
func Param(r *http.Request, key string) string {
	val := r.PostForm.Get(key)
	if len(val) != 0 {
		return val
	}

	return r.URL.Query().Get(key)
}
<commit_msg>Parse the form before attempting to pull out parameters in the router<commit_after>// Defines the Route interface, and registers routes to a server
package router

import (
	"net/http"
	"time"

	"github.com/KyleBanks/go-kit/log"
)

// Interface for the provided server to comply with
type Server interface {
	HandleFunc(string, func(http.ResponseWriter, *http.Request))
}

// Defines an executable Route
type Route struct {
	Path   string // The URL path to listen for (i.e. "/api")
	Handle func(w http.ResponseWriter, r *http.Request)
}

// Register registers each Route with the Server provided.
//
// Each Route will be wrapped in a middleware function that adds trace logging.
func Register(s Server, routes []Route) {
	for _, route := range routes {
		log.Info("Registering route:", route.Path)
		s.HandleFunc(route.Path, handleWrapper(route))
	}
}

// handleWrapper returns a request handling function that wraps the provided route.
func handleWrapper(route Route) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		// Enable CORS
		w.Header().Set("Access-Control-Allow-Origin", "*")

		start := time.Now().Unix()

		log.Info("START:", r.URL.Path, r.URL.RawQuery, r.PostForm)
		route.Handle(w, r)
		log.Info("END:", r.URL.Path, r.URL.RawQuery, r.PostForm, time.Now().Unix()-start)
	}
}

// Param returns a POST/GET parameter from the request.
//
// If the parameter is found in the POST and the GET parameter set, the POST parameter
// will be given priority.
func Param(r *http.Request, key string) string {
	r.ParseForm()

	val := r.PostForm.Get(key)
	if len(val) != 0 {
		return val
	}

	return r.URL.Query().Get(key)
}
<|endoftext|>


<commit_before>package router

import (
	"bytes"
	"fmt"
	"sync"
)

type Routes struct {
	sync.RWMutex
	ourself      *Peer
	peers        *Peers
	unicast      map[PeerName]PeerName
	unicastAll   map[PeerName]PeerName // [1]
	broadcast    map[PeerName][]PeerName
	broadcastAll map[PeerName][]PeerName // [1]
	recalculate  chan<- *struct{}
	wait         chan<- chan struct{}
	// [1] based on *all* connections, not just established &
	// symmetric ones
}

func NewRoutes(ourself *Peer, peers *Peers) *Routes {
	routes := &Routes{
		ourself:      ourself,
		peers:        peers,
		unicast:      make(map[PeerName]PeerName),
		unicastAll:   make(map[PeerName]PeerName),
		broadcast:    make(map[PeerName][]PeerName),
		broadcastAll: make(map[PeerName][]PeerName)}
	routes.unicast[ourself.Name] = UnknownPeerName
	routes.unicastAll[ourself.Name] = UnknownPeerName
	routes.broadcast[ourself.Name] = []PeerName{}
	routes.broadcastAll[ourself.Name] = []PeerName{}
	return routes
}

func (routes *Routes) Start() {
	recalculate := make(chan *struct{}, 1)
	wait := make(chan chan struct{})
	routes.recalculate = recalculate
	routes.wait = wait
	go routes.run(recalculate, wait)
}

func (routes *Routes) Unicast(name PeerName) (PeerName, bool) {
	routes.RLock()
	defer routes.RUnlock()
	hop, found := routes.unicast[name]
	return hop, found
}

func (routes *Routes) UnicastAll(name PeerName) (PeerName, bool) {
	routes.RLock()
	defer routes.RUnlock()
	hop, found := routes.unicastAll[name]
	return hop, found
}

func (routes *Routes) Broadcast(name PeerName) []PeerName {
	routes.RLock()
	defer routes.RUnlock()
	hops, found := routes.broadcast[name]
	if !found {
		return []PeerName{}
	}
	return hops
}

func (routes *Routes) BroadcastAll(name PeerName) []PeerName {
	routes.RLock()
	defer routes.RUnlock()
	hops, found := routes.broadcastAll[name]
	if !found {
		return []PeerName{}
	}
	return hops
}

func (routes *Routes) String() string {
	var buf bytes.Buffer
	routes.RLock()
	defer routes.RUnlock()
	fmt.Fprintln(&buf, "unicast:")
	for name, hop := range routes.unicast {
		fmt.Fprintf(&buf, "%s -> %s\n", name, hop)
	}
	fmt.Fprintln(&buf, "broadcast:")
	for name, hops := range routes.broadcast {
		fmt.Fprintf(&buf, "%s -> %v\n", name, hops)
	}
	// We don't include the 'all' routes here since they are of
	// limited utility in troubleshooting
	return buf.String()
}

// Request recalculation of the routing table. This is async but can
// effectively be made synchronous with a subsequent call to
// EnsureRecalculated.
func (routes *Routes) Recalculate() {
	// The use of a 1-capacity channel in combination with the
	// non-blocking send is an optimisation that results in multiple
	// requests being coalesced.
	select {
	case routes.recalculate <- nil:
	default:
	}
}

// Wait for any preceding Recalculate requests to be processed.
func (routes *Routes) EnsureRecalculated() {
	done := make(chan struct{})
	routes.wait <- done
	<-done
}

func (routes *Routes) run(recalculate <-chan *struct{}, wait <-chan chan struct{}) {
	for {
		select {
		case <-recalculate:
			routes.calculate()
		case done := <-wait:
			select {
			case <-recalculate:
				routes.calculate()
			default:
			}
			close(done)
		}
	}
}

func (routes *Routes) calculate() {
	var (
		unicast      = routes.calculateUnicast(true)
		unicastAll   = routes.calculateUnicast(false)
		broadcast    = routes.calculateBroadcast(true)
		broadcastAll = routes.calculateBroadcast(false)
	)
	routes.Lock()
	routes.unicast = unicast
	routes.unicastAll = unicastAll
	routes.broadcast = broadcast
	routes.broadcastAll = broadcastAll
	routes.Unlock()
}

// Calculate all the routes for the question: if *we* want to send a
// packet to Peer X, what is the next hop?
//
// When we sniff a packet, we determine the destination peer
// ourself. Consequently, we can relay the packet via any
// arbitrary peers - the intermediate peers do not have to have
// any knowledge of the MAC address at all. Thus there's no need
// to exchange knowledge of MAC addresses, nor any constraints on
// the routes that we construct.
func (routes *Routes) calculateUnicast(establishedAndSymmetric bool) map[PeerName]PeerName {
	_, unicast := routes.ourself.Routes(nil, establishedAndSymmetric)
	return unicast
}

// Calculate all the routes for the question: if we receive a
// broadcast originally from Peer X, which peers should we pass the
// frames on to?
//
// When the topology is stable, and thus all peers perform route
// calculations based on the same data, the algorithm ensures that
// broadcasts reach every peer exactly once.
//
// This is largely due to properties of the Peer.Routes algorithm. In
// particular:
//
// ForAll X,Y,Z in Peers.
//     X.Routes(Y) <= X.Routes(Z) \/
//     X.Routes(Z) <= X.Routes(Y)
// ForAll X,Y,Z in Peers.
//     Y =/= Z /\ X.Routes(Y) <= X.Routes(Z) =>
//     X.Routes(Y) u [P | Y.HasSymmetricConnectionTo(P)] <= X.Routes(Z)
// where <= is the subset relationship on keys of the returned map.
func (routes *Routes) calculateBroadcast(establishedAndSymmetric bool) map[PeerName][]PeerName {
	broadcast := make(map[PeerName][]PeerName)
	ourself := routes.ourself

	routes.peers.ForEach(func(peer *Peer) {
		hops := []PeerName{}
		if found, reached := peer.Routes(ourself, establishedAndSymmetric); found {
			// This is rather similar to the inner loop on
			// peer.Routes(...); the main difference is in the
			// locking.
			for conn := range ourself.Connections() {
				if establishedAndSymmetric && !conn.Established() {
					continue
				}
				remoteName := conn.Remote().Name
				if _, found := reached[remoteName]; found {
					continue
				}
				if remoteConn, found := conn.Remote().ConnectionTo(ourself.Name); !establishedAndSymmetric || (found && remoteConn.Established()) {
					hops = append(hops, remoteName)
				}
			}
		}
		broadcast[peer.Name] = hops
	})
	return broadcast
}
<commit_msg>optimise broadcast route calculation ...by obtaining our connection set just once. This also guards against potential locking issues, since we no longer obtain the required lock on our local peer while holding a lock on Peers.<commit_after>package router

import (
	"bytes"
	"fmt"
	"sync"
)

type Routes struct {
	sync.RWMutex
	ourself      *Peer
	peers        *Peers
	unicast      map[PeerName]PeerName
	unicastAll   map[PeerName]PeerName // [1]
	broadcast    map[PeerName][]PeerName
	broadcastAll map[PeerName][]PeerName // [1]
	recalculate  chan<- *struct{}
	wait         chan<- chan struct{}
	// [1] based on *all* connections, not just established &
	// symmetric ones
}

func NewRoutes(ourself *Peer, peers *Peers) *Routes {
	routes := &Routes{
		ourself:      ourself,
		peers:        peers,
		unicast:      make(map[PeerName]PeerName),
		unicastAll:   make(map[PeerName]PeerName),
		broadcast:    make(map[PeerName][]PeerName),
		broadcastAll: make(map[PeerName][]PeerName)}
	routes.unicast[ourself.Name] = UnknownPeerName
	routes.unicastAll[ourself.Name] = UnknownPeerName
	routes.broadcast[ourself.Name] = []PeerName{}
	routes.broadcastAll[ourself.Name] = []PeerName{}
	return routes
}

func (routes *Routes) Start() {
	recalculate := make(chan *struct{}, 1)
	wait := make(chan chan struct{})
	routes.recalculate = recalculate
	routes.wait = wait
	go routes.run(recalculate, wait)
}

func (routes *Routes) Unicast(name PeerName) (PeerName, bool) {
	routes.RLock()
	defer routes.RUnlock()
	hop, found := routes.unicast[name]
	return hop, found
}

func (routes *Routes) UnicastAll(name PeerName) (PeerName, bool) {
	routes.RLock()
	defer routes.RUnlock()
	hop, found := routes.unicastAll[name]
	return hop, found
}

func (routes *Routes) Broadcast(name PeerName) []PeerName {
	routes.RLock()
	defer routes.RUnlock()
	hops, found := routes.broadcast[name]
	if !found {
		return []PeerName{}
	}
	return hops
}

func (routes *Routes) BroadcastAll(name PeerName) []PeerName {
	routes.RLock()
	defer routes.RUnlock()
	hops, found := routes.broadcastAll[name]
	if !found {
		return []PeerName{}
	}
	return hops
}

func (routes *Routes) String() string {
	var buf bytes.Buffer
	routes.RLock()
	defer routes.RUnlock()
	fmt.Fprintln(&buf, "unicast:")
	for name, hop := range routes.unicast {
		fmt.Fprintf(&buf, "%s -> %s\n", name, hop)
	}
	fmt.Fprintln(&buf, "broadcast:")
	for name, hops := range routes.broadcast {
		fmt.Fprintf(&buf, "%s -> %v\n", name, hops)
	}
	// We don't include the 'all' routes here since they are of
	// limited utility in troubleshooting
	return buf.String()
}

// Request recalculation of the routing table. This is async but can
// effectively be made synchronous with a subsequent call to
// EnsureRecalculated.
func (routes *Routes) Recalculate() {
	// The use of a 1-capacity channel in combination with the
	// non-blocking send is an optimisation that results in multiple
	// requests being coalesced.
	select {
	case routes.recalculate <- nil:
	default:
	}
}

// Wait for any preceding Recalculate requests to be processed.
func (routes *Routes) EnsureRecalculated() {
	done := make(chan struct{})
	routes.wait <- done
	<-done
}

func (routes *Routes) run(recalculate <-chan *struct{}, wait <-chan chan struct{}) {
	for {
		select {
		case <-recalculate:
			routes.calculate()
		case done := <-wait:
			select {
			case <-recalculate:
				routes.calculate()
			default:
			}
			close(done)
		}
	}
}

func (routes *Routes) calculate() {
	var (
		unicast      = routes.calculateUnicast(true)
		unicastAll   = routes.calculateUnicast(false)
		broadcast    = routes.calculateBroadcast(true)
		broadcastAll = routes.calculateBroadcast(false)
	)
	routes.Lock()
	routes.unicast = unicast
	routes.unicastAll = unicastAll
	routes.broadcast = broadcast
	routes.broadcastAll = broadcastAll
	routes.Unlock()
}

// Calculate all the routes for the question: if *we* want to send a
// packet to Peer X, what is the next hop?
//
// When we sniff a packet, we determine the destination peer
// ourself. Consequently, we can relay the packet via any
// arbitrary peers - the intermediate peers do not have to have
// any knowledge of the MAC address at all. Thus there's no need
// to exchange knowledge of MAC addresses, nor any constraints on
// the routes that we construct.
func (routes *Routes) calculateUnicast(establishedAndSymmetric bool) map[PeerName]PeerName {
	_, unicast := routes.ourself.Routes(nil, establishedAndSymmetric)
	return unicast
}

// Calculate all the routes for the question: if we receive a
// broadcast originally from Peer X, which peers should we pass the
// frames on to?
//
// When the topology is stable, and thus all peers perform route
// calculations based on the same data, the algorithm ensures that
// broadcasts reach every peer exactly once.
//
// This is largely due to properties of the Peer.Routes algorithm. In
// particular:
//
// ForAll X,Y,Z in Peers.
//     X.Routes(Y) <= X.Routes(Z) \/
//     X.Routes(Z) <= X.Routes(Y)
// ForAll X,Y,Z in Peers.
//     Y =/= Z /\ X.Routes(Y) <= X.Routes(Z) =>
//     X.Routes(Y) u [P | Y.HasSymmetricConnectionTo(P)] <= X.Routes(Z)
// where <= is the subset relationship on keys of the returned map.
func (routes *Routes) calculateBroadcast(establishedAndSymmetric bool) map[PeerName][]PeerName {
	broadcast := make(map[PeerName][]PeerName)
	ourself := routes.ourself
	ourConnections := ourself.Connections()

	routes.peers.ForEach(func(peer *Peer) {
		hops := []PeerName{}
		if found, reached := peer.Routes(ourself, establishedAndSymmetric); found {
			// This is rather similar to the inner loop on
			// peer.Routes(...); the main difference is in the
			// locking.
			for conn := range ourConnections {
				if establishedAndSymmetric && !conn.Established() {
					continue
				}
				remoteName := conn.Remote().Name
				if _, found := reached[remoteName]; found {
					continue
				}
				if remoteConn, found := conn.Remote().ConnectionTo(ourself.Name); !establishedAndSymmetric || (found && remoteConn.Established()) {
					hops = append(hops, remoteName)
				}
			}
		}
		broadcast[peer.Name] = hops
	})
	return broadcast
}
<|endoftext|>


<commit_before>package negotiator

import (
	"net/http"
	"strings"
)

const MimeJSON = "application/json"
const MimeXML = "application/xml"

// Encoder is an interface for a struct that can encode data into []byte
type Encoder interface {
	Encode(data interface{}) ([]byte, error)
	ContentType() string
}

// A Negotiator can Negotiate to determine what content type to convert
// a struct into for client consumption
type Negotiator interface {
	Negotiate(req *http.Request, data interface{}) ([]byte, error)
}

// ContentNegotiator is a Neotiator that supports pretty printing
// and a fallback/default encoder as well as dynamically adding
// encoders
type ContentNegotiator struct {
	PrettyPrint    bool
	DefaultEncoder Encoder
	ResponseWriter http.ResponseWriter
	encoderMap     map[string]Encoder
}

func NewJsonXmlContentNegotiator(prettyPrint bool, defaultEncoder Encoder, responseWriter http.ResponseWriter) ContentNegotiator {
	result := ContentNegotiator{prettyPrint, defaultEncoder, responseWriter}
	result.AddEncoder(MimeJSON, JsonEncoder{prettyPrint})
	result.AddEncoder(MimeXML, XmlEncoder{prettyPrint})
	return result
}

// Negotiate inspects the request for the accept header and
// encodes the response appropriately.
func (cn ContentNegotiator) Negotiate(req *http.Request, data interface{}) ([]byte, error) {
	if len(cn.encoderMap) <= 0 {
		panic("No Encoders present. Please add them using ContentNegotiator.AddEncoder()")
	}
	var e = cn.getEncoder(req)
	cn.ResponseWriter.Header().Set("Content-Type", e.ContentType())
	return e.Encode(data)
}

// AddEncoder registers a mimetype and its encoder to be used if a client
// requests that mimetype
func (cn ContentNegotiator) AddEncoder(mimeType string, enc Encoder) {
	cn.encoderMap[mimeType] = enc
}

// getEncoder parses the Accept header an returns the appropriate encoder to use
func (cn ContentNegotiator) getEncoder(req *http.Request) Encoder {
	var result = cn.DefaultEncoder
	accept := req.Header.Get("Accept")

	for k, v := range cn.encoderMap {
		if strings.Contains(accept, k) {
			result = v
			break
		}
	}
	return result
}

// Check for an error and panic
func Must(data []byte, err error) []byte {
	if err != nil {
		panic(err)
	}
	return data
}
<commit_msg>Adding basic NewContentNegotiator helper<commit_after>package negotiator

import (
	"net/http"
	"strings"
)

const MimeJSON = "application/json"
const MimeXML = "application/xml"

// Encoder is an interface for a struct that can encode data into []byte
type Encoder interface {
	Encode(data interface{}) ([]byte, error)
	ContentType() string
}

// A Negotiator can Negotiate to determine what content type to convert
// a struct into for client consumption
type Negotiator interface {
	Negotiate(req *http.Request, data interface{}) ([]byte, error)
}

// ContentNegotiator is a Neotiator that supports a fallback/default
// encoder as well as dynamically adding encoders
type ContentNegotiator struct {
	DefaultEncoder Encoder
	ResponseWriter http.ResponseWriter
	encoderMap     map[string]Encoder
}

// NewContentNegotiator creates a basic ContentNegotiator with out any attached
// encoders
func NewContentNegotiator(defaultEncoder Encoder, responseWriter http.ResponseWriter) ContentNegotiator {
	result := ContentNegotiator{}
	result.DefaultEncoder = defaultEncoder
	result.ResponseWriter = responseWriter
	return result
}

// NewJsonXmlContentNegotiator creates a basic ContentNegotiator and attaches
// a JSON and an XML encoder to it.
func NewJsonXmlContentNegotiator(defaultEncoder Encoder, responseWriter http.ResponseWriter, prettyPrint bool) ContentNegotiator {
	result := NewContentNegotiator(defaultEncoder, responseWriter)
	result.AddEncoder(MimeJSON, JsonEncoder{prettyPrint})
	result.AddEncoder(MimeXML, XmlEncoder{prettyPrint})
	return result
}

// Negotiate inspects the request for the accept header and
// encodes the response appropriately.
func (cn ContentNegotiator) Negotiate(req *http.Request, data interface{}) ([]byte, error) {
	if len(cn.encoderMap) <= 0 {
		panic("No Encoders present. Please add them using ContentNegotiator.AddEncoder()")
	}
	var e = cn.getEncoder(req)
	cn.ResponseWriter.Header().Set("Content-Type", e.ContentType())
	return e.Encode(data)
}

// AddEncoder registers a mimetype and its encoder to be used if a client
// requests that mimetype
func (cn ContentNegotiator) AddEncoder(mimeType string, enc Encoder) {
	cn.encoderMap[mimeType] = enc
}

// getEncoder parses the Accept header an returns the appropriate encoder to use
func (cn ContentNegotiator) getEncoder(req *http.Request) Encoder {
	var result = cn.DefaultEncoder
	accept := req.Header.Get("Accept")

	for k, v := range cn.encoderMap {
		if strings.Contains(accept, k) {
			result = v
			break
		}
	}
	return result
}

// Check for an error and panic
func Must(data []byte, err error) []byte {
	if err != nil {
		panic(err)
	}
	return data
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package logger

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"

	keybase1 "github.com/keybase/client/go/protocol"
	logging "github.com/op/go-logging"
	"golang.org/x/net/context"
)

const permDir os.FileMode = 0700

var initLoggingBackendOnce sync.Once
var logRotateMutex sync.Mutex

// CtxStandardLoggerKey is a type defining context keys used by the
// Standard logger.
type CtxStandardLoggerKey int

const (
	// CtxLogTags defines a context key that can hold a slice of context
	// keys, the value of which should be logged by a Standard logger if
	// one of those keys is seen in a context during a log call.
	CtxLogTagsKey CtxStandardLoggerKey = iota
)

type CtxLogTags map[interface{}]string

// NewContext returns a new Context that carries adds the given log
// tag mappings (context key -> display string).
func NewContextWithLogTags(
	ctx context.Context, logTagsToAdd CtxLogTags) context.Context {
	currTags, ok := LogTagsFromContext(ctx)
	if !ok {
		currTags = make(CtxLogTags)
	}
	for key, tag := range logTagsToAdd {
		currTags[key] = tag
	}
	return context.WithValue(ctx, CtxLogTagsKey, currTags)
}

// LogTagsFromContext returns the log tags being passed along with the
// given context.
func LogTagsFromContext(ctx context.Context) (CtxLogTags, bool) {
	logTags, ok := ctx.Value(CtxLogTagsKey).(CtxLogTags)
	return logTags, ok
}

type ExternalLogger interface {
	Log(level keybase1.LogLevel, format string, args []interface{})
}

type Standard struct {
	internal       *logging.Logger
	filename       string
	configureMutex sync.Mutex
	module         string

	externalLoggers      map[uint64]ExternalLogger
	externalLoggersCount uint64
	externalLogLevel     keybase1.LogLevel
	externalLoggersMutex sync.RWMutex
}

// New creates a new Standard logger for module.
func New(module string) *Standard {
	return NewWithCallDepth(module, 0)
}

// Verify Standard fully implements the Logger interface.
var _ Logger = (*Standard)(nil)

// NewWithCallDepth creates a new Standard logger for module, and when
// printing file names and line numbers, it goes extraCallDepth up the
// stack from where logger was invoked.
func NewWithCallDepth(module string, extraCallDepth int) *Standard {
	log := logging.MustGetLogger(module)
	log.ExtraCalldepth = 1 + extraCallDepth
	ret := &Standard{
		internal:             log,
		module:               module,
		externalLoggers:      make(map[uint64]ExternalLogger),
		externalLoggersCount: 0,
		externalLogLevel:     keybase1.LogLevel_INFO,
	}
	ret.initLogging()
	return ret
}

func (log *Standard) initLogging() {
	// Logging is always done to stderr. It's the responsibility of the
	// launcher (like launchd on OSX, or the autoforking code) to set up stderr
	// to point to the appropriate log file.
	initLoggingBackendOnce.Do(func() {
		logBackend := logging.NewLogBackend(os.Stderr, "", 0)
		logging.SetBackend(logBackend)
		logging.SetLevel(logging.INFO, log.module)
	})
}

func (log *Standard) prepareString(
	ctx context.Context, fmts string) string {
	if ctx == nil {
		return fmts
	}
	logTags, ok := LogTagsFromContext(ctx)
	if !ok || len(logTags) == 0 {
		return fmts
	}
	var tags []string
	for key, tag := range logTags {
		if v := ctx.Value(key); v != nil {
			tags = append(tags, fmt.Sprintf("%s=%s", tag, v))
		}
	}
	return fmts + " [tags:" + strings.Join(tags, ",") + "]"
}

func (log *Standard) Debug(fmt string, arg ...interface{}) {
	log.internal.Debug(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_DEBUG, fmt, arg)
}

func (log *Standard) CDebugf(ctx context.Context, fmt string,
	arg ...interface{}) {
	if log.internal.IsEnabledFor(logging.DEBUG) {
		log.Debug(log.prepareString(ctx, fmt), arg...)
	}
}

func (log *Standard) Info(fmt string, arg ...interface{}) {
	log.internal.Info(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_INFO, fmt, arg)
}

func (log *Standard) CInfof(ctx context.Context, fmt string,
	arg ...interface{}) {
	if log.internal.IsEnabledFor(logging.INFO) {
		log.Info(log.prepareString(ctx, fmt), arg...)
	}
}

func (log *Standard) Notice(fmt string, arg ...interface{}) {
	log.internal.Notice(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_NOTICE, fmt, arg)
}

func (log *Standard) CNoticef(ctx context.Context, fmt string,
	arg ...interface{}) {
	if log.internal.IsEnabledFor(logging.NOTICE) {
		log.Notice(log.prepareString(ctx, fmt), arg...)
	}
}

func (log *Standard) Warning(fmt string, arg ...interface{}) {
	log.internal.Warning(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_WARN, fmt, arg)
}

func (log *Standard) CWarningf(ctx context.Context, fmt string,
	arg ...interface{}) {
	if log.internal.IsEnabledFor(logging.WARNING) {
		log.Warning(log.prepareString(ctx, fmt), arg...)
	}
}

func (log *Standard) Error(fmt string, arg ...interface{}) {
	log.internal.Error(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_ERROR, fmt, arg)
}

func (log *Standard) Errorf(fmt string, arg ...interface{}) {
	log.Error(fmt, arg...)
}

func (log *Standard) CErrorf(ctx context.Context, fmt string,
	arg ...interface{}) {
	if log.internal.IsEnabledFor(logging.ERROR) {
		log.Error(log.prepareString(ctx, fmt), arg...)
	}
}

func (log *Standard) Critical(fmt string, arg ...interface{}) {
	log.internal.Critical(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_CRITICAL, fmt, arg)
}

func (log *Standard) CCriticalf(ctx context.Context, fmt string,
	arg ...interface{}) {
	if log.internal.IsEnabledFor(logging.CRITICAL) {
		log.Critical(log.prepareString(ctx, fmt), arg...)
	}
}

func (log *Standard) Fatalf(fmt string, arg ...interface{}) {
	log.internal.Fatalf(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_FATAL, fmt, arg)
}

func (log *Standard) CFatalf(ctx context.Context, fmt string,
	arg ...interface{}) {
	log.Fatalf(log.prepareString(ctx, fmt), arg...)
}

func (log *Standard) Profile(fmts string, arg ...interface{}) {
	log.Debug(fmts, arg...)
}

func (log *Standard) Configure(style string, debug bool, filename string) {
	log.configureMutex.Lock()
	defer log.configureMutex.Unlock()

	log.filename = filename

	var logfmt string
	if debug {
		logfmt = fancyFormat
	} else {
		logfmt = defaultFormat
	}

	// Override the format above if an explicit style was specified.
	switch style {
	case "default":
		logfmt = defaultFormat // Default
	case "plain":
		logfmt = plainFormat // Plain
	case "file":
		logfmt = fileFormat // Good for logging to files
	case "fancy":
		logfmt = fancyFormat // Fancy, good for terminals with color
	}

	if debug {
		logging.SetLevel(logging.DEBUG, log.module)
	}

	logging.SetFormatter(logging.MustStringFormatter(logfmt))
}

func OpenLogFile(filename string) (name string, file *os.File, err error) {
	name = filename
	if err = MakeParentDirs(name); err != nil {
		return
	}
	file, err = os.OpenFile(name, (os.O_APPEND | os.O_WRONLY | os.O_CREATE), 0600)
	if err != nil {
		return
	}
	return
}

func FileExists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return false, err
}

func MakeParentDirs(filename string) error {
	dir, _ := filepath.Split(filename)
	exists, err := FileExists(dir)
	if err != nil {
		return err
	}

	if !exists {
		err = os.MkdirAll(dir, permDir)
		if err != nil {
			return err
		}
	}
	return nil
}

func PickFirstError(errors ...error) error {
	for _, e := range errors {
		if e != nil {
			return e
		}
	}
	return nil
}

func (log *Standard) AddExternalLogger(externalLogger ExternalLogger) uint64 {
	log.externalLoggersMutex.Lock()
	defer log.externalLoggersMutex.Unlock()

	handle := log.externalLoggersCount
	log.externalLoggersCount++
	log.externalLoggers[handle] = externalLogger
	return handle
}

func (log *Standard) RemoveExternalLogger(handle uint64) {
	log.externalLoggersMutex.Lock()
	defer log.externalLoggersMutex.Unlock()

	delete(log.externalLoggers, handle)
}

func (log *Standard) logToExternalLoggers(level keybase1.LogLevel, format string, args []interface{}) {
	log.externalLoggersMutex.RLock()
	defer log.externalLoggersMutex.RUnlock()

	// Short circuit logs that are more verbose than the current external log
	// level.
	if level < log.externalLogLevel {
		return
	}

	for _, externalLogger := range log.externalLoggers {
		externalLogger.Log(level, format, args)
	}
}

func (log *Standard) SetExternalLogLevel(level keybase1.LogLevel) {
	log.externalLoggersMutex.Lock()
	defer log.externalLoggersMutex.Unlock()

	log.externalLogLevel = level
}
<commit_msg>Back out #1365; found a deadlock in synchronous logging<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package logger

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"

	keybase1 "github.com/keybase/client/go/protocol"
	logging "github.com/op/go-logging"
	"golang.org/x/net/context"
)

const permDir os.FileMode = 0700

var initLoggingBackendOnce sync.Once
var logRotateMutex sync.Mutex

// CtxStandardLoggerKey is a type defining context keys used by the
// Standard logger.
type CtxStandardLoggerKey int

const (
	// CtxLogTags defines a context key that can hold a slice of context
	// keys, the value of which should be logged by a Standard logger if
	// one of those keys is seen in a context during a log call.
	CtxLogTagsKey CtxStandardLoggerKey = iota
)

type CtxLogTags map[interface{}]string

// NewContext returns a new Context that carries adds the given log
// tag mappings (context key -> display string).
func NewContextWithLogTags(
	ctx context.Context, logTagsToAdd CtxLogTags) context.Context {
	currTags, ok := LogTagsFromContext(ctx)
	if !ok {
		currTags = make(CtxLogTags)
	}
	for key, tag := range logTagsToAdd {
		currTags[key] = tag
	}
	return context.WithValue(ctx, CtxLogTagsKey, currTags)
}

// LogTagsFromContext returns the log tags being passed along with the
// given context.
func LogTagsFromContext(ctx context.Context) (CtxLogTags, bool) {
	logTags, ok := ctx.Value(CtxLogTagsKey).(CtxLogTags)
	return logTags, ok
}

type ExternalLogger interface {
	Log(level keybase1.LogLevel, format string, args []interface{})
}

type Standard struct {
	internal       *logging.Logger
	filename       string
	configureMutex sync.Mutex
	module         string

	externalLoggers      map[uint64]ExternalLogger
	externalLoggersCount uint64
	externalLogLevel     keybase1.LogLevel
	externalLoggersMutex sync.RWMutex
}

// New creates a new Standard logger for module.
func New(module string) *Standard {
	return NewWithCallDepth(module, 0)
}

// Verify Standard fully implements the Logger interface.
var _ Logger = (*Standard)(nil)

// NewWithCallDepth creates a new Standard logger for module, and when
// printing file names and line numbers, it goes extraCallDepth up the
// stack from where logger was invoked.
func NewWithCallDepth(module string, extraCallDepth int) *Standard {
	log := logging.MustGetLogger(module)
	log.ExtraCalldepth = 1 + extraCallDepth
	ret := &Standard{
		internal:             log,
		module:               module,
		externalLoggers:      make(map[uint64]ExternalLogger),
		externalLoggersCount: 0,
		externalLogLevel:     keybase1.LogLevel_INFO,
	}
	ret.initLogging()
	return ret
}

func (log *Standard) initLogging() {
	// Logging is always done to stderr. It's the responsibility of the
	// launcher (like launchd on OSX, or the autoforking code) to set up stderr
	// to point to the appropriate log file.
	initLoggingBackendOnce.Do(func() {
		logBackend := logging.NewLogBackend(os.Stderr, "", 0)
		logging.SetBackend(logBackend)
		logging.SetLevel(logging.INFO, log.module)
	})
}

func (log *Standard) prepareString(
	ctx context.Context, fmts string) string {
	if ctx == nil {
		return fmts
	}
	logTags, ok := LogTagsFromContext(ctx)
	if !ok || len(logTags) == 0 {
		return fmts
	}
	var tags []string
	for key, tag := range logTags {
		if v := ctx.Value(key); v != nil {
			tags = append(tags, fmt.Sprintf("%s=%s", tag, v))
		}
	}
	return fmts + " [tags:" + strings.Join(tags, ",") + "]"
}

func (log *Standard) Debug(fmt string, arg ...interface{}) {
	log.internal.Debug(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_DEBUG, fmt, arg)
}

func (log *Standard) CDebugf(ctx context.Context, fmt string,
	arg ...interface{}) {
	if log.internal.IsEnabledFor(logging.DEBUG) {
		log.Debug(log.prepareString(ctx, fmt), arg...)
	}
}

func (log *Standard) Info(fmt string, arg ...interface{}) {
	log.internal.Info(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_INFO, fmt, arg)
}

func (log *Standard) CInfof(ctx context.Context, fmt string,
	arg ...interface{}) {
	if log.internal.IsEnabledFor(logging.INFO) {
		log.Info(log.prepareString(ctx, fmt), arg...)
	}
}

func (log *Standard) Notice(fmt string, arg ...interface{}) {
	log.internal.Notice(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_NOTICE, fmt, arg)
}

func (log *Standard) CNoticef(ctx context.Context, fmt string,
	arg ...interface{}) {
	if log.internal.IsEnabledFor(logging.NOTICE) {
		log.Notice(log.prepareString(ctx, fmt), arg...)
	}
}

func (log *Standard) Warning(fmt string, arg ...interface{}) {
	log.internal.Warning(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_WARN, fmt, arg)
}

func (log *Standard) CWarningf(ctx context.Context, fmt string,
	arg ...interface{}) {
	if log.internal.IsEnabledFor(logging.WARNING) {
		log.Warning(log.prepareString(ctx, fmt), arg...)
	}
}

func (log *Standard) Error(fmt string, arg ...interface{}) {
	log.internal.Error(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_ERROR, fmt, arg)
}

func (log *Standard) Errorf(fmt string, arg ...interface{}) {
	log.Error(fmt, arg...)
}

func (log *Standard) CErrorf(ctx context.Context, fmt string,
	arg ...interface{}) {
	if log.internal.IsEnabledFor(logging.ERROR) {
		log.Error(log.prepareString(ctx, fmt), arg...)
	}
}

func (log *Standard) Critical(fmt string, arg ...interface{}) {
	log.internal.Critical(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_CRITICAL, fmt, arg)
}

func (log *Standard) CCriticalf(ctx context.Context, fmt string,
	arg ...interface{}) {
	if log.internal.IsEnabledFor(logging.CRITICAL) {
		log.Critical(log.prepareString(ctx, fmt), arg...)
	}
}

func (log *Standard) Fatalf(fmt string, arg ...interface{}) {
	log.internal.Fatalf(fmt, arg...)
	log.logToExternalLoggers(keybase1.LogLevel_FATAL, fmt, arg)
}

func (log *Standard) CFatalf(ctx context.Context, fmt string,
	arg ...interface{}) {
	log.Fatalf(log.prepareString(ctx, fmt), arg...)
}

func (log *Standard) Profile(fmts string, arg ...interface{}) {
	log.Debug(fmts, arg...)
}

func (log *Standard) Configure(style string, debug bool, filename string) {
	log.configureMutex.Lock()
	defer log.configureMutex.Unlock()

	log.filename = filename

	var logfmt string
	if debug {
		logfmt = fancyFormat
	} else {
		logfmt = defaultFormat
	}

	// Override the format above if an explicit style was specified.
	switch style {
	case "default":
		logfmt = defaultFormat // Default
	case "plain":
		logfmt = plainFormat // Plain
	case "file":
		logfmt = fileFormat // Good for logging to files
	case "fancy":
		logfmt = fancyFormat // Fancy, good for terminals with color
	}

	if debug {
		logging.SetLevel(logging.DEBUG, log.module)
	}

	logging.SetFormatter(logging.MustStringFormatter(logfmt))
}

func OpenLogFile(filename string) (name string, file *os.File, err error) {
	name = filename
	if err = MakeParentDirs(name); err != nil {
		return
	}
	file, err = os.OpenFile(name, (os.O_APPEND | os.O_WRONLY | os.O_CREATE), 0600)
	if err != nil {
		return
	}
	return
}

func FileExists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return false, err
}

func MakeParentDirs(filename string) error {
	dir, _ := filepath.Split(filename)
	exists, err := FileExists(dir)
	if err != nil {
		return err
	}

	if !exists {
		err = os.MkdirAll(dir, permDir)
		if err != nil {
			return err
		}
	}
	return nil
}

func PickFirstError(errors ...error) error {
	for _, e := range errors {
		if e != nil {
			return e
		}
	}
	return nil
}

func (log *Standard) AddExternalLogger(externalLogger ExternalLogger) uint64 {
	log.externalLoggersMutex.Lock()
	defer log.externalLoggersMutex.Unlock()

	handle := log.externalLoggersCount
	log.externalLoggersCount++
	log.externalLoggers[handle] = externalLogger
	return handle
}

func (log *Standard) RemoveExternalLogger(handle uint64) {
	log.externalLoggersMutex.Lock()
	defer log.externalLoggersMutex.Unlock()

	delete(log.externalLoggers, handle)
}

func (log *Standard) logToExternalLoggers(level keybase1.LogLevel, format string, args []interface{}) {
	log.externalLoggersMutex.RLock()
	defer log.externalLoggersMutex.RUnlock()

	// Short circuit logs that are more verbose than the current external log
	// level.
	if level < log.externalLogLevel {
		return
	}

	for _, externalLogger := range log.externalLoggers {
		go externalLogger.Log(level, format, args)
	}
}

func (log *Standard) SetExternalLogLevel(level keybase1.LogLevel) {
	log.externalLoggersMutex.Lock()
	defer log.externalLoggersMutex.Unlock()

	log.externalLogLevel = level
}
<|endoftext|>


<commit_before>package player

import "github.com/lean-poker/poker-player-go/leanpoker"

const VERSION = "Default Go folding player"

func BetRequest(state *leanpoker.Game) int {
	return 1
}

func Showdown(state *leanpoker.Game) {

}

func Version() string {
	return VERSION
}
<commit_msg>Return 1000<commit_after>package player

import "github.com/lean-poker/poker-player-go/leanpoker"

const VERSION = "Default Go folding player"

func BetRequest(state *leanpoker.Game) int {
	return 1000
}

func Showdown(state *leanpoker.Game) {

}

func Version() string {
	return VERSION
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package keybase1

import (
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"strings"
	"time"

	jsonw "github.com/keybase/go-jsonw"
)

const (
	UID_LEN          = 16
	UID_SUFFIX       = 0x00
	UID_SUFFIX_2     = 0x19
	UID_SUFFIX_HEX   = "00"
	UID_SUFFIX_2_HEX = "19"
	PUBLIC_UID       = "ffffffffffffffffffffffffffffff00"
)

// UID for the special "public" user.
var PublicUID = UID(PUBLIC_UID)

const (
	SIG_ID_LEN         = 32
	SIG_ID_SUFFIX      = 0x0f
	SIG_SHORT_ID_BYTES = 27
	SigIDQueryMin      = 16
)

const (
	DeviceIDLen       = 16
	DeviceIDSuffix    = 0x18
	DeviceIDSuffixHex = "18"
)

const (
	KidLen     = 35   // bytes
	KidSuffix  = 0x0a // a byte
	KidVersion = 0x1
)

func Unquote(data []byte) string {
	return strings.Trim(string(data), "\"")
}

func Quote(s string) []byte {
	return []byte("\"" + s + "\"")
}

func KIDFromSlice(b []byte) KID {
	return KID(hex.EncodeToString(b))
}

func KIDFromStringChecked(s string) (KID, error) {

	// It's OK to have a 0-length KID. That means, no such key
	// (or NULL kid).
	if len(s) == 0 {
		return KID(""), nil
	}

	b, err := hex.DecodeString(s)
	if err != nil {
		return KID(""), err
	}

	if len(b) != KidLen {
		return KID(""), fmt.Errorf("KID wrong length; wanted %d but got %d bytes",
			KidLen, len(b))
	}
	if b[len(b)-1] != KidSuffix {
		return KID(""), fmt.Errorf("Bad KID suffix: got 0x%02x, wanted 0x%02x",
			b[len(b)-1], KidSuffix)
	}
	if b[0] != KidVersion {
		return KID(""), fmt.Errorf("Bad KID version; got 0x%02x but wanted 0x%02x",
			b[0], KidVersion)
	}
	return KID(s), nil
}

func KIDFromString(s string) KID {
	// there are no validations for KIDs (length, suffixes)
	return KID(s)
}

func (k KID) IsValid() bool {
	return len(k) > 0
}

func (k KID) String() string {
	return string(k)
}

func (k KID) IsNil() bool {
	return len(k) == 0
}

func (k KID) Exists() bool {
	return !k.IsNil()
}

func (k KID) Equal(v KID) bool {
	return k == v
}

func (k KID) NotEqual(v KID) bool {
	return !k.Equal(v)
}

func (k KID) Match(q string, exact bool) bool {
	if k.IsNil() {
		return false
	}

	if exact {
		return strings.ToLower(k.String()) == strings.ToLower(q)
	}

	if strings.HasPrefix(k.String(), strings.ToLower(q)) {
		return true
	}
	if strings.HasPrefix(k.ToShortIDString(), q) {
		return true
	}
	return false
}

func (k KID) ToBytes() []byte {
	b, err := hex.DecodeString(string(k))
	if err != nil {
		return nil
	}
	return b
}

func (k KID) GetKeyType() byte {
	raw := k.ToBytes()
	if len(raw) < 2 {
		return 0
	}
	return raw[1]
}

func (k KID) ToShortIDString() string {
	return encode(k.ToBytes()[0:12])
}

func (k KID) ToJsonw() *jsonw.Wrapper {
	if k.IsNil() {
		return jsonw.NewNil()
	}
	return jsonw.NewString(string(k))
}

func DeviceIDFromBytes(b [DeviceIDLen]byte) DeviceID {
	return DeviceID(hex.EncodeToString(b[:]))
}

func (d DeviceID) ToBytes(out []byte) error {
	tmp, err := hex.DecodeString(string(d))
	if err != nil {
		return err
	}
	if len(tmp) != DeviceIDLen {
		return fmt.Errorf("Bad device ID; wanted %d bytes but got %d", DeviceIDLen, len(tmp))
	}
	if len(out) != DeviceIDLen {
		return fmt.Errorf("Need to output to a slice with %d bytes", DeviceIDLen)
	}
	copy(out[:], tmp)
	return nil
}

func DeviceIDFromSlice(b []byte) (DeviceID, error) {
	if len(b) != DeviceIDLen {
		return "", fmt.Errorf("invalid byte slice for DeviceID: len == %d, expected %d", len(b), DeviceIDLen)
	}
	var x [DeviceIDLen]byte
	copy(x[:], b)
	return DeviceIDFromBytes(x), nil
}

func DeviceIDFromString(s string) (DeviceID, error) {
	if len(s) != hex.EncodedLen(DeviceIDLen) {
		return "", fmt.Errorf("Bad Device ID length: %d", len(s))
	}
	suffix := s[len(s)-2:]
	if suffix != DeviceIDSuffixHex {
		return "", fmt.Errorf("Bad suffix byte: %s", suffix)
	}
	return DeviceID(s), nil
}

func (d DeviceID) String() string {
	return string(d)
}

func (d DeviceID) IsNil() bool {
	return len(d) == 0
}

func (d DeviceID) Exists() bool {
	return !d.IsNil()
}

func (d DeviceID) Eq(d2 DeviceID) bool {
	return d.Eq(d2)
}

func UIDFromString(s string) (UID, error) {
	if len(s) != hex.EncodedLen(UID_LEN) {
		return "", fmt.Errorf("Bad UID '%s'; must be %d bytes long", s, UID_LEN)
	}
	suffix := s[len(s)-2:]
	if suffix != UID_SUFFIX_HEX && suffix != UID_SUFFIX_2_HEX {
		return "", fmt.Errorf("Bad UID '%s': must end in 0x%x or 0x%x", s, UID_SUFFIX, UID_SUFFIX_2)
	}
	return UID(s), nil
}

// Used by unit tests.
func MakeTestUID(n uint32) UID {
	b := make([]byte, 8)
	binary.LittleEndian.PutUint32(b, n)
	s := hex.EncodeToString(b)
	c := 2*UID_LEN - len(UID_SUFFIX_HEX) - len(s)
	s += strings.Repeat("0", c) + UID_SUFFIX_HEX
	uid, err := UIDFromString(s)
	if err != nil {
		panic(err)
	}
	return uid
}

func (u UID) String() string {
	return string(u)
}

func (u UID) ToBytes() []byte {
	b, err := hex.DecodeString(string(u))
	if err != nil {
		return nil
	}
	return b
}

func (u UID) IsNil() bool {
	return len(u) == 0
}

func (u UID) Exists() bool {
	return !u.IsNil()
}

func (u UID) Equal(v UID) bool {
	return u == v
}

func (u UID) NotEqual(v UID) bool {
	return !u.Equal(v)
}

func (u UID) Less(v UID) bool {
	return u < v
}

// Returns a number in [0, shardCount) which can be treated as roughly
// uniformly distributed. Used for things that need to shard by user.
func (u UID) GetShard(shardCount int) (int, error) {
	bytes, err := hex.DecodeString(string(u))
	if err != nil {
		return 0, err
	}
	n := binary.LittleEndian.Uint32(bytes)
	return int(n % uint32(shardCount)), nil
}

func (s SigID) IsNil() bool {
	return len(s) == 0
}

func (s SigID) Exists() bool {
	return !s.IsNil()
}

func (s SigID) Equal(t SigID) bool {
	return s == t
}

func (s SigID) Match(q string, exact bool) bool {
	if s.IsNil() {
		return false
	}

	if exact {
		return strings.ToLower(s.ToString(true)) == strings.ToLower(q)
	}

	if strings.HasPrefix(s.ToString(true), strings.ToLower(q)) {
		return true
	}

	return false
}

func (s SigID) NotEqual(t SigID) bool {
	return !s.Equal(t)
}

func (s SigID) ToDisplayString(verbose bool) string {
	if verbose {
		return string(s)
	}
	return fmt.Sprintf("%s...", s[0:16])
}

func (s SigID) ToString(suffix bool) string {
	if len(s) == 0 {
		return ""
	}
	if suffix {
		return string(s)
	}
	return string(s[0 : len(s)-2])
}

func SigIDFromString(s string, suffix bool) (SigID, error) {
	blen := SIG_ID_LEN
	if suffix {
		blen++
	}
	if len(s) != hex.EncodedLen(blen) {
		return "", fmt.Errorf("Invalid SigID string length: %d, expected %d (suffix = %v)", len(s), hex.EncodedLen(blen), suffix)
	}
	if suffix {
		return SigID(s), nil
	}
	return SigID(fmt.Sprintf("%s%02x", s, SIG_ID_SUFFIX)), nil
}

func SigIDFromBytes(b [SIG_ID_LEN]byte) SigID {
	s := hex.EncodeToString(b[:])
	return SigID(fmt.Sprintf("%s%02x", s, SIG_ID_SUFFIX))
}

func SigIDFromSlice(b []byte) (SigID, error) {
	if len(b) != SIG_ID_LEN {
		return "", fmt.Errorf("invalid byte slice for SigID: len == %d, expected %d", len(b), SIG_ID_LEN)
	}
	var x [SIG_ID_LEN]byte
	copy(x[:], b)
	return SigIDFromBytes(x), nil
}

func (s SigID) toBytes() []byte {
	b, err := hex.DecodeString(string(s))
	if err != nil {
		return nil
	}
	return b[0:SIG_ID_LEN]
}

func (s SigID) ToMediumID() string {
	return encode(s.toBytes())
}

func (s SigID) ToShortID() string {
	return encode(s.toBytes()[0:SIG_SHORT_ID_BYTES])
}

func encode(b []byte) string {
	return strings.TrimRight(base64.URLEncoding.EncodeToString(b), "=")
}

func FromTime(t Time) time.Time {
	return time.Unix(0, int64(t)*1000000)
}

func ToTime(t time.Time) Time {
	return Time(t.UnixNano() / 1000000)
}

func TimeFromSeconds(seconds int64) Time {
	return Time(seconds * 1000)
}

func (t Time) IsZero() bool        { return t == 0 }
func (t Time) After(t2 Time) bool  { return t > t2 }
func (t Time) Before(t2 Time) bool { return t < t2 }

func FormatTime(t Time) string {
	layout := "2006-01-02 15:04:05 MST"
	return FromTime(t).Format(layout)
}

func (s Status) Error() string {
	if s.Code == 0 {
		return ""
	}
	return fmt.Sprintf("%s (%s/%d)", s.Desc, s.Name, s.Code)
}

func (s InstallStatus) String() string {
	switch s {
	case InstallStatus_UNKNOWN:
		return "Unknown"
	case InstallStatus_ERROR:
		return "Error"
	case InstallStatus_NOT_INSTALLED:
		return "Not Installed"
	case InstallStatus_INSTALLED:
		return "Installed"
	}
	return ""
}

func (s InstallAction) String() string {
	switch s {
	case InstallAction_UNKNOWN:
		return "Unknown"
	case InstallAction_NONE:
		return "None"
	case InstallAction_UPGRADE:
		return "Upgrade"
	case InstallAction_REINSTALL:
		return "Re-Install"
	case InstallAction_INSTALL:
		return "Install"
	}
	return ""
}

func (s ServiceStatus) NeedsInstall() bool {
	return s.InstallAction == InstallAction_INSTALL ||
		s.InstallAction == InstallAction_REINSTALL ||
		s.InstallAction == InstallAction_UPGRADE
}

func (k *KID) UnmarshalJSON(b []byte) error {
	kid, err := KIDFromStringChecked(Unquote(b))
	if err != nil {
		return err
	}
	*k = KID(kid)
	return nil
}

func (k *KID) MarshalJSON() ([]byte, error) {
	return Quote(k.String()), nil
}

func (f Folder) ToString() string {
	prefix := "public/"
	if f.Private {
		prefix = "private/"
	}
	return prefix + f.Name
}

func (t TrackToken) String() string {
	return string(t)
}

func KIDFromRawKey(b []byte, keyType byte) KID {
	tmp := []byte{KidVersion, keyType}
	tmp = append(tmp, b...)
	tmp = append(tmp, byte(KidSuffix))
	return KIDFromSlice(tmp)
}

type APIStatus interface {
	Status() Status
}

type Error struct {
	code    StatusCode
	message string
}

func NewError(code StatusCode, message string) *Error {
	if code == StatusCode_SCOk {
		return nil
	}
	return &Error{code: code, message: message}
}

func FromError(err error) *Error {
	return &Error{code: StatusCode_SCGeneric, message: err.Error()}
}

func StatusOK() Status {
	return Status{Code: int(StatusCode_SCOk), Name: "OK", Desc: "OK"}
}

func StatusFromCode(code StatusCode, message string) Status {
	if code == StatusCode_SCOk {
		return StatusOK()
	}
	return NewError(code, message).Status()
}

func (e *Error) Error() string {
	return e.message
}

func (e *Error) Status() Status {
	return Status{Code: int(e.code), Name: "ERROR", Desc: e.message}
}
<commit_msg>Use SigIDQueryMin for ToDisplayString size<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package keybase1

import (
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"strings"
	"time"

	jsonw "github.com/keybase/go-jsonw"
)

const (
	UID_LEN          = 16
	UID_SUFFIX       = 0x00
	UID_SUFFIX_2     = 0x19
	UID_SUFFIX_HEX   = "00"
	UID_SUFFIX_2_HEX = "19"
	PUBLIC_UID       = "ffffffffffffffffffffffffffffff00"
)

// UID for the special "public" user.
var PublicUID = UID(PUBLIC_UID)

const (
	SIG_ID_LEN         = 32
	SIG_ID_SUFFIX      = 0x0f
	SIG_SHORT_ID_BYTES = 27
	SigIDQueryMin      = 16
)

const (
	DeviceIDLen       = 16
	DeviceIDSuffix    = 0x18
	DeviceIDSuffixHex = "18"
)

const (
	KidLen     = 35   // bytes
	KidSuffix  = 0x0a // a byte
	KidVersion = 0x1
)

func Unquote(data []byte) string {
	return strings.Trim(string(data), "\"")
}

func Quote(s string) []byte {
	return []byte("\"" + s + "\"")
}

func KIDFromSlice(b []byte) KID {
	return KID(hex.EncodeToString(b))
}

func KIDFromStringChecked(s string) (KID, error) {

	// It's OK to have a 0-length KID. That means, no such key
	// (or NULL kid).
	if len(s) == 0 {
		return KID(""), nil
	}

	b, err := hex.DecodeString(s)
	if err != nil {
		return KID(""), err
	}

	if len(b) != KidLen {
		return KID(""), fmt.Errorf("KID wrong length; wanted %d but got %d bytes",
			KidLen, len(b))
	}
	if b[len(b)-1] != KidSuffix {
		return KID(""), fmt.Errorf("Bad KID suffix: got 0x%02x, wanted 0x%02x",
			b[len(b)-1], KidSuffix)
	}
	if b[0] != KidVersion {
		return KID(""), fmt.Errorf("Bad KID version; got 0x%02x but wanted 0x%02x",
			b[0], KidVersion)
	}
	return KID(s), nil
}

func KIDFromString(s string) KID {
	// there are no validations for KIDs (length, suffixes)
	return KID(s)
}

func (k KID) IsValid() bool {
	return len(k) > 0
}

func (k KID) String() string {
	return string(k)
}

func (k KID) IsNil() bool {
	return len(k) == 0
}

func (k KID) Exists() bool {
	return !k.IsNil()
}

func (k KID) Equal(v KID) bool {
	return k == v
}

func (k KID) NotEqual(v KID) bool {
	return !k.Equal(v)
}

func (k KID) Match(q string, exact bool) bool {
	if k.IsNil() {
		return false
	}

	if exact {
		return strings.ToLower(k.String()) == strings.ToLower(q)
	}

	if strings.HasPrefix(k.String(), strings.ToLower(q)) {
		return true
	}
	if strings.HasPrefix(k.ToShortIDString(), q) {
		return true
	}
	return false
}

func (k KID) ToBytes() []byte {
	b, err := hex.DecodeString(string(k))
	if err != nil {
		return nil
	}
	return b
}

func (k KID) GetKeyType() byte {
	raw := k.ToBytes()
	if len(raw) < 2 {
		return 0
	}
	return raw[1]
}

func (k KID) ToShortIDString() string {
	return encode(k.ToBytes()[0:12])
}

func (k KID) ToJsonw() *jsonw.Wrapper {
	if k.IsNil() {
		return jsonw.NewNil()
	}
	return jsonw.NewString(string(k))
}

func DeviceIDFromBytes(b [DeviceIDLen]byte) DeviceID {
	return DeviceID(hex.EncodeToString(b[:]))
}

func (d DeviceID) ToBytes(out []byte) error {
	tmp, err := hex.DecodeString(string(d))
	if err != nil {
		return err
	}
	if len(tmp) != DeviceIDLen {
		return fmt.Errorf("Bad device ID; wanted %d bytes but got %d", DeviceIDLen, len(tmp))
	}
	if len(out) != DeviceIDLen {
		return fmt.Errorf("Need to output to a slice with %d bytes", DeviceIDLen)
	}
	copy(out[:], tmp)
	return nil
}

func DeviceIDFromSlice(b []byte) (DeviceID, error) {
	if len(b) != DeviceIDLen {
		return "", fmt.Errorf("invalid byte slice for DeviceID: len == %d, expected %d", len(b), DeviceIDLen)
	}
	var x [DeviceIDLen]byte
	copy(x[:], b)
	return DeviceIDFromBytes(x), nil
}

func DeviceIDFromString(s string) (DeviceID, error) {
	if len(s) != hex.EncodedLen(DeviceIDLen) {
		return "", fmt.Errorf("Bad Device ID length: %d", len(s))
	}
	suffix := s[len(s)-2:]
	if suffix != DeviceIDSuffixHex {
		return "", fmt.Errorf("Bad suffix byte: %s", suffix)
	}
	return DeviceID(s), nil
}

func (d DeviceID) String() string {
	return string(d)
}

func (d DeviceID) IsNil() bool {
	return len(d) == 0
}

func (d DeviceID) Exists() bool {
	return !d.IsNil()
}

func (d DeviceID) Eq(d2 DeviceID) bool {
	return d.Eq(d2)
}

func UIDFromString(s string) (UID, error) {
	if len(s) != hex.EncodedLen(UID_LEN) {
		return "", fmt.Errorf("Bad UID '%s'; must be %d bytes long", s, UID_LEN)
	}
	suffix := s[len(s)-2:]
	if suffix != UID_SUFFIX_HEX && suffix != UID_SUFFIX_2_HEX {
		return "", fmt.Errorf("Bad UID '%s': must end in 0x%x or 0x%x", s, UID_SUFFIX, UID_SUFFIX_2)
	}
	return UID(s), nil
}

// Used by unit tests.
func MakeTestUID(n uint32) UID {
	b := make([]byte, 8)
	binary.LittleEndian.PutUint32(b, n)
	s := hex.EncodeToString(b)
	c := 2*UID_LEN - len(UID_SUFFIX_HEX) - len(s)
	s += strings.Repeat("0", c) + UID_SUFFIX_HEX
	uid, err := UIDFromString(s)
	if err != nil {
		panic(err)
	}
	return uid
}

func (u UID) String() string {
	return string(u)
}

func (u UID) ToBytes() []byte {
	b, err := hex.DecodeString(string(u))
	if err != nil {
		return nil
	}
	return b
}

func (u UID) IsNil() bool {
	return len(u) == 0
}

func (u UID) Exists() bool {
	return !u.IsNil()
}

func (u UID) Equal(v UID) bool {
	return u == v
}

func (u UID) NotEqual(v UID) bool {
	return !u.Equal(v)
}

func (u UID) Less(v UID) bool {
	return u < v
}

// Returns a number in [0, shardCount) which can be treated as roughly
// uniformly distributed. Used for things that need to shard by user.
func (u UID) GetShard(shardCount int) (int, error) {
	bytes, err := hex.DecodeString(string(u))
	if err != nil {
		return 0, err
	}
	n := binary.LittleEndian.Uint32(bytes)
	return int(n % uint32(shardCount)), nil
}

func (s SigID) IsNil() bool {
	return len(s) == 0
}

func (s SigID) Exists() bool {
	return !s.IsNil()
}

func (s SigID) Equal(t SigID) bool {
	return s == t
}

func (s SigID) Match(q string, exact bool) bool {
	if s.IsNil() {
		return false
	}

	if exact {
		return strings.ToLower(s.ToString(true)) == strings.ToLower(q)
	}

	if strings.HasPrefix(s.ToString(true), strings.ToLower(q)) {
		return true
	}

	return false
}

func (s SigID) NotEqual(t SigID) bool {
	return !s.Equal(t)
}

func (s SigID) ToDisplayString(verbose bool) string {
	if verbose {
		return string(s)
	}
	return fmt.Sprintf("%s...", s[0:SigIDQueryMin])
}

func (s SigID) ToString(suffix bool) string {
	if len(s) == 0 {
		return ""
	}
	if suffix {
		return string(s)
	}
	return string(s[0 : len(s)-2])
}

func SigIDFromString(s string, suffix bool) (SigID, error) {
	blen := SIG_ID_LEN
	if suffix {
		blen++
	}
	if len(s) != hex.EncodedLen(blen) {
		return "", fmt.Errorf("Invalid SigID string length: %d, expected %d (suffix = %v)", len(s), hex.EncodedLen(blen), suffix)
	}
	if suffix {
		return SigID(s), nil
	}
	return SigID(fmt.Sprintf("%s%02x", s, SIG_ID_SUFFIX)), nil
}

func SigIDFromBytes(b [SIG_ID_LEN]byte) SigID {
	s := hex.EncodeToString(b[:])
	return SigID(fmt.Sprintf("%s%02x", s, SIG_ID_SUFFIX))
}

func SigIDFromSlice(b []byte) (SigID, error) {
	if len(b) != SIG_ID_LEN {
		return "", fmt.Errorf("invalid byte slice for SigID: len == %d, expected %d", len(b), SIG_ID_LEN)
	}
	var x [SIG_ID_LEN]byte
	copy(x[:], b)
	return SigIDFromBytes(x), nil
}

func (s SigID) toBytes() []byte {
	b, err := hex.DecodeString(string(s))
	if err != nil {
		return nil
	}
	return b[0:SIG_ID_LEN]
}

func (s SigID) ToMediumID() string {
	return encode(s.toBytes())
}

func (s SigID) ToShortID() string {
	return encode(s.toBytes()[0:SIG_SHORT_ID_BYTES])
}

func encode(b []byte) string {
	return strings.TrimRight(base64.URLEncoding.EncodeToString(b), "=")
}

func FromTime(t Time) time.Time {
	return time.Unix(0, int64(t)*1000000)
}

func ToTime(t time.Time) Time {
	return Time(t.UnixNano() / 1000000)
}

func TimeFromSeconds(seconds int64) Time {
	return Time(seconds * 1000)
}

func (t Time) IsZero() bool        { return t == 0 }
func (t Time) After(t2 Time) bool  { return t > t2 }
func (t Time) Before(t2 Time) bool { return t < t2 }

func FormatTime(t Time) string {
	layout := "2006-01-02 15:04:05 MST"
	return FromTime(t).Format(layout)
}

func (s Status) Error() string {
	if s.Code == 0 {
		return ""
	}
	return fmt.Sprintf("%s (%s/%d)", s.Desc, s.Name, s.Code)
}

func (s InstallStatus) String() string {
	switch s {
	case InstallStatus_UNKNOWN:
		return "Unknown"
	case InstallStatus_ERROR:
		return "Error"
	case InstallStatus_NOT_INSTALLED:
		return "Not Installed"
	case InstallStatus_INSTALLED:
		return "Installed"
	}
	return ""
}

func (s InstallAction) String() string {
	switch s {
	case InstallAction_UNKNOWN:
		return "Unknown"
	case InstallAction_NONE:
		return "None"
	case InstallAction_UPGRADE:
		return "Upgrade"
	case InstallAction_REINSTALL:
		return "Re-Install"
	case InstallAction_INSTALL:
		return "Install"
	}
	return ""
}

func (s ServiceStatus) NeedsInstall() bool {
	return s.InstallAction == InstallAction_INSTALL ||
		s.InstallAction == InstallAction_REINSTALL ||
		s.InstallAction == InstallAction_UPGRADE
}

func (k *KID) UnmarshalJSON(b []byte) error {
	kid, err := KIDFromStringChecked(Unquote(b))
	if err != nil {
		return err
	}
	*k = KID(kid)
	return nil
}

func (k *KID) MarshalJSON() ([]byte, error) {
	return Quote(k.String()), nil
}

func (f Folder) ToString() string {
	prefix := "public/"
	if f.Private {
		prefix = "private/"
	}
	return prefix + f.Name
}

func (t TrackToken) String() string {
	return string(t)
}

func KIDFromRawKey(b []byte, keyType byte) KID {
	tmp := []byte{KidVersion, keyType}
	tmp = append(tmp, b...)
	tmp = append(tmp, byte(KidSuffix))
	return KIDFromSlice(tmp)
}

type APIStatus interface {
	Status() Status
}

type Error struct {
	code    StatusCode
	message string
}

func NewError(code StatusCode, message string) *Error {
	if code == StatusCode_SCOk {
		return nil
	}
	return &Error{code: code, message: message}
}

func FromError(err error) *Error {
	return &Error{code: StatusCode_SCGeneric, message: err.Error()}
}

func StatusOK() Status {
	return Status{Code: int(StatusCode_SCOk), Name: "OK", Desc: "OK"}
}

func StatusFromCode(code StatusCode, message string) Status {
	if code == StatusCode_SCOk {
		return StatusOK()
	}
	return NewError(code, message).Status()
}

func (e *Error) Error() string {
	return e.message
}

func (e *Error) Status() Status {
	return Status{Code: int(e.code), Name: "ERROR", Desc: e.message}
}
<|endoftext|>


<commit_before>package RethinkDBStorage

import (
	"github.com/RangelReale/osin"
	"github.com/mitchellh/mapstructure"
	r "gopkg.in/dancannon/gorethink.v1"
)

const (
	clientsTable      = "oauth_clients"
	authorizeTable    = "oauth_authorize_data"
	accessTable       = "oauth_access_data"
	accessTokenField  = "AccessToken"
	refreshTokenField = "RefreshToken"
)

// RethinkDBStorage implements storage for osin
type RethinkDBStorage struct {
	session *r.Session
}

// New initializes and returns a new RethinkDBStorage
func New(session *r.Session) *RethinkDBStorage {
	storage := &RethinkDBStorage{session}
	return storage
}

// Clone the storage if needed.
func (s *RethinkDBStorage) Clone() osin.Storage {
	return s
}

// Close the resources the Storage potentially holds
func (s *RethinkDBStorage) Close() {}

// CreateClient inserts a new client
func (s *RethinkDBStorage) CreateClient(c osin.Client) error {
	_, err := r.Table(clientsTable).Insert(c).RunWrite(s.session)
	return err
}

// GetClient returns client with given ID
func (s *RethinkDBStorage) GetClient(clientID string) (osin.Client, error) {
	result, err := r.Table(clientsTable).Filter(r.Row.Field("Id").Eq(clientID)).Run(s.session)
	if err != nil {
		return nil, err
	}
	defer result.Close()

	var clientMap map[string]interface{}
	err = result.One(&clientMap)
	if err != nil {
		return nil, err
	}

	var clientStruct osin.DefaultClient
	err = mapstructure.Decode(clientMap, &clientStruct)
	if err != nil {
		return nil, err
	}

	return &clientStruct, nil
}

// UpdateClient updates given client
func (s *RethinkDBStorage) UpdateClient(c osin.Client) error {
	result, err := r.Table(clientsTable).Filter(r.Row.Field("Id").Eq(c.GetId())).Run(s.session)
	if err != nil {
		return err
	}
	defer result.Close()

	var clientMap map[string]interface{}
	err = result.One(&clientMap)
	if err != nil {
		return err
	}

	_, err = r.Table(clientsTable).Get(clientMap["id"]).Update(c).RunWrite(s.session)
	return err
}

// DeleteClient deletes given client
func (s *RethinkDBStorage) DeleteClient(c osin.Client) error {
	result, err := r.Table(clientsTable).Filter(r.Row.Field("Id").Eq(c.GetId())).Run(s.session)
	if err != nil {
		return err
	}
	defer result.Close()

	var clientMap map[string]interface{}
	err = result.One(&clientMap)
	if err != nil {
		return err
	}

	_, err = r.Table(clientsTable).Get(clientMap["id"]).Delete().RunWrite(s.session)
	return err
}

// SaveAuthorize creates a new authorization
func (s *RethinkDBStorage) SaveAuthorize(data *osin.AuthorizeData) error {
	_, err := r.Table(authorizeTable).Insert(data).RunWrite(s.session)
	return err
}

// LoadAuthorize gets authorization data with given code
func (s *RethinkDBStorage) LoadAuthorize(code string) (*osin.AuthorizeData, error) {
	result, err := r.Table(authorizeTable).Filter(r.Row.Field("Code").Eq(code)).Run(s.session)
	if err != nil {
		return nil, err
	}
	defer result.Close()

	var dataMap map[string]interface{}
	err = result.One(&dataMap)
	if err != nil {
		return nil, err
	}

	var client osin.Client
	clientID := dataMap["Client"].(map[string]interface{})["Id"].(string)
	client, err = s.GetClient(clientID)
	if err != nil {
		return nil, err
	}
	dataMap["Client"] = client

	var dataStruct osin.AuthorizeData
	err = mapstructure.Decode(dataMap, &dataStruct)
	if err != nil {
		return nil, err
	}

	return &dataStruct, nil
}

// RemoveAuthorize deletes given authorization
func (s *RethinkDBStorage) RemoveAuthorize(code string) error {
	result, err := r.Table(authorizeTable).Filter(r.Row.Field("Code").Eq(code)).Run(s.session)
	if err != nil {
		return err
	}
	defer result.Close()

	var dataMap map[string]interface{}
	err = result.One(&dataMap)
	if err != nil {
		return err
	}

	_, err = r.Table(authorizeTable).Get(dataMap["id"]).Delete().RunWrite(s.session)
	return err
}

// SaveAccess creates a new access data
func (s *RethinkDBStorage) SaveAccess(data *osin.AccessData) error {
	// limit nested AccessData to one level
	if data.AccessData != nil {
		data.AccessData.AccessData = nil
	}

	_, err := r.Table(accessTable).Insert(data).RunWrite(s.session)
	return err
}

// LoadAccess gets access data with given access token
func (s *RethinkDBStorage) LoadAccess(accessToken string) (*osin.AccessData, error) {
	return s.getAccessData(accessTokenField, accessToken)
}

// RemoveAccess deletes AccessData with given access token
func (s *RethinkDBStorage) RemoveAccess(accessToken string) error {
	return s.removeAccessData(accessTokenField, accessToken)
}

// LoadRefresh gets access data with given refresh token
func (s *RethinkDBStorage) LoadRefresh(refreshToken string) (*osin.AccessData, error) {
	return s.getAccessData(refreshTokenField, refreshToken)
}

// RemoveRefresh deletes AccessData with given refresh token
func (s *RethinkDBStorage) RemoveRefresh(refreshToken string) error {
	return s.removeAccessData(refreshTokenField, refreshToken)
}

// getAccessData is a common function to get AccessData by field
func (s *RethinkDBStorage) getAccessData(fieldName, token string) (*osin.AccessData, error) {
	result, err := r.Table(accessTable).Filter(r.Row.Field(fieldName).Eq(token)).Run(s.session)
	if err != nil {
		return nil, err
	}
	defer result.Close()

	var dataMap map[string]interface{}
	err = result.One(&dataMap)
	if err != nil {
		return nil, err
	}

	var client osin.Client
	clientID := dataMap["Client"].(map[string]interface{})["Id"].(string)
	client, err = s.GetClient(clientID)
	if err != nil {
		return nil, err
	}
	dataMap["Client"] = client

	if authorizeData := dataMap["AuthorizeData"]; authorizeData != nil {
		if authorizeDataClient := authorizeData.(map[string]interface{})["Client"]; authorizeDataClient != nil {
			var authorizeDataClientStruct osin.Client
			if authorizeDataClientID := authorizeDataClient.(map[string]interface{})["Id"]; authorizeDataClientID != nil {
				authorizeDataClientStruct, err = s.GetClient(authorizeDataClientID.(string))
				if err != nil {
					return nil, err
				}
				dataMap["AuthorizeData"].(map[string]interface{})["Client"] = authorizeDataClientStruct
			}
		}
	}

	if accessData := dataMap["AccessData"]; accessData != nil {

		if accessDataClient := accessData.(map[string]interface{})["Client"]; accessDataClient != nil {
			var accessDataClientStruct osin.Client
			if accessDataClientID := accessDataClient.(map[string]interface{})["Id"]; accessDataClientID != nil {
				accessDataClientStruct, err = s.GetClient(accessDataClientID.(string))
				if err != nil {
					return nil, err
				}
				dataMap["AccessData"].(map[string]interface{})["Client"] = accessDataClientStruct
			}
		}
	}

	var dataStruct osin.AccessData
	err = mapstructure.Decode(dataMap, &dataStruct)
	if err != nil {
		return nil, err
	}

	return &dataStruct, nil
}

// removeAccessData is a common function to remove AccessData by field
func (s *RethinkDBStorage) removeAccessData(fieldName, token string) error {
	result, err := r.Table(accessTable).Filter(r.Row.Field(fieldName).Eq(token)).Run(s.session)
	if err != nil {
		return err
	}
	defer result.Close()

	var dataMap map[string]interface{}
	err = result.One(&dataMap)
	if err != nil {
		return err
	}

	_, err = r.Table(accessTable).Get(dataMap["id"]).Delete().RunWrite(s.session)
	return err
}
<commit_msg>allow external packages to define their own osin.Client implementation<commit_after>package RethinkDBStorage

import (
	"github.com/RangelReale/osin"
	"github.com/mitchellh/mapstructure"
	r "gopkg.in/dancannon/gorethink.v1"
)

const (
	clientsTable      = "oauth_clients"
	authorizeTable    = "oauth_authorize_data"
	accessTable       = "oauth_access_data"
	accessTokenField  = "AccessToken"
	refreshTokenField = "RefreshToken"
)

// ClientGetter returns a function that returns an osin.Client allowing
// other packages to define their own Clients
var ClientGetter = func() osin.Client {
	return &osin.DefaultClient{}
}

// RethinkDBStorage implements storage for osin
type RethinkDBStorage struct {
	session *r.Session
}

// New initializes and returns a new RethinkDBStorage
func New(session *r.Session) *RethinkDBStorage {
	storage := &RethinkDBStorage{session}
	return storage
}

// Clone the storage if needed.
func (s *RethinkDBStorage) Clone() osin.Storage {
	return s
}

// Close the resources the Storage potentially holds
func (s *RethinkDBStorage) Close() {}

// CreateClient inserts a new client
func (s *RethinkDBStorage) CreateClient(c osin.Client) error {
	_, err := r.Table(clientsTable).Insert(c).RunWrite(s.session)
	return err
}

// GetClient returns client with given ID
func (s *RethinkDBStorage) GetClient(clientID string) (osin.Client, error) {
	result, err := r.Table(clientsTable).Filter(r.Row.Field("Id").Eq(clientID)).Run(s.session)
	if err != nil {
		return nil, err
	}
	defer result.Close()

	client := ClientGetter()
	err = result.One(client)
	if err != nil {
		return nil, err
	}

	return client, nil
}

// UpdateClient updates given client
func (s *RethinkDBStorage) UpdateClient(c osin.Client) error {
	result, err := r.Table(clientsTable).Filter(r.Row.Field("Id").Eq(c.GetId())).Run(s.session)
	if err != nil {
		return err
	}
	defer result.Close()

	var clientMap map[string]interface{}
	err = result.One(&clientMap)
	if err != nil {
		return err
	}

	_, err = r.Table(clientsTable).Get(clientMap["id"]).Update(c).RunWrite(s.session)
	return err
}

// DeleteClient deletes given client
func (s *RethinkDBStorage) DeleteClient(c osin.Client) error {
	result, err := r.Table(clientsTable).Filter(r.Row.Field("Id").Eq(c.GetId())).Run(s.session)
	if err != nil {
		return err
	}
	defer result.Close()

	var clientMap map[string]interface{}
	err = result.One(&clientMap)
	if err != nil {
		return err
	}

	_, err = r.Table(clientsTable).Get(clientMap["id"]).Delete().RunWrite(s.session)
	return err
}

// SaveAuthorize creates a new authorization
func (s *RethinkDBStorage) SaveAuthorize(data *osin.AuthorizeData) error {
	_, err := r.Table(authorizeTable).Insert(data).RunWrite(s.session)
	return err
}

// LoadAuthorize gets authorization data with given code
func (s *RethinkDBStorage) LoadAuthorize(code string) (*osin.AuthorizeData, error) {
	result, err := r.Table(authorizeTable).Filter(r.Row.Field("Code").Eq(code)).Run(s.session)
	if err != nil {
		return nil, err
	}
	defer result.Close()

	var dataMap map[string]interface{}
	err = result.One(&dataMap)
	if err != nil {
		return nil, err
	}

	var client osin.Client
	clientID := dataMap["Client"].(map[string]interface{})["Id"].(string)
	client, err = s.GetClient(clientID)
	if err != nil {
		return nil, err
	}
	dataMap["Client"] = client

	var dataStruct osin.AuthorizeData
	err = mapstructure.Decode(dataMap, &dataStruct)
	if err != nil {
		return nil, err
	}

	return &dataStruct, nil
}

// RemoveAuthorize deletes given authorization
func (s *RethinkDBStorage) RemoveAuthorize(code string) error {
	result, err := r.Table(authorizeTable).Filter(r.Row.Field("Code").Eq(code)).Run(s.session)
	if err != nil {
		return err
	}
	defer result.Close()

	var dataMap map[string]interface{}
	err = result.One(&dataMap)
	if err != nil {
		return err
	}

	_, err = r.Table(authorizeTable).Get(dataMap["id"]).Delete().RunWrite(s.session)
	return err
}

// SaveAccess creates a new access data
func (s *RethinkDBStorage) SaveAccess(data *osin.AccessData) error {
	// limit nested AccessData to one level
	if data.AccessData != nil {
		data.AccessData.AccessData = nil
	}

	_, err := r.Table(accessTable).Insert(data).RunWrite(s.session)
	return err
}

// LoadAccess gets access data with given access token
func (s *RethinkDBStorage) LoadAccess(accessToken string) (*osin.AccessData, error) {
	return s.getAccessData(accessTokenField, accessToken)
}

// RemoveAccess deletes AccessData with given access token
func (s *RethinkDBStorage) RemoveAccess(accessToken string) error {
	return s.removeAccessData(accessTokenField, accessToken)
}

// LoadRefresh gets access data with given refresh token
func (s *RethinkDBStorage) LoadRefresh(refreshToken string) (*osin.AccessData, error) {
	return s.getAccessData(refreshTokenField, refreshToken)
}

// RemoveRefresh deletes AccessData with given refresh token
func (s *RethinkDBStorage) RemoveRefresh(refreshToken string) error {
	return s.removeAccessData(refreshTokenField, refreshToken)
}

// getAccessData is a common function to get AccessData by field
func (s *RethinkDBStorage) getAccessData(fieldName, token string) (*osin.AccessData, error) {
	result, err := r.Table(accessTable).Filter(r.Row.Field(fieldName).Eq(token)).Run(s.session)
	if err != nil {
		return nil, err
	}
	defer result.Close()

	var dataMap map[string]interface{}
	err = result.One(&dataMap)
	if err != nil {
		return nil, err
	}

	var client osin.Client
	clientID := dataMap["Client"].(map[string]interface{})["Id"].(string)
	client, err = s.GetClient(clientID)
	if err != nil {
		return nil, err
	}
	dataMap["Client"] = client

	if authorizeData := dataMap["AuthorizeData"]; authorizeData != nil {
		if authorizeDataClient := authorizeData.(map[string]interface{})["Client"]; authorizeDataClient != nil {
			var authorizeDataClientStruct osin.Client
			if authorizeDataClientID := authorizeDataClient.(map[string]interface{})["Id"]; authorizeDataClientID != nil {
				authorizeDataClientStruct, err = s.GetClient(authorizeDataClientID.(string))
				if err != nil {
					return nil, err
				}
				dataMap["AuthorizeData"].(map[string]interface{})["Client"] = authorizeDataClientStruct
			}
		}
	}

	if accessData := dataMap["AccessData"]; accessData != nil {

		if accessDataClient := accessData.(map[string]interface{})["Client"]; accessDataClient != nil {
			var accessDataClientStruct osin.Client
			if accessDataClientID := accessDataClient.(map[string]interface{})["Id"]; accessDataClientID != nil {
				accessDataClientStruct, err = s.GetClient(accessDataClientID.(string))
				if err != nil {
					return nil, err
				}
				dataMap["AccessData"].(map[string]interface{})["Client"] = accessDataClientStruct
			}
		}
	}

	var dataStruct osin.AccessData
	err = mapstructure.Decode(dataMap, &dataStruct)
	if err != nil {
		return nil, err
	}

	return &dataStruct, nil
}

// removeAccessData is a common function to remove AccessData by field
func (s *RethinkDBStorage) removeAccessData(fieldName, token string) error {
	result, err := r.Table(accessTable).Filter(r.Row.Field(fieldName).Eq(token)).Run(s.session)
	if err != nil {
		return err
	}
	defer result.Close()

	var dataMap map[string]interface{}
	err = result.One(&dataMap)
	if err != nil {
		return err
	}

	_, err = r.Table(accessTable).Get(dataMap["id"]).Delete().RunWrite(s.session)
	return err
}
<|endoftext|>