<commit_before>package main

import (
	"crypto/sha256"
	"errors"
	"fmt"
	"github.com/nu7hatch/gouuid"
	"math/big"
	"net"
	"strings"
	"time"
)

func stringToBig(str string) *big.Int {
	sha := sha256.New()
	sha.Write([]byte(str))
	b := sha.Sum(nil)

	ret := new(big.Int).SetBytes(b)
	return ret
}

type Bully struct {
	myId     *big.Int
	cmdChan  chan *command
	ctrlChan chan *control
	ln       net.Listener
}

type Candidate struct {
	Id   *big.Int
	Addr net.Addr
}

var ErrUnknownError = errors.New("Unknown")

func (self *Bully) MyId() *big.Int {
	return self.myId
}

func (self *Bully) AddCandidate(addrStr string, id *big.Int, timeout time.Duration) error {
	addr, err := net.ResolveTCPAddr("tcp", addrStr)
	if err != nil {
		return err
	}
	if addr == nil {
		return ErrUnknownError
	}
	if timeout <= 1*time.Second {
		timeout = 2 * time.Second
	}
	ctrl := new(control)
	ctrl.addr = addr
	ctrl.id = id
	ctrl.cmd = ctrlADD
	ctrl.timeout = timeout
	replyChan := make(chan *controlReply)
	ctrl.replyChan = replyChan

	self.ctrlChan <- ctrl
	reply := <-replyChan
	if reply == nil {
		return ErrUnknownError
	}
	if reply.err != nil {
		return reply.err
	}
	return nil
}

func (self *Bully) CandidateList() []*Candidate {
	ctrl := new(control)
	ctrl.cmd = ctrlQUERY_CANDY
	replyChan := make(chan *controlReply)
	ctrl.replyChan = replyChan

	self.ctrlChan <- ctrl

	ret := make([]*Candidate, 0, 10)
	for reply := range replyChan {
		c := new(Candidate)
		c.Id = reply.id
		c.Addr = reply.addr
		ret = append(ret, c)
	}
	return ret
}

func commandCollector(src *big.Int, conn net.Conn, cmdChan chan<- *command, timeout time.Duration) {
	defer func() {
		conn.Close()
		cmd := new(command)
		cmd.src = src
		cmd.Cmd = cmdBYE
		cmdChan <- cmd
	}()
	for {
		cmd, err := readCommand(conn)
		if err != nil {
			return
		}
		if cmd.Cmd == cmdITSME || cmd.Cmd == cmdBYE {
			return
		}
		cmd.src = src
		cmd.replyWriter = conn
		select {
		case cmdChan <- cmd:
			continue
		case <-time.After(timeout):
			return
		}
	}
}

func NewBully(ln net.Listener, myId *big.Int) *Bully {
	ret := new(Bully)
	if myId == nil {
		uu, _ := uuid.NewV4()
		ret.myId = stringToBig(uu.String())
	} else {
		ret.myId = myId
	}
	ret.cmdChan = make(chan *command)
	ret.ctrlChan = make(chan *control)
	ret.ln = ln
	go ret.listen(ln)
	go ret.process()
	return ret
}

type node struct {
	id   *big.Int
	conn net.Conn
}

func insertNode(l []*node, id *big.Int, conn net.Conn) ([]*node, bool) {
	n := findNode(l, id)
	if nil != n {
		return l, false
	}
	n = new(node)
	n.id = id
	n.conn = conn
	return append(l, n), true
}

func removeNode(l []*node, id *big.Int) []*node {
	idx := -1
	for i, n := range l {
		if n.id.Cmp(id) == 0 {
			idx = i
		}
	}
	if idx >= 0 {
		l[idx] = l[len(l)-1]
		l[len(l)-1] = nil
		l = l[:len(l)-1]
	}
	return l
}

func findNode(l []*node, id *big.Int) *node {
	for _, n := range l {
		if n.id.Cmp(id) == 0 {
			return n
		}
	}
	return nil
}

const (
	ctrlADD = iota
	ctrlQUERY_LEADER
	ctrlQUERY_CANDY
)

type controlReply struct {
	addr net.Addr
	id   *big.Int
	err  error
}

type control struct {
	cmd       int
	addr      net.Addr
	id        *big.Int
	timeout   time.Duration
	replyChan chan<- *controlReply
}

var ErrUnmatchedId = errors.New("Unmatched node id")

func (self *Bully) handshake(addr net.Addr, id *big.Int, candy []*node, timeout time.Duration) ([]*node, error) {
	if id != nil {
		cmp := id.Cmp(self.myId)
		if cmp > 0 {
			n := findNode(candy, id)
			if n != nil {
				return candy, nil
			}
		} else if cmp < 0 {
			// If we know the id of a node,
			// then we only connect to the nodes with higher id,
			// and let the nodes with lower id connect us.
			return candy, nil
		} else {
			// It is ourselves, don't need to add it.
			return candy, nil
		}
	}
	conn, err := net.DialTimeout("tcp", addr.String(), timeout)
	if err != nil {
		return candy, err
	}
	cmd := new(command)
	cmd.Cmd = cmdHELLO
	cmd.Body = self.myId.Bytes()
	err = writeCommand(conn, cmd)
	if err != nil {
		return candy, err
	}
	reply, err := readCommand(conn)
	if err != nil {
		return candy, err
	}
	if reply.Cmd != cmdHELLO_REPLY {
		fmt.Printf("Not a good move\n")
		conn.Close()
		return candy, nil
	}
	rId := new(big.Int).SetBytes(reply.Body)
	if id != nil {
		if rId.Cmp(id) != 0 {
			return candy, ErrUnmatchedId
		}
	}
	candy, _ = insertNode(candy, rId, conn)
	go commandCollector(rId, conn, self.cmdChan, 10*time.Second)

	return candy, nil
}

var ErrNeedNewElection = errors.New("Need another round")

func (self *Bully) elect(candy []*node, timeout time.Duration) (leader *node, err error) {
	higherCandy := make([]*node, 0, len(candy))
	for _, c := range candy {
		if c.id.Cmp(self.myId) > 0 {
			cmd := new(command)
			cmd.Cmd = cmdELECT
			err := writeCommand(c.conn, cmd)
			if err == nil {
				higherCandy = append(higherCandy, c)
			}
		}
	}

	// No one is higher than me.
	// I am the leader.
	if len(higherCandy) <= 0 {
		leader.conn = nil
		leader.id = self.myId
		for _, c := range candy {
			cmd := new(command)
			cmd.Cmd = cmdCOORDIN
			writeCommand(c.conn, cmd)
		}
		return
	}
	slaved := false
	for {
		select {
		case cmd := <-self.cmdChan:
			switch cmd.Cmd {
			case cmdHELLO:
				reply := new(command)
				reply.Cmd = cmdTRY_LATER
				writeCommand(cmd.replyWriter, reply)
			case cmdELECT:
				reply := new(command)
				reply.Cmd = cmdELECT_OK
				writeCommand(cmd.replyWriter, reply)
			case cmdELECT_OK:
				n := findNode(higherCandy, cmd.src)
				if n == nil {
					continue
				}
				slaved = true
			case cmdCOORDIN:
				n := findNode(candy, cmd.src)
				if n == nil {
					err = ErrNeedNewElection
					return
				}
				if n.id.Cmp(self.myId) < 0 {
					err = ErrNeedNewElection
					return
				}
				leader = n
				return
			}
		case <-time.After(timeout):
			break
		}
	}

	// No one replied within time out.
	// I am the leader.
	if !slaved {
		leader.conn = nil
		leader.id = self.myId
		for _, c := range candy {
			cmd := new(command)
			cmd.Cmd = cmdCOORDIN
			writeCommand(c.conn, cmd)
		}
		return
	}
	err = ErrNeedNewElection
	return
}

func (self *Bully) myAddress() net.Addr {
	addrStr := self.ln.Addr().String()
	ae := strings.Split(addrStr, ":")
	addrStr = fmt.Sprintf("127.0.0.1:%v", ae[len(ae)-1])
	ret, _ := net.ResolveTCPAddr("tcp", addrStr)
	return ret
}

func (self *Bully) process() {
	candy := make([]*node, 0, 10)
	var leader *node
	leaderTimeout := 10 * time.Second
	for {
		select {
		case cmd := <-self.cmdChan:
			fmt.Printf("Command: %v; myid=%v\n", cmd.Cmd, self.myId)
			switch cmd.Cmd {
			case cmdHELLO:
				if cmd.src.Cmp(self.myId) == 0 {
					reply := new(command)
					reply.Cmd = cmdITSME
					err := writeCommand(cmd.replyWriter, reply)
					if err != nil {
						cmd.replyWriter.Close()
					}
					continue
				}
				n := findNode(candy, cmd.src)
				if n == nil {
					reply := new(command)
					reply.Cmd = cmdHELLO_REPLY
					reply.Body = self.myId.Bytes()
					err := writeCommand(cmd.replyWriter, reply)
					if err != nil {
						cmd.replyWriter.Close()
						continue
					}
					candy, _ = insertNode(candy, cmd.src, cmd.replyWriter)
				} else {
					reply := new(command)
					reply.Cmd = cmdDUP_CONN
					writeCommand(cmd.replyWriter, reply)
					cmd.replyWriter.Close()
				}
			case cmdBYE:
				candy = removeNode(candy, cmd.src)
			case cmdELECT:
				reply := new(command)
				reply.Cmd = cmdELECT_OK
				err := writeCommand(cmd.replyWriter, reply)
				if err != nil {
					continue
				}
				leader, err = self.elect(candy, leaderTimeout)
				// Try until we get a leader
				for leader == nil || err != nil {
					leader, err = self.elect(candy, leaderTimeout)
				}
			}
		case ctrl := <-self.ctrlChan:
			switch ctrl.cmd {
			case ctrlADD:
				var err error
				candy, err = self.handshake(ctrl.addr, ctrl.id, candy, ctrl.timeout)
				reply := new(controlReply)
				if err != nil {
					reply.err = err
					ctrl.replyChan <- reply
					continue
				}
				ctrl.replyChan <- reply
			case ctrlQUERY_CANDY:
				reply := new(controlReply)
				reply.addr = self.myAddress()
				reply.id = self.myId
				ctrl.replyChan <- reply
				for _, node := range candy {
					reply := new(controlReply)
					reply.addr = node.conn.RemoteAddr()
					reply.id = node.id
					ctrl.replyChan <- reply
				}
				close(ctrl.replyChan)
			case ctrlQUERY_LEADER:
				if leader == nil {
					var err error
					leader, err = self.elect(candy, leaderTimeout)
					// Try until we get a leader
					for leader == nil || err != nil {
						leader, err = self.elect(candy, leaderTimeout)
					}
				}
				reply := new(controlReply)
				if leader.conn != nil {
					reply.addr = leader.conn.RemoteAddr()
				} else {
					reply.addr = self.myAddress()
				}
				reply.id = leader.id
				ctrl.replyChan <- reply
			}
		}
	}
}

func (self *Bully) replyHandshake(conn net.Conn) {
	cmd, err := readCommand(conn)
	if err != nil {
		conn.Close()
		return
	}
	if cmd.Cmd != cmdHELLO {
		conn.Close()
		return
	}
	if len(cmd.Body) == 0 {
		conn.Close()
		return
	}

	rId := new(big.Int).SetBytes(cmd.Body)
	if rId.Cmp(self.myId) == 0 {
		reply := new(command)
		reply.Cmd = cmdITSME
		writeCommand(conn, reply)
		conn.Close()
		return
	}
	cmd.src = rId
	cmd.replyWriter = conn
	go commandCollector(rId, conn, self.cmdChan, 10*time.Second)
	select {
	case self.cmdChan <- cmd:
	case <-time.After(10 * time.Second):
	}
	return
}

func (self *Bully) listen(ln net.Listener) {
	for {
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		if conn == nil {
			continue
		}
		go self.replyHandshake(conn)
	}
}

<commit_msg>[Mod] Check the reply message on handshake<commit_after>package main

import (
	"crypto/sha256"
	"errors"
	"fmt"
	"github.com/nu7hatch/gouuid"
	"math/big"
	"net"
	"strings"
	"time"
)

func stringToBig(str string) *big.Int {
	sha := sha256.New()
	sha.Write([]byte(str))
	b := sha.Sum(nil)

	ret := new(big.Int).SetBytes(b)
	return ret
}

type Bully struct {
	myId     *big.Int
	cmdChan  chan *command
	ctrlChan chan *control
	ln       net.Listener
}

type Candidate struct {
	Id   *big.Int
	Addr net.Addr
}

var ErrUnknownError = errors.New("Unknown")

func (self *Bully) MyId() *big.Int {
	return self.myId
}

func (self *Bully) AddCandidate(addrStr string, id *big.Int, timeout time.Duration) error {
	addr, err := net.ResolveTCPAddr("tcp", addrStr)
	if err != nil {
		return err
	}
	if addr == nil {
		return ErrUnknownError
	}
	if timeout <= 1*time.Second {
		timeout = 2 * time.Second
	}
	ctrl := new(control)
	ctrl.addr = addr
	ctrl.id = id
	ctrl.cmd = ctrlADD
	ctrl.timeout = timeout
	replyChan := make(chan *controlReply)
	ctrl.replyChan = replyChan

	self.ctrlChan <- ctrl
	reply := <-replyChan
	if reply == nil {
		return ErrUnknownError
	}
	if reply.err != nil {
		return reply.err
	}
	return nil
}

func (self *Bully) CandidateList() []*Candidate {
	ctrl := new(control)
	ctrl.cmd = ctrlQUERY_CANDY
	replyChan := make(chan *controlReply)
	ctrl.replyChan = replyChan

	self.ctrlChan <- ctrl

	ret := make([]*Candidate, 0, 10)
	for reply := range replyChan {
		c := new(Candidate)
		c.Id = reply.id
		c.Addr = reply.addr
		ret = append(ret, c)
	}
	return ret
}

func commandCollector(src *big.Int, conn net.Conn, cmdChan chan<- *command, timeout time.Duration) {
	defer func() {
		conn.Close()
		cmd := new(command)
		cmd.src = src
		cmd.Cmd = cmdBYE
		cmdChan <- cmd
	}()
	for {
		cmd, err := readCommand(conn)
		if err != nil {
			return
		}
		if cmd.Cmd == cmdITSME || cmd.Cmd == cmdBYE {
			return
		}
		cmd.src = src
		cmd.replyWriter = conn
		select {
		case cmdChan <- cmd:
			continue
		case <-time.After(timeout):
			return
		}
	}
}

func NewBully(ln net.Listener, myId *big.Int) *Bully {
	ret := new(Bully)
	if myId == nil {
		uu, _ := uuid.NewV4()
		ret.myId = stringToBig(uu.String())
	} else {
		ret.myId = myId
	}
	ret.cmdChan = make(chan *command)
	ret.ctrlChan = make(chan *control)
	ret.ln = ln
	go ret.listen(ln)
	go ret.process()
	return ret
}

type node struct {
	id   *big.Int
	conn net.Conn
}

func insertNode(l []*node, id *big.Int, conn net.Conn) ([]*node, bool) {
	n := findNode(l, id)
	if nil != n {
		return l, false
	}
	n = new(node)
	n.id = id
	n.conn = conn
	return append(l, n), true
}

func removeNode(l []*node, id *big.Int) []*node {
	idx := -1
	for i, n := range l {
		if n.id.Cmp(id) == 0 {
			idx = i
		}
	}
	if idx >= 0 {
		l[idx] = l[len(l)-1]
		l[len(l)-1] = nil
		l = l[:len(l)-1]
	}
	return l
}

func findNode(l []*node, id *big.Int) *node {
	for _, n := range l {
		if n.id.Cmp(id) == 0 {
			return n
		}
	}
	return nil
}

const (
	ctrlADD = iota
	ctrlQUERY_LEADER
	ctrlQUERY_CANDY
)

type controlReply struct {
	addr net.Addr
	id   *big.Int
	err  error
}

type control struct {
	cmd       int
	addr      net.Addr
	id        *big.Int
	timeout   time.Duration
	replyChan chan<- *controlReply
}

var ErrUnmatchedId = errors.New("Unmatched node id")
var ErrTryLater = errors.New("Try it later")
var ErrBadProtoImpl = errors.New("Bad protocol implementation")

func (self *Bully) handshake(addr net.Addr, id *big.Int, candy []*node, timeout time.Duration) ([]*node, error) {
	if id != nil {
		cmp := id.Cmp(self.myId)
		if cmp > 0 {
			n := findNode(candy, id)
			if n != nil {
				return candy, nil
			}
		} else if cmp < 0 {
			// If we know the id of a node,
			// then we only connect to the nodes with higher id,
			// and let the nodes with lower id connect us.
			return candy, nil
		} else {
			// It is ourselves, don't need to add it.
			return candy, nil
		}
	}
	conn, err := net.DialTimeout("tcp", addr.String(), timeout)
	if err != nil {
		return candy, err
	}
	cmd := new(command)
	cmd.Cmd = cmdHELLO
	cmd.Body = self.myId.Bytes()
	err = writeCommand(conn, cmd)
	if err != nil {
		return candy, err
	}
	reply, err := readCommand(conn)
	if err != nil {
		return candy, err
	}
	if reply.Cmd != cmdHELLO_REPLY {
		conn.Close()
		switch reply.Cmd {
		case cmdTRY_LATER:
			return candy, ErrTryLater
		case cmdDUP_CONN:
			return candy, nil
		case cmdITSME:
			return candy, nil
		}
		return candy, ErrBadProtoImpl
	}
	rId := new(big.Int).SetBytes(reply.Body)
	if id != nil {
		if rId.Cmp(id) != 0 {
			return candy, ErrUnmatchedId
		}
	}
	candy, _ = insertNode(candy, rId, conn)
	go commandCollector(rId, conn, self.cmdChan, 10*time.Second)

	return candy, nil
}

func (self *Bully) electUntilDie(candy []*node, timeout time.Duration) *node {
	leader, err := self.elect(candy, timeout)
	// Try until we get a leader
	for leader == nil || err != nil {
		leader, err = self.elect(candy, timeout)
	}
	return leader
}

var ErrNeedNewElection = errors.New("Need another round")

func (self *Bully) elect(candy []*node, timeout time.Duration) (leader *node, err error) {
	higherCandy := make([]*node, 0, len(candy))
	for _, c := range candy {
		if c.id.Cmp(self.myId) > 0 {
			cmd := new(command)
			cmd.Cmd = cmdELECT
			err := writeCommand(c.conn, cmd)
			if err == nil {
				higherCandy = append(higherCandy, c)
			}
		}
	}

	// No one is higher than me.
	// I am the leader.
	if len(higherCandy) <= 0 {
		leader.conn = nil
		leader.id = self.myId
		for _, c := range candy {
			cmd := new(command)
			cmd.Cmd = cmdCOORDIN
			writeCommand(c.conn, cmd)
		}
		return
	}
	slaved := false
	for {
		select {
		case cmd := <-self.cmdChan:
			switch cmd.Cmd {
			case cmdHELLO:
				reply := new(command)
				reply.Cmd = cmdTRY_LATER
				writeCommand(cmd.replyWriter, reply)
			case cmdELECT:
				reply := new(command)
				reply.Cmd = cmdELECT_OK
				writeCommand(cmd.replyWriter, reply)
			case cmdELECT_OK:
				n := findNode(higherCandy, cmd.src)
				if n == nil {
					continue
				}
				slaved = true
			case cmdCOORDIN:
				n := findNode(candy, cmd.src)
				if n == nil {
					err = ErrNeedNewElection
					return
				}
				if n.id.Cmp(self.myId) < 0 {
					err = ErrNeedNewElection
					return
				}
				leader = n
				return
			}
		case <-time.After(timeout):
			break
		}
	}

	// No one replied within time out.
	// I am the leader.
	if !slaved {
		leader.conn = nil
		leader.id = self.myId
		for _, c := range candy {
			cmd := new(command)
			cmd.Cmd = cmdCOORDIN
			writeCommand(c.conn, cmd)
		}
		return
	}
	err = ErrNeedNewElection
	return
}

func (self *Bully) myAddress() net.Addr {
	addrStr := self.ln.Addr().String()
	ae := strings.Split(addrStr, ":")
	addrStr = fmt.Sprintf("127.0.0.1:%v", ae[len(ae)-1])
	ret, _ := net.ResolveTCPAddr("tcp", addrStr)
	return ret
}

func (self *Bully) process() {
	candy := make([]*node, 0, 10)
	var leader *node
	leaderTimeout := 10 * time.Second
	for {
		select {
		case cmd := <-self.cmdChan:
			fmt.Printf("Command: %v; myid=%v\n", cmd.Cmd, self.myId)
			switch cmd.Cmd {
			case cmdHELLO:
				if cmd.src.Cmp(self.myId) == 0 {
					reply := new(command)
					reply.Cmd = cmdITSME
					err := writeCommand(cmd.replyWriter, reply)
					if err != nil {
						cmd.replyWriter.Close()
					}
					continue
				}
				n := findNode(candy, cmd.src)
				if n == nil {
					reply := new(command)
					reply.Cmd = cmdHELLO_REPLY
					reply.Body = self.myId.Bytes()
					err := writeCommand(cmd.replyWriter, reply)
					if err != nil {
						cmd.replyWriter.Close()
						continue
					}
					candy, _ = insertNode(candy, cmd.src, cmd.replyWriter)
				} else {
					reply := new(command)
					reply.Cmd = cmdDUP_CONN
					writeCommand(cmd.replyWriter, reply)
					cmd.replyWriter.Close()
				}
			case cmdBYE:
				candy = removeNode(candy, cmd.src)
			case cmdELECT:
				reply := new(command)
				reply.Cmd = cmdELECT_OK
				err := writeCommand(cmd.replyWriter, reply)
				if err != nil {
					continue
				}
				leader = self.electUntilDie(candy, leaderTimeout)
			case cmdCOORDIN:
				n := findNode(candy, cmd.src)
				if n == nil {
					leader = self.electUntilDie(candy, leaderTimeout)
				} else if n.id.Cmp(self.myId) < 0 {
					leader = self.electUntilDie(candy, leaderTimeout)
				} else {
					leader = n
				}
			}
		case ctrl := <-self.ctrlChan:
			switch ctrl.cmd {
			case ctrlADD:
				var err error
				candy, err = self.handshake(ctrl.addr, ctrl.id, candy, ctrl.timeout)
				reply := new(controlReply)
				if err != nil {
					reply.err = err
					ctrl.replyChan <- reply
					continue
				}
				ctrl.replyChan <- reply
			case ctrlQUERY_CANDY:
				reply := new(controlReply)
				reply.addr = self.myAddress()
				reply.id = self.myId
				ctrl.replyChan <- reply
				for _, node := range candy {
					reply := new(controlReply)
					reply.addr = node.conn.RemoteAddr()
					reply.id = node.id
					ctrl.replyChan <- reply
				}
				close(ctrl.replyChan)
			case ctrlQUERY_LEADER:
				if leader == nil {
					leader = self.electUntilDie(candy, leaderTimeout)
				}
				reply := new(controlReply)
				if leader.conn != nil {
					reply.addr = leader.conn.RemoteAddr()
				} else {
					reply.addr = self.myAddress()
				}
				reply.id = leader.id
				ctrl.replyChan <- reply
			}
		}
	}
}

func (self *Bully) replyHandshake(conn net.Conn) {
	cmd, err := readCommand(conn)
	if err != nil {
		conn.Close()
		return
	}
	if cmd.Cmd != cmdHELLO {
		conn.Close()
		return
	}
	if len(cmd.Body) == 0 {
		conn.Close()
		return
	}

	rId := new(big.Int).SetBytes(cmd.Body)
	if rId.Cmp(self.myId) == 0 {
		reply := new(command)
		reply.Cmd = cmdITSME
		writeCommand(conn, reply)
		conn.Close()
		return
	}
	cmd.src = rId
	cmd.replyWriter = conn
	go commandCollector(rId, conn, self.cmdChan, 10*time.Second)
	select {
	case self.cmdChan <- cmd:
	case <-time.After(10 * time.Second):
	}
	return
}

func (self *Bully) listen(ln net.Listener) {
	for {
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		if conn == nil {
			continue
		}
		go self.replyHandshake(conn)
	}
}
<|endoftext|>


<commit_before>// Package kontrol provides an implementation for the name service kite.
// It can be queried to get the list of running kites.
package kontrol

import (
	"errors"
	"math/rand"
	"strings"
	"sync"
	"time"

	"github.com/dgrijalva/jwt-go"
	"github.com/koding/kite"
	"github.com/koding/kite/config"
	"github.com/koding/kite/kitekey"
	kontrolprotocol "github.com/koding/kite/kontrol/protocol"
	"github.com/nu7hatch/gouuid"
)

const (
	KontrolVersion = "0.0.4"
	KitesPrefix    = "/kites"
	TokenLeeway    = 1 * time.Minute
)

var (
	TokenTTL    = 48 * time.Hour
	DefaultPort = 4000

	tokenCache   = make(map[string]string)
	tokenCacheMu sync.Mutex

	// HeartbeatInterval is the interval in which kites are sending heartbeats
	HeartbeatInterval = time.Second * 10

	// HeartbeatDelay is the compensation interval which is added to the
	// heartbeat to avoid network delays
	HeartbeatDelay = time.Second * 10

	// UpdateInterval is the interval in which the key gets updated
	// periodically. Keeping it low increase the write load to the storage, so
	// be cautious when changing it.
	UpdateInterval = time.Second * 60

	// KeyTLL is the timeout in which a key expires. Each storage
	// implementation needs to set keys according to this Key. If a storage
	// doesn't support TTL mechanism (such as PostgreSQL), it should use a
	// background cleaner which cleans up keys that are KeyTTL old.
	KeyTTL = time.Second * 90
)

type Kontrol struct {
	Kite *kite.Kite

	// MachineAuthenticate is used to authenticate the request in the
	// "handleMachine" method.  The reason for a separate auth function is, the
	// request must not be authenticated because clients do not have a kite.key
	// before they register to this machine.
	MachineAuthenticate func(r *kite.Request) error

	// RSA keys
	publicKey  string // for validating tokens
	privateKey string // for signing tokens

	clientLocks *IdLock

	// storage defines the storage of the kites.
	storage Storage

	// RegisterURL defines the URL that is used to self register when adding
	// itself to the storage backend
	RegisterURL string

	log kite.Logger
}

// New creates a new kontrol instance with the given version and config
// instance. Publickey is used for validating tokens and privateKey is used for
// signing tokens.
//
// Public and private keys are RSA pem blocks that can be generated with the
// following command:
//     openssl genrsa -out testkey.pem 2048
//     openssl rsa -in testkey.pem -pubout > testkey_pub.pem
//
func New(conf *config.Config, version, publicKey, privateKey string) *Kontrol {
	k := kite.New("kontrol", version)
	k.Config = conf

	// Listen on 4000 by default
	if k.Config.Port == 0 {
		k.Config.Port = DefaultPort
	}

	kontrol := &Kontrol{
		Kite:        k,
		publicKey:   publicKey,
		privateKey:  privateKey,
		log:         k.Log,
		clientLocks: NewIdlock(),
	}

	k.HandleFunc("register", kontrol.handleRegister)
	k.HandleFunc("registerMachine", kontrol.handleMachine).DisableAuthentication()
	k.HandleFunc("getKites", kontrol.handleGetKites)
	k.HandleFunc("getToken", kontrol.handleGetToken)

	return kontrol
}

func (k *Kontrol) AddAuthenticator(keyType string, fn func(*kite.Request) error) {
	k.Kite.Authenticators[keyType] = fn
}

func (k *Kontrol) Run() {
	rand.Seed(time.Now().UnixNano())

	if k.storage == nil {
		panic("kontrol storage is not set")
	}

	// now go and register ourself
	go k.registerSelf()

	k.Kite.Run()
}

// SetStorage sets the backend storage that kontrol is going to use to store
// kites
func (k *Kontrol) SetStorage(storage Storage) {
	k.storage = storage
}

// Close stops kontrol and closes all connections
func (k *Kontrol) Close() {
	k.Kite.Close()
}

// InitializeSelf registers his host by writing a key to ~/.kite/kite.key
func (k *Kontrol) InitializeSelf() error {
	key, err := k.registerUser(k.Kite.Config.Username)
	if err != nil {
		return err
	}
	return kitekey.Write(key)
}

func (k *Kontrol) registerUser(username string) (kiteKey string, err error) {
	// Only accept requests of type machine
	tknID, err := uuid.NewV4()
	if err != nil {
		return "", errors.New("cannot generate a token")
	}

	token := jwt.New(jwt.GetSigningMethod("RS256"))

	token.Claims = map[string]interface{}{
		"iss":        k.Kite.Kite().Username,         // Issuer
		"sub":        username,                       // Subject
		"iat":        time.Now().UTC().Unix(),        // Issued At
		"jti":        tknID.String(),                 // JWT ID
		"kontrolURL": k.Kite.Config.KontrolURL,       // Kontrol URL
		"kontrolKey": strings.TrimSpace(k.publicKey), // Public key of kontrol
	}

	k.Kite.Log.Info("Registered machine on user: %s", username)

	return token.SignedString([]byte(k.privateKey))
}

// registerSelf adds Kontrol itself to the storage as a kite.
func (k *Kontrol) registerSelf() {
	value := &kontrolprotocol.RegisterValue{
		URL: k.Kite.Config.KontrolURL,
	}

	// change if the user wants something different
	if k.RegisterURL != "" {
		value.URL = k.RegisterURL
	}

	// Register first by adding the value to the storage. We don't return any
	// error because we need to know why kontrol doesn't register itself
	if err := k.storage.Add(k.Kite.Kite(), value); err != nil {
		k.log.Error(err.Error())
	}

	for {
		if err := k.storage.Update(k.Kite.Kite(), value); err != nil {
			k.log.Error(err.Error())
			time.Sleep(time.Second)
			continue
		}

		time.Sleep(HeartbeatDelay + HeartbeatInterval)
	}
}

// generateToken returns a JWT token string. Please see the URL for details:
// http://tools.ietf.org/html/draft-ietf-oauth-json-web-token-13#section-4.1
func generateToken(aud, username, issuer, privateKey string) (string, error) {
	tokenCacheMu.Lock()
	defer tokenCacheMu.Unlock()

	uniqKey := aud + username + issuer // neglect privateKey, its always the same
	signed, ok := tokenCache[uniqKey]
	if ok {
		return signed, nil
	}

	tknID, err := uuid.NewV4()
	if err != nil {
		return "", errors.New("Server error: Cannot generate a token")
	}

	// Identifies the expiration time after which the JWT MUST NOT be accepted
	// for processing.
	ttl := TokenTTL

	// Implementers MAY provide for some small leeway, usually no more than
	// a few minutes, to account for clock skew.
	leeway := TokenLeeway

	tkn := jwt.New(jwt.GetSigningMethod("RS256"))
	tkn.Claims["iss"] = issuer                                       // Issuer
	tkn.Claims["sub"] = username                                     // Subject
	tkn.Claims["aud"] = aud                                          // Audience
	tkn.Claims["exp"] = time.Now().UTC().Add(ttl).Add(leeway).Unix() // Expiration Time
	tkn.Claims["nbf"] = time.Now().UTC().Add(-leeway).Unix()         // Not Before
	tkn.Claims["iat"] = time.Now().UTC().Unix()                      // Issued At
	tkn.Claims["jti"] = tknID.String()                               // JWT ID

	signed, err = tkn.SignedString([]byte(privateKey))
	if err != nil {
		return "", errors.New("Server error: Cannot generate a token")
	}

	// cache our token
	tokenCache[uniqKey] = signed

	// cache invalidation, because we cache the token in tokenCache we need to
	// invalidate it expiration time. This was handled usually within JWT, but
	// now we have to do it manually for our own cache.
	time.AfterFunc(TokenTTL, func() {
		tokenCacheMu.Lock()
		defer tokenCacheMu.Unlock()

		delete(tokenCache, uniqKey)
	})

	return signed, nil
}
<commit_msg>kontrol/kontrol.go: increase delay for heavy network based environments<commit_after>// Package kontrol provides an implementation for the name service kite.
// It can be queried to get the list of running kites.
package kontrol

import (
	"errors"
	"math/rand"
	"strings"
	"sync"
	"time"

	"github.com/dgrijalva/jwt-go"
	"github.com/koding/kite"
	"github.com/koding/kite/config"
	"github.com/koding/kite/kitekey"
	kontrolprotocol "github.com/koding/kite/kontrol/protocol"
	"github.com/nu7hatch/gouuid"
)

const (
	KontrolVersion = "0.0.4"
	KitesPrefix    = "/kites"
	TokenLeeway    = 1 * time.Minute
)

var (
	TokenTTL    = 48 * time.Hour
	DefaultPort = 4000

	tokenCache   = make(map[string]string)
	tokenCacheMu sync.Mutex

	// HeartbeatInterval is the interval in which kites are sending heartbeats
	HeartbeatInterval = time.Second * 10

	// HeartbeatDelay is the compensation interval which is added to the
	// heartbeat to avoid network delays
	HeartbeatDelay = time.Second * 20

	// UpdateInterval is the interval in which the key gets updated
	// periodically. Keeping it low increase the write load to the storage, so
	// be cautious when changing it.
	UpdateInterval = time.Second * 60

	// KeyTLL is the timeout in which a key expires. Each storage
	// implementation needs to set keys according to this Key. If a storage
	// doesn't support TTL mechanism (such as PostgreSQL), it should use a
	// background cleaner which cleans up keys that are KeyTTL old.
	KeyTTL = time.Second * 90
)

type Kontrol struct {
	Kite *kite.Kite

	// MachineAuthenticate is used to authenticate the request in the
	// "handleMachine" method.  The reason for a separate auth function is, the
	// request must not be authenticated because clients do not have a kite.key
	// before they register to this machine.
	MachineAuthenticate func(r *kite.Request) error

	// RSA keys
	publicKey  string // for validating tokens
	privateKey string // for signing tokens

	clientLocks *IdLock

	// storage defines the storage of the kites.
	storage Storage

	// RegisterURL defines the URL that is used to self register when adding
	// itself to the storage backend
	RegisterURL string

	log kite.Logger
}

// New creates a new kontrol instance with the given version and config
// instance. Publickey is used for validating tokens and privateKey is used for
// signing tokens.
//
// Public and private keys are RSA pem blocks that can be generated with the
// following command:
//     openssl genrsa -out testkey.pem 2048
//     openssl rsa -in testkey.pem -pubout > testkey_pub.pem
//
func New(conf *config.Config, version, publicKey, privateKey string) *Kontrol {
	k := kite.New("kontrol", version)
	k.Config = conf

	// Listen on 4000 by default
	if k.Config.Port == 0 {
		k.Config.Port = DefaultPort
	}

	kontrol := &Kontrol{
		Kite:        k,
		publicKey:   publicKey,
		privateKey:  privateKey,
		log:         k.Log,
		clientLocks: NewIdlock(),
	}

	k.HandleFunc("register", kontrol.handleRegister)
	k.HandleFunc("registerMachine", kontrol.handleMachine).DisableAuthentication()
	k.HandleFunc("getKites", kontrol.handleGetKites)
	k.HandleFunc("getToken", kontrol.handleGetToken)

	return kontrol
}

func (k *Kontrol) AddAuthenticator(keyType string, fn func(*kite.Request) error) {
	k.Kite.Authenticators[keyType] = fn
}

func (k *Kontrol) Run() {
	rand.Seed(time.Now().UnixNano())

	if k.storage == nil {
		panic("kontrol storage is not set")
	}

	// now go and register ourself
	go k.registerSelf()

	k.Kite.Run()
}

// SetStorage sets the backend storage that kontrol is going to use to store
// kites
func (k *Kontrol) SetStorage(storage Storage) {
	k.storage = storage
}

// Close stops kontrol and closes all connections
func (k *Kontrol) Close() {
	k.Kite.Close()
}

// InitializeSelf registers his host by writing a key to ~/.kite/kite.key
func (k *Kontrol) InitializeSelf() error {
	key, err := k.registerUser(k.Kite.Config.Username)
	if err != nil {
		return err
	}
	return kitekey.Write(key)
}

func (k *Kontrol) registerUser(username string) (kiteKey string, err error) {
	// Only accept requests of type machine
	tknID, err := uuid.NewV4()
	if err != nil {
		return "", errors.New("cannot generate a token")
	}

	token := jwt.New(jwt.GetSigningMethod("RS256"))

	token.Claims = map[string]interface{}{
		"iss":        k.Kite.Kite().Username,         // Issuer
		"sub":        username,                       // Subject
		"iat":        time.Now().UTC().Unix(),        // Issued At
		"jti":        tknID.String(),                 // JWT ID
		"kontrolURL": k.Kite.Config.KontrolURL,       // Kontrol URL
		"kontrolKey": strings.TrimSpace(k.publicKey), // Public key of kontrol
	}

	k.Kite.Log.Info("Registered machine on user: %s", username)

	return token.SignedString([]byte(k.privateKey))
}

// registerSelf adds Kontrol itself to the storage as a kite.
func (k *Kontrol) registerSelf() {
	value := &kontrolprotocol.RegisterValue{
		URL: k.Kite.Config.KontrolURL,
	}

	// change if the user wants something different
	if k.RegisterURL != "" {
		value.URL = k.RegisterURL
	}

	// Register first by adding the value to the storage. We don't return any
	// error because we need to know why kontrol doesn't register itself
	if err := k.storage.Add(k.Kite.Kite(), value); err != nil {
		k.log.Error(err.Error())
	}

	for {
		if err := k.storage.Update(k.Kite.Kite(), value); err != nil {
			k.log.Error(err.Error())
			time.Sleep(time.Second)
			continue
		}

		time.Sleep(HeartbeatDelay + HeartbeatInterval)
	}
}

// generateToken returns a JWT token string. Please see the URL for details:
// http://tools.ietf.org/html/draft-ietf-oauth-json-web-token-13#section-4.1
func generateToken(aud, username, issuer, privateKey string) (string, error) {
	tokenCacheMu.Lock()
	defer tokenCacheMu.Unlock()

	uniqKey := aud + username + issuer // neglect privateKey, its always the same
	signed, ok := tokenCache[uniqKey]
	if ok {
		return signed, nil
	}

	tknID, err := uuid.NewV4()
	if err != nil {
		return "", errors.New("Server error: Cannot generate a token")
	}

	// Identifies the expiration time after which the JWT MUST NOT be accepted
	// for processing.
	ttl := TokenTTL

	// Implementers MAY provide for some small leeway, usually no more than
	// a few minutes, to account for clock skew.
	leeway := TokenLeeway

	tkn := jwt.New(jwt.GetSigningMethod("RS256"))
	tkn.Claims["iss"] = issuer                                       // Issuer
	tkn.Claims["sub"] = username                                     // Subject
	tkn.Claims["aud"] = aud                                          // Audience
	tkn.Claims["exp"] = time.Now().UTC().Add(ttl).Add(leeway).Unix() // Expiration Time
	tkn.Claims["nbf"] = time.Now().UTC().Add(-leeway).Unix()         // Not Before
	tkn.Claims["iat"] = time.Now().UTC().Unix()                      // Issued At
	tkn.Claims["jti"] = tknID.String()                               // JWT ID

	signed, err = tkn.SignedString([]byte(privateKey))
	if err != nil {
		return "", errors.New("Server error: Cannot generate a token")
	}

	// cache our token
	tokenCache[uniqKey] = signed

	// cache invalidation, because we cache the token in tokenCache we need to
	// invalidate it expiration time. This was handled usually within JWT, but
	// now we have to do it manually for our own cache.
	time.AfterFunc(TokenTTL, func() {
		tokenCacheMu.Lock()
		defer tokenCacheMu.Unlock()

		delete(tokenCache, uniqKey)
	})

	return signed, nil
}
<|endoftext|>


<commit_before>package main

import
//	"fmt"

(
	"fmt"
	"os/exec"
)

func main() {
	done := make(chan error, 1) // buffer because goroutine might write to done<- before main goroutine starts to wait on <-done
	done <- fmt.Errorf("2")
	cmd := exec.Command("uname", "-a")
	go func() {
		err := cmd.Wait()
		select {
		case done <- err:
		default:
		}
	}()
	// don't panic if buffered
	err := <-done
	fmt.Println(err)
}
<commit_msg>buffer example 	modified:   main.go<commit_after>package main

import (
	"fmt"
	"os/exec"
)

func main() {
	done := make(chan error, 1) // buffer because goroutine might write to done<- before main goroutine starts to wait on <-done
	cmd := exec.Command("ls")
	go func() {
		select {
		case done <- cmd.Wait():
		default:
		}
	}()
	cmd.Start()
	fmt.Printf("done = %+v\n", <-done)
}
<|endoftext|>


<commit_before>package search

import (
	"context"
	"fmt"
	"sync"

	"github.com/keybase/client/go/chat/globals"
	"github.com/keybase/client/go/chat/storage"
	"github.com/keybase/client/go/chat/utils"
	"github.com/keybase/client/go/encrypteddb"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
)

type store struct {
	sync.Mutex
	globals.Contextified
	utils.DebugLabeler
	encryptedDB *encrypteddb.EncryptedDB
}

// store keeps an encrypted index of chat messages for all conversations to
// enable full inbox search locally.
// Data is stored in an encrypted leveldb in the form:
// (uid,convID) -> {
//                token: { msgID,...},
//                ...
//             },
//     ...       ->        ...
// NOTE: as a performance optimization we may want to split the metadata from
// the index itself so we can quickly operate on the metadata separately from
// the index and have less bytes to encrypt/decrypt on reads (metadata only
// contains only msg ids and no user content).
func newStore(g *globals.Context) *store {
	keyFn := func(ctx context.Context) ([32]byte, error) {
		return storage.GetSecretBoxKey(ctx, g.ExternalG(), storage.DefaultSecretUI)
	}
	dbFn := func(g *libkb.GlobalContext) *libkb.JSONLocalDb {
		return g.LocalChatDb
	}
	return &store{
		Contextified: globals.NewContextified(g),
		DebugLabeler: utils.NewDebugLabeler(g.GetLog(), "Search.store", false),
		encryptedDB:  encrypteddb.New(g.ExternalG(), dbFn, keyFn),
	}
}

func (s *store) dbKey(convID chat1.ConversationID, uid gregor1.UID) libkb.DbKey {
	return libkb.DbKey{
		Typ: libkb.DBChatIndex,
		Key: fmt.Sprintf("idx:%s:%s", uid, convID),
	}
}

func (s *store) getLocked(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID) (ret *chat1.ConversationIndex, err error) {
	defer func() {
		// return a blank index
		if err == nil && ret == nil {
			ret = &chat1.ConversationIndex{}
			ret.Index = map[string]map[chat1.MessageID]chat1.EmptyStruct{}
			ret.Alias = map[string]map[string]chat1.EmptyStruct{}
			ret.Metadata.SeenIDs = map[chat1.MessageID]chat1.EmptyStruct{}
		}
		if err != nil {
			if derr := s.deleteLocked(ctx, convID, uid); derr != nil {
				s.Debug(ctx, "unable to delete: %v", derr)
			}
		}
	}()
	dbKey := s.dbKey(convID, uid)
	var entry chat1.ConversationIndex
	found, err := s.encryptedDB.Get(ctx, dbKey, &entry)
	if err != nil || !found {
		return nil, err
	}
	if entry.Metadata.Version != IndexVersion {
		// drop the whole index for this conv
		err = s.deleteLocked(ctx, convID, uid)
		return nil, err
	}
	return &entry, nil
}

func (s *store) putLocked(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID, entry *chat1.ConversationIndex) error {
	if entry == nil {
		return nil
	}
	dbKey := s.dbKey(convID, uid)
	entry.Metadata.Version = IndexVersion
	return s.encryptedDB.Put(ctx, dbKey, *entry)
}

func (s *store) deleteLocked(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID) error {
	dbKey := s.dbKey(convID, uid)
	return s.encryptedDB.Delete(ctx, dbKey)
}

func (s *store) getConvIndex(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID) (entry *chat1.ConversationIndex, err error) {
	s.Lock()
	defer s.Unlock()
	return s.getLocked(ctx, convID, uid)
}

// addTokensLocked add the given tokens to the index under the given message
// id, when ingesting EDIT messages the msgID is of the superseded msg but the
// tokens are from the EDIT itself.
func (s *store) addTokensLocked(entry *chat1.ConversationIndex, tokens tokenMap, msgID chat1.MessageID) {
	for token, aliases := range tokens {
		msgIDs, ok := entry.Index[token]
		if !ok {
			msgIDs = map[chat1.MessageID]chat1.EmptyStruct{}
		}
		msgIDs[msgID] = chat1.EmptyStruct{}
		entry.Index[token] = msgIDs
		for alias := range aliases {
			atoken, ok := entry.Alias[alias]
			if !ok {
				atoken = map[string]chat1.EmptyStruct{}
			}
			atoken[token] = chat1.EmptyStruct{}
			entry.Alias[alias] = atoken
		}
	}
}

func (s *store) addMsgLocked(entry *chat1.ConversationIndex, msg chat1.MessageUnboxed) {
	tokens := tokensFromMsg(msg)
	s.addTokensLocked(entry, tokens, msg.GetMessageID())
}

func (s *store) removeMsgLocked(entry *chat1.ConversationIndex, msg chat1.MessageUnboxed) {
	// find the msgID that the index stores
	var msgID chat1.MessageID
	switch msg.GetMessageType() {
	case chat1.MessageType_EDIT, chat1.MessageType_ATTACHMENTUPLOADED:
		superIDs, err := utils.GetSupersedes(msg)
		if err != nil || len(superIDs) != 1 {
			return
		}
		msgID = superIDs[0]
	default:
		msgID = msg.GetMessageID()
	}

	for token, aliases := range tokensFromMsg(msg) {
		msgIDs := entry.Index[token]
		delete(msgIDs, msgID)
		if len(msgIDs) == 0 {
			delete(entry.Index, token)
		}
		for alias := range aliases {
			for atoken := range entry.Alias[alias] {
				_, ok := entry.Index[atoken]
				if !ok {
					delete(entry.Alias[alias], atoken)
				}
			}
			if len(entry.Alias[alias]) == 0 {
				delete(entry.Alias, alias)
			}
		}
	}
}

func (s *store) add(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID,
	msgs []chat1.MessageUnboxed) (err error) {
	s.Lock()
	defer s.Unlock()

	entry, err := s.getLocked(ctx, convID, uid)
	if err != nil {
		return err
	}

	fetchSupersededMsgs := func(msg chat1.MessageUnboxed) []chat1.MessageUnboxed {
		superIDs, err := utils.GetSupersedes(msg)
		if err != nil {
			s.Debug(ctx, "unable to get supersedes: %v", err)
			return nil
		}
		reason := chat1.GetThreadReason_INDEXED_SEARCH
		supersededMsgs, err := s.G().ChatHelper.GetMessages(ctx, uid, convID, superIDs,
			false /* resolveSupersedes*/, &reason)
		if err != nil {
			// Log but ignore error
			s.Debug(ctx, "unable to get fetch messages: %v", err)
			return nil
		}
		return supersededMsgs
	}

	for _, msg := range msgs {
		seenIDs := entry.Metadata.SeenIDs
		// Don't add if we've seen
		if _, ok := seenIDs[msg.GetMessageID()]; ok {
			continue
		}
		seenIDs[msg.GetMessageID()] = chat1.EmptyStruct{}
		// NOTE DELETE and DELETEHISTORY are handled through calls to `remove`,
		// other messages will be added if there is any content that can be
		// indexed.
		switch msg.GetMessageType() {
		case chat1.MessageType_ATTACHMENTUPLOADED:
			supersededMsgs := fetchSupersededMsgs(msg)
			for _, sm := range supersededMsgs {
				seenIDs[sm.GetMessageID()] = chat1.EmptyStruct{}
				s.addMsgLocked(entry, sm)
			}
		case chat1.MessageType_EDIT:
			tokens := tokensFromMsg(msg)
			supersededMsgs := fetchSupersededMsgs(msg)
			// remove the original message text and replace it with the edited
			// contents (using the original id in the index)
			for _, sm := range supersededMsgs {
				seenIDs[sm.GetMessageID()] = chat1.EmptyStruct{}
				s.removeMsgLocked(entry, sm)
				s.addTokensLocked(entry, tokens, sm.GetMessageID())
			}
		default:
			s.addMsgLocked(entry, msg)
		}
	}
	err = s.putLocked(ctx, convID, uid, entry)
	return err
}

// Remove tokenizes the message content and updates/removes index keys for each token.
func (s *store) remove(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID,
	msgs []chat1.MessageUnboxed) (err error) {
	s.Lock()
	defer s.Unlock()

	entry, err := s.getLocked(ctx, convID, uid)
	if err != nil {
		return err
	}

	seenIDs := entry.Metadata.SeenIDs
	for _, msg := range msgs {
		// Don't remove if we haven't seen
		if _, ok := entry.Metadata.SeenIDs[msg.GetMessageID()]; !ok {
			continue
		}
		seenIDs[msg.GetMessageID()] = chat1.EmptyStruct{}
		s.removeMsgLocked(entry, msg)
	}
	err = s.putLocked(ctx, convID, uid, entry)
	return err
}
<commit_msg>Update comment about index storage (#16930)<commit_after>package search

import (
	"context"
	"fmt"
	"sync"

	"github.com/keybase/client/go/chat/globals"
	"github.com/keybase/client/go/chat/storage"
	"github.com/keybase/client/go/chat/utils"
	"github.com/keybase/client/go/encrypteddb"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
)

type store struct {
	sync.Mutex
	globals.Contextified
	utils.DebugLabeler
	encryptedDB *encrypteddb.EncryptedDB
}

// store keeps an encrypted index of chat messages for all conversations
// to enable full inbox search locally.  Index data for each conversation
// is stored in an encrypted leveldb in the form:
// (uid,convID) -> {
//  Index: {
//     token: { msgID,... },
//     ...
//  },
//  Alias: {
//    alias: { token,... },
//    ...
//  },
//  Metadata: chat1.ConversationIndexMetadata
//}
// NOTE: as a performance optimization we may want to split the metadata
// from the index itself so we can quickly operate on the metadata
// separately from the index and have less bytes to encrypt/decrypt on
// reads (metadata only contains only msg ids and no user content).
func newStore(g *globals.Context) *store {
	keyFn := func(ctx context.Context) ([32]byte, error) {
		return storage.GetSecretBoxKey(ctx, g.ExternalG(), storage.DefaultSecretUI)
	}
	dbFn := func(g *libkb.GlobalContext) *libkb.JSONLocalDb {
		return g.LocalChatDb
	}
	return &store{
		Contextified: globals.NewContextified(g),
		DebugLabeler: utils.NewDebugLabeler(g.GetLog(), "Search.store", false),
		encryptedDB:  encrypteddb.New(g.ExternalG(), dbFn, keyFn),
	}
}

func (s *store) dbKey(convID chat1.ConversationID, uid gregor1.UID) libkb.DbKey {
	return libkb.DbKey{
		Typ: libkb.DBChatIndex,
		Key: fmt.Sprintf("idx:%s:%s", uid, convID),
	}
}

func (s *store) getLocked(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID) (ret *chat1.ConversationIndex, err error) {
	defer func() {
		// return a blank index
		if err == nil && ret == nil {
			ret = &chat1.ConversationIndex{}
			ret.Index = map[string]map[chat1.MessageID]chat1.EmptyStruct{}
			ret.Alias = map[string]map[string]chat1.EmptyStruct{}
			ret.Metadata.SeenIDs = map[chat1.MessageID]chat1.EmptyStruct{}
		}
		if err != nil {
			if derr := s.deleteLocked(ctx, convID, uid); derr != nil {
				s.Debug(ctx, "unable to delete: %v", derr)
			}
		}
	}()
	dbKey := s.dbKey(convID, uid)
	var entry chat1.ConversationIndex
	found, err := s.encryptedDB.Get(ctx, dbKey, &entry)
	if err != nil || !found {
		return nil, err
	}
	if entry.Metadata.Version != IndexVersion {
		// drop the whole index for this conv
		err = s.deleteLocked(ctx, convID, uid)
		return nil, err
	}
	return &entry, nil
}

func (s *store) putLocked(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID, entry *chat1.ConversationIndex) error {
	if entry == nil {
		return nil
	}
	dbKey := s.dbKey(convID, uid)
	entry.Metadata.Version = IndexVersion
	return s.encryptedDB.Put(ctx, dbKey, *entry)
}

func (s *store) deleteLocked(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID) error {
	dbKey := s.dbKey(convID, uid)
	return s.encryptedDB.Delete(ctx, dbKey)
}

func (s *store) getConvIndex(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID) (entry *chat1.ConversationIndex, err error) {
	s.Lock()
	defer s.Unlock()
	return s.getLocked(ctx, convID, uid)
}

// addTokensLocked add the given tokens to the index under the given message
// id, when ingesting EDIT messages the msgID is of the superseded msg but the
// tokens are from the EDIT itself.
func (s *store) addTokensLocked(entry *chat1.ConversationIndex, tokens tokenMap, msgID chat1.MessageID) {
	for token, aliases := range tokens {
		msgIDs, ok := entry.Index[token]
		if !ok {
			msgIDs = map[chat1.MessageID]chat1.EmptyStruct{}
		}
		msgIDs[msgID] = chat1.EmptyStruct{}
		entry.Index[token] = msgIDs
		for alias := range aliases {
			atoken, ok := entry.Alias[alias]
			if !ok {
				atoken = map[string]chat1.EmptyStruct{}
			}
			atoken[token] = chat1.EmptyStruct{}
			entry.Alias[alias] = atoken
		}
	}
}

func (s *store) addMsgLocked(entry *chat1.ConversationIndex, msg chat1.MessageUnboxed) {
	tokens := tokensFromMsg(msg)
	s.addTokensLocked(entry, tokens, msg.GetMessageID())
}

func (s *store) removeMsgLocked(entry *chat1.ConversationIndex, msg chat1.MessageUnboxed) {
	// find the msgID that the index stores
	var msgID chat1.MessageID
	switch msg.GetMessageType() {
	case chat1.MessageType_EDIT, chat1.MessageType_ATTACHMENTUPLOADED:
		superIDs, err := utils.GetSupersedes(msg)
		if err != nil || len(superIDs) != 1 {
			return
		}
		msgID = superIDs[0]
	default:
		msgID = msg.GetMessageID()
	}

	for token, aliases := range tokensFromMsg(msg) {
		msgIDs := entry.Index[token]
		delete(msgIDs, msgID)
		if len(msgIDs) == 0 {
			delete(entry.Index, token)
		}
		for alias := range aliases {
			for atoken := range entry.Alias[alias] {
				_, ok := entry.Index[atoken]
				if !ok {
					delete(entry.Alias[alias], atoken)
				}
			}
			if len(entry.Alias[alias]) == 0 {
				delete(entry.Alias, alias)
			}
		}
	}
}

func (s *store) add(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID,
	msgs []chat1.MessageUnboxed) (err error) {
	s.Lock()
	defer s.Unlock()

	entry, err := s.getLocked(ctx, convID, uid)
	if err != nil {
		return err
	}

	fetchSupersededMsgs := func(msg chat1.MessageUnboxed) []chat1.MessageUnboxed {
		superIDs, err := utils.GetSupersedes(msg)
		if err != nil {
			s.Debug(ctx, "unable to get supersedes: %v", err)
			return nil
		}
		reason := chat1.GetThreadReason_INDEXED_SEARCH
		supersededMsgs, err := s.G().ChatHelper.GetMessages(ctx, uid, convID, superIDs,
			false /* resolveSupersedes*/, &reason)
		if err != nil {
			// Log but ignore error
			s.Debug(ctx, "unable to get fetch messages: %v", err)
			return nil
		}
		return supersededMsgs
	}

	for _, msg := range msgs {
		seenIDs := entry.Metadata.SeenIDs
		// Don't add if we've seen
		if _, ok := seenIDs[msg.GetMessageID()]; ok {
			continue
		}
		seenIDs[msg.GetMessageID()] = chat1.EmptyStruct{}
		// NOTE DELETE and DELETEHISTORY are handled through calls to `remove`,
		// other messages will be added if there is any content that can be
		// indexed.
		switch msg.GetMessageType() {
		case chat1.MessageType_ATTACHMENTUPLOADED:
			supersededMsgs := fetchSupersededMsgs(msg)
			for _, sm := range supersededMsgs {
				seenIDs[sm.GetMessageID()] = chat1.EmptyStruct{}
				s.addMsgLocked(entry, sm)
			}
		case chat1.MessageType_EDIT:
			tokens := tokensFromMsg(msg)
			supersededMsgs := fetchSupersededMsgs(msg)
			// remove the original message text and replace it with the edited
			// contents (using the original id in the index)
			for _, sm := range supersededMsgs {
				seenIDs[sm.GetMessageID()] = chat1.EmptyStruct{}
				s.removeMsgLocked(entry, sm)
				s.addTokensLocked(entry, tokens, sm.GetMessageID())
			}
		default:
			s.addMsgLocked(entry, msg)
		}
	}
	err = s.putLocked(ctx, convID, uid, entry)
	return err
}

// Remove tokenizes the message content and updates/removes index keys for each token.
func (s *store) remove(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID,
	msgs []chat1.MessageUnboxed) (err error) {
	s.Lock()
	defer s.Unlock()

	entry, err := s.getLocked(ctx, convID, uid)
	if err != nil {
		return err
	}

	seenIDs := entry.Metadata.SeenIDs
	for _, msg := range msgs {
		// Don't remove if we haven't seen
		if _, ok := entry.Metadata.SeenIDs[msg.GetMessageID()]; !ok {
			continue
		}
		seenIDs[msg.GetMessageID()] = chat1.EmptyStruct{}
		s.removeMsgLocked(entry, msg)
	}
	err = s.putLocked(ctx, convID, uid, entry)
	return err
}
<|endoftext|>


<commit_before>package vtctl

import (
	"testing"

	"github.com/google/go-cmp/cmp"
)

func TestGenerateShardRanges(t *testing.T) {
	type args struct {
		shards int
	}

	tests := []struct {
		name    string
		args    args
		want    []string
		wantErr bool
	}{
		{
			"errors for shards less than 0",
			args{0},
			nil,
			true,
		},
		{
			"errors for shards more than 65536",
			args{65537},
			nil,
			true,
		},
		{
			"works for a single shard",
			args{1},
			[]string{"-"},
			false,
		},
		{
			"works for more than one shard",
			args{2},
			[]string{"-80", "80-"},
			false,
		},
		{
			"works for an odd number of shards",
			args{7},
			[]string{"-24", "24-49", "49-6d", "6d-92", "92-b6", "b6-db", "db-"},
			false,
		},
		{
			"works for large number of shards",
			args{256},
			[]string{"-01", "01-02", "02-03", "03-04", "04-05", "05-06", "06-07", "07-08", "08-09", "09-0a", "0a-0b", "0b-0c", "0c-0d", "0d-0e", "0e-0f", "0f-10", "10-11", "11-12", "12-13", "13-14", "14-15", "15-16", "16-17", "17-18", "18-19", "19-1a", "1a-1b", "1b-1c", "1c-1d", "1d-1e", "1e-1f", "1f-20", "20-21", "21-22", "22-23", "23-24", "24-25", "25-26", "26-27", "27-28", "28-29", "29-2a", "2a-2b", "2b-2c", "2c-2d", "2d-2e", "2e-2f", "2f-30", "30-31", "31-32", "32-33", "33-34", "34-35", "35-36", "36-37", "37-38", "38-39", "39-3a", "3a-3b", "3b-3c", "3c-3d", "3d-3e", "3e-3f", "3f-40", "40-41", "41-42", "42-43", "43-44", "44-45", "45-46", "46-47", "47-48", "48-49", "49-4a", "4a-4b", "4b-4c", "4c-4d", "4d-4e", "4e-4f", "4f-50", "50-51", "51-52", "52-53", "53-54", "54-55", "55-56", "56-57", "57-58", "58-59", "59-5a", "5a-5b", "5b-5c", "5c-5d", "5d-5e", "5e-5f", "5f-60", "60-61", "61-62", "62-63", "63-64", "64-65", "65-66", "66-67", "67-68", "68-69", "69-6a", "6a-6b", "6b-6c", "6c-6d", "6d-6e", "6e-6f", "6f-70", "70-71", "71-72", "72-73", "73-74", "74-75", "75-76", "76-77", "77-78", "78-79", "79-7a", "7a-7b", "7b-7c", "7c-7d", "7d-7e", "7e-7f", "7f-80", "80-81", "81-82", "82-83", "83-84", "84-85", "85-86", "86-87", "87-88", "88-89", "89-8a", "8a-8b", "8b-8c", "8c-8d", "8d-8e", "8e-8f", "8f-90", "90-91", "91-92", "92-93", "93-94", "94-95", "95-96", "96-97", "97-98", "98-99", "99-9a", "9a-9b", "9b-9c", "9c-9d", "9d-9e", "9e-9f", "9f-a0", "a0-a1", "a1-a2", "a2-a3", "a3-a4", "a4-a5", "a5-a6", "a6-a7", "a7-a8", "a8-a9", "a9-aa", "aa-ab", "ab-ac", "ac-ad", "ad-ae", "ae-af", "af-b0", "b0-b1", "b1-b2", "b2-b3", "b3-b4", "b4-b5", "b5-b6", "b6-b7", "b7-b8", "b8-b9", "b9-ba", "ba-bb", "bb-bc", "bc-bd", "bd-be", "be-bf", "bf-c0", "c0-c1", "c1-c2", "c2-c3", "c3-c4", "c4-c5", "c5-c6", "c6-c7", "c7-c8", "c8-c9", "c9-ca", "ca-cb", "cb-cc", "cc-cd", "cd-ce", "ce-cf", "cf-d0", "d0-d1", "d1-d2", "d2-d3", "d3-d4", "d4-d5", "d5-d6", "d6-d7", "d7-d8", "d8-d9", "d9-da", "da-db", "db-dc", "dc-dd", "dd-de", "de-df", "df-e0", "e0-e1", "e1-e2", "e2-e3", "e3-e4", "e4-e5", "e5-e6", "e6-e7", "e7-e8", "e8-e9", "e9-ea", "ea-eb", "eb-ec", "ec-ed", "ed-ee", "ee-ef", "ef-f0", "f0-f1", "f1-f2", "f2-f3", "f3-f4", "f4-f5", "f5-f6", "f6-f7", "f7-f8", "f8-f9", "f9-fa", "fa-fb", "fb-fc", "fc-fd", "fd-fe", "fe-ff", "ff-"},
			false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := generateShardRanges(tt.args.shards)
			if (err != nil) != tt.wantErr {
				t.Errorf("listShardRanges() error = %v, wantErr %v", err, tt.wantErr)

				return
			}

			if !cmp.Equal(got, tt.want) {
				t.Errorf("listShardRanges() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestShardCalculatorForShardsGreaterThan512(t *testing.T) {
	got, err := generateShardRanges(512)
	if err != nil {
		t.Errorf("listShardRanges() error = %v", err)
	}

	want := "ff80-"

	if got[511] != want {
		t.Errorf("Invalid mapping for a 512 shard keyspace. Expected %v, want %v", got[511], want)
	}
}
<commit_msg>Use require/assert in tests<commit_after>package vtctl

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestGenerateShardRanges(t *testing.T) {
	type args struct {
		shards int
	}

	tests := []struct {
		name    string
		args    args
		want    []string
		wantErr bool
	}{
		{
			"errors for shards less than 0",
			args{0},
			nil,
			true,
		},
		{
			"errors for shards more than 65536",
			args{65537},
			nil,
			true,
		},
		{
			"works for a single shard",
			args{1},
			[]string{"-"},
			false,
		},
		{
			"works for more than one shard",
			args{2},
			[]string{"-80", "80-"},
			false,
		},
		{
			"works for an odd number of shards",
			args{7},
			[]string{"-24", "24-49", "49-6d", "6d-92", "92-b6", "b6-db", "db-"},
			false,
		},
		{
			"works for large number of shards",
			args{256},
			[]string{"-01", "01-02", "02-03", "03-04", "04-05", "05-06", "06-07", "07-08", "08-09", "09-0a", "0a-0b", "0b-0c", "0c-0d", "0d-0e", "0e-0f", "0f-10", "10-11", "11-12", "12-13", "13-14", "14-15", "15-16", "16-17", "17-18", "18-19", "19-1a", "1a-1b", "1b-1c", "1c-1d", "1d-1e", "1e-1f", "1f-20", "20-21", "21-22", "22-23", "23-24", "24-25", "25-26", "26-27", "27-28", "28-29", "29-2a", "2a-2b", "2b-2c", "2c-2d", "2d-2e", "2e-2f", "2f-30", "30-31", "31-32", "32-33", "33-34", "34-35", "35-36", "36-37", "37-38", "38-39", "39-3a", "3a-3b", "3b-3c", "3c-3d", "3d-3e", "3e-3f", "3f-40", "40-41", "41-42", "42-43", "43-44", "44-45", "45-46", "46-47", "47-48", "48-49", "49-4a", "4a-4b", "4b-4c", "4c-4d", "4d-4e", "4e-4f", "4f-50", "50-51", "51-52", "52-53", "53-54", "54-55", "55-56", "56-57", "57-58", "58-59", "59-5a", "5a-5b", "5b-5c", "5c-5d", "5d-5e", "5e-5f", "5f-60", "60-61", "61-62", "62-63", "63-64", "64-65", "65-66", "66-67", "67-68", "68-69", "69-6a", "6a-6b", "6b-6c", "6c-6d", "6d-6e", "6e-6f", "6f-70", "70-71", "71-72", "72-73", "73-74", "74-75", "75-76", "76-77", "77-78", "78-79", "79-7a", "7a-7b", "7b-7c", "7c-7d", "7d-7e", "7e-7f", "7f-80", "80-81", "81-82", "82-83", "83-84", "84-85", "85-86", "86-87", "87-88", "88-89", "89-8a", "8a-8b", "8b-8c", "8c-8d", "8d-8e", "8e-8f", "8f-90", "90-91", "91-92", "92-93", "93-94", "94-95", "95-96", "96-97", "97-98", "98-99", "99-9a", "9a-9b", "9b-9c", "9c-9d", "9d-9e", "9e-9f", "9f-a0", "a0-a1", "a1-a2", "a2-a3", "a3-a4", "a4-a5", "a5-a6", "a6-a7", "a7-a8", "a8-a9", "a9-aa", "aa-ab", "ab-ac", "ac-ad", "ad-ae", "ae-af", "af-b0", "b0-b1", "b1-b2", "b2-b3", "b3-b4", "b4-b5", "b5-b6", "b6-b7", "b7-b8", "b8-b9", "b9-ba", "ba-bb", "bb-bc", "bc-bd", "bd-be", "be-bf", "bf-c0", "c0-c1", "c1-c2", "c2-c3", "c3-c4", "c4-c5", "c5-c6", "c6-c7", "c7-c8", "c8-c9", "c9-ca", "ca-cb", "cb-cc", "cc-cd", "cd-ce", "ce-cf", "cf-d0", "d0-d1", "d1-d2", "d2-d3", "d3-d4", "d4-d5", "d5-d6", "d6-d7", "d7-d8", "d8-d9", "d9-da", "da-db", "db-dc", "dc-dd", "dd-de", "de-df", "df-e0", "e0-e1", "e1-e2", "e2-e3", "e3-e4", "e4-e5", "e5-e6", "e6-e7", "e7-e8", "e8-e9", "e9-ea", "ea-eb", "eb-ec", "ec-ed", "ed-ee", "ee-ef", "ef-f0", "f0-f1", "f1-f2", "f2-f3", "f3-f4", "f4-f5", "f5-f6", "f6-f7", "f7-f8", "f8-f9", "f9-fa", "fa-fb", "fb-fc", "fc-fd", "fd-fe", "fe-ff", "ff-"},
			false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := generateShardRanges(tt.args.shards)
			if tt.wantErr {
				require.Error(t, err)
			} else {
				require.NoError(t, err)
			}

			assert.Equal(t, got, tt.want)
		})
	}
}

func TestShardCalculatorForShardsGreaterThan512(t *testing.T) {
	got, err := generateShardRanges(512)
	if err != nil {
		t.Errorf("listShardRanges() error = %v", err)
	}

	want := "ff80-"

	if got[511] != want {
		t.Errorf("Invalid mapping for a 512 shard keyspace. Expected %v, want %v", got[511], want)
	}
}
<|endoftext|>


<commit_before>/*******************************************************************************
The MIT License (MIT)

Copyright (c) 2015 Hajime Nakagami

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*******************************************************************************/

package main

import (
	"sort"
	"sync"
)

// ------------------------------ Topic ---------------------------------------

type Topics struct {
	m map[string]map[string]bool
	sync.RWMutex
}

func NewTopics() *Topics {
	return &Topics{
		m: make(map[string]map[string]bool),
	}
}

func (t *Topics) Add(topicName string, clientID string) {
	t.Lock()
	defer t.Unlock()
	topic, ok := t.m[topicName]
	if !ok {
		topic = make(map[string]bool)
		t.m[topicName] = topic
	}
	topic[clientID] = true
}

func (t *Topics) Remove(topicName string, clientID string) {
	t.Lock()
	defer t.Unlock()
	topic, ok := t.m[topicName]
	if ok {
		delete(topic, clientID)
	}
}

func (t *Topics) TopicList() []string {
	t.RLock()
	defer t.RUnlock()
	list := make([]string, 0)
	for k, _ := range t.m {
		list = append(list, k)
	}
	sort.Strings(list)
	return list
}

func (t *Topics) List(topicName string) []string {
	t.RLock()
	defer t.RUnlock()
	list := make([]string, 0)
	topic, ok := t.m[topicName]
	if ok {
		for item := range topic {
			list = append(list, item)
		}
	}
	return list
}
<commit_msg>remove comment line<commit_after>/*******************************************************************************
The MIT License (MIT)

Copyright (c) 2015 Hajime Nakagami

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*******************************************************************************/

package main

import (
	"sort"
	"sync"
)

type Topics struct {
	m map[string]map[string]bool
	sync.RWMutex
}

func NewTopics() *Topics {
	return &Topics{
		m: make(map[string]map[string]bool),
	}
}

func (t *Topics) Add(topicName string, clientID string) {
	t.Lock()
	defer t.Unlock()
	topic, ok := t.m[topicName]
	if !ok {
		topic = make(map[string]bool)
		t.m[topicName] = topic
	}
	topic[clientID] = true
}

func (t *Topics) Remove(topicName string, clientID string) {
	t.Lock()
	defer t.Unlock()
	topic, ok := t.m[topicName]
	if ok {
		delete(topic, clientID)
	}
}

func (t *Topics) TopicList() []string {
	t.RLock()
	defer t.RUnlock()
	list := make([]string, 0)
	for k, _ := range t.m {
		list = append(list, k)
	}
	sort.Strings(list)
	return list
}

func (t *Topics) List(topicName string) []string {
	t.RLock()
	defer t.RUnlock()
	list := make([]string, 0)
	topic, ok := t.m[topicName]
	if ok {
		for item := range topic {
			list = append(list, item)
		}
	}
	return list
}
<|endoftext|>


<commit_before>package commands

import (
	"cf"
	"cf/api"
	"cf/requirements"
	"cf/terminal"
	"errors"
	"fmt"
	"github.com/codegangsta/cli"
	"strings"
)

type CreateService struct {
	ui          terminal.UI
	serviceRepo api.ServiceRepository
}

func NewCreateService(ui terminal.UI, sR api.ServiceRepository) (cmd CreateService) {
	cmd.ui = ui
	cmd.serviceRepo = sR
	return
}

func (cmd CreateService) GetRequirements(reqFactory requirements.Factory, c *cli.Context) (reqs []requirements.Requirement, err error) {
	offeringName := c.String("offering")
	parameterList := c.String("parameters")

	if offeringName == "user-provided" && parameterList == "" {
		err = errors.New("Incorrect Usage")
		cmd.ui.FailWithUsage(c, "create-service")
		return
	}

	return
}

func (cmd CreateService) Run(c *cli.Context) {
	name := c.String("name")
	offeringName := c.String("offering")

	if offeringName == "user-provided" {
		params := c.String("parameters")
		cmd.createUserProvidedService(name, params)
	} else {
		planName := c.String("plan")
		cmd.createService(name, offeringName, planName)
	}
}

func (cmd CreateService) createUserProvidedService(name string, params string) {
	paramsMap := make(map[string]string)
	params = strings.Trim(params, `"`)

	for _, param := range strings.Split(params, ",") {
		param = strings.Trim(param, " ")
		paramsMap[param] = cmd.ui.Ask("%s%s", param, terminal.PromptColor(">"))
	}

	cmd.ui.Say("Creating service...")
	err := cmd.serviceRepo.CreateUserProvidedServiceInstance(name, paramsMap)
	if err != nil {
		cmd.ui.Failed(err.Error())
		return
	}

	cmd.ui.Ok()
	return
}

func (cmd CreateService) createService(name string, offeringName string, planName string) {
	offerings, apiErr := cmd.serviceRepo.GetServiceOfferings()
	if apiErr != nil {
		cmd.ui.Failed(apiErr.Error())
		return
	}

	offering, err := findOffering(offerings, offeringName)
	if err != nil {
		cmd.ui.Failed("Offering not found")
		return
	}

	plan, err := findPlan(offering.Plans, planName)
	if err != nil {
		cmd.ui.Failed("Plan not found")
		return
	}

	cmd.ui.Say("Creating service %s", terminal.EntityNameColor(name))

	var alreadyExists bool
	alreadyExists, apiErr = cmd.serviceRepo.CreateServiceInstance(name, plan)
	if apiErr != nil {
		cmd.ui.Failed(apiErr.Error())
		return
	}

	cmd.ui.Ok()

	if alreadyExists {
		cmd.ui.Warn("Service %s already exists", name)
	}
	return
}

func findOffering(offerings []cf.ServiceOffering, name string) (offering cf.ServiceOffering, err error) {
	for _, offering := range offerings {
		if name == offering.Label {
			return offering, nil
		}
	}

	err = errors.New(fmt.Sprintf("Could not find offering with name %s", name))
	return
}

func findPlan(plans []cf.ServicePlan, name string) (plan cf.ServicePlan, err error) {
	for _, plan := range plans {
		if name == plan.Name {
			return plan, nil
		}
	}

	err = errors.New(fmt.Sprintf("Could not find plan with name %s", name))
	return
}
<commit_msg>write out actual error message instead of guessing<commit_after>package commands

import (
	"cf"
	"cf/api"
	"cf/requirements"
	"cf/terminal"
	"errors"
	"fmt"
	"github.com/codegangsta/cli"
	"strings"
)

type CreateService struct {
	ui          terminal.UI
	serviceRepo api.ServiceRepository
}

func NewCreateService(ui terminal.UI, sR api.ServiceRepository) (cmd CreateService) {
	cmd.ui = ui
	cmd.serviceRepo = sR
	return
}

func (cmd CreateService) GetRequirements(reqFactory requirements.Factory, c *cli.Context) (reqs []requirements.Requirement, err error) {
	offeringName := c.String("offering")
	parameterList := c.String("parameters")

	if offeringName == "user-provided" && parameterList == "" {
		err = errors.New("Incorrect Usage")
		cmd.ui.FailWithUsage(c, "create-service")
		return
	}

	return
}

func (cmd CreateService) Run(c *cli.Context) {
	name := c.String("name")
	offeringName := c.String("offering")

	if offeringName == "user-provided" {
		params := c.String("parameters")
		cmd.createUserProvidedService(name, params)
	} else {
		planName := c.String("plan")
		cmd.createService(name, offeringName, planName)
	}
}

func (cmd CreateService) createUserProvidedService(name string, params string) {
	paramsMap := make(map[string]string)
	params = strings.Trim(params, `"`)

	for _, param := range strings.Split(params, ",") {
		param = strings.Trim(param, " ")
		paramsMap[param] = cmd.ui.Ask("%s%s", param, terminal.PromptColor(">"))
	}

	cmd.ui.Say("Creating service...")
	err := cmd.serviceRepo.CreateUserProvidedServiceInstance(name, paramsMap)
	if err != nil {
		cmd.ui.Failed(err.Error())
		return
	}

	cmd.ui.Ok()
	return
}

func (cmd CreateService) createService(name string, offeringName string, planName string) {
	offerings, apiErr := cmd.serviceRepo.GetServiceOfferings()
	if apiErr != nil {
		cmd.ui.Failed(apiErr.Error())
		return
	}

	offering, err := findOffering(offerings, offeringName)
	if err != nil {
		cmd.ui.Failed(err.Error())
		return
	}

	plan, err := findPlan(offering.Plans, planName)
	if err != nil {
		cmd.ui.Failed(err.Error())
		return
	}

	cmd.ui.Say("Creating service %s", terminal.EntityNameColor(name))

	var alreadyExists bool
	alreadyExists, apiErr = cmd.serviceRepo.CreateServiceInstance(name, plan)
	if apiErr != nil {
		cmd.ui.Failed(apiErr.Error())
		return
	}

	cmd.ui.Ok()

	if alreadyExists {
		cmd.ui.Warn("Service %s already exists", name)
	}
	return
}

func findOffering(offerings []cf.ServiceOffering, name string) (offering cf.ServiceOffering, err error) {
	for _, offering := range offerings {
		if name == offering.Label {
			return offering, nil
		}
	}

	err = errors.New(fmt.Sprintf("Could not find offering with name %s", name))
	return
}

func findPlan(plans []cf.ServicePlan, name string) (plan cf.ServicePlan, err error) {
	for _, plan := range plans {
		if name == plan.Name {
			return plan, nil
		}
	}

	err = errors.New(fmt.Sprintf("Could not find plan with name %s", name))
	return
}
<|endoftext|>


<commit_before>package app

import (
	//"errors"
	// /"github.com/paulmach/go.geojson"
	"encoding/json"
	"fmt"
	"log"
	"strings"
	"testing"
	"time"
)

// go test -bench=.
// go test -bench=. -test.benchmem

// const (
// 	testDbFile         string = "./test.db"
// 	testCustomerApikey string = "testKey"
// 	testDatasource     string = "testLayer"
// )

//var testDb Database
var testTcpServer TcpServer

func init() {
	COMMIT_LOG_FILE = "./test_commit.log"
	DB = Database{File: testDbFile}
	DB.Init()
	enable_test_logging()
	testTcpServer = TcpServer{Host: "localhost", Port: "3333"}
	testTcpServer.Start()
}

/*

{"method":"assign_datasource","apikey":"70P78vbIeKex","datasource":"bf1f964abdab49aea6739bf7f6b32867"}

{"method":"import_file","file":"springfield_projects_edit.geojson"}
{"method":"import_file","file":"test.geojson"}

{"method":"edit_feature"}
{"method":"edit_feature","data":{"datasource":"bf1f964abdab49aea6739bf7f6b32867","geo_id":"0"}}
{"method":"insert_feature","data":{"datasource":"bf1f964abdab49aea6739bf7f6b32867","feature":{}}}
{"method":"insert_feature","data":{"datasource":"bf1f964abdab49aea6739bf7f6b32867","feature":{"type":"Feature","geometry":{"type":"Point","coordinates":[47.279229,47.27922900257082]},"properties":{"date_created":1.487653451e+09,"date_modified":1.487653451e+09,"geo_id":"1487653451","is_active":true,"is_deleted":false,"name":"Dot"}}}}

*/

func parseRequest(message string) TcpMessage {
	req := TcpMessage{}
	err := json.Unmarshal([]byte(message), &req)
	if err != nil {
		log.Println(err)
	}
	return req
}

func parseResponse(message string) map[string]interface {
	data := map[string]interface{}
	//var data interface{}
	err := json.Umarshal(&data)
	if err != nil {
		log.Println(err)
	}
	return data
}

func TestTCPCreateApikeySuccess(t *testing.T) {
	req := parseRequest(`{"method": "create_apikey"}`)
	resp := testTcpServer.create_apikey(req)
	// check for error in response
	if strings.Contains(resp, `"status": "error"`) {
		t.Error(resp)
	}
	// check if "apikey" in response
	if !strings.Contains(resp, `"apikey":`) {
		t.Error(resp)
	}
}

/*
func TestTCPMethodError(t *testing.T) {
	req := parseRequest(`{"method": "this_is_not_a_supported_method"}`)
	resp := testTcpServer.create_apikey(req)

	log.Println(resp)

	// check for error in response
	if strings.Contains(resp, `"status": "error"`) {
		t.Error(resp)
	}
}
*/

func TestTCPInsertExportApikeySuccess(t *testing.T) {
	// create and insert apikey
	now := time.Now().Second()
	test_apikey := fmt.Sprintf("test_apikey_%s", now)
	req := parseRequest(`{"method": "insert_apikey", "data": { "apikey": "` + test_apikey + `" } }`)
	resp := testTcpServer.insert_apikey(req)
	// check for error in response
	if !strings.Contains(resp, `"status": "ok"`) {
		t.Error(resp)
	}
	// check if "test_apikey" in response
	if !strings.Contains(resp, test_apikey) {
		t.Error(resp)
	}
	//
	req = parseRequest(`{"method": "export_apikey", "apikey": "` + test_apikey + `" }`)
	resp = testTcpServer.export_apikey(req)
	// check if "test_apikey" in response
	if !strings.Contains(resp, test_apikey) {
		t.Error(resp)
	}
	//
	req = parseRequest(`{"method": "export_apikeys"}`)
	resp = testTcpServer.export_apikeys(req)
	// check if "test_apikey" in response
	if !strings.Contains(resp, test_apikey) {
		t.Error(resp)
	}
}

func TestTCPCreateExportDatasources(t *testing.T) {
	req := parseRequest(`{"method":"create_datasource"}`)
	resp := testTcpServer.create_datasource(req)
	// check for error in response
	if strings.Contains(resp, `"status": "error"`) {
		t.Error(resp)
	}
	// // check if "apikey" in response
	// if !strings.Contains(resp, `"apikey":`) {
	// 	t.Error(resp)
	// }
	log.Println(parseResponse(resp))
}

// {"method":"export_datasource","datasource":"bf1f964abdab49aea6739bf7f6b32867"}
// {"method":"export_datasources"}
<commit_msg>tcp tests<commit_after>package app

import (
	//"errors"
	// /"github.com/paulmach/go.geojson"
	"encoding/json"
	"fmt"
	"log"
	"strings"
	"testing"
	"time"
)

// go test -bench=.
// go test -bench=. -test.benchmem

// const (
// 	testDbFile         string = "./test.db"
// 	testCustomerApikey string = "testKey"
// 	testDatasource     string = "testLayer"
// )

//var testDb Database
var testTcpServer TcpServer

func init() {
	COMMIT_LOG_FILE = "./test_commit.log"
	DB = Database{File: testDbFile}
	DB.Init()
	enable_test_logging()
	testTcpServer = TcpServer{Host: "localhost", Port: "3333"}
	testTcpServer.Start()
}

/*

{"method":"assign_datasource","apikey":"70P78vbIeKex","datasource":"bf1f964abdab49aea6739bf7f6b32867"}

{"method":"import_file","file":"springfield_projects_edit.geojson"}
{"method":"import_file","file":"test.geojson"}

{"method":"edit_feature"}
{"method":"edit_feature","data":{"datasource":"bf1f964abdab49aea6739bf7f6b32867","geo_id":"0"}}
{"method":"insert_feature","data":{"datasource":"bf1f964abdab49aea6739bf7f6b32867","feature":{}}}
{"method":"insert_feature","data":{"datasource":"bf1f964abdab49aea6739bf7f6b32867","feature":{"type":"Feature","geometry":{"type":"Point","coordinates":[47.279229,47.27922900257082]},"properties":{"date_created":1.487653451e+09,"date_modified":1.487653451e+09,"geo_id":"1487653451","is_active":true,"is_deleted":false,"name":"Dot"}}}}

*/

func parseRequest(message string) TcpMessage {
	req := TcpMessage{}
	err := json.Unmarshal([]byte(message), &req)
	if err != nil {
		log.Println(err)
	}
	return req
}

func parseResponse(message string) map[string]interface {
	data := map[string]interface{}
	//var data interface{}
	err := json.Unmarshal(message &data)
	if err != nil {
		log.Println(err)
	}
	return data
}

func TestTCPCreateApikeySuccess(t *testing.T) {
	req := parseRequest(`{"method": "create_apikey"}`)
	resp := testTcpServer.create_apikey(req)
	// check for error in response
	if strings.Contains(resp, `"status": "error"`) {
		t.Error(resp)
	}
	// check if "apikey" in response
	if !strings.Contains(resp, `"apikey":`) {
		t.Error(resp)
	}
}

/*
func TestTCPMethodError(t *testing.T) {
	req := parseRequest(`{"method": "this_is_not_a_supported_method"}`)
	resp := testTcpServer.create_apikey(req)

	log.Println(resp)

	// check for error in response
	if strings.Contains(resp, `"status": "error"`) {
		t.Error(resp)
	}
}
*/

func TestTCPInsertExportApikeySuccess(t *testing.T) {
	// create and insert apikey
	now := time.Now().Second()
	test_apikey := fmt.Sprintf("test_apikey_%s", now)
	req := parseRequest(`{"method": "insert_apikey", "data": { "apikey": "` + test_apikey + `" } }`)
	resp := testTcpServer.insert_apikey(req)
	// check for error in response
	if !strings.Contains(resp, `"status": "ok"`) {
		t.Error(resp)
	}
	// check if "test_apikey" in response
	if !strings.Contains(resp, test_apikey) {
		t.Error(resp)
	}
	//
	req = parseRequest(`{"method": "export_apikey", "apikey": "` + test_apikey + `" }`)
	resp = testTcpServer.export_apikey(req)
	// check if "test_apikey" in response
	if !strings.Contains(resp, test_apikey) {
		t.Error(resp)
	}
	//
	req = parseRequest(`{"method": "export_apikeys"}`)
	resp = testTcpServer.export_apikeys(req)
	// check if "test_apikey" in response
	if !strings.Contains(resp, test_apikey) {
		t.Error(resp)
	}
}

func TestTCPCreateExportDatasources(t *testing.T) {
	req := parseRequest(`{"method":"create_datasource"}`)
	resp := testTcpServer.create_datasource(req)
	// check for error in response
	if strings.Contains(resp, `"status": "error"`) {
		t.Error(resp)
	}
	// // check if "apikey" in response
	// if !strings.Contains(resp, `"apikey":`) {
	// 	t.Error(resp)
	// }
	log.Println(parseResponse(resp))
}

// {"method":"export_datasource","datasource":"bf1f964abdab49aea6739bf7f6b32867"}
// {"method":"export_datasources"}
<|endoftext|>


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

package source

import (
	sources "knative.dev/eventing-kafka/pkg/apis/sources/v1beta1"
)

type SecretLocator struct {
	*sources.KafkaSource
}

func (ks *SecretLocator) SecretName() (string, bool) {
	if !hasAuthConfig(ks.KafkaSource) {
		return "", false
	}
	return getName(ks.KafkaSource), hasAuthConfig(ks.KafkaSource)
}

func (ks *SecretLocator) SecretNamespace() (string, bool) {
	return ks.Namespace, hasAuthConfig(ks.KafkaSource)
}

func getName(ks *sources.KafkaSource) string {
	return ""
}

func hasAuthConfig(ks *sources.KafkaSource) bool {
	// TODO handle secrets
	//return (ks.Spec.KafkaAuthSpec.Net.TLS.Enable || ks.Spec.KafkaAuthSpec.Net.SASL.Enable) &&
	//	(ks.Spec.KafkaAuthSpec.Net.TLS.Cert.SecretKeyRef != nil ||
	//		ks.Spec.KafkaAuthSpec.Net.SASL.User.SecretKeyRef != nil ||
	//		ks.Spec.KafkaAuthSpec.Net.SASL.Password.SecretKeyRef != nil ||
	//		ks.Spec.KafkaAuthSpec.Net.SASL.Type.SecretKeyRef != nil)
	return false
}
<commit_msg>Enable KafkaSource auth support (#1403)<commit_after>/*
 * Copyright 2021 The Knative Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package source

import (
	sources "knative.dev/eventing-kafka/pkg/apis/sources/v1beta1"
)

type SecretLocator struct {
	*sources.KafkaSource
}

func (ks *SecretLocator) SecretName() (string, bool) {
	if !hasAuthConfig(ks.KafkaSource) {
		return "", false
	}
	// KafkaSource uses a secret provider `security.NetSpecSecretProviderFunc` that ignores name and namespace.
	// so there is no need to return a name here.
	return "", hasAuthConfig(ks.KafkaSource)
}

func (ks *SecretLocator) SecretNamespace() (string, bool) {
	return ks.Namespace, hasAuthConfig(ks.KafkaSource)
}

func hasAuthConfig(ks *sources.KafkaSource) bool {
	return (ks.Spec.KafkaAuthSpec.Net.TLS.Enable || ks.Spec.KafkaAuthSpec.Net.SASL.Enable) &&
		(ks.Spec.KafkaAuthSpec.Net.TLS.Cert.SecretKeyRef != nil ||
			ks.Spec.KafkaAuthSpec.Net.TLS.CACert.SecretKeyRef != nil ||
			ks.Spec.KafkaAuthSpec.Net.TLS.Key.SecretKeyRef != nil ||
			ks.Spec.KafkaAuthSpec.Net.SASL.User.SecretKeyRef != nil ||
			ks.Spec.KafkaAuthSpec.Net.SASL.Password.SecretKeyRef != nil ||
			ks.Spec.KafkaAuthSpec.Net.SASL.Type.SecretKeyRef != nil)
}
<|endoftext|>


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

// Package modcmd implements the ``go mod'' command.
package modcmd

import (
	"cmd/go/internal/base"
	"cmd/go/internal/cfg"
)

var CmdMod = &base.Command{
	UsageLine: "go mod",
	Short:     "module maintenance",
	Long: `Go mod provides access to operations on modules.

Note that support for modules is built into all the go commands,
not just 'go mod'. For example, day-to-day adding, removing, upgrading,
and downgrading of dependencies should be done using 'go get'.
See 'go help modules' for an overview of module functionality.
	`,

	Commands: []*base.Command{
		cmdDownload,
		cmdEdit,
		cmdGraph,
		cmdInit,
		cmdTidy,
		cmdVendor,
		cmdVerify,
		cmdWhy,
	},
}

func addModFlags(cmd *base.Command) {
	cmd.Flag.StringVar(&cfg.ModFile, "modfile", "", "")
}
<commit_msg>cmd/go/internal/modcmd: remove dead function addModFlag<commit_after>// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package modcmd implements the ``go mod'' command.
package modcmd

import (
	"cmd/go/internal/base"
)

var CmdMod = &base.Command{
	UsageLine: "go mod",
	Short:     "module maintenance",
	Long: `Go mod provides access to operations on modules.

Note that support for modules is built into all the go commands,
not just 'go mod'. For example, day-to-day adding, removing, upgrading,
and downgrading of dependencies should be done using 'go get'.
See 'go help modules' for an overview of module functionality.
	`,

	Commands: []*base.Command{
		cmdDownload,
		cmdEdit,
		cmdGraph,
		cmdInit,
		cmdTidy,
		cmdVendor,
		cmdVerify,
		cmdWhy,
	},
}
<|endoftext|>