<commit_before>package control

import (
	"path/filepath"
	"testing"

	"bazil.org/bazil/server"
	"bazil.org/bazil/server/control/wire"
	"bazil.org/bazil/tokens"
	"bazil.org/bazil/util/grpcunix"
	"bazil.org/bazil/util/tempdir"
	"github.com/agl/ed25519"
	"github.com/boltdb/bolt"
	"golang.org/x/net/context"
	"google.golang.org/grpc/codes"
)

func TestPeerAdd(t *testing.T) {
	tmp := tempdir.New(t)
	defer tmp.Cleanup()
	app, err := server.New(tmp.Path)
	if err != nil {
		t.Fatal(err)
	}
	defer controlListenAndServe(t, app)()

	pub := [ed25519.PublicKeySize]byte{1, 2, 3, 4, 5}
	addReq := &wire.PeerAddRequest{
		Pub: pub[:],
	}

	rpcConn, err := grpcunix.Dial(filepath.Join(app.DataDir, "control"))
	if err != nil {
		t.Fatal(err)
	}
	defer rpcConn.Close()
	rpcClient := wire.NewControlClient(rpcConn)

	ctx := context.Background()
	if _, err := rpcClient.PeerAdd(ctx, addReq); err != nil {
		t.Fatalf("adding peer failed: %v", err)
	}

	p, err := app.GetPeer(&pub)
	if err != nil {
		t.Fatalf("checking stored peer: %v", err)
	}
	if g, e := *p.Pub, pub; g != e {
		t.Errorf("wrong public key stored: %x != %x", g, e)
	}
}

func TestPeerAddBadPubLong(t *testing.T) {
	tmp := tempdir.New(t)
	defer tmp.Cleanup()
	app, err := server.New(tmp.Path)
	if err != nil {
		t.Fatal(err)
	}
	defer controlListenAndServe(t, app)()

	tooLong := make([]byte, 33)
	addReq := &wire.PeerAddRequest{
		Pub: tooLong,
	}

	rpcConn, err := grpcunix.Dial(filepath.Join(app.DataDir, "control"))
	if err != nil {
		t.Fatal(err)
	}
	defer rpcConn.Close()
	rpcClient := wire.NewControlClient(rpcConn)

	ctx := context.Background()
	_, err = rpcClient.PeerAdd(ctx, addReq)
	if err == nil {
		t.Fatalf("expected error from PeerAdd with too long public key")
	}
	if err := checkRPCError(err, codes.InvalidArgument, "peer public key must be exactly 32 bytes"); err != nil {
		t.Error(err)
	}

	check := func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte(tokens.BucketPeer))
		val := bucket.Get([]byte("foo"))
		if g := val; g != nil {
			t.Errorf("public key stored even on error: %x", g)
		}
		return nil
	}
	if err := app.DB.View(check); err != nil {
		t.Error(err)
	}
}

func TestPeerAddBadPubShort(t *testing.T) {
	tmp := tempdir.New(t)
	defer tmp.Cleanup()
	app, err := server.New(tmp.Path)
	if err != nil {
		t.Fatal(err)
	}
	defer controlListenAndServe(t, app)()

	tooShort := make([]byte, 33)
	addReq := &wire.PeerAddRequest{
		Pub: tooShort,
	}

	rpcConn, err := grpcunix.Dial(filepath.Join(app.DataDir, "control"))
	if err != nil {
		t.Fatal(err)
	}
	defer rpcConn.Close()
	rpcClient := wire.NewControlClient(rpcConn)

	ctx := context.Background()
	_, err = rpcClient.PeerAdd(ctx, addReq)
	if err == nil {
		t.Fatalf("expected error from PeerAdd with too short public key")
	}
	if err := checkRPCError(err, codes.InvalidArgument, "peer public key must be exactly 32 bytes"); err != nil {
		t.Error(err)
	}

	check := func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte(tokens.BucketPeer))
		val := bucket.Get([]byte("foo"))
		if g := val; g != nil {
			t.Errorf("public key stored even on error: %x", g)
		}
		return nil
	}
	if err := app.DB.View(check); err != nil {
		t.Error(err)
	}
}
<commit_msg>server/control test: Better database post-condition validation<commit_after>package control

import (
	"fmt"
	"path/filepath"
	"testing"

	"bazil.org/bazil/server"
	"bazil.org/bazil/server/control/wire"
	"bazil.org/bazil/tokens"
	"bazil.org/bazil/util/grpcunix"
	"bazil.org/bazil/util/tempdir"
	"github.com/agl/ed25519"
	"github.com/boltdb/bolt"
	"golang.org/x/net/context"
	"google.golang.org/grpc/codes"
)

func checkHasNoPeers(app *server.App) error {
	check := func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(tokens.BucketPeer))
		c := b.Cursor()
		k, _ := c.First()
		if k != nil {
			return fmt.Errorf("did not expect stored public key: %x", k)
		}
		return nil
	}
	if err := app.DB.View(check); err != nil {
		return err
	}
	return nil
}

func TestPeerAdd(t *testing.T) {
	tmp := tempdir.New(t)
	defer tmp.Cleanup()
	app, err := server.New(tmp.Path)
	if err != nil {
		t.Fatal(err)
	}
	defer controlListenAndServe(t, app)()

	pub := [ed25519.PublicKeySize]byte{1, 2, 3, 4, 5}
	addReq := &wire.PeerAddRequest{
		Pub: pub[:],
	}

	rpcConn, err := grpcunix.Dial(filepath.Join(app.DataDir, "control"))
	if err != nil {
		t.Fatal(err)
	}
	defer rpcConn.Close()
	rpcClient := wire.NewControlClient(rpcConn)

	ctx := context.Background()
	if _, err := rpcClient.PeerAdd(ctx, addReq); err != nil {
		t.Fatalf("adding peer failed: %v", err)
	}

	p, err := app.GetPeer(&pub)
	if err != nil {
		t.Fatalf("checking stored peer: %v", err)
	}
	if g, e := *p.Pub, pub; g != e {
		t.Errorf("wrong public key stored: %x != %x", g, e)
	}
}

func TestPeerAddBadPubLong(t *testing.T) {
	tmp := tempdir.New(t)
	defer tmp.Cleanup()
	app, err := server.New(tmp.Path)
	if err != nil {
		t.Fatal(err)
	}
	defer controlListenAndServe(t, app)()

	tooLong := make([]byte, 33)
	addReq := &wire.PeerAddRequest{
		Pub: tooLong,
	}

	rpcConn, err := grpcunix.Dial(filepath.Join(app.DataDir, "control"))
	if err != nil {
		t.Fatal(err)
	}
	defer rpcConn.Close()
	rpcClient := wire.NewControlClient(rpcConn)

	ctx := context.Background()
	_, err = rpcClient.PeerAdd(ctx, addReq)
	if err == nil {
		t.Fatalf("expected error from PeerAdd with too long public key")
	}
	if err := checkRPCError(err, codes.InvalidArgument, "peer public key must be exactly 32 bytes"); err != nil {
		t.Error(err)
	}
	if err := checkHasNoPeers(app); err != nil {
		t.Error(err)
	}
}

func TestPeerAddBadPubShort(t *testing.T) {
	tmp := tempdir.New(t)
	defer tmp.Cleanup()
	app, err := server.New(tmp.Path)
	if err != nil {
		t.Fatal(err)
	}
	defer controlListenAndServe(t, app)()

	tooShort := make([]byte, 33)
	addReq := &wire.PeerAddRequest{
		Pub: tooShort,
	}

	rpcConn, err := grpcunix.Dial(filepath.Join(app.DataDir, "control"))
	if err != nil {
		t.Fatal(err)
	}
	defer rpcConn.Close()
	rpcClient := wire.NewControlClient(rpcConn)

	ctx := context.Background()
	_, err = rpcClient.PeerAdd(ctx, addReq)
	if err == nil {
		t.Fatalf("expected error from PeerAdd with too short public key")
	}
	if err := checkRPCError(err, codes.InvalidArgument, "peer public key must be exactly 32 bytes"); err != nil {
		t.Error(err)
	}
	if err := checkHasNoPeers(app); err != nil {
		t.Error(err)
	}
}
<|endoftext|>


<commit_before>package dexcom

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/ecc1/ble"
)

type bleConn struct {
	tx ble.Characteristic
	rx chan byte
}

const (
	// maximum size of writes to GATT characteristics
	gattMTU = 20
)

func (conn *bleConn) Send(data []byte) error {
	// Dexcom G4 Share expects each BLE message to start with two 01 bytes.
	data = append([]byte{0x01, 0x01}, data...)
	for {
		n := len(data)
		if n == 0 {
			return nil
		}
		if n > gattMTU {
			n = gattMTU
		}
		err := conn.tx.WriteValue(data[:n])
		if err != nil {
			return err
		}
		data = data[n:]
	}
}

func (conn *bleConn) Receive(data []byte) error {
	for i := 0; i < len(data); i++ {
		b, ok := <-conn.rx
		if !ok {
			return fmt.Errorf("input channel closed")
		}
		data[i] = b
	}
	return nil
}

var (
	receiverService = dexcomUUID(0xa0b1)
)

func connect() error {
	device, err := ble.Discover(time.Minute, receiverService)
	if err != nil {
		return err
	}
	if !device.Connected() {
		err = device.Connect()
		if err != nil {
			return err
		}
		log.Printf("%s: connected", device.Name())
	} else {
		log.Printf("%s: already connected", device.Name())
	}
	if !device.Paired() {
		err = device.Pair()
		if err != nil {
			return err
		}
		log.Printf("%s: paired", device.Name())
	} else {
		log.Printf("%s: already paired", device.Name())
	}
	err = ble.Update()
	if err != nil {
		return err
	}
	return authenticate(device)
}

var (
	authentication = dexcomUUID(0xacac)
	authEnvVar     = "DEXCOM_CGM_ID"
	authCode       = []byte{}
)

func initAuthCode() error {
	if len(authCode) != 0 {
		return nil
	}
	id := os.Getenv(authEnvVar)
	if len(id) == 0 {
		return fmt.Errorf("%s environment variable is not set", authEnvVar)
	}
	if len(id) != 10 {
		return fmt.Errorf("%s environment variable must be 2 letters followed by 8 digits", authEnvVar)
	}
	authCode = []byte(id + "000000")
	return nil
}

func authenticate(device ble.Device) error {
	err := initAuthCode()
	if err != nil {
		return err
	}
	auth, err := ble.GetCharacteristic(authentication)
	if err != nil {
		return err
	}
	data, err := auth.ReadValue()
	if err != nil {
		return err
	}
	if bytes.Equal(data, authCode) {
		log.Printf("%s: already authenticated", device.Name())
		return nil
	}
	err = auth.WriteValue(authCode)
	if err != nil {
		return err
	}
	log.Printf("%s: authenticated", device.Name())
	return nil
}

var (
	heartbeat   = dexcomUUID(0x2b18)
	sendData    = dexcomUUID(0xb20a)
	receiveData = dexcomUUID(0xb20b)
)

func OpenBLE() (Connection, error) {
	err := connect()
	if err != nil {
		return nil, err
	}

	// We need to enable heartbeat notifications
	// or else we won't get any receiveData responses.
	err = ble.HandleNotify(heartbeat, func(data []byte) {})
	if err != nil {
		return nil, err
	}

	rx := make(chan byte, 1600)
	err = ble.HandleNotify(receiveData, func(data []byte) {
		for _, b := range data {
			rx <- b
		}
	})
	if err != nil {
		return nil, err
	}

	tx, err := ble.GetCharacteristic(sendData)
	if err != nil {
		return nil, err
	}

	return &bleConn{tx: tx, rx: rx}, nil
}

func dexcomUUID(id uint16) string {
	return fmt.Sprintf("f0ac%04x-ebfa-f96f-28da-076c35a521db", id)
}

func (*bleConn) Close() {}
<commit_msg>Use connection returned from ble.Open<commit_after>package dexcom

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/ecc1/ble"
)

type bleConn struct {
	conn *ble.Connection
	tx   ble.Characteristic
	rx   chan byte
}

const (
	// maximum size of writes to GATT characteristics
	gattMTU = 20
)

func (conn *bleConn) Send(data []byte) error {
	// Dexcom G4 Share expects each BLE message to start with two 01 bytes.
	data = append([]byte{0x01, 0x01}, data...)
	for {
		n := len(data)
		if n == 0 {
			return nil
		}
		if n > gattMTU {
			n = gattMTU
		}
		err := conn.tx.WriteValue(data[:n])
		if err != nil {
			return err
		}
		data = data[n:]
	}
}

func (conn *bleConn) Receive(data []byte) error {
	for i := 0; i < len(data); i++ {
		b, ok := <-conn.rx
		if !ok {
			return fmt.Errorf("input channel closed")
		}
		data[i] = b
	}
	return nil
}

var (
	receiverService = dexcomUUID(0xa0b1)
)

func connect(conn *ble.Connection) error {
	device, err := conn.Discover(time.Minute, receiverService)
	if err != nil {
		return err
	}
	if !device.Connected() {
		err = device.Connect()
		if err != nil {
			return err
		}
		log.Printf("%s: connected", device.Name())
	} else {
		log.Printf("%s: already connected", device.Name())
	}
	if !device.Paired() {
		err = device.Pair()
		if err != nil {
			return err
		}
		log.Printf("%s: paired", device.Name())
	} else {
		log.Printf("%s: already paired", device.Name())
	}
	err = conn.Update()
	if err != nil {
		return err
	}
	return authenticate(device)
}

var (
	authentication = dexcomUUID(0xacac)
	authEnvVar     = "DEXCOM_CGM_ID"
	authCode       = []byte{}
)

func initAuthCode() error {
	if len(authCode) != 0 {
		return nil
	}
	id := os.Getenv(authEnvVar)
	if len(id) == 0 {
		return fmt.Errorf("%s environment variable is not set", authEnvVar)
	}
	if len(id) != 10 {
		return fmt.Errorf("%s environment variable must be 2 letters followed by 8 digits", authEnvVar)
	}
	authCode = []byte(id + "000000")
	return nil
}

func authenticate(device ble.Device) error {
	err := initAuthCode()
	if err != nil {
		return err
	}
	auth, err := device.Conn().GetCharacteristic(authentication)
	if err != nil {
		return err
	}
	data, err := auth.ReadValue()
	if err != nil {
		return err
	}
	if bytes.Equal(data, authCode) {
		log.Printf("%s: already authenticated", device.Name())
		return nil
	}
	err = auth.WriteValue(authCode)
	if err != nil {
		return err
	}
	log.Printf("%s: authenticated", device.Name())
	return nil
}

var (
	heartbeat   = dexcomUUID(0x2b18)
	sendData    = dexcomUUID(0xb20a)
	receiveData = dexcomUUID(0xb20b)
)

func OpenBLE() (Connection, error) {
	conn, err := ble.Open()
	if err != nil {
		return nil, err
	}

	err = connect(conn)
	if err != nil {
		conn.Close()
		return nil, err
	}

	// We need to enable heartbeat notifications
	// or else we won't get any receiveData responses.
	err = conn.HandleNotify(heartbeat, func(data []byte) {})
	if err != nil {
		conn.Close()
		return nil, err
	}

	rx := make(chan byte, 1600)
	err = conn.HandleNotify(receiveData, func(data []byte) {
		for _, b := range data {
			rx <- b
		}
	})
	if err != nil {
		conn.Close()
		return nil, err
	}

	tx, err := conn.GetCharacteristic(sendData)
	if err != nil {
		conn.Close()
		return nil, err
	}

	return &bleConn{conn: conn, tx: tx, rx: rx}, nil
}

func dexcomUUID(id uint16) string {
	return fmt.Sprintf("f0ac%04x-ebfa-f96f-28da-076c35a521db", id)
}

func (conn *bleConn) Close() {
	conn.conn.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"io"
	"os"
	"fmt"
)

type shell struct {
	appinfo
}

func (app shell) run() {
	fmt.Printf("Welcome to the LDS Scriptures interactive shell.\n")
	cin := bufio.NewReader(os.Stdin)

	for {
		app.handleLine(cin)
	}
}

func (app shell) handleLine(cin *bufio.Reader) {
	fmt.Printf("> ");
	line, isPrefix, err := cin.ReadLine()
	if err != nil {
		panic(err)
	}
}<commit_msg>Made shell.go compile<commit_after>package main

import (
	"bufio"
	_ "io"
	"os"
	"fmt"
)

type shell struct {
	appinfo
}

func (app shell) run() {
	fmt.Printf("Welcome to the LDS Scriptures interactive shell.\n")
	cin := bufio.NewReader(os.Stdin)

	for {
		app.handleLine(cin)
	}
}

func (app shell) handleLine(cin *bufio.Reader) {
	fmt.Printf("> ");
	line, isPrefix, err := cin.ReadLine()
	_ = line
	_ = isPrefix
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/evolsnow/robot/conn"
	"github.com/go-telegram-bot-api/telegram-bot-api"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"sync"
	"time"
)

func (rb *Robot) Start(update tgbotapi.Update) string {
	user := update.Message.Chat.UserName
	go conn.SetUserChatId(user, update.Message.Chat.ID)
	return "welcome: " + user
}

func (rb *Robot) Help(update tgbotapi.Update) string {
	helpMsg := `
/alarm - set a reminder
/evolve	- self evolution of me
/movie - find movie download links
/trans - translate words between english and chinese
/help - show this help message
`
	return helpMsg
}

func (rb *Robot) Evolve(update tgbotapi.Update) {
	if update.Message.Chat.FirstName != "Evol" || update.Message.Chat.LastName != "Gan" {
		rb.bot.Send(tgbotapi.NewMessage(update.Message.Chat.ID, "sorry, unauthorized"))
		return
	}
	<-saidGoodBye
	close(saidGoodBye)
	cmd := exec.Command("bash", "/root/evolve_"+rb.nickName)
	cmd.Start()
	os.Exit(1)
}

func (rb *Robot) Translate(update tgbotapi.Update) string {
	var info string
	if string(update.Message.Text[0]) == "/" {
		raw := strings.SplitAfterN(update.Message.Text, " ", 2)
		if len(raw) < 2 {
			return "what do you want to translate, try '/trans cat'?"
		} else {
			info = "翻译" + raw[1]
		}
	} else {
		info = update.Message.Text
	}

	return qinAI(info)

}
func (rb *Robot) Talk(update tgbotapi.Update) string {
	info := update.Message.Text
	if strings.Contains(info, rb.name) {
		if strings.Contains(info, "闭嘴") || strings.Contains(info, "别说话") {
			rb.shutUp = true
		} else if rb.shutUp && strings.Contains(info, "说话") {
			rb.shutUp = false
			return fmt.Sprintf("%s终于可以说话啦", rb.nickName)
		}
		info = strings.Replace(info, fmt.Sprintf("@%s", rb.name), "", -1)
	}

	if rb.shutUp {
		return ""
	}
	log.Printf(info)

	if rb.nickName == "samaritan" {
		if chinese(info) {
			return tlAI(info)
		} else {
			return mitAI(info)
		}
	} else { //jarvis use another AI
		return qinAI(info)
	}
	//	return response
}

func (rb *Robot) SetReminder(update tgbotapi.Update, step int) string {
	user := update.Message.Chat.UserName
	switch step {
	case 0:
		//known issue of go, you can not just assign update.Message.Chat.ID to userTask[user].ChatId
		tmpTask := userTask[user]
		tmpTask.ChatId = update.Message.Chat.ID
		tmpTask.Owner = update.Message.Chat.UserName
		userTask[user] = tmpTask

		tmpAction := userAction[user]
		tmpAction.ActionStep++
		userAction[user] = tmpAction
		return "Ok, what should I remind you to do?"
	case 1:
		//save thing
		tmpTask := userTask[user]
		tmpTask.Desc = update.Message.Text
		userTask[user] = tmpTask

		tmpAction := userAction[user]
		tmpAction.ActionStep++
		userAction[user] = tmpAction
		return "When or how much time after?\n" +
			"You can type:\n" +
			"'*2/14 11:30*' means 11:30 at 2/14 \n" + //first format
			"'*11:30*' means  11:30 today\n" + //second format
			"'*5m10s*' means 5 minutes 10 seconds later" //third format
	case 2:
		defer delete(userAction, user)
		//save time duration
		text := update.Message.Text
		text = strings.Replace(text, "：", ":", -1)
		firstFormat := "1/02 15:04"
		secondFormat := "15:04"
		thirdFormat := "15:04:05"
		var showTime string
		var scheduledTime time.Time
		var nowTime = time.Now()
		var du time.Duration
		var err error
		if strings.Contains(text, ":") {
			scheduledTime, err = time.Parse(firstFormat, text)
			nowTime, _ = time.Parse(firstFormat, nowTime.Format(firstFormat))
			showTime = scheduledTime.Format(firstFormat)
			if err != nil { //try to parse with first format
				scheduledTime, err = time.Parse(secondFormat, text)
				nowTime, _ = time.Parse(secondFormat, nowTime.Format(secondFormat))
				showTime = scheduledTime.Format(secondFormat)
				if err != nil {
					return "wrong format, try '2/14 11:30' or '11:30'?"
				}
			}
			du = scheduledTime.Sub(nowTime)
		} else {

			du, err = time.ParseDuration(text)
			scheduledTime = nowTime.Add(du)
			showTime = scheduledTime.Format(thirdFormat)
			if err != nil {
				return "wrong format, try '1h2m3s'?"
			}
		}
		//		tmpTask := userTask[user]
		//		tmpTask.When = scheduledTime
		//		userTask[user] = tmpTask
		go func(rb *Robot, ts Task) {
			timer := time.NewTimer(du)
			rawMsg := fmt.Sprintf("Hi %s, maybe it's time to:\n*%s*", ts.Owner, ts.Desc)
			msg := tgbotapi.NewMessage(ts.ChatId, rawMsg)
			msg.ParseMode = "markdown"
			<-timer.C
			_, err := rb.bot.Send(msg)
			if err != nil {
				rb.bot.Send(tgbotapi.NewMessage(conn.GetUserChatId(ts.Owner), rawMsg))
			}
			delete(userTask, user)
		}(rb, userTask[user])

		//		delete(userAction, user)
		return fmt.Sprintf("Ok, I will remind you that\n*%s* - *%s*", showTime, userTask[user].Desc)
	}
	return ""
}

func (rb *Robot) DownloadMovie(update tgbotapi.Update, step int, results chan string) (ret string) {
	user := update.Message.Chat.UserName
	switch step {
	case 0:
		//known issue of go, you can not just assign update.Message.Chat.ID to userTask[user].ChatId
		tmpAction := userAction[user]
		tmpAction.ActionStep++
		userAction[user] = tmpAction
		ret = "Ok, which movie do you want to download?"
	case 1:
		defer func() {
			delete(userAction, user)
			results <- "done"
		}()
		results <- "Searching..."
		movie := update.Message.Text
		var wg sync.WaitGroup
		wg.Add(2)
		go getMovieFromZmz(movie, results, &wg)
		go getMovieFromLbl(movie, results, &wg)
		wg.Wait()
	}
	return
}

func getMovieFromZmz(movie string, results chan string, wg *sync.WaitGroup) {
	defer wg.Done()
	return
}

func getMovieFromLbl(movie string, results chan string, wg *sync.WaitGroup) {
	defer wg.Done()
	var id string
	resp, _ := http.Get("http://www.lbldy.com/search/" + movie)
	defer resp.Body.Close()
	body, _ := ioutil.ReadAll(resp.Body)
	re, _ := regexp.Compile("
")
	//find first match case
	firstId := re.FindSubmatch(body)
	if len(firstId) == 0 {
		results <- fmt.Sprintf("no answer for *%s* from lbl", movie)
		return
	} else {
		id = string(firstId[1])
		resp, _ = http.Get("http://www.lbldy.com/movie/" + id + ".html")
		defer resp.Body.Close()
		re, _ = regexp.Compile(`<p><a href="(.*?)"( target="_blank">|>)(.*?)</a></p>`)
		body, _ := ioutil.ReadAll(resp.Body)
		//go does not support (?!)
		body = []byte(strings.Replace(string(body), `<a href="/xunlei/"`, "", -1))
		downloads := re.FindAllSubmatch(body, -1)
		if len(downloads) == 0 {
			results <- fmt.Sprintf("no answer for *%s* from lbl", movie)
			return
		} else {
			results <- "Results from lbl:\n\n"
			ret := ""
			for i := range downloads {
				ret += "[" + string(downloads[i][3]) + "]\n" + string(downloads[i][1]) + "\n\n"
			}
			results <- ret
		}
	}
}
<commit_msg>format<commit_after>package main

import (
	"fmt"
	"github.com/evolsnow/robot/conn"
	"github.com/go-telegram-bot-api/telegram-bot-api"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"sync"
	"time"
)

func (rb *Robot) Start(update tgbotapi.Update) string {
	user := update.Message.Chat.UserName
	go conn.SetUserChatId(user, update.Message.Chat.ID)
	return "welcome: " + user
}

func (rb *Robot) Help(update tgbotapi.Update) string {
	helpMsg := `
/alarm - set a reminder
/evolve	- self evolution of me
/movie - find movie download links
/trans - translate words between english and chinese
/help - show this help message
`
	return helpMsg
}

func (rb *Robot) Evolve(update tgbotapi.Update) {
	if update.Message.Chat.FirstName != "Evol" || update.Message.Chat.LastName != "Gan" {
		rb.bot.Send(tgbotapi.NewMessage(update.Message.Chat.ID, "sorry, unauthorized"))
		return
	}
	<-saidGoodBye
	close(saidGoodBye)
	cmd := exec.Command("bash", "/root/evolve_"+rb.nickName)
	cmd.Start()
	os.Exit(1)
}

func (rb *Robot) Translate(update tgbotapi.Update) string {
	var info string
	if string(update.Message.Text[0]) == "/" {
		raw := strings.SplitAfterN(update.Message.Text, " ", 2)
		if len(raw) < 2 {
			return "what do you want to translate, try '/trans cat'?"
		} else {
			info = "翻译" + raw[1]
		}
	} else {
		info = update.Message.Text
	}

	return qinAI(info)

}
func (rb *Robot) Talk(update tgbotapi.Update) string {
	info := update.Message.Text
	if strings.Contains(info, rb.name) {
		if strings.Contains(info, "闭嘴") || strings.Contains(info, "别说话") {
			rb.shutUp = true
		} else if rb.shutUp && strings.Contains(info, "说话") {
			rb.shutUp = false
			return fmt.Sprintf("%s终于可以说话啦", rb.nickName)
		}
		info = strings.Replace(info, fmt.Sprintf("@%s", rb.name), "", -1)
	}

	if rb.shutUp {
		return ""
	}
	log.Printf(info)

	if rb.nickName == "samaritan" {
		if chinese(info) {
			return tlAI(info)
		} else {
			return mitAI(info)
		}
	} else { //jarvis use another AI
		return qinAI(info)
	}
	//	return response
}

func (rb *Robot) SetReminder(update tgbotapi.Update, step int) string {
	user := update.Message.Chat.UserName
	switch step {
	case 0:
		//known issue of go, you can not just assign update.Message.Chat.ID to userTask[user].ChatId
		tmpTask := userTask[user]
		tmpTask.ChatId = update.Message.Chat.ID
		tmpTask.Owner = update.Message.Chat.UserName
		userTask[user] = tmpTask

		tmpAction := userAction[user]
		tmpAction.ActionStep++
		userAction[user] = tmpAction
		return "Ok, what should I remind you to do?"
	case 1:
		//save thing
		tmpTask := userTask[user]
		tmpTask.Desc = update.Message.Text
		userTask[user] = tmpTask

		tmpAction := userAction[user]
		tmpAction.ActionStep++
		userAction[user] = tmpAction
		return "When or how much time after?\n" +
			"You can type:\n" +
			"'*2/14 11:30*' means 11:30 at 2/14 \n" + //first format
			"'*11:30*' means  11:30 today\n" + //second format
			"'*5m10s*' means 5 minutes 10 seconds later" //third format
	case 2:
		defer delete(userAction, user)
		//save time duration
		text := update.Message.Text
		text = strings.Replace(text, "：", ":", -1)
		firstFormat := "1/02 15:04"
		secondFormat := "15:04"
		thirdFormat := "15:04:05"
		var showTime string
		var scheduledTime time.Time
		var nowTime = time.Now()
		var du time.Duration
		var err error
		if strings.Contains(text, ":") {
			scheduledTime, err = time.Parse(firstFormat, text)
			nowTime, _ = time.Parse(firstFormat, nowTime.Format(firstFormat))
			showTime = scheduledTime.Format(firstFormat)
			if err != nil { //try to parse with first format
				scheduledTime, err = time.Parse(secondFormat, text)
				nowTime, _ = time.Parse(secondFormat, nowTime.Format(secondFormat))
				showTime = scheduledTime.Format(secondFormat)
				if err != nil {
					return "wrong format, try '2/14 11:30' or '11:30'?"
				}
			}
			du = scheduledTime.Sub(nowTime)
		} else {

			du, err = time.ParseDuration(text)
			scheduledTime = nowTime.Add(du)
			showTime = scheduledTime.Format(thirdFormat)
			if err != nil {
				return "wrong format, try '1h2m3s'?"
			}
		}
		//		tmpTask := userTask[user]
		//		tmpTask.When = scheduledTime
		//		userTask[user] = tmpTask
		go func(rb *Robot, ts Task) {
			timer := time.NewTimer(du)
			rawMsg := fmt.Sprintf("Hi %s, maybe it's time to:\n*%s*", ts.Owner, ts.Desc)
			msg := tgbotapi.NewMessage(ts.ChatId, rawMsg)
			msg.ParseMode = "markdown"
			<-timer.C
			_, err := rb.bot.Send(msg)
			if err != nil {
				rb.bot.Send(tgbotapi.NewMessage(conn.GetUserChatId(ts.Owner), rawMsg))
			}
			delete(userTask, user)
		}(rb, userTask[user])

		//		delete(userAction, user)
		return fmt.Sprintf("Ok, I will remind you that\n*%s* - *%s*", showTime, userTask[user].Desc)
	}
	return ""
}

func (rb *Robot) DownloadMovie(update tgbotapi.Update, step int, results chan string) (ret string) {
	user := update.Message.Chat.UserName
	switch step {
	case 0:
		//known issue of go, you can not just assign update.Message.Chat.ID to userTask[user].ChatId
		tmpAction := userAction[user]
		tmpAction.ActionStep++
		userAction[user] = tmpAction
		ret = "Ok, which movie do you want to download?"
	case 1:
		defer func() {
			delete(userAction, user)
			results <- "done"
		}()
		results <- "Searching..."
		movie := update.Message.Text
		var wg sync.WaitGroup
		wg.Add(2)
		go getMovieFromZmz(movie, results, &wg)
		go getMovieFromLbl(movie, results, &wg)
		wg.Wait()
	}
	return
}

func getMovieFromZmz(movie string, results chan string, wg *sync.WaitGroup) {
	defer wg.Done()
	return
}

func getMovieFromLbl(movie string, results chan string, wg *sync.WaitGroup) {
	defer wg.Done()
	var id string
	resp, _ := http.Get("http://www.lbldy.com/search/" + movie)
	defer resp.Body.Close()
	body, _ := ioutil.ReadAll(resp.Body)
	re, _ := regexp.Compile("
")
	//find first match case
	firstId := re.FindSubmatch(body)
	if len(firstId) == 0 {
		results <- fmt.Sprintf("no answer for *%s* from lbl", movie)
		return
	} else {
		id = string(firstId[1])
		resp, _ = http.Get("http://www.lbldy.com/movie/" + id + ".html")
		defer resp.Body.Close()
		re, _ = regexp.Compile(`<p><a href="(.*?)"( target="_blank">|>)(.*?)</a></p>`)
		body, _ := ioutil.ReadAll(resp.Body)
		//go does not support (?!)
		body = []byte(strings.Replace(string(body), `<a href="/xunlei/"`, "", -1))
		downloads := re.FindAllSubmatch(body, -1)
		if len(downloads) == 0 {
			results <- fmt.Sprintf("no answer for *%s* from lbl", movie)
			return
		} else {
			results <- "Results from lbl:\n\n"
			ret := ""
			for i := range downloads {
				ret += fmt.Sprintf("[%s]\n%s\n\n", string(downloads[i][3]), string(downloads[i][1]))
			}
			results <- ret
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"log"

	"github.com/nlopes/slack"
)

type user struct {
	fullName string
	email    string
}

func generateMessageAttachment(api *slack.Client, messageTimestamp, username, channel, decisionVerb string) slack.Attachment {

	history, err := api.GetChannelHistory(channel, slack.HistoryParameters{
		Count:     1,
		Inclusive: true,
		Latest:    messageTimestamp,
	})
	if err != nil {
		log.Println(err)
	}

	messageAttachment := history.Messages[0].Msg.Attachments[0]
	messageAttachment.Fields[0] = slack.AttachmentField{
		Value: fmt.Sprintf("%s by %s", decisionVerb, username),
		Short: false,
	}

	return messageAttachment
}

func updateMessage(api *slack.Client, attachment slack.Attachment, messageTimestamp, channel string) {
	//We use SendMessage here instead of UpdateMessage because the latter does not support updating attachments
	// https://github.com/nlopes/slack/pull/121
	api.SendMessage(channel, slack.MsgOptionUpdate(messageTimestamp), slack.MsgOptionAttachments(attachment))
}

func getMessageCallbackInfo(api *slack.Client, messageTimestamp, channel string) map[string]interface{} {
	history, err := api.GetChannelHistory(channel, slack.HistoryParameters{
		Count:     1,
		Inclusive: true,
		Latest:    messageTimestamp,
	})
	if err != nil {
		log.Println(err)
	}

	messageInfo := history.Messages[0].Msg.Attachments[0].CallbackID

	var callbackInfo map[string]interface{}
	json.Unmarshal([]byte(messageInfo), &callbackInfo)

	return callbackInfo

}

func getUser(api *slack.Client, userID string) (*user, error) {
	userInfo, err := api.GetUserInfo(userID)
	if err != nil {
		log.Printf("%s\n", err)
		return nil, err
	}

	return &user{userInfo.Profile.RealName, userInfo.Profile.Email}, nil

}

func processDecision(api *slack.Client, lgts *lgts, event *slack.ReactionAddedEvent) {

	userID := event.User
	channel := event.Item.Channel
	messageTimestamp := event.Item.Timestamp
	emojiUsed := event.Reaction

	//Check if the emoji used is one of the predefined emojis
	var isApproved bool
	var decisionVerb string

	switch {
	case lgts.isApprovalEmoji(emojiUsed):
		isApproved = true
		decisionVerb = "approved"
	case lgts.isRejectionEmoji(emojiUsed):
		isApproved = false
		decisionVerb = "rejected"
	default:
		return
	}

	//Get information about the message
	callbackInfo := getMessageCallbackInfo(api, messageTimestamp, channel)

	if _, ok := callbackInfo["message_id"]; !ok {
		log.Printf("message_id parameter missing from callback id string")
		return
	}
	if _, ok := callbackInfo["app_name"]; !ok {
		log.Printf("app_name parameter missing from callback id string")
		return
	}

	//Use information to recieve proper app/message object and check for existence
	message, present := lgts.Messages[callbackInfo["message_id"].(string)]
	if !present {
		log.Println("Message processed but not found in queue")
		return
	}

	app, present := lgts.Apps[message.AppName]
	if !present {
		log.Println("Message processed but app not registered")
		return
	}

	//Check if user who used the emoji is part of approved list
	userInfo, err := getUser(api, userID)
	if err != nil {
		log.Printf("Cannot find slack user: %v", err)
		return
	}

	if !app.isAuthorizedUser(userInfo.email) {
		log.Println("User not authorized to approve")
		return
	}

	//User, emoji, and messageid check out
	// update slack message and send callback url proper message

	err = app.sendMessageApproval(callbackInfo, userInfo.email, isApproved)
	if err != nil {
		//return
		log.Printf("Couldn't send proper request: %v", err)
	}

	attachment := generateMessageAttachment(api, messageTimestamp, userInfo.fullName, channel, decisionVerb)
	updateMessage(api, attachment, messageTimestamp, channel)

	log.Printf("Message ID %s was %s by slack user %s", message.ID, decisionVerb, userInfo.fullName)

	delete(lgts.Messages, message.ID)

}

func runrtm(lgts *lgts, slackToken string) {

	api := slack.New(slackToken)
	//api.SetDebug(true)

	rtm := api.NewRTM()
	go rtm.ManageConnection()

	log.Println("Starting slack event reader")

	for msg := range rtm.IncomingEvents {
		switch ev := msg.Data.(type) {
		case *slack.ReactionAddedEvent:
			processDecision(api, lgts, ev)

		case *slack.RTMError:
			log.Printf("Error: %s\n", ev.Error())

		case *slack.InvalidAuthEvent:
			log.Printf("Invalid credentials")
			return

		default:
			// Ignore other events..
		}
	}
}
<commit_msg>Removing testing return skip<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"log"

	"github.com/nlopes/slack"
)

type user struct {
	fullName string
	email    string
}

func generateMessageAttachment(api *slack.Client, messageTimestamp, username, channel, decisionVerb string) slack.Attachment {

	history, err := api.GetChannelHistory(channel, slack.HistoryParameters{
		Count:     1,
		Inclusive: true,
		Latest:    messageTimestamp,
	})
	if err != nil {
		log.Println(err)
	}

	messageAttachment := history.Messages[0].Msg.Attachments[0]
	messageAttachment.Fields[0] = slack.AttachmentField{
		Value: fmt.Sprintf("%s by %s", decisionVerb, username),
		Short: false,
	}

	return messageAttachment
}

func updateMessage(api *slack.Client, attachment slack.Attachment, messageTimestamp, channel string) {
	//We use SendMessage here instead of UpdateMessage because the latter does not support updating attachments
	// https://github.com/nlopes/slack/pull/121
	api.SendMessage(channel, slack.MsgOptionUpdate(messageTimestamp), slack.MsgOptionAttachments(attachment))
}

func getMessageCallbackInfo(api *slack.Client, messageTimestamp, channel string) map[string]interface{} {
	history, err := api.GetChannelHistory(channel, slack.HistoryParameters{
		Count:     1,
		Inclusive: true,
		Latest:    messageTimestamp,
	})
	if err != nil {
		log.Println(err)
	}

	messageInfo := history.Messages[0].Msg.Attachments[0].CallbackID

	var callbackInfo map[string]interface{}
	json.Unmarshal([]byte(messageInfo), &callbackInfo)

	return callbackInfo

}

func getUser(api *slack.Client, userID string) (*user, error) {
	userInfo, err := api.GetUserInfo(userID)
	if err != nil {
		log.Printf("%s\n", err)
		return nil, err
	}

	return &user{userInfo.Profile.RealName, userInfo.Profile.Email}, nil

}

func processDecision(api *slack.Client, lgts *lgts, event *slack.ReactionAddedEvent) {

	userID := event.User
	channel := event.Item.Channel
	messageTimestamp := event.Item.Timestamp
	emojiUsed := event.Reaction

	//Check if the emoji used is one of the predefined emojis
	var isApproved bool
	var decisionVerb string

	switch {
	case lgts.isApprovalEmoji(emojiUsed):
		isApproved = true
		decisionVerb = "approved"
	case lgts.isRejectionEmoji(emojiUsed):
		isApproved = false
		decisionVerb = "rejected"
	default:
		return
	}

	//Get information about the message
	callbackInfo := getMessageCallbackInfo(api, messageTimestamp, channel)

	if _, ok := callbackInfo["message_id"]; !ok {
		log.Printf("message_id parameter missing from callback id string")
		return
	}
	if _, ok := callbackInfo["app_name"]; !ok {
		log.Printf("app_name parameter missing from callback id string")
		return
	}

	//Use information to recieve proper app/message object and check for existence
	message, present := lgts.Messages[callbackInfo["message_id"].(string)]
	if !present {
		log.Println("Message processed but not found in queue")
		return
	}

	app, present := lgts.Apps[message.AppName]
	if !present {
		log.Println("Message processed but app not registered")
		return
	}

	//Check if user who used the emoji is part of approved list
	userInfo, err := getUser(api, userID)
	if err != nil {
		log.Printf("Cannot find slack user: %v", err)
		return
	}

	if !app.isAuthorizedUser(userInfo.email) {
		log.Println("User not authorized to approve")
		return
	}

	//User, emoji, and messageid check out
	// update slack message and send callback url proper message

	err = app.sendMessageApproval(callbackInfo, userInfo.email, isApproved)
	if err != nil {
		log.Printf("Couldn't send proper request: %v", err)
		return
	}

	attachment := generateMessageAttachment(api, messageTimestamp, userInfo.fullName, channel, decisionVerb)
	updateMessage(api, attachment, messageTimestamp, channel)

	log.Printf("Message ID %s was %s by slack user %s", message.ID, decisionVerb, userInfo.fullName)

	delete(lgts.Messages, message.ID)

}

func runrtm(lgts *lgts, slackToken string) {

	api := slack.New(slackToken)
	//api.SetDebug(true)

	rtm := api.NewRTM()
	go rtm.ManageConnection()

	log.Println("Starting slack event reader")

	for msg := range rtm.IncomingEvents {
		switch ev := msg.Data.(type) {
		case *slack.ReactionAddedEvent:
			processDecision(api, lgts, ev)

		case *slack.RTMError:
			log.Printf("Error: %s\n", ev.Error())

		case *slack.InvalidAuthEvent:
			log.Printf("Invalid credentials")
			return

		default:
			// Ignore other events..
		}
	}
}
<|endoftext|>


<commit_before>package organisations

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/Financial-Times/go-fthealth/v1a"
	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/mux"
)

// OrganisationDriver for cypher queries
var OrganisationDriver Driver
var CacheControlHeader string

// HealthCheck does something
func HealthCheck() v1a.Check {
	return v1a.Check{
		BusinessImpact:   "Unable to respond to Public Organisations api requests",
		Name:             "Check connectivity to Neo4j - neoUrl is a parameter in hieradata for this service",
		PanicGuide:       "https://sites.google.com/a/ft.com/ft-technology-service-transition/home/run-book-library/organisation-read-api",
		Severity:         1,
		TechnicalSummary: "Cannot connect to Neo4j a instance with at least one organisation loaded in it",
		Checker:          Checker,
	}
}

// Checker does more stuff
func Checker() (string, error) {
	err := OrganisationDriver.CheckConnectivity()
	if err == nil {
		return "Connectivity to neo4j is ok", err
	}
	return "Error connecting to neo4j", err
}

// Ping says pong
func Ping(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "pong")
}

// BuildInfoHandler - This is a stop gap and will be added to when we can define what we should display here
func BuildInfoHandler(w http.ResponseWriter, req *http.Request) {
	fmt.Fprintf(w, "build-info")
}

// MethodNotAllowedHandler does stuff
func MethodNotAllowedHandler(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusMethodNotAllowed)
	return
}

// GetOrganisation is the public API
func GetOrganisation(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	uuid := vars["uuid"]

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	if uuid == "" {
		http.Error(w, "uuid required", http.StatusBadRequest)
		return
	}
	organisation, found, err := OrganisationDriver.Read(uuid)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte(`{"message": "` + err.Error() + `"}`))
		return
	}
	if !found {
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte(`{"message":"Organisation not found."}`))
		return
	}
	Jason, _ := json.Marshal(organisation)
	log.Debugf("Organisation(uuid:%s): %s\n", Jason)
	w.Header().Set("Cache-Control", CacheControlHeader)
	w.WriteHeader(http.StatusOK)
	json.NewEncoder(w).Encode(organisation)
}
<commit_msg>Removing double unmarshalling.<commit_after>package organisations

import (
	"encoding/json"
	"fmt"
	"github.com/Financial-Times/go-fthealth/v1a"
	"github.com/gorilla/mux"
	"net/http"
)

// OrganisationDriver for cypher queries
var OrganisationDriver Driver
var CacheControlHeader string

// HealthCheck does something
func HealthCheck() v1a.Check {
	return v1a.Check{
		BusinessImpact:   "Unable to respond to Public Organisations api requests",
		Name:             "Check connectivity to Neo4j - neoUrl is a parameter in hieradata for this service",
		PanicGuide:       "https://sites.google.com/a/ft.com/ft-technology-service-transition/home/run-book-library/organisation-read-api",
		Severity:         1,
		TechnicalSummary: "Cannot connect to Neo4j a instance with at least one organisation loaded in it",
		Checker:          Checker,
	}
}

// Checker does more stuff
func Checker() (string, error) {
	err := OrganisationDriver.CheckConnectivity()
	if err == nil {
		return "Connectivity to neo4j is ok", err
	}
	return "Error connecting to neo4j", err
}

// Ping says pong
func Ping(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "pong")
}

// BuildInfoHandler - This is a stop gap and will be added to when we can define what we should display here
func BuildInfoHandler(w http.ResponseWriter, req *http.Request) {
	fmt.Fprintf(w, "build-info")
}

// MethodNotAllowedHandler does stuff
func MethodNotAllowedHandler(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusMethodNotAllowed)
	return
}

// GetOrganisation is the public API
func GetOrganisation(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	uuid := vars["uuid"]

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	if uuid == "" {
		http.Error(w, "uuid required", http.StatusBadRequest)
		return
	}
	organisation, found, err := OrganisationDriver.Read(uuid)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte(`{"message": "` + err.Error() + `"}`))
		return
	}
	if !found {
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte(`{"message":"Organisation not found."}`))
		return
	}

	w.Header().Set("Cache-Control", CacheControlHeader)
	w.WriteHeader(http.StatusOK)
	err = json.NewEncoder(w).Encode(organisation)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte(`{"message":"Organisation could not be marshelled, err=` + err.Error() + `"}`))
	}
}
<|endoftext|>


<commit_before>package ping

import (
	"bytes"
	"context"
	"errors"
	"io"
	"time"

	u "github.com/ipfs/go-ipfs-util"
	logging "github.com/ipfs/go-log"
	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
)

var log = logging.Logger("ping")

const PingSize = 32

const ID = "/ipfs/ping/1.0.0"

const pingTimeout = time.Second * 60

type PingService struct {
	Host host.Host
}

func NewPingService(h host.Host) *PingService {
	ps := &PingService{h}
	h.SetStreamHandler(ID, ps.PingHandler)
	return ps
}

func (p *PingService) PingHandler(s inet.Stream) {
	buf := make([]byte, PingSize)

	errCh := make(chan error, 1)
	defer close(errCh)
	timer := time.NewTimer(pingTimeout)
	defer timer.Stop()

	go func() {
		select {
		case <-timer.C:
			log.Debug("ping timeout")
		case err, ok := <-errCh:
			if ok {
				log.Debug(err)
			} else {
				log.Error("ping loop failed without error")
			}
		}
		s.Reset()
	}()

	for {
		_, err := io.ReadFull(s, buf)
		if err != nil {
			errCh <- err
			return
		}

		_, err = s.Write(buf)
		if err != nil {
			errCh <- err
			return
		}

		timer.Reset(pingTimeout)
	}
}

// Result is a result of a ping attempt, either an RTT or an error.
type Result struct {
	RTT   time.Duration
	Error error
}

func (ps *PingService) Ping(ctx context.Context, p peer.ID) <-chan Result {
	return Ping(ctx, ps.Host, p)
}

// Ping pings the remote peer until the context is canceled, returning a stream
// of RTTs or errors.
func Ping(ctx context.Context, h host.Host, p peer.ID) <-chan Result {
	s, err := h.NewStream(ctx, p, ID)
	if err != nil {
		ch := make(chan Result, 1)
		ch <- Result{Error: err}
		close(ch)
		return ch
	}

	ctx, cancel := context.WithCancel(ctx)

	out := make(chan Result)
	go func() {
		defer close(out)
		defer cancel()

		for {
			select {
			case <-ctx.Done():
				return
			default:
				var res Result
				res.RTT, res.Error = ping(s)

				// canceled, ignore everything.
				if ctx.Err() != nil {
					return
				}

				// No error, record the RTT.
				if res.Error == nil {
					h.Peerstore().RecordLatency(p, res.RTT)
				}

				select {
				case out <- res:
				case <-ctx.Done():
					return
				}
			}
		}
	}()
	go func() {
		// forces the ping to abort.
		<-ctx.Done()
		s.Reset()
	}()

	return out
}

func ping(s inet.Stream) (time.Duration, error) {
	buf := make([]byte, PingSize)
	u.NewTimeSeededRand().Read(buf)

	before := time.Now()
	_, err := s.Write(buf)
	if err != nil {
		return 0, err
	}

	rbuf := make([]byte, PingSize)
	_, err = io.ReadFull(s, rbuf)
	if err != nil {
		return 0, err
	}

	if !bytes.Equal(buf, rbuf) {
		return 0, errors.New("ping packet was incorrect!")
	}

	return time.Since(before), nil
}
<commit_msg>ping: simplify ping loop<commit_after>package ping

import (
	"bytes"
	"context"
	"errors"
	"io"
	"time"

	u "github.com/ipfs/go-ipfs-util"
	logging "github.com/ipfs/go-log"
	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
)

var log = logging.Logger("ping")

const PingSize = 32

const ID = "/ipfs/ping/1.0.0"

const pingTimeout = time.Second * 60

type PingService struct {
	Host host.Host
}

func NewPingService(h host.Host) *PingService {
	ps := &PingService{h}
	h.SetStreamHandler(ID, ps.PingHandler)
	return ps
}

func (p *PingService) PingHandler(s inet.Stream) {
	buf := make([]byte, PingSize)

	errCh := make(chan error, 1)
	defer close(errCh)
	timer := time.NewTimer(pingTimeout)
	defer timer.Stop()

	go func() {
		select {
		case <-timer.C:
			log.Debug("ping timeout")
		case err, ok := <-errCh:
			if ok {
				log.Debug(err)
			} else {
				log.Error("ping loop failed without error")
			}
		}
		s.Reset()
	}()

	for {
		_, err := io.ReadFull(s, buf)
		if err != nil {
			errCh <- err
			return
		}

		_, err = s.Write(buf)
		if err != nil {
			errCh <- err
			return
		}

		timer.Reset(pingTimeout)
	}
}

// Result is a result of a ping attempt, either an RTT or an error.
type Result struct {
	RTT   time.Duration
	Error error
}

func (ps *PingService) Ping(ctx context.Context, p peer.ID) <-chan Result {
	return Ping(ctx, ps.Host, p)
}

// Ping pings the remote peer until the context is canceled, returning a stream
// of RTTs or errors.
func Ping(ctx context.Context, h host.Host, p peer.ID) <-chan Result {
	s, err := h.NewStream(ctx, p, ID)
	if err != nil {
		ch := make(chan Result, 1)
		ch <- Result{Error: err}
		close(ch)
		return ch
	}

	ctx, cancel := context.WithCancel(ctx)

	out := make(chan Result)
	go func() {
		defer close(out)
		defer cancel()

		for ctx.Err() == nil {
			var res Result
			res.RTT, res.Error = ping(s)

			// canceled, ignore everything.
			if ctx.Err() != nil {
				return
			}

			// No error, record the RTT.
			if res.Error == nil {
				h.Peerstore().RecordLatency(p, res.RTT)
			}

			select {
			case out <- res:
			case <-ctx.Done():
				return
			}
		}
	}()
	go func() {
		// forces the ping to abort.
		<-ctx.Done()
		s.Reset()
	}()

	return out
}

func ping(s inet.Stream) (time.Duration, error) {
	buf := make([]byte, PingSize)
	u.NewTimeSeededRand().Read(buf)

	before := time.Now()
	_, err := s.Write(buf)
	if err != nil {
		return 0, err
	}

	rbuf := make([]byte, PingSize)
	_, err = io.ReadFull(s, rbuf)
	if err != nil {
		return 0, err
	}

	if !bytes.Equal(buf, rbuf) {
		return 0, errors.New("ping packet was incorrect!")
	}

	return time.Since(before), nil
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package poll

import (
	"fmt"
	"strings"
	"time"

	"github.com/ava-labs/avalanchego/utils/linkedhashmap"

	"github.com/prometheus/client_golang/prometheus"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/ava-labs/avalanchego/utils/metric"
)

type pollHolder interface {
	GetPoll() Poll
	StartTime() time.Time
}

type poll struct {
	Poll
	start time.Time
}

func (p poll) GetPoll() Poll {
	return p
}

func (p poll) StartTime() time.Time {
	return p.start
}

type set struct {
	log      logging.Logger
	numPolls prometheus.Gauge
	durPolls prometheus.Histogram
	factory  Factory
	// maps requestID -> poll
	polls linkedhashmap.LinkedHashmap
}

// NewSet returns a new empty set of polls
func NewSet(
	factory Factory,
	log logging.Logger,
	namespace string,
	registerer prometheus.Registerer,
) Set {
	numPolls := prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: namespace,
		Name:      "polls",
		Help:      "Number of pending network polls",
	})
	if err := registerer.Register(numPolls); err != nil {
		log.Error("failed to register polls statistics due to %s", err)
	}

	durPolls := prometheus.NewHistogram(prometheus.HistogramOpts{
		Namespace: namespace,
		Name:      "poll_duration",
		Help:      "Length of time the poll existed in milliseconds",
		Buckets:   metric.MillisecondsBuckets,
	})
	if err := registerer.Register(durPolls); err != nil {
		log.Error("failed to register poll_duration statistics due to %s", err)
	}

	return &set{
		log:      log,
		numPolls: numPolls,
		durPolls: durPolls,
		factory:  factory,
		polls:    linkedhashmap.New(),
	}
}

// Add to the current set of polls
// Returns true if the poll was registered correctly and the network sample
//         should be made.
func (s *set) Add(requestID uint32, vdrs ids.ShortBag) bool {
	if _, exists := s.polls.Get(requestID); exists {
		s.log.Debug("dropping poll due to duplicated requestID: %d", requestID)
		return false
	}

	s.log.Verbo("creating poll with requestID %d and validators %s",
		requestID,
		&vdrs)

	s.polls.Put(requestID, poll{
		Poll:  s.factory.New(vdrs), // create the new poll
		start: time.Now(),
	})
	s.numPolls.Inc() // increase the metrics
	return true
}

// Vote registers the connections response to a query for [id]. If there was no
// query, or the response has already be registered, nothing is performed.
func (s *set) Vote(requestID uint32, vdr ids.ShortID, vote ids.ID) ([]ids.Bag, bool) {
	pollHolderIntf, exists := s.polls.Get(requestID)
	if !exists {
		s.log.Verbo("dropping vote from %s to an unknown poll with requestID: %d",
			vdr,
			requestID)
		return []ids.Bag{}, false
	}

	holder := pollHolderIntf.(pollHolder)
	p := holder.GetPoll()

	s.log.Verbo("processing vote from %s in the poll with requestID: %d with the vote %s",
		vdr,
		requestID,
		vote)

	p.Vote(vdr, vote)
	if !p.Finished() {
		return []ids.Bag{}, false
	}

	s.log.Verbo("poll with requestID %d finished as %s", requestID, p)
	s.durPolls.Observe(float64(time.Since(holder.StartTime()).Milliseconds()))
	s.numPolls.Dec() // decrease the metrics

	var results []ids.Bag
	// check if previous poll(s) have finished
	if oldestRequestID, _, exists := s.polls.Oldest(); exists && oldestRequestID == requestID {
		// this is the oldest poll that has just finished
		// iterate from oldest to newest
		iter := s.polls.NewIterator()
		var keysToDelete []interface{}
		for iter.Next() {
			holder := iter.Value().(pollHolder)
			p = holder.GetPoll()
			if !p.Finished() {
				// since we're iterating from oldest to newest, if the next poll has not finished,
				// we can break and return what we have so far
				break
			}

			results = append(results, p.Result())
			keysToDelete = append(keysToDelete, iter.Key())
		}

		// delete the keys to be deleted
		for _, keyToDelete := range keysToDelete {
			s.polls.Delete(keyToDelete)
		}
	}

	// only gets here if the poll has finished
	return results, len(results) > 0
}

// Drop registers the connections response to a query for [id]. If there was no
// query, or the response has already be registered, nothing is performed.
func (s *set) Drop(requestID uint32, vdr ids.ShortID) ([]ids.Bag, bool) {
	pollHolderIntf, exists := s.polls.Get(requestID)
	if !exists {
		s.log.Verbo("dropping vote from %s to an unknown poll with requestID: %d",
			vdr,
			requestID)
		return []ids.Bag{}, false
	}

	s.log.Verbo("processing dropped vote from %s in the poll with requestID: %d",
		vdr,
		requestID)

	pollHolder := pollHolderIntf.(pollHolder)
	poll := pollHolder.GetPoll()

	poll.Drop(vdr)
	if !poll.Finished() {
		return []ids.Bag{}, false
	}

	s.log.Verbo("poll with requestID %d finished as %s", requestID, poll)

	s.polls.Delete(requestID) // remove the poll from the current set
	s.durPolls.Observe(float64(time.Since(pollHolder.StartTime()).Milliseconds()))
	s.numPolls.Dec() // decrease the metrics
	return []ids.Bag{poll.Result()}, true
}

// Len returns the number of outstanding polls
func (s *set) Len() int { return s.polls.Len() }

func (s *set) String() string {
	sb := strings.Builder{}
	sb.WriteString(fmt.Sprintf("current polls: (Size = %d)", s.polls.Len()))
	iter := s.polls.NewIterator()
	for iter.Next() {
		requestID := iter.Key()
		poll := iter.Value().(Poll)
		sb.WriteString(fmt.Sprintf("\n    %d: %s", requestID, poll.PrefixedString("    ")))
	}
	return sb.String()
}
<commit_msg>add comment<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package poll

import (
	"fmt"
	"strings"
	"time"

	"github.com/ava-labs/avalanchego/utils/linkedhashmap"

	"github.com/prometheus/client_golang/prometheus"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/ava-labs/avalanchego/utils/metric"
)

type pollHolder interface {
	GetPoll() Poll
	StartTime() time.Time
}

type poll struct {
	Poll
	start time.Time
}

func (p poll) GetPoll() Poll {
	return p
}

func (p poll) StartTime() time.Time {
	return p.start
}

type set struct {
	log      logging.Logger
	numPolls prometheus.Gauge
	durPolls prometheus.Histogram
	factory  Factory
	// maps requestID -> poll
	polls linkedhashmap.LinkedHashmap
}

// NewSet returns a new empty set of polls
func NewSet(
	factory Factory,
	log logging.Logger,
	namespace string,
	registerer prometheus.Registerer,
) Set {
	numPolls := prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: namespace,
		Name:      "polls",
		Help:      "Number of pending network polls",
	})
	if err := registerer.Register(numPolls); err != nil {
		log.Error("failed to register polls statistics due to %s", err)
	}

	durPolls := prometheus.NewHistogram(prometheus.HistogramOpts{
		Namespace: namespace,
		Name:      "poll_duration",
		Help:      "Length of time the poll existed in milliseconds",
		Buckets:   metric.MillisecondsBuckets,
	})
	if err := registerer.Register(durPolls); err != nil {
		log.Error("failed to register poll_duration statistics due to %s", err)
	}

	return &set{
		log:      log,
		numPolls: numPolls,
		durPolls: durPolls,
		factory:  factory,
		polls:    linkedhashmap.New(),
	}
}

// Add to the current set of polls
// Returns true if the poll was registered correctly and the network sample
//         should be made.
func (s *set) Add(requestID uint32, vdrs ids.ShortBag) bool {
	if _, exists := s.polls.Get(requestID); exists {
		s.log.Debug("dropping poll due to duplicated requestID: %d", requestID)
		return false
	}

	s.log.Verbo("creating poll with requestID %d and validators %s",
		requestID,
		&vdrs)

	s.polls.Put(requestID, poll{
		Poll:  s.factory.New(vdrs), // create the new poll
		start: time.Now(),
	})
	s.numPolls.Inc() // increase the metrics
	return true
}

// Vote registers the connections response to a query for [id]. If there was no
// query, or the response has already be registered, nothing is performed.
func (s *set) Vote(requestID uint32, vdr ids.ShortID, vote ids.ID) ([]ids.Bag, bool) {
	pollHolderIntf, exists := s.polls.Get(requestID)
	if !exists {
		s.log.Verbo("dropping vote from %s to an unknown poll with requestID: %d",
			vdr,
			requestID)
		return []ids.Bag{}, false
	}

	holder := pollHolderIntf.(pollHolder)
	p := holder.GetPoll()

	s.log.Verbo("processing vote from %s in the poll with requestID: %d with the vote %s",
		vdr,
		requestID,
		vote)

	p.Vote(vdr, vote)
	if !p.Finished() {
		return []ids.Bag{}, false
	}

	s.log.Verbo("poll with requestID %d finished as %s", requestID, p)
	s.durPolls.Observe(float64(time.Since(holder.StartTime()).Milliseconds()))
	s.numPolls.Dec() // decrease the metrics

	var results []ids.Bag
	// check if previous poll(s) have finished
	if oldestRequestID, _, exists := s.polls.Oldest(); exists && oldestRequestID == requestID {
		// this is the oldest poll that has just finished
		// iterate from oldest to newest
		iter := s.polls.NewIterator()
		var keysToDelete []interface{}
		for iter.Next() {
			holder := iter.Value().(pollHolder)
			p = holder.GetPoll()
			if !p.Finished() {
				// since we're iterating from oldest to newest, if the next poll has not finished,
				// we can break and return what we have so far
				break
			}

			results = append(results, p.Result())
			keysToDelete = append(keysToDelete, iter.Key())
		}

		// delete the keys to be deleted
		for _, keyToDelete := range keysToDelete {
			s.polls.Delete(keyToDelete)
		}
	}

	// only gets here if the poll has finished
	// results will have values if this and other newer polls have finished
	return results, len(results) > 0
}

// Drop registers the connections response to a query for [id]. If there was no
// query, or the response has already be registered, nothing is performed.
func (s *set) Drop(requestID uint32, vdr ids.ShortID) ([]ids.Bag, bool) {
	pollHolderIntf, exists := s.polls.Get(requestID)
	if !exists {
		s.log.Verbo("dropping vote from %s to an unknown poll with requestID: %d",
			vdr,
			requestID)
		return []ids.Bag{}, false
	}

	s.log.Verbo("processing dropped vote from %s in the poll with requestID: %d",
		vdr,
		requestID)

	pollHolder := pollHolderIntf.(pollHolder)
	poll := pollHolder.GetPoll()

	poll.Drop(vdr)
	if !poll.Finished() {
		return []ids.Bag{}, false
	}

	s.log.Verbo("poll with requestID %d finished as %s", requestID, poll)

	s.polls.Delete(requestID) // remove the poll from the current set
	s.durPolls.Observe(float64(time.Since(pollHolder.StartTime()).Milliseconds()))
	s.numPolls.Dec() // decrease the metrics
	return []ids.Bag{poll.Result()}, true
}

// Len returns the number of outstanding polls
func (s *set) Len() int { return s.polls.Len() }

func (s *set) String() string {
	sb := strings.Builder{}
	sb.WriteString(fmt.Sprintf("current polls: (Size = %d)", s.polls.Len()))
	iter := s.polls.NewIterator()
	for iter.Next() {
		requestID := iter.Key()
		poll := iter.Value().(Poll)
		sb.WriteString(fmt.Sprintf("\n    %d: %s", requestID, poll.PrefixedString("    ")))
	}
	return sb.String()
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package common

import (
	"fmt"
	"time"

	stdmath "math"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/utils/math"
)

const (
	// MaxContainersPerMultiPut is the maximum number of containers that can be
	// sent in a MultiPut message
	MaxContainersPerMultiPut = 2000

	// StatusUpdateFrequency is how many containers should be processed between
	// logs
	StatusUpdateFrequency = 2500

	// MaxOutstandingRequests is the maximum number of GetAncestors sent but not
	// responded to/failed
	MaxOutstandingRequests = 8

	// MaxTimeFetchingAncestors is the maximum amount of time to spend fetching
	// vertices during a call to GetAncestors
	MaxTimeFetchingAncestors = 50 * time.Millisecond
)

// Bootstrapper implements the Engine interface.
type Bootstrapper struct {
	Config

	RequestID uint32

	// IDs of validators we have requested the accepted frontier from but haven't
	// received a reply from
	pendingAcceptedFrontier ids.ShortSet
	acceptedFrontier        ids.Set

	pendingAccepted ids.ShortSet
	acceptedVotes   map[ids.ID]uint64

	// current weight
	started bool
	weight  uint64

	bootstrapAttempts int
	// validators that failed to respond with their frontiers
	failedAcceptedFrontierVdrs ids.ShortSet
}

// Initialize implements the Engine interface.
func (b *Bootstrapper) Initialize(config Config) error {
	b.Config = config

	beacons, err := b.Beacons.Sample(config.SampleK)
	if err != nil {
		return err
	}

	for _, vdr := range beacons {
		vdrID := vdr.ID()
		b.pendingAcceptedFrontier.Add(vdrID)
	}

	for _, vdr := range b.Beacons.List() {
		vdrID := vdr.ID()
		b.pendingAccepted.Add(vdrID)
	}

	b.acceptedVotes = make(map[ids.ID]uint64)
	if b.Config.StartupAlpha > 0 {
		return nil
	}
	return b.Startup()
}

// Startup implements the Engine interface.
func (b *Bootstrapper) Startup() error {

	b.bootstrapAttempts++
	b.started = true
	if b.pendingAcceptedFrontier.Len() == 0 {
		b.Ctx.Log.Info("Bootstrapping skipped due to no provided bootstraps")
		return b.Bootstrapable.ForceAccepted(nil)
	}

	// Ask each of the bootstrap validators to send their accepted frontier
	vdrs := ids.ShortSet{}
	vdrs.Union(b.pendingAcceptedFrontier)

	b.RequestID++
	b.Sender.GetAcceptedFrontier(vdrs, b.RequestID)
	return nil
}

// GetAcceptedFrontier implements the Engine interface.
func (b *Bootstrapper) GetAcceptedFrontier(validatorID ids.ShortID, requestID uint32) error {
	b.Sender.AcceptedFrontier(validatorID, requestID, b.Bootstrapable.CurrentAcceptedFrontier())
	return nil
}

// GetAcceptedFrontierFailed implements the Engine interface.
func (b *Bootstrapper) GetAcceptedFrontierFailed(validatorID ids.ShortID, requestID uint32) error {
	// If we can't get a response from [validatorID], act as though they said their accepted frontier is empty
	// and we add the validator to the failed list
	b.failedAcceptedFrontierVdrs.Add(validatorID)
	return b.AcceptedFrontier(validatorID, requestID, nil)
}

// AcceptedFrontier implements the Engine interface.
func (b *Bootstrapper) AcceptedFrontier(validatorID ids.ShortID, requestID uint32, containerIDs []ids.ID) error {
	if !b.pendingAcceptedFrontier.Contains(validatorID) {
		b.Ctx.Log.Debug("Received an AcceptedFrontier message from %s unexpectedly", validatorID)
		return nil
	}
	// Mark that we received a response from [validatorID]
	b.pendingAcceptedFrontier.Remove(validatorID)

	// Union the reported accepted frontier from [validatorID] with the accepted frontier we got from others
	b.acceptedFrontier.Add(containerIDs...)

	// still waiting on requests
	if b.pendingAcceptedFrontier.Len() != 0 {
		return nil
	}

	// We've received the accepted frontier from every bootstrap validator
	// Ask each bootstrap validator to filter the list of containers that we were
	// told are on the accepted frontier such that the list only contains containers
	// they think are accepted
	var err error
	totalWeight := uint64(0)

	// measure total weight
	for _, beacon := range b.Beacons.List() {
		if !b.failedAcceptedFrontierVdrs.Contains(beacon.ID()) {
			totalWeight, err = math.Add64(totalWeight, beacon.Weight())
			if err != nil {
				totalWeight = stdmath.MaxUint64
			}
		}
	}

	// restart the bootstrap
	if totalWeight < b.Alpha {
		b.Ctx.Log.Info("Didn't receive enough AcceptedFrontier to bootstrap - bootstrap attempt: %d", b.bootstrapAttempts)
		return b.restartBootstrap()
	}

	vdrs := ids.ShortSet{}
	vdrs.Union(b.pendingAccepted)

	b.RequestID++
	b.Sender.GetAccepted(vdrs, b.RequestID, b.acceptedFrontier.List())

	return nil
}

// GetAccepted implements the Engine interface.
func (b *Bootstrapper) GetAccepted(validatorID ids.ShortID, requestID uint32, containerIDs []ids.ID) error {
	b.Sender.Accepted(validatorID, requestID, b.Bootstrapable.FilterAccepted(containerIDs))
	return nil
}

// GetAcceptedFailed implements the Engine interface.
func (b *Bootstrapper) GetAcceptedFailed(validatorID ids.ShortID, requestID uint32) error {
	// If we can't get a response from [validatorID], act as though they said
	// that they think none of the containers we sent them in GetAccepted are accepted
	return b.Accepted(validatorID, requestID, nil)
}

// Accepted implements the Engine interface.
func (b *Bootstrapper) Accepted(validatorID ids.ShortID, requestID uint32, containerIDs []ids.ID) error {
	if !b.pendingAccepted.Contains(validatorID) {
		b.Ctx.Log.Debug("Received an Accepted message from %s unexpectedly", validatorID)
		return nil
	}
	// Mark that we received a response from [validatorID]
	b.pendingAccepted.Remove(validatorID)

	weight := uint64(0)
	if w, ok := b.Beacons.GetWeight(validatorID); ok {
		weight = w
	}

	for _, containerID := range containerIDs {
		previousWeight := b.acceptedVotes[containerID]
		newWeight, err := math.Add64(weight, previousWeight)
		if err != nil {
			newWeight = stdmath.MaxUint64
		}
		b.acceptedVotes[containerID] = newWeight
	}

	// wait on pending responses
	if b.pendingAccepted.Len() != 0 {
		return nil
	}

	// We've received the filtered accepted frontier from every bootstrap validator
	// Accept all containers that have a sufficient weight behind them
	accepted := make([]ids.ID, 0, len(b.acceptedVotes))
	for containerID, weight := range b.acceptedVotes {
		if weight >= b.Alpha {
			accepted = append(accepted, containerID)
		}
	}

	// if we don't have enough weight for the bootstrap to be accepted then restart or fail the bootstrap
	size := len(accepted)

	if size == 0 && b.Beacons.Len() > 0 {
		b.Ctx.Log.Info("Bootstrapping finished with no accepted frontier. This is likely a result of failing to "+
			"be able to connect to the specified bootstraps, or no transactions have been issued on this chain yet"+
			" - Number of beacons: %d - bootstrap attempt: %d", b.Beacons.Len(), b.bootstrapAttempts)
		return b.restartBootstrap()
	}

	b.Ctx.Log.Info("Bootstrapping started syncing with %d vertices in the accepted frontier", size)

	return b.Bootstrapable.ForceAccepted(accepted)
}

// Connected implements the Engine interface.
func (b *Bootstrapper) Connected(validatorID ids.ShortID) error {
	if b.started {
		return nil
	}
	weight, ok := b.Beacons.GetWeight(validatorID)
	if !ok {
		return nil
	}
	weight, err := math.Add64(weight, b.weight)
	if err != nil {
		return err
	}
	b.weight = weight
	if b.weight < b.StartupAlpha {
		return nil
	}
	return b.Startup()
}

// Disconnected implements the Engine interface.
func (b *Bootstrapper) Disconnected(validatorID ids.ShortID) error {
	if weight, ok := b.Beacons.GetWeight(validatorID); ok {
		// TODO: Account for weight changes in a more robust manner.

		// Sub64 should rarely error since only validators that have added their
		// weight can become disconnected. Because it is possible that there are
		// changes to the validators set, we utilize that Sub64 returns 0 on
		// error.
		b.weight, _ = math.Sub64(b.weight, weight)
	}
	return nil
}

func (b *Bootstrapper) restartBootstrap() error {
	if b.bootstrapAttempts > 3 {
		return fmt.Errorf("failed to boostrap the chain after %d attempts", b.bootstrapAttempts)
	}

	// reset the failed frontier responses
	b.failedAcceptedFrontierVdrs = ids.ShortSet{}

	// ensure it doesn't repeat non-stop
	time.Sleep(1 * time.Second)

	beacons, err := b.Beacons.Sample(b.Config.SampleK)
	if err != nil {
		return err
	}

	for _, vdr := range beacons {
		vdrID := vdr.ID()
		b.pendingAcceptedFrontier.Add(vdrID)
	}

	for _, vdr := range b.Beacons.List() {
		vdrID := vdr.ID()
		b.pendingAccepted.Add(vdrID)
	}

	b.acceptedVotes = make(map[ids.ID]uint64)
	return b.Startup()
}
<commit_msg>addressing PR comments<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package common

import (
	"fmt"
	"time"

	stdmath "math"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/utils/math"
)

const (
	// MaxContainersPerMultiPut is the maximum number of containers that can be
	// sent in a MultiPut message
	MaxContainersPerMultiPut = 2000

	// StatusUpdateFrequency is how many containers should be processed between
	// logs
	StatusUpdateFrequency = 2500

	// MaxOutstandingRequests is the maximum number of GetAncestors sent but not
	// responded to/failed
	MaxOutstandingRequests = 8

	// MaxTimeFetchingAncestors is the maximum amount of time to spend fetching
	// vertices during a call to GetAncestors
	MaxTimeFetchingAncestors = 50 * time.Millisecond
)

// Bootstrapper implements the Engine interface.
type Bootstrapper struct {
	Config

	RequestID uint32

	// IDs of validators we have requested the accepted frontier from but haven't
	// received a reply from
	pendingAcceptedFrontier ids.ShortSet
	acceptedFrontier        ids.Set

	pendingAccepted ids.ShortSet
	acceptedVotes   map[ids.ID]uint64

	// current weight
	started bool
	weight  uint64

	// number of times the bootstrap was attempted
	bootstrapAttempts int
	// validators that failed to respond with their frontiers
	failedAcceptedFrontierVdrs ids.ShortSet
}

// Initialize implements the Engine interface.
func (b *Bootstrapper) Initialize(config Config) error {
	b.Config = config

	beacons, err := b.Beacons.Sample(config.SampleK)
	if err != nil {
		return err
	}

	for _, vdr := range beacons {
		vdrID := vdr.ID()
		b.pendingAcceptedFrontier.Add(vdrID)
	}

	for _, vdr := range b.Beacons.List() {
		vdrID := vdr.ID()
		b.pendingAccepted.Add(vdrID)
	}

	b.acceptedVotes = make(map[ids.ID]uint64)
	if b.Config.StartupAlpha > 0 {
		return nil
	}
	return b.Startup()
}

// Startup implements the Engine interface.
func (b *Bootstrapper) Startup() error {
	b.bootstrapAttempts++
	b.started = true
	if b.pendingAcceptedFrontier.Len() == 0 {
		b.Ctx.Log.Info("Bootstrapping skipped due to no provided bootstraps")
		return b.Bootstrapable.ForceAccepted(nil)
	}

	// Ask each of the bootstrap validators to send their accepted frontier
	vdrs := ids.ShortSet{}
	vdrs.Union(b.pendingAcceptedFrontier)

	b.RequestID++
	b.Sender.GetAcceptedFrontier(vdrs, b.RequestID)
	return nil
}

// GetAcceptedFrontier implements the Engine interface.
func (b *Bootstrapper) GetAcceptedFrontier(validatorID ids.ShortID, requestID uint32) error {
	b.Sender.AcceptedFrontier(validatorID, requestID, b.Bootstrapable.CurrentAcceptedFrontier())
	return nil
}

// GetAcceptedFrontierFailed implements the Engine interface.
func (b *Bootstrapper) GetAcceptedFrontierFailed(validatorID ids.ShortID, requestID uint32) error {
	// If we can't get a response from [validatorID], act as though they said their accepted frontier is empty
	// and we add the validator to the failed list
	b.failedAcceptedFrontierVdrs.Add(validatorID)
	return b.AcceptedFrontier(validatorID, requestID, nil)
}

// AcceptedFrontier implements the Engine interface.
func (b *Bootstrapper) AcceptedFrontier(validatorID ids.ShortID, requestID uint32, containerIDs []ids.ID) error {
	if !b.pendingAcceptedFrontier.Contains(validatorID) {
		b.Ctx.Log.Debug("Received an AcceptedFrontier message from %s unexpectedly", validatorID)
		return nil
	}
	// Mark that we received a response from [validatorID]
	b.pendingAcceptedFrontier.Remove(validatorID)

	// Union the reported accepted frontier from [validatorID] with the accepted frontier we got from others
	b.acceptedFrontier.Add(containerIDs...)

	// still waiting on requests
	if b.pendingAcceptedFrontier.Len() != 0 {
		return nil
	}

	// We've received the accepted frontier from every bootstrap validator
	// Ask each bootstrap validator to filter the list of containers that we were
	// told are on the accepted frontier such that the list only contains containers
	// they think are accepted
	var err error
	totalWeight := uint64(0)

	// measure total weight of accepted frontiers from the validators
	for _, beacon := range b.Beacons.List() {
		if b.failedAcceptedFrontierVdrs.Contains(beacon.ID()) {
			continue
		}

		totalWeight, err = math.Add64(totalWeight, beacon.Weight())
		if err != nil {
			b.Ctx.Log.Error("Error calculating the AcceptedFrontier beacons weight - totalWeight: %v, beacon.Weight: %v", totalWeight, beacon.Weight())
			totalWeight = stdmath.MaxUint64
		}
	}

	// restart the bootstrap
	if totalWeight < b.Alpha {
		b.Ctx.Log.Info("Didn't receive enough AcceptedFrontier to bootstrap - bootstrap attempt: %d", b.bootstrapAttempts)
		return b.restartBootstrap()
	}

	vdrs := ids.ShortSet{}
	vdrs.Union(b.pendingAccepted)

	b.RequestID++
	b.Sender.GetAccepted(vdrs, b.RequestID, b.acceptedFrontier.List())

	return nil
}

// GetAccepted implements the Engine interface.
func (b *Bootstrapper) GetAccepted(validatorID ids.ShortID, requestID uint32, containerIDs []ids.ID) error {
	b.Sender.Accepted(validatorID, requestID, b.Bootstrapable.FilterAccepted(containerIDs))
	return nil
}

// GetAcceptedFailed implements the Engine interface.
func (b *Bootstrapper) GetAcceptedFailed(validatorID ids.ShortID, requestID uint32) error {
	// If we can't get a response from [validatorID], act as though they said
	// that they think none of the containers we sent them in GetAccepted are accepted
	return b.Accepted(validatorID, requestID, nil)
}

// Accepted implements the Engine interface.
func (b *Bootstrapper) Accepted(validatorID ids.ShortID, requestID uint32, containerIDs []ids.ID) error {
	if !b.pendingAccepted.Contains(validatorID) {
		b.Ctx.Log.Debug("Received an Accepted message from %s unexpectedly", validatorID)
		return nil
	}
	// Mark that we received a response from [validatorID]
	b.pendingAccepted.Remove(validatorID)

	weight := uint64(0)
	if w, ok := b.Beacons.GetWeight(validatorID); ok {
		weight = w
	}

	for _, containerID := range containerIDs {
		previousWeight := b.acceptedVotes[containerID]
		newWeight, err := math.Add64(weight, previousWeight)
		if err != nil {
			b.Ctx.Log.Error("Error calculating the Accepted votes - weight: %v, previousWeight: %v", weight, previousWeight)
			newWeight = stdmath.MaxUint64
		}
		b.acceptedVotes[containerID] = newWeight
	}

	// wait on pending responses
	if b.pendingAccepted.Len() != 0 {
		return nil
	}

	// We've received the filtered accepted frontier from every bootstrap validator
	// Accept all containers that have a sufficient weight behind them
	accepted := make([]ids.ID, 0, len(b.acceptedVotes))
	for containerID, weight := range b.acceptedVotes {
		if weight >= b.Alpha {
			accepted = append(accepted, containerID)
		}
	}

	// if we don't have enough weight for the bootstrap to be accepted then restart or fail the bootstrap
	size := len(accepted)

	if size == 0 && b.Beacons.Len() > 0 {
		b.Ctx.Log.Info("Bootstrapping finished with no accepted frontier. This is likely a result of failing to "+
			"be able to connect to the specified bootstraps, or no transactions have been issued on this chain yet"+
			" - Number of beacons: %d - bootstrap attempt: %d", b.Beacons.Len(), b.bootstrapAttempts)
		return b.restartBootstrap()
	}

	b.Ctx.Log.Info("Bootstrapping started syncing with %d vertices in the accepted frontier", size)

	return b.Bootstrapable.ForceAccepted(accepted)
}

// Connected implements the Engine interface.
func (b *Bootstrapper) Connected(validatorID ids.ShortID) error {
	if b.started {
		return nil
	}
	weight, ok := b.Beacons.GetWeight(validatorID)
	if !ok {
		return nil
	}
	weight, err := math.Add64(weight, b.weight)
	if err != nil {
		return err
	}
	b.weight = weight
	if b.weight < b.StartupAlpha {
		return nil
	}
	return b.Startup()
}

// Disconnected implements the Engine interface.
func (b *Bootstrapper) Disconnected(validatorID ids.ShortID) error {
	if weight, ok := b.Beacons.GetWeight(validatorID); ok {
		// TODO: Account for weight changes in a more robust manner.

		// Sub64 should rarely error since only validators that have added their
		// weight can become disconnected. Because it is possible that there are
		// changes to the validators set, we utilize that Sub64 returns 0 on
		// error.
		b.weight, _ = math.Sub64(b.weight, weight)
	}
	return nil
}

func (b *Bootstrapper) restartBootstrap() error {
	if b.bootstrapAttempts > 3 {
		return fmt.Errorf("failed to boostrap the chain after %d attempts", b.bootstrapAttempts)
	}

	// reset the failed frontier responses
	b.failedAcceptedFrontierVdrs = ids.ShortSet{}

	// ensure it doesn't repeat non-stop
	time.Sleep(1 * time.Second)

	beacons, err := b.Beacons.Sample(b.Config.SampleK)
	if err != nil {
		return err
	}

	for _, vdr := range beacons {
		vdrID := vdr.ID()
		b.pendingAcceptedFrontier.Add(vdrID)
	}

	for _, vdr := range b.Beacons.List() {
		vdrID := vdr.ID()
		b.pendingAccepted.Add(vdrID)
	}

	b.acceptedVotes = make(map[ids.ID]uint64)
	return b.Startup()
}
<|endoftext|>


<commit_before>// Copyright (c) 2020 Tigera, Inc. All rights reserved.
//
// 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.

// Copyright (c) 2020  All rights reserved.

package tc

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"

	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"

	"github.com/projectcalico/libcalico-go/lib/set"

	"github.com/projectcalico/felix/bpf"
)

type AttachPoint struct {
	Type       EndpointType
	ToOrFrom   ToOrFromEp
	Hook       Hook
	Iface      string
	LogLevel   string
	IP         net.IP
	FIB        bool
	ToHostDrop bool
	DSR        bool
	TunnelMTU  uint16
}

var tcLock sync.Mutex

type ErrAttachFailed struct {
	ExitCode int
	Stderr   string
}

func (e ErrAttachFailed) Error() string {
	return fmt.Sprintf("tc failed with exit code %d; stderr=%v", e.ExitCode, e.Stderr)
}

// AttachProgram attaches a BPF program from a file to the TC attach point
func (ap AttachPoint) AttachProgram() error {
	// FIXME we use this lock so that two copies of tc running in parallel don't re-use the same jump map.
	// This can happen if tc incorrectly decides the two programs are identical (when in fact they differ by attach
	// point).
	logCxt := log.WithField("attachPoint", ap)
	logCxt.Debug("AttachProgram waiting for lock...")
	tcLock.Lock()
	defer tcLock.Unlock()
	logCxt.Debug("AttachProgram got lock.")

	// Work around tc map name collision: when we load two identical BPF programs onto different interfaces, tc
	// pins object-local maps to a namespace based on the hash of the BPF program, which is the same for both
	// interfaces.  Since we want one map per interface instead, we search for such maps and rename them before we
	// release the tc lock.
	//
	// For our purposes, it should work to simply delete the map.  However, when we tried that, the contents of the
	// map get deleted even though it is in use by a BPF program.
	defer repinJumpMaps()

	tempDir, err := ioutil.TempDir("", "calico-tc")
	if err != nil {
		return errors.Wrap(err, "failed to create temporary directory")
	}
	defer func() {
		_ = os.RemoveAll(tempDir)
	}()

	filename := ap.FileName()
	preCompiledBinary := path.Join(bpf.ObjectDir, filename)
	tempBinary := path.Join(tempDir, filename)

	err = ap.patchBinary(preCompiledBinary, tempBinary)
	if err != nil {
		logCxt.WithError(err).Error("Failed to patch binary")
		return err
	}

	tcCmd := exec.Command("tc",
		"filter", "add", "dev", ap.Iface,
		string(ap.Hook),
		"bpf", "da", "obj", tempBinary,
		"sec", SectionName(ap.Type, ap.ToOrFrom))

	out, err := tcCmd.CombinedOutput()
	if err != nil {
		out := string(out)
		if strings.Contains(out, "Cannot find device") {
			// Avoid a big, spammy log when the issue is that the interface isn't present.
			logCxt.WithField("iface", ap.Iface).Info(
				"Failed to attach BPF program; interface not found.  Will retry if it show up.")
			return nil
		}
		logCxt.WithError(err).WithFields(log.Fields{"out": out}).
			WithField("command", tcCmd).Error("Failed to attach BPF program")
		if err, ok := err.(*exec.ExitError); ok {
			// ExitError is really unhelpful dumped to the log, swap it for a custom one.
			return ErrAttachFailed{
				ExitCode: err.ExitCode(),
				Stderr:   string(err.Stderr),
			}
		}
		return errors.Wrap(err, "failed to attach TC program")
	}

	return nil
}

func (ap AttachPoint) patchBinary(ifile, ofile string) error {
	b, err := bpf.BinaryFromFile(ifile)
	if err != nil {
		return errors.Wrap(err, "failed to read pre-compiled BPF binary")
	}

	log.WithField("ap", ap).WithField("ip", ap.IP).Debug("Patching in IP")
	err = b.PatchIPv4(ap.IP)
	if err != nil {
		return errors.WithMessage(err, "patching in IPv4")
	}

	b.PatchLogPrefix(ap.Iface)
	b.PatchTunnelMTU(ap.TunnelMTU)

	err = b.WriteToFile(ofile)
	if err != nil {
		return errors.Wrap(err, "failed to write patched BPF binary")
	}

	return nil
}

// ProgramName returnt the name of the program associated with this AttachPoint
func (ap AttachPoint) ProgramName() string {
	return SectionName(ap.Type, ap.ToOrFrom)
}

// FileName returnthe file the AttachPoint will load the program from
func (ap AttachPoint) FileName() string {
	return ProgFilename(ap.Type, ap.ToOrFrom, ap.ToHostDrop, ap.FIB, ap.DSR, ap.LogLevel)
}

func repinJumpMaps() {
	// Find the maps we care about by walking the BPF filesystem.
	err := filepath.Walk("/sys/fs/bpf/tc", func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.Name() == "cali_jump" {
			log.WithField("path", path).Debug("Queueing deletion of map")

			if log.GetLevel() >= log.DebugLevel {
				out, err := exec.Command("bpftool", "map", "dump", "pinned", path).Output()
				if err != nil {
					log.WithError(err).Panic("Failed to dump map")
				}
				log.WithField("dump", string(out)).Debug("Map dump before deletion")
			}

			out, err := exec.Command("bpftool", "map", "show", "pinned", path).Output()
			if err != nil {
				log.WithError(err).Panic("Failed to show map")
			}
			log.WithField("dump", string(out)).Debug("Map show before deletion")
			id := string(bytes.Split(out, []byte(":"))[0])

			newPath := path + fmt.Sprint(rand.Uint32())
			out, err = exec.Command("bpftool", "map", "pin", "id", id, newPath).Output()
			if err != nil {
				log.WithError(err).Panic("Failed to repin map")
			}
			log.WithField("dump", string(out)).Debug("Repin output")

			err = os.Remove(path)
			if err != nil {
				log.WithError(err).Panic("Failed to remove old map pin")
			}

			if log.GetLevel() >= log.DebugLevel {
				out, err = exec.Command("bpftool", "map", "dump", "pinned", newPath).Output()
				if err != nil {
					log.WithError(err).Panic("Failed to show map")
				}
				log.WithField("dump", string(out)).Debug("Map show after repin")
			}
		}
		return nil
	})
	if os.IsNotExist(err) {
		log.WithError(err).Warn("tc directory missing from BPF file system?")
		return
	}
	if err != nil {
		log.WithError(err).Panic("Failed to walk BPF filesystem")
	}
	log.Debug("Finished moving map pins that we don't need.")
}

// tcDirRegex matches tc's auto-created directory names so we can clean them up when removing maps without accidentally
// removing other user-created dirs..
var tcDirRegex = regexp.MustCompile(`[0-9a-f]{40}`)

// CleanUpJumpMaps scans for cali_jump maps that are still pinned to the filesystem but no longer referenced by
// our BPF programs.
func CleanUpJumpMaps() {
	// So that we serialise with AttachProgram()
	log.Debug("CleanUpJumpMaps waiting for lock...")
	tcLock.Lock()
	defer tcLock.Unlock()
	log.Debug("CleanUpJumpMaps got lock, cleaning up...")

	// Find the maps we care about by walking the BPF filesystem.
	mapIDToPath := make(map[int]string)
	err := filepath.Walk("/sys/fs/bpf/tc", func(p string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if strings.HasPrefix(info.Name(), "cali_jump") {
			log.WithField("path", p).Debug("Examining map")

			out, err := exec.Command("bpftool", "map", "show", "pinned", p).Output()
			if err != nil {
				log.WithError(err).Panic("Failed to show map")
			}
			log.WithField("dump", string(out)).Debug("Map show before deletion")
			idStr := string(bytes.Split(out, []byte(":"))[0])
			id, err := strconv.Atoi(idStr)
			if err != nil {
				log.WithError(err).WithField("dump", string(out)).Error("Failed to parse bpftool output.")
				return err
			}
			mapIDToPath[id] = p
		}
		return nil
	})
	if os.IsNotExist(err) {
		log.WithError(err).Warn("tc directory missing from BPF file system?")
		return
	}
	if err != nil {
		log.WithError(err).Error("Error while looking for maps.")
	}

	// Find all the programs that are attached to interfaces.
	out, err := exec.Command("bpftool", "net", "-j").Output()
	if err != nil {
		log.WithError(err).Panic("Failed to list attached bpf programs")
	}
	log.WithField("dump", string(out)).Debug("Attached BPF programs")

	var attached []struct {
		TC []struct {
			DevName string `json:"devname"`
			ID      int    `json:"id"`
		} `json:"tc"`
	}
	err = json.Unmarshal(out, &attached)
	if err != nil {
		log.WithError(err).WithField("dump", string(out)).Error("Failed to parse list of attached BPF programs")
	}
	attachedProgs := set.New()
	for _, prog := range attached[0].TC {
		log.WithField("prog", prog).Debug("Adding prog to attached set")
		attachedProgs.Add(prog.ID)
	}

	// Find all the maps that the attached programs refer to and remove them from consideration.
	progsJSON, err := exec.Command("bpftool", "prog", "list", "--json").Output()
	if err != nil {
		log.WithError(err).Info("Failed to list BPF programs, assuming there's nothing to clean up.")
		return
	}
	var progs []struct {
		ID   int    `json:"id"`
		Name string `json:"name"`
		Maps []int  `json:"map_ids"`
	}
	err = json.Unmarshal(progsJSON, &progs)
	if err != nil {
		log.WithError(err).Info("Failed to parse bpftool output.  Assuming nothing to clean up.")
		return
	}
	for _, p := range progs {
		if !attachedProgs.Contains(p.ID) {
			log.WithField("prog", p).Debug("Prog is not in the attached set, skipping")
			continue
		}
		for _, id := range p.Maps {
			log.WithField("mapID", id).WithField("prog", p).Debug("Map is still in use")
			delete(mapIDToPath, id)
		}
	}

	// Remove the pins.
	for id, p := range mapIDToPath {
		log.WithFields(log.Fields{"id": id, "path": p}).Debug("Removing stale BPF map pin.")
		err := os.Remove(p)
		if err != nil {
			log.WithError(err).Warn("Removed stale BPF map pin.")
		}
		log.WithFields(log.Fields{"id": id, "path": p}).Info("Removed stale BPF map pin.")
	}

	// Look for empty dirs.
	emptyAutoDirs := set.New()
	err = filepath.Walk("/sys/fs/bpf/tc", func(p string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() && tcDirRegex.MatchString(info.Name()) {
			p := path.Clean(p)
			log.WithField("path", p).Debug("Found tc auto-created dir.")
			emptyAutoDirs.Add(p)
		} else {
			dirPath := path.Clean(path.Dir(p))
			log.WithField("path", dirPath).Debug("tc dir is not empty.")
			emptyAutoDirs.Discard(dirPath)
		}
		return nil
	})
	if os.IsNotExist(err) {
		log.WithError(err).Warn("tc directory missing from BPF file system?")
		return
	}
	if err != nil {
		log.WithError(err).Error("Error while looking for maps.")
	}

	emptyAutoDirs.Iter(func(item interface{}) error {
		p := item.(string)
		log.WithField("path", p).Debug("Removing empty dir.")
		err := os.Remove(p)
		if err != nil {
			log.WithError(err).Error("Error while removing empty dir.")
		}
		return nil
	})
}

// EnsureQdisc makes sure that qdisc is attached to the given interface
func EnsureQdisc(ifaceName string) {
	// FIXME Avoid flapping the tc program and qdisc
	cmd := exec.Command("tc", "qdisc", "del", "dev", ifaceName, "clsact")
	_ = cmd.Run()
	cmd = exec.Command("tc", "qdisc", "add", "dev", ifaceName, "clsact")
	_ = cmd.Run()
}
<commit_msg>Markups.<commit_after>// Copyright (c) 2020 Tigera, Inc. All rights reserved.
//
// 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.

// Copyright (c) 2020  All rights reserved.

package tc

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"

	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"

	"github.com/projectcalico/libcalico-go/lib/set"

	"github.com/projectcalico/felix/bpf"
)

type AttachPoint struct {
	Type       EndpointType
	ToOrFrom   ToOrFromEp
	Hook       Hook
	Iface      string
	LogLevel   string
	IP         net.IP
	FIB        bool
	ToHostDrop bool
	DSR        bool
	TunnelMTU  uint16
}

var tcLock sync.Mutex

type ErrAttachFailed struct {
	ExitCode int
	Stderr   string
}

func (e ErrAttachFailed) Error() string {
	return fmt.Sprintf("tc failed with exit code %d; stderr=%v", e.ExitCode, e.Stderr)
}

// AttachProgram attaches a BPF program from a file to the TC attach point
func (ap AttachPoint) AttachProgram() error {
	// FIXME we use this lock so that two copies of tc running in parallel don't re-use the same jump map.
	// This can happen if tc incorrectly decides the two programs are identical (when in fact they differ by attach
	// point).
	logCxt := log.WithField("attachPoint", ap)
	logCxt.Debug("AttachProgram waiting for lock...")
	tcLock.Lock()
	defer tcLock.Unlock()
	logCxt.Debug("AttachProgram got lock.")

	// Work around tc map name collision: when we load two identical BPF programs onto different interfaces, tc
	// pins object-local maps to a namespace based on the hash of the BPF program, which is the same for both
	// interfaces.  Since we want one map per interface instead, we search for such maps and rename them before we
	// release the tc lock.
	//
	// For our purposes, it should work to simply delete the map.  However, when we tried that, the contents of the
	// map get deleted even though it is in use by a BPF program.
	defer repinJumpMaps()

	tempDir, err := ioutil.TempDir("", "calico-tc")
	if err != nil {
		return errors.Wrap(err, "failed to create temporary directory")
	}
	defer func() {
		_ = os.RemoveAll(tempDir)
	}()

	filename := ap.FileName()
	preCompiledBinary := path.Join(bpf.ObjectDir, filename)
	tempBinary := path.Join(tempDir, filename)

	err = ap.patchBinary(logCxt, preCompiledBinary, tempBinary)
	if err != nil {
		logCxt.WithError(err).Error("Failed to patch binary")
		return err
	}

	tcCmd := exec.Command("tc",
		"filter", "add", "dev", ap.Iface,
		string(ap.Hook),
		"bpf", "da", "obj", tempBinary,
		"sec", SectionName(ap.Type, ap.ToOrFrom))

	out, err := tcCmd.CombinedOutput()
	if err != nil {
		out := string(out)
		if strings.Contains(out, "Cannot find device") {
			// Avoid a big, spammy log when the issue is that the interface isn't present.
			logCxt.WithField("iface", ap.Iface).Info(
				"Failed to attach BPF program; interface not found.  Will retry if it show up.")
			return nil
		}
		logCxt.WithError(err).WithFields(log.Fields{"out": out}).
			WithField("command", tcCmd).Error("Failed to attach BPF program")
		if err, ok := err.(*exec.ExitError); ok {
			// ExitError is really unhelpful dumped to the log, swap it for a custom one.
			return ErrAttachFailed{
				ExitCode: err.ExitCode(),
				Stderr:   string(err.Stderr),
			}
		}
		return errors.Wrap(err, "failed to attach TC program")
	}

	return nil
}

func (ap AttachPoint) patchBinary(logCtx *log.Entry, ifile, ofile string) error {
	b, err := bpf.BinaryFromFile(ifile)
	if err != nil {
		return errors.Wrap(err, "failed to read pre-compiled BPF binary")
	}

	logCtx.WithField("ip", ap.IP).Debug("Patching in IP")
	err = b.PatchIPv4(ap.IP)
	if err != nil {
		return errors.WithMessage(err, "patching in IPv4")
	}

	b.PatchLogPrefix(ap.Iface)
	b.PatchTunnelMTU(ap.TunnelMTU)

	err = b.WriteToFile(ofile)
	if err != nil {
		return errors.Wrap(err, "failed to write patched BPF binary")
	}

	return nil
}

// ProgramName returnt the name of the program associated with this AttachPoint
func (ap AttachPoint) ProgramName() string {
	return SectionName(ap.Type, ap.ToOrFrom)
}

// FileName returnthe file the AttachPoint will load the program from
func (ap AttachPoint) FileName() string {
	return ProgFilename(ap.Type, ap.ToOrFrom, ap.ToHostDrop, ap.FIB, ap.DSR, ap.LogLevel)
}

func repinJumpMaps() {
	// Find the maps we care about by walking the BPF filesystem.
	err := filepath.Walk("/sys/fs/bpf/tc", func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.Name() == "cali_jump" {
			log.WithField("path", path).Debug("Queueing deletion of map")

			if log.GetLevel() >= log.DebugLevel {
				out, err := exec.Command("bpftool", "map", "dump", "pinned", path).Output()
				if err != nil {
					log.WithError(err).Panic("Failed to dump map")
				}
				log.WithField("dump", string(out)).Debug("Map dump before deletion")
			}

			out, err := exec.Command("bpftool", "map", "show", "pinned", path).Output()
			if err != nil {
				log.WithError(err).Panic("Failed to show map")
			}
			log.WithField("dump", string(out)).Debug("Map show before deletion")
			id := string(bytes.Split(out, []byte(":"))[0])

			newPath := path + fmt.Sprint(rand.Uint32())
			out, err = exec.Command("bpftool", "map", "pin", "id", id, newPath).Output()
			if err != nil {
				log.WithError(err).Panic("Failed to repin map")
			}
			log.WithField("dump", string(out)).Debug("Repin output")

			err = os.Remove(path)
			if err != nil {
				log.WithError(err).Panic("Failed to remove old map pin")
			}

			if log.GetLevel() >= log.DebugLevel {
				out, err = exec.Command("bpftool", "map", "dump", "pinned", newPath).Output()
				if err != nil {
					log.WithError(err).Panic("Failed to show map")
				}
				log.WithField("dump", string(out)).Debug("Map show after repin")
			}
		}
		return nil
	})
	if os.IsNotExist(err) {
		log.WithError(err).Warn("tc directory missing from BPF file system?")
		return
	}
	if err != nil {
		log.WithError(err).Panic("Failed to walk BPF filesystem")
	}
	log.Debug("Finished moving map pins that we don't need.")
}

// tcDirRegex matches tc's auto-created directory names so we can clean them up when removing maps without accidentally
// removing other user-created dirs..
var tcDirRegex = regexp.MustCompile(`[0-9a-f]{40}`)

// CleanUpJumpMaps scans for cali_jump maps that are still pinned to the filesystem but no longer referenced by
// our BPF programs.
func CleanUpJumpMaps() {
	// So that we serialise with AttachProgram()
	log.Debug("CleanUpJumpMaps waiting for lock...")
	tcLock.Lock()
	defer tcLock.Unlock()
	log.Debug("CleanUpJumpMaps got lock, cleaning up...")

	// Find the maps we care about by walking the BPF filesystem.
	mapIDToPath := make(map[int]string)
	err := filepath.Walk("/sys/fs/bpf/tc", func(p string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if strings.HasPrefix(info.Name(), "cali_jump") {
			log.WithField("path", p).Debug("Examining map")

			out, err := exec.Command("bpftool", "map", "show", "pinned", p).Output()
			if err != nil {
				log.WithError(err).Panic("Failed to show map")
			}
			log.WithField("dump", string(out)).Debug("Map show before deletion")
			idStr := string(bytes.Split(out, []byte(":"))[0])
			id, err := strconv.Atoi(idStr)
			if err != nil {
				log.WithError(err).WithField("dump", string(out)).Error("Failed to parse bpftool output.")
				return err
			}
			mapIDToPath[id] = p
		}
		return nil
	})
	if os.IsNotExist(err) {
		log.WithError(err).Warn("tc directory missing from BPF file system?")
		return
	}
	if err != nil {
		log.WithError(err).Error("Error while looking for maps.")
	}

	// Find all the programs that are attached to interfaces.
	out, err := exec.Command("bpftool", "net", "-j").Output()
	if err != nil {
		log.WithError(err).Panic("Failed to list attached bpf programs")
	}
	log.WithField("dump", string(out)).Debug("Attached BPF programs")

	var attached []struct {
		TC []struct {
			DevName string `json:"devname"`
			ID      int    `json:"id"`
		} `json:"tc"`
	}
	err = json.Unmarshal(out, &attached)
	if err != nil {
		log.WithError(err).WithField("dump", string(out)).Error("Failed to parse list of attached BPF programs")
	}
	attachedProgs := set.New()
	for _, prog := range attached[0].TC {
		log.WithField("prog", prog).Debug("Adding prog to attached set")
		attachedProgs.Add(prog.ID)
	}

	// Find all the maps that the attached programs refer to and remove them from consideration.
	progsJSON, err := exec.Command("bpftool", "prog", "list", "--json").Output()
	if err != nil {
		log.WithError(err).Info("Failed to list BPF programs, assuming there's nothing to clean up.")
		return
	}
	var progs []struct {
		ID   int    `json:"id"`
		Name string `json:"name"`
		Maps []int  `json:"map_ids"`
	}
	err = json.Unmarshal(progsJSON, &progs)
	if err != nil {
		log.WithError(err).Info("Failed to parse bpftool output.  Assuming nothing to clean up.")
		return
	}
	for _, p := range progs {
		if !attachedProgs.Contains(p.ID) {
			log.WithField("prog", p).Debug("Prog is not in the attached set, skipping")
			continue
		}
		for _, id := range p.Maps {
			log.WithField("mapID", id).WithField("prog", p).Debug("Map is still in use")
			delete(mapIDToPath, id)
		}
	}

	// Remove the pins.
	for id, p := range mapIDToPath {
		log.WithFields(log.Fields{"id": id, "path": p}).Debug("Removing stale BPF map pin.")
		err := os.Remove(p)
		if err != nil {
			log.WithError(err).Warn("Removed stale BPF map pin.")
		}
		log.WithFields(log.Fields{"id": id, "path": p}).Info("Removed stale BPF map pin.")
	}

	// Look for empty dirs.
	emptyAutoDirs := set.New()
	err = filepath.Walk("/sys/fs/bpf/tc", func(p string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() && tcDirRegex.MatchString(info.Name()) {
			p := path.Clean(p)
			log.WithField("path", p).Debug("Found tc auto-created dir.")
			emptyAutoDirs.Add(p)
		} else {
			dirPath := path.Clean(path.Dir(p))
			log.WithField("path", dirPath).Debug("tc dir is not empty.")
			emptyAutoDirs.Discard(dirPath)
		}
		return nil
	})
	if os.IsNotExist(err) {
		log.WithError(err).Warn("tc directory missing from BPF file system?")
		return
	}
	if err != nil {
		log.WithError(err).Error("Error while looking for maps.")
	}

	emptyAutoDirs.Iter(func(item interface{}) error {
		p := item.(string)
		log.WithField("path", p).Debug("Removing empty dir.")
		err := os.Remove(p)
		if err != nil {
			log.WithError(err).Error("Error while removing empty dir.")
		}
		return nil
	})
}

// EnsureQdisc makes sure that qdisc is attached to the given interface
func EnsureQdisc(ifaceName string) {
	// FIXME Avoid flapping the tc program and qdisc
	cmd := exec.Command("tc", "qdisc", "del", "dev", ifaceName, "clsact")
	_ = cmd.Run()
	cmd = exec.Command("tc", "qdisc", "add", "dev", ifaceName, "clsact")
	_ = cmd.Run()
}
<|endoftext|>