<commit_before>package controllers

import (
	"concierge/config"
	"concierge/database"
	"concierge/models"
	"concierge/pkg"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	log "github.com/sirupsen/logrus"
)

type info struct {
	Userinfo  models.Users
	Leaseinfo models.Leases
}

//ShowAllowedIngress ...
func ShowAllowedIngress(c *gin.Context) {
	clientset := config.KubeClient.ClientSet
	User, _ := c.Get("User")

	myclientset := pkg.MyClientSet{Clientset: clientset}
	ns := ""
	ns = c.Query("ns")
	log.Infof("Listing ingress in namespace %s for user %s\n", ns, User.(*models.Users).Email)
	data, err := myclientset.GetIngresses(ns)
	if err != nil {
		log.Error("Error", err)
		c.HTML(http.StatusOK, "showingresslist.gohtml", gin.H{
			"data": data,
			"user": User,
			"message": map[string]string{
				"class":   "Danger",
				"message": err.Error(),
			},
		})
		return
	}
	c.HTML(http.StatusOK, "showingresslist.gohtml", gin.H{
		"data": data,
		"user": User,
	})
}

//WhiteListIP ...
func WhiteListIP(c *gin.Context) {
	clientset := config.KubeClient.ClientSet

	User, _ := c.Get("User")
	myclientset := pkg.MyClientSet{Clientset: clientset}
	ns := c.Param("ns")
	name := c.Param("name")
	expiry, _ := strconv.Atoi(c.PostForm("expiry"))
	var leases []models.Leases
	leases = GetActiveLeases(ns, name)
	data, err := myclientset.GetIngress(ns, name)
	if err != nil {
		log.Error("Error", err)
		c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
			"data": data,
			"user": User,
			"message": map[string]string{
				"class":   "Danger",
				"message": err.Error(),
			},
			"activeLeases": leases,
		})
		return
	}
	ips := c.Request.Header["X-Forwarded-For"][0]
	ip := strings.Split(ips, ",")[0]
	ip = ip + "/32"
	updateStatus, err := myclientset.WhiteListIP(ns, name, ip)

	if err != nil {
		log.Error("Error", err)
		c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
			"data": data,
			"user": User,
			"message": map[string]string{
				"class":   "Danger",
				"message": err.Error(),
			},
			"activeLeases": leases,
		})
		return
	}
	if updateStatus {
		msgInfo := "Whitelisted IP " + ip + " to ingress " + name + " in namespace " + ns + " for user " + User.(*models.Users).Email
		slackNotification(msgInfo, User.(*models.Users).Email)
		log.Info(msgInfo)
		if database.DB == nil {
			database.Conn()
		}

		lease := models.Leases{
			UserID:    User.(*models.Users).ID,
			LeaseIP:   ip,
			LeaseType: "Ingress",
			GroupID:   ns + ":" + name,
			Expiry:    uint(expiry),
		}

		database.DB.Create(&lease)
		leases = GetActiveLeases(ns, name)

		c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
			"data": data,
			"user": User,
			"message": map[string]string{
				"class":   "Success",
				"message": "Lease is successfully taken",
			},
			"activeLeases": leases,
		})
		return
	}

	c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
		"data": data,
		"user": User,
		"message": map[string]string{
			"class":   "Danger",
			"message": "Your IP is already present",
		},
		"activeLeases": leases,
	})
}

//DeleteIPFromIngress ...
func DeleteIPFromIngress(c *gin.Context) {
	clientset := config.KubeClient.ClientSet

	User, _ := c.Get("User")

	myclientset := pkg.MyClientSet{Clientset: clientset}
	ns := c.Param("ns")
	name := c.Param("name")
	leaseID, err := strconv.Atoi(c.Param("id"))
	ID := uint(leaseID)
	leases := GetActiveLeases(ns, name)

	data, err := myclientset.GetIngress(ns, name)
	if err != nil {
		log.Error("Error", err)
		c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
			"data": data,
			"user": User,
			"message": map[string]string{
				"class":   "Danger",
				"message": err.Error(),
			},
			"activeLeases": leases,
		})
		return
	}
	ips := c.Request.Header["X-Forwarded-For"][0]
	ip := strings.Split(ips, ",")[0]
	ip = ip + "/32"
	updateStatus, err := DeleteLeases(ns, name, ip, ID)
	if err != nil {
		log.Error("Error", err)
		c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
			"data": data,
			"user": User,
			"message": map[string]string{
				"class":   "Danger",
				"message": err.Error(),
			},
			"activeLeases": leases,
		})
		return
	}
	if updateStatus {
		msgInfo := "Removed IP " + ip + " from ingress " + name + " in namespace " + ns + " for user " + User.(*models.Users).Email
		slackNotification(msgInfo, User.(*models.Users).Email)
		log.Info(msgInfo)
		leases = GetActiveLeases(ns, name)
		c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
			"data":         data,
			"user":         User,
			"activeLeases": leases,
			"message": map[string]string{
				"class":   "Success",
				"message": "Lease is successfully deleted",
			},
		})
		return
	}
}

//IngressDetails ...
func IngressDetails(c *gin.Context) {
	clientset := config.KubeClient.ClientSet

	User, _ := c.Get("User")
	myclientset := pkg.MyClientSet{Clientset: clientset}
	ns := c.Param("ns")
	name := c.Param("name")
	leases := GetActiveLeases(ns, name)
	data, err := myclientset.GetIngress(ns, name)
	if err != nil {
		log.Error("Error", err)
		c.HTML(http.StatusNotFound, "manageingress.gohtml", gin.H{
			"message": map[string]string{
				"class":   "Danger",
				"message": err.Error(),
			},
			"user":         User,
			"activeLeases": leases,
		})
		return
	}
	c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
		"data":         data,
		"user":         User,
		"activeLeases": leases,
	})
}

//GetActiveLeases ...
func GetActiveLeases(ns string, name string) []models.Leases {
	if database.DB == nil {
		database.Conn()
	}

	leases := []models.Leases{}
	if ns == "" && name == "" {
		database.DB.Preload("User").Where(models.Leases{
			LeaseType: "Ingress",
		}).Find(&leases)
	} else {
		database.DB.Preload("User").Where(models.Leases{
			LeaseType: "Ingress",
			GroupID:   ns + ":" + name,
		}).Find(&leases)
	}
	myleases := []models.Leases{}
	for i, lease := range leases {
		splitGroupID := strings.Split(lease.GroupID, ":")
		ns = splitGroupID[0]
		name = splitGroupID[1]
		t := uint(lease.CreatedAt.Unix()) + lease.Expiry
		if t < uint(time.Now().Unix()) {
			leases[i].Expiry = uint(0)
			_, err := DeleteLeases(ns, name, lease.LeaseIP, lease.ID)
			if err != nil {
				log.Error("Error", err)
			}
			log.Infof("Removed expired IP %s from ingress %s in namespace %s for User %s\n", lease.LeaseIP, name, ns, lease.User.Email)
		} else {
			leases[i].Expiry = t - uint(time.Now().Unix())
			myleases = append(myleases, leases[i])
		}
	}
	return myleases
}

//DeleteLeases ...
func DeleteLeases(ns string, name string, ip string, ID uint) (bool, error) {
	clientset := config.KubeClient.ClientSet

	myclientset := pkg.MyClientSet{Clientset: clientset}
	if database.DB == nil {
		database.Conn()
	}

	database.DB.Delete(models.Leases{
		ID: ID,
	})
	updateStatus, err := myclientset.RemoveIngressIP(ns, name, ip)
	return updateStatus, err
}

//ClearExpiredLeases ...
func ClearExpiredLeases(c *gin.Context) {
	GetActiveLeases("", "")
	c.String(200, "Done")
}

func slackNotification(msg string, user string) {
	slackWebhookURL := os.Getenv("SLACK_WEBHOOK_URL")
	if slackWebhookURL == "" {
		return
	}
	payload := pkg.Payload{
		Title:      "Concierge",
		Pretext:    msg,
		Text:       msg,
		Color:      "#36a64f",
		AuthorName: user,
		TitleLink:  "",
		Footer:     "Concierge",
		Timestamp:  strconv.FormatInt(time.Now().Unix(), 10),
	}
	payloads := pkg.Payloads{
		Attachments: map[string][]pkg.Payload{
			"attachments": []pkg.Payload{
				payload,
			},
		},
	}
	payloads.SlackNotification(slackWebhookURL)
}
<commit_msg>Removing leases from db if removed from ingress<commit_after>package controllers

import (
	"concierge/config"
	"concierge/database"
	"concierge/models"
	"concierge/pkg"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	log "github.com/sirupsen/logrus"
)

type info struct {
	Userinfo  models.Users
	Leaseinfo models.Leases
}

//ShowAllowedIngress ...
func ShowAllowedIngress(c *gin.Context) {
	clientset := config.KubeClient.ClientSet
	User, _ := c.Get("User")

	myclientset := pkg.MyClientSet{Clientset: clientset}
	ns := ""
	ns = c.Query("ns")
	log.Infof("Listing ingress in namespace %s for user %s\n", ns, User.(*models.Users).Email)
	data, err := myclientset.GetIngresses(ns)
	if err != nil {
		log.Error("Error", err)
		c.HTML(http.StatusOK, "showingresslist.gohtml", gin.H{
			"data": data,
			"user": User,
			"message": map[string]string{
				"class":   "Danger",
				"message": err.Error(),
			},
		})
		return
	}
	c.HTML(http.StatusOK, "showingresslist.gohtml", gin.H{
		"data": data,
		"user": User,
	})
}

//WhiteListIP ...
func WhiteListIP(c *gin.Context) {
	clientset := config.KubeClient.ClientSet

	User, _ := c.Get("User")
	myclientset := pkg.MyClientSet{Clientset: clientset}
	ns := c.Param("ns")
	name := c.Param("name")
	expiry, _ := strconv.Atoi(c.PostForm("expiry"))
	var leases []models.Leases
	leases = GetActiveLeases(ns, name)
	data, err := myclientset.GetIngress(ns, name)
	if err != nil {
		log.Error("Error", err)
		c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
			"data": data,
			"user": User,
			"message": map[string]string{
				"class":   "Danger",
				"message": err.Error(),
			},
			"activeLeases": leases,
		})
		return
	}
	ips := c.Request.Header["X-Forwarded-For"][0]
	ip := strings.Split(ips, ",")[0]
	ip = ip + "/32"
	updateStatus, err := myclientset.WhiteListIP(ns, name, ip)

	if err != nil {
		log.Error("Error", err)
		c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
			"data": data,
			"user": User,
			"message": map[string]string{
				"class":   "Danger",
				"message": err.Error(),
			},
			"activeLeases": leases,
		})
		return
	}
	if updateStatus {
		msgInfo := "Whitelisted IP " + ip + " to ingress " + name + " in namespace " + ns + " for user " + User.(*models.Users).Email
		slackNotification(msgInfo, User.(*models.Users).Email)
		log.Info(msgInfo)
		if database.DB == nil {
			database.Conn()
		}

		lease := models.Leases{
			UserID:    User.(*models.Users).ID,
			LeaseIP:   ip,
			LeaseType: "Ingress",
			GroupID:   ns + ":" + name,
			Expiry:    uint(expiry),
		}

		database.DB.Create(&lease)
		leases = GetActiveLeases(ns, name)

		c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
			"data": data,
			"user": User,
			"message": map[string]string{
				"class":   "Success",
				"message": "Lease is successfully taken",
			},
			"activeLeases": leases,
		})
		return
	}

	c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
		"data": data,
		"user": User,
		"message": map[string]string{
			"class":   "Danger",
			"message": "Your IP is already present",
		},
		"activeLeases": leases,
	})
}

//DeleteIPFromIngress ...
func DeleteIPFromIngress(c *gin.Context) {
	clientset := config.KubeClient.ClientSet

	User, _ := c.Get("User")

	myclientset := pkg.MyClientSet{Clientset: clientset}
	ns := c.Param("ns")
	name := c.Param("name")
	leaseID, err := strconv.Atoi(c.Param("id"))
	ID := uint(leaseID)
	leases := GetActiveLeases(ns, name)

	data, err := myclientset.GetIngress(ns, name)
	if err != nil {
		log.Error("Error", err)
		c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
			"data": data,
			"user": User,
			"message": map[string]string{
				"class":   "Danger",
				"message": err.Error(),
			},
			"activeLeases": leases,
		})
		return
	}
	ips := c.Request.Header["X-Forwarded-For"][0]
	ip := strings.Split(ips, ",")[0]
	ip = ip + "/32"
	updateStatus, err := DeleteLeases(ns, name, ip, ID)
	if err != nil {
		log.Error("Error", err)
		c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
			"data": data,
			"user": User,
			"message": map[string]string{
				"class":   "Danger",
				"message": err.Error(),
			},
			"activeLeases": leases,
		})
		return
	}
	if updateStatus {
		msgInfo := "Removed IP " + ip + " from ingress " + name + " in namespace " + ns + " for user " + User.(*models.Users).Email
		slackNotification(msgInfo, User.(*models.Users).Email)
		log.Info(msgInfo)
		leases = GetActiveLeases(ns, name)
		c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
			"data":         data,
			"user":         User,
			"activeLeases": leases,
			"message": map[string]string{
				"class":   "Success",
				"message": "Lease is successfully deleted",
			},
		})
		return
	}
}

//IngressDetails ...
func IngressDetails(c *gin.Context) {
	clientset := config.KubeClient.ClientSet

	User, _ := c.Get("User")
	myclientset := pkg.MyClientSet{Clientset: clientset}
	ns := c.Param("ns")
	name := c.Param("name")
	leases := GetActiveLeases(ns, name)
	data, err := myclientset.GetIngress(ns, name)
	if err != nil {
		log.Error("Error", err)
		c.HTML(http.StatusNotFound, "manageingress.gohtml", gin.H{
			"message": map[string]string{
				"class":   "Danger",
				"message": err.Error(),
			},
			"user":         User,
			"activeLeases": leases,
		})
		return
	}
	c.HTML(http.StatusOK, "manageingress.gohtml", gin.H{
		"data":         data,
		"user":         User,
		"activeLeases": leases,
	})
}

//GetActiveLeases ...
func GetActiveLeases(ns string, name string) []models.Leases {
	if database.DB == nil {
		database.Conn()
	}

	leases := []models.Leases{}
	if ns == "" && name == "" {
		database.DB.Preload("User").Where(models.Leases{
			LeaseType: "Ingress",
		}).Find(&leases)
	} else {
		database.DB.Preload("User").Where(models.Leases{
			LeaseType: "Ingress",
			GroupID:   ns + ":" + name,
		}).Find(&leases)
	}
	myleases := []models.Leases{}
	for i, lease := range leases {
		splitGroupID := strings.Split(lease.GroupID, ":")
		ns = splitGroupID[0]
		name = splitGroupID[1]
		t := uint(lease.CreatedAt.Unix()) + lease.Expiry
		if t < uint(time.Now().Unix()) {
			leases[i].Expiry = uint(0)
			_, err := DeleteLeases(ns, name, lease.LeaseIP, lease.ID)
			if err != nil {
				log.Error("Error", err)
			}
			log.Infof("Removed expired IP %s from ingress %s in namespace %s for User %s\n", lease.LeaseIP, name, ns, lease.User.Email)
		} else {
			leases[i].Expiry = t - uint(time.Now().Unix())
			myleases = append(myleases, leases[i])
		}
	}
	return myleases
}

//DeleteLeases ...
func DeleteLeases(ns string, name string, ip string, ID uint) (bool, error) {
	clientset := config.KubeClient.ClientSet

	myclientset := pkg.MyClientSet{Clientset: clientset}
	if database.DB == nil {
		database.Conn()
	}
	updateStatus, err := myclientset.RemoveIngressIP(ns, name, ip)
	if updateStatus {
		database.DB.Delete(models.Leases{
			ID: ID,
		})
	}
	return updateStatus, err
}

//ClearExpiredLeases ...
func ClearExpiredLeases(c *gin.Context) {
	GetActiveLeases("", "")
	c.String(200, "Done")
}

func slackNotification(msg string, user string) {
	slackWebhookURL := os.Getenv("SLACK_WEBHOOK_URL")
	if slackWebhookURL == "" {
		return
	}
	payload := pkg.Payload{
		Title:      "Concierge",
		Pretext:    msg,
		Text:       msg,
		Color:      "#36a64f",
		AuthorName: user,
		TitleLink:  "",
		Footer:     "Concierge",
		Timestamp:  strconv.FormatInt(time.Now().Unix(), 10),
	}
	payloads := pkg.Payloads{
		Attachments: map[string][]pkg.Payload{
			"attachments": []pkg.Payload{
				payload,
			},
		},
	}
	payloads.SlackNotification(slackWebhookURL)
}
<|endoftext|>


<commit_before>package h2spec

import (
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	"github.com/bradfitz/http2"
	"github.com/bradfitz/http2/hpack"
	"io"
	"net"
	"os"
	"strings"
	"time"
)

var TIMEOUT = errors.New("Timeout")

type RunMode int

const (
	ModeAll       RunMode = 0
	ModeGroupOnly RunMode = 1
	ModeSkip      RunMode = 2
)

type Context struct {
	Port      int
	Host      string
	Tls       bool
	TlsConfig *tls.Config
	Sections  map[string]bool
	Timeout   time.Duration
	Settings  map[http2.SettingID]uint32
}

func (ctx *Context) Authority() string {
	return fmt.Sprintf("%s:%d", ctx.Host, ctx.Port)
}

func (ctx *Context) GetRunMode(section string) RunMode {
	if ctx.Sections == nil {
		return ModeAll
	}

	val, ok := ctx.Sections[section]
	if !ok {
		return ModeSkip
	}
	if !val {
		return ModeGroupOnly
	}

	return ModeAll
}

type Test interface {
	Run(*Context, int)
}

type TestGroup struct {
	Section    string
	Name       string
	testGroups []TestGroup
	testCases  []TestCase
}

func (tg *TestGroup) Run(ctx *Context, level int) {
	runMode := ctx.GetRunMode(tg.Section)

	if runMode == ModeSkip {
		return
	}

	tg.PrintHeader(level)
	if runMode == ModeAll {
		for _, testCase := range tg.testCases {
			testCase.Run(ctx, level+1)
		}
		tg.PrintFooter(level)
	}

	for _, testGroup := range tg.testGroups {
		testGroup.Run(ctx, level+1)
	}
}

func (tg *TestGroup) AddTestCase(testCase *TestCase) {
	tg.testCases = append(tg.testCases, *testCase)
}

func (tg *TestGroup) AddTestGroup(testGroup *TestGroup) {
	tg.testGroups = append(tg.testGroups, *testGroup)
}

func (tg *TestGroup) PrintHeader(level int) {
	indent := strings.Repeat("  ", level)
	fmt.Printf("%s%s. %s\n", indent, tg.Section, tg.Name)
}

func (tg *TestGroup) PrintFooter(level int) {
	if len(tg.testCases) == 0 {
		return
	}
	fmt.Println("")
}

type TestCase struct {
	Desc    string
	Spec    string
	handler func(*Context) ([]Result, Result)
}

func (tc *TestCase) Run(ctx *Context, level int) {
	expected, actual := tc.handler(ctx)
	if tc.evaluateResult(expected, actual) {
		tc.PrintResult(level)
	} else {
		tc.PrintError(expected, actual, level)
	}
}

func (tc *TestCase) evaluateResult(expected []Result, actual Result) bool {
	actualStr := actual.String()
	for _, exp := range expected {
		if exp.String() == actualStr {
			return true
		}
	}

	return false
}

func (tc *TestCase) HandleFunc(handler func(*Context) ([]Result, Result)) {
	tc.handler = handler
}

func (tc *TestCase) PrintResult(level int) {
	mark := "✓"
	indent := strings.Repeat("  ", level)
	fmt.Printf("%s\x1b[32m%s\x1b[0m \x1b[90m%s\x1b[0m\n", indent, mark, tc.Desc)
}

func (tc *TestCase) PrintError(expected []Result, actual Result, level int) {
	mark := "×"
	indent := strings.Repeat("  ", level)

	fmt.Printf("\x1b[31m")
	fmt.Printf("%s%s %s\n", indent, mark, tc.Desc)
	fmt.Printf("%s  - %s\n", indent, tc.Spec)
	fmt.Printf("\x1b[33m")
	for i, exp := range expected {
		var lavel string
		if i == 0 {
			lavel = "Expected:"
		} else {
			lavel = strings.Repeat(" ", 9)
		}
		fmt.Printf("%s    %s %s\n", indent, lavel, exp)
	}
	fmt.Printf("\x1b[32m")
	fmt.Printf("%s      Actual: %s\n", indent, actual)
	fmt.Printf("\x1b[0m")
}

func NewTestGroup(section, name string) *TestGroup {
	return &TestGroup{section, name, nil, nil}
}

func NewTestCase(desc, spec string, handler func(*Context) ([]Result, Result)) *TestCase {
	return &TestCase{desc, spec, handler}
}

var FlagDefault http2.Flags = 0x0
var ErrCodeDefault http2.ErrCode = 0xff

type Result interface {
	String() string
}

type ResultFrame struct {
	Type  http2.FrameType
	Flags http2.Flags
	Code  http2.ErrCode
}

func (rf *ResultFrame) String() string {
	parts := []string{}

	if rf.Flags != FlagDefault {
		parts = append(parts, fmt.Sprintf("Flags: %d", rf.Flags))
	}
	if rf.Code != ErrCodeDefault {
		parts = append(parts, fmt.Sprintf("Error Code: %s", rf.Code.String()))
	}

	res := fmt.Sprintf("%s frame", rf.Type.String())
	if len(parts) > 0 {
		res += fmt.Sprintf(" (%s)", strings.Join(parts, ", "))
	}

	return res
}

type ResultConnectionClose struct{}

func (rcc *ResultConnectionClose) String() string {
	return "Connection close"
}

type ResultTestTimeout struct{}

func (ttr *ResultTestTimeout) String() string {
	return "Test timeout"
}

type ResultError struct {
	Error error
}

func (re *ResultError) String() string {
	return fmt.Sprintf("Error: %s", re.Error)
}

type TcpConn struct {
	conn   net.Conn
	dataCh chan []byte
	errCh  chan error
}

type Http2Conn struct {
	conn           net.Conn
	dataCh         chan http2.Frame
	errCh          chan error
	fr             *http2.Framer
	HpackEncoder   *hpack.Encoder
	HeaderWriteBuf bytes.Buffer
}

// ReadFrame reads a complete HTTP/2 frame from underlying connection.
// This function blocks until a complete frame is received or timeout
// t is expired.  The returned http2.Frame must not be used after next
// ReadFrame call.
func (h2Conn *Http2Conn) ReadFrame(t time.Duration) (http2.Frame, error) {
	go func() {
		f, err := h2Conn.fr.ReadFrame()
		if err != nil {
			h2Conn.errCh <- err
			return
		}
		h2Conn.dataCh <- f
	}()

	select {
	case f := <-h2Conn.dataCh:
		return f, nil
	case err := <-h2Conn.errCh:
		return nil, err
	case <-time.After(t):
		return nil, TIMEOUT
	}
}

// EncodeHeader encodes header and returns encoded bytes.  h2Conn
// retains encoding context and next call of EncodeHeader will be
// performed using the same encoding context.
func (h2Conn *Http2Conn) EncodeHeader(header []hpack.HeaderField) []byte {
	h2Conn.HeaderWriteBuf.Reset()

	for _, hf := range header {
		_ = h2Conn.HpackEncoder.WriteField(hf)
	}

	dst := make([]byte, h2Conn.HeaderWriteBuf.Len())
	copy(dst, h2Conn.HeaderWriteBuf.Bytes())

	return dst
}

func connectTls(ctx *Context) (net.Conn, error) {
	if ctx.TlsConfig == nil {
		ctx.TlsConfig = new(tls.Config)
	}

	if ctx.TlsConfig.NextProtos == nil {
		ctx.TlsConfig.NextProtos = append(ctx.TlsConfig.NextProtos, "h2-14", "h2-15", "h2-16")
	}

	conn, err := tls.Dial("tcp", ctx.Authority(), ctx.TlsConfig)
	if err != nil {
		return nil, err
	}

	cs := conn.ConnectionState()
	if !cs.NegotiatedProtocolIsMutual {
		return nil, fmt.Errorf("HTTP/2 protocol was not negotiated")
	}

	return conn, err
}

func CreateTcpConn(ctx *Context) *TcpConn {
	var conn net.Conn
	var err error

	if ctx.Tls {
		conn, err = connectTls(ctx)
	} else {
		conn, err = net.Dial("tcp", ctx.Authority())
	}

	if err != nil {
		fmt.Printf("Unable to connect to the target server: %v\n", err)
		os.Exit(1)
	}

	dataCh := make(chan []byte)
	errCh := make(chan error, 1)

	tcpConn := &TcpConn{
		conn:   conn,
		dataCh: dataCh,
		errCh:  errCh,
	}

	go func() {
		for {
			buf := make([]byte, 512)
			_, err := conn.Read(buf)
			dataCh <- buf
			if err != nil {
				errCh <- err
				return
			}
		}
	}()

	return tcpConn
}

func CreateHttp2Conn(ctx *Context, sn bool) *Http2Conn {
	var conn net.Conn
	var err error

	if ctx.Tls {
		conn, err = connectTls(ctx)
	} else {
		conn, err = net.Dial("tcp", ctx.Authority())
	}

	if err != nil {
		fmt.Printf("Unable to connect to the target server: %v\n", err)
		os.Exit(1)
	}

	fmt.Fprintf(conn, "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n")

	fr := http2.NewFramer(conn, conn)

	if sn {
		doneCh := make(chan bool, 1)
		errCh := make(chan error, 1)
		fr.WriteSettings()

		go func() {
			local := false
			remote := false

			for {
				f, err := fr.ReadFrame()
				if err != nil {
					errCh <- err
					return
				}

				switch f := f.(type) {
				case *http2.SettingsFrame:
					if f.IsAck() {
						local = true
					} else {
						ctx.Settings = map[http2.SettingID]uint32{}
						f.ForeachSetting(func(setting http2.Setting) error {
							ctx.Settings[setting.ID] = setting.Val
							return nil
						})
						fr.WriteSettingsAck()
						remote = true
					}
				}

				if local && remote {
					doneCh <- true
					return
				}
			}
		}()

		select {
		case <-doneCh:
			// Nothing to. do
		case <-errCh:
			fmt.Println("HTTP/2 settings negotiation failed")
			os.Exit(1)
		case <-time.After(ctx.Timeout):
			fmt.Println("HTTP/2 settings negotiation timeout")
			os.Exit(1)
		}
	}

	fr.AllowIllegalWrites = true
	dataCh := make(chan http2.Frame)
	errCh := make(chan error, 1)

	http2Conn := &Http2Conn{
		conn:   conn,
		fr:     fr,
		dataCh: dataCh,
		errCh:  errCh,
	}

	http2Conn.HpackEncoder = hpack.NewEncoder(&http2Conn.HeaderWriteBuf)

	return http2Conn
}

func TestConnectionError(ctx *Context, http2Conn *Http2Conn, codes []http2.ErrCode) (expected []Result, actual Result) {
	for _, code := range codes {
		expected = append(expected, &ResultFrame{http2.FrameGoAway, FlagDefault, code})
	}
	expected = append(expected, &ResultConnectionClose{})

loop:
	for {
		f, err := http2Conn.ReadFrame(ctx.Timeout)
		if err != nil {
			if err == io.EOF {
				actual = &ResultConnectionClose{}
			} else if err == TIMEOUT {
				if actual == nil {
					actual = &ResultTestTimeout{}
				}
			} else {
				actual = &ResultError{err}
			}
			break loop
		}

		switch f := f.(type) {
		case *http2.GoAwayFrame:
			actual = &ResultFrame{f.Header().Type, FlagDefault, f.ErrCode}
			if TestErrorCode(f.ErrCode, codes) {
				break loop
			}
		default:
			actual = &ResultFrame{f.Header().Type, FlagDefault, ErrCodeDefault}
		}
	}

	return expected, actual
}

func TestStreamError(ctx *Context, http2Conn *Http2Conn, codes []http2.ErrCode) (expected []Result, actual Result) {
	for _, code := range codes {
		expected = append(expected, &ResultFrame{http2.FrameGoAway, FlagDefault, code})
		expected = append(expected, &ResultFrame{http2.FrameRSTStream, FlagDefault, code})
	}
	expected = append(expected, &ResultConnectionClose{})

loop:
	for {
		f, err := http2Conn.ReadFrame(ctx.Timeout)
		if err != nil {
			if err == io.EOF {
				actual = &ResultConnectionClose{}
			} else if err == TIMEOUT {
				if actual == nil {
					actual = &ResultTestTimeout{}
				}
			} else {
				actual = &ResultError{err}
			}
			break loop
		}

		switch f := f.(type) {
		case *http2.GoAwayFrame:
			actual = &ResultFrame{f.Header().Type, FlagDefault, f.ErrCode}
			if TestErrorCode(f.ErrCode, codes) {
				break loop
			}
		case *http2.RSTStreamFrame:
			actual = &ResultFrame{f.Header().Type, FlagDefault, f.ErrCode}
			if TestErrorCode(f.ErrCode, codes) {
				break loop
			}
		default:
			actual = &ResultFrame{f.Header().Type, FlagDefault, ErrCodeDefault}
		}
	}

	return expected, actual
}

func TestErrorCode(code http2.ErrCode, expected []http2.ErrCode) bool {
	for _, exp := range expected {
		if code == exp {
			return true
		}
	}
	return false
}

func dummyData(num int) string {
	var data string
	for i := 0; i < num; i++ {
		data += "x"
	}
	return data
}

func pair(name, value string) hpack.HeaderField {
	return hpack.HeaderField{Name: name, Value: value}
}

func Run(ctx *Context) {
	groups := []*TestGroup{
		Http2ConnectionPrefaceTestGroup(),
		FrameSizeTestGroup(),
		HeaderCompressionAndDecompressionTestGroup(),
		StreamStatesTestGroup(),
		ErrorHandlingTestGroup(),
		ExtendingHttp2TestGroup(),
		DataTestGroup(),
		HeadersTestGroup(),
		PriorityTestGroup(),
		RstStreamTestGroup(),
		SettingsTestGroup(),
		PingTestGroup(),
		GoawayTestGroup(),
		WindowUpdateTestGroup(),
		ContinuationTestGroup(),
		HttpRequestResponseExchangeTestGroup(),
		ServerPushTestGroup(),
	}

	for _, group := range groups {
		group.Run(ctx, 1)
	}
}
<commit_msg>Update error string color<commit_after>package h2spec

import (
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	"github.com/bradfitz/http2"
	"github.com/bradfitz/http2/hpack"
	"io"
	"net"
	"os"
	"strings"
	"time"
)

var TIMEOUT = errors.New("Timeout")

type RunMode int

const (
	ModeAll       RunMode = 0
	ModeGroupOnly RunMode = 1
	ModeSkip      RunMode = 2
)

type Context struct {
	Port      int
	Host      string
	Tls       bool
	TlsConfig *tls.Config
	Sections  map[string]bool
	Timeout   time.Duration
	Settings  map[http2.SettingID]uint32
}

func (ctx *Context) Authority() string {
	return fmt.Sprintf("%s:%d", ctx.Host, ctx.Port)
}

func (ctx *Context) GetRunMode(section string) RunMode {
	if ctx.Sections == nil {
		return ModeAll
	}

	val, ok := ctx.Sections[section]
	if !ok {
		return ModeSkip
	}
	if !val {
		return ModeGroupOnly
	}

	return ModeAll
}

type Test interface {
	Run(*Context, int)
}

type TestGroup struct {
	Section    string
	Name       string
	testGroups []TestGroup
	testCases  []TestCase
}

func (tg *TestGroup) Run(ctx *Context, level int) {
	runMode := ctx.GetRunMode(tg.Section)

	if runMode == ModeSkip {
		return
	}

	tg.PrintHeader(level)
	if runMode == ModeAll {
		for _, testCase := range tg.testCases {
			testCase.Run(ctx, level+1)
		}
		tg.PrintFooter(level)
	}

	for _, testGroup := range tg.testGroups {
		testGroup.Run(ctx, level+1)
	}
}

func (tg *TestGroup) AddTestCase(testCase *TestCase) {
	tg.testCases = append(tg.testCases, *testCase)
}

func (tg *TestGroup) AddTestGroup(testGroup *TestGroup) {
	tg.testGroups = append(tg.testGroups, *testGroup)
}

func (tg *TestGroup) PrintHeader(level int) {
	indent := strings.Repeat("  ", level)
	fmt.Printf("%s%s. %s\n", indent, tg.Section, tg.Name)
}

func (tg *TestGroup) PrintFooter(level int) {
	if len(tg.testCases) == 0 {
		return
	}
	fmt.Println("")
}

type TestCase struct {
	Desc    string
	Spec    string
	handler func(*Context) ([]Result, Result)
}

func (tc *TestCase) Run(ctx *Context, level int) {
	expected, actual := tc.handler(ctx)
	if tc.evaluateResult(expected, actual) {
		tc.PrintResult(level)
	} else {
		tc.PrintError(expected, actual, level)
	}
}

func (tc *TestCase) evaluateResult(expected []Result, actual Result) bool {
	actualStr := actual.String()
	for _, exp := range expected {
		if exp.String() == actualStr {
			return true
		}
	}

	return false
}

func (tc *TestCase) HandleFunc(handler func(*Context) ([]Result, Result)) {
	tc.handler = handler
}

func (tc *TestCase) PrintResult(level int) {
	mark := "✓"
	indent := strings.Repeat("  ", level)
	fmt.Printf("%s\x1b[32m%s\x1b[0m \x1b[90m%s\x1b[0m\n", indent, mark, tc.Desc)
}

func (tc *TestCase) PrintError(expected []Result, actual Result, level int) {
	mark := "×"
	indent := strings.Repeat("  ", level)

	fmt.Printf("\x1b[31m")
	fmt.Printf("%s%s %s\n", indent, mark, tc.Desc)
	fmt.Printf("%s  - %s\n", indent, tc.Spec)
	fmt.Printf("\x1b[32m")
	for i, exp := range expected {
		var lavel string
		if i == 0 {
			lavel = "Expected:"
		} else {
			lavel = strings.Repeat(" ", 9)
		}
		fmt.Printf("%s    %s %s\n", indent, lavel, exp)
	}
	fmt.Printf("\x1b[33m")
	fmt.Printf("%s      Actual: %s\n", indent, actual)
	fmt.Printf("\x1b[0m")
}

func NewTestGroup(section, name string) *TestGroup {
	return &TestGroup{section, name, nil, nil}
}

func NewTestCase(desc, spec string, handler func(*Context) ([]Result, Result)) *TestCase {
	return &TestCase{desc, spec, handler}
}

var FlagDefault http2.Flags = 0x0
var ErrCodeDefault http2.ErrCode = 0xff

type Result interface {
	String() string
}

type ResultFrame struct {
	Type  http2.FrameType
	Flags http2.Flags
	Code  http2.ErrCode
}

func (rf *ResultFrame) String() string {
	parts := []string{}

	if rf.Flags != FlagDefault {
		parts = append(parts, fmt.Sprintf("Flags: %d", rf.Flags))
	}
	if rf.Code != ErrCodeDefault {
		parts = append(parts, fmt.Sprintf("Error Code: %s", rf.Code.String()))
	}

	res := fmt.Sprintf("%s frame", rf.Type.String())
	if len(parts) > 0 {
		res += fmt.Sprintf(" (%s)", strings.Join(parts, ", "))
	}

	return res
}

type ResultConnectionClose struct{}

func (rcc *ResultConnectionClose) String() string {
	return "Connection close"
}

type ResultTestTimeout struct{}

func (ttr *ResultTestTimeout) String() string {
	return "Test timeout"
}

type ResultError struct {
	Error error
}

func (re *ResultError) String() string {
	return fmt.Sprintf("Error: %s", re.Error)
}

type TcpConn struct {
	conn   net.Conn
	dataCh chan []byte
	errCh  chan error
}

type Http2Conn struct {
	conn           net.Conn
	dataCh         chan http2.Frame
	errCh          chan error
	fr             *http2.Framer
	HpackEncoder   *hpack.Encoder
	HeaderWriteBuf bytes.Buffer
}

// ReadFrame reads a complete HTTP/2 frame from underlying connection.
// This function blocks until a complete frame is received or timeout
// t is expired.  The returned http2.Frame must not be used after next
// ReadFrame call.
func (h2Conn *Http2Conn) ReadFrame(t time.Duration) (http2.Frame, error) {
	go func() {
		f, err := h2Conn.fr.ReadFrame()
		if err != nil {
			h2Conn.errCh <- err
			return
		}
		h2Conn.dataCh <- f
	}()

	select {
	case f := <-h2Conn.dataCh:
		return f, nil
	case err := <-h2Conn.errCh:
		return nil, err
	case <-time.After(t):
		return nil, TIMEOUT
	}
}

// EncodeHeader encodes header and returns encoded bytes.  h2Conn
// retains encoding context and next call of EncodeHeader will be
// performed using the same encoding context.
func (h2Conn *Http2Conn) EncodeHeader(header []hpack.HeaderField) []byte {
	h2Conn.HeaderWriteBuf.Reset()

	for _, hf := range header {
		_ = h2Conn.HpackEncoder.WriteField(hf)
	}

	dst := make([]byte, h2Conn.HeaderWriteBuf.Len())
	copy(dst, h2Conn.HeaderWriteBuf.Bytes())

	return dst
}

func connectTls(ctx *Context) (net.Conn, error) {
	if ctx.TlsConfig == nil {
		ctx.TlsConfig = new(tls.Config)
	}

	if ctx.TlsConfig.NextProtos == nil {
		ctx.TlsConfig.NextProtos = append(ctx.TlsConfig.NextProtos, "h2-14", "h2-15", "h2-16")
	}

	conn, err := tls.Dial("tcp", ctx.Authority(), ctx.TlsConfig)
	if err != nil {
		return nil, err
	}

	cs := conn.ConnectionState()
	if !cs.NegotiatedProtocolIsMutual {
		return nil, fmt.Errorf("HTTP/2 protocol was not negotiated")
	}

	return conn, err
}

func CreateTcpConn(ctx *Context) *TcpConn {
	var conn net.Conn
	var err error

	if ctx.Tls {
		conn, err = connectTls(ctx)
	} else {
		conn, err = net.Dial("tcp", ctx.Authority())
	}

	if err != nil {
		fmt.Printf("Unable to connect to the target server: %v\n", err)
		os.Exit(1)
	}

	dataCh := make(chan []byte)
	errCh := make(chan error, 1)

	tcpConn := &TcpConn{
		conn:   conn,
		dataCh: dataCh,
		errCh:  errCh,
	}

	go func() {
		for {
			buf := make([]byte, 512)
			_, err := conn.Read(buf)
			dataCh <- buf
			if err != nil {
				errCh <- err
				return
			}
		}
	}()

	return tcpConn
}

func CreateHttp2Conn(ctx *Context, sn bool) *Http2Conn {
	var conn net.Conn
	var err error

	if ctx.Tls {
		conn, err = connectTls(ctx)
	} else {
		conn, err = net.Dial("tcp", ctx.Authority())
	}

	if err != nil {
		fmt.Printf("Unable to connect to the target server: %v\n", err)
		os.Exit(1)
	}

	fmt.Fprintf(conn, "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n")

	fr := http2.NewFramer(conn, conn)

	if sn {
		doneCh := make(chan bool, 1)
		errCh := make(chan error, 1)
		fr.WriteSettings()

		go func() {
			local := false
			remote := false

			for {
				f, err := fr.ReadFrame()
				if err != nil {
					errCh <- err
					return
				}

				switch f := f.(type) {
				case *http2.SettingsFrame:
					if f.IsAck() {
						local = true
					} else {
						ctx.Settings = map[http2.SettingID]uint32{}
						f.ForeachSetting(func(setting http2.Setting) error {
							ctx.Settings[setting.ID] = setting.Val
							return nil
						})
						fr.WriteSettingsAck()
						remote = true
					}
				}

				if local && remote {
					doneCh <- true
					return
				}
			}
		}()

		select {
		case <-doneCh:
			// Nothing to. do
		case <-errCh:
			fmt.Println("HTTP/2 settings negotiation failed")
			os.Exit(1)
		case <-time.After(ctx.Timeout):
			fmt.Println("HTTP/2 settings negotiation timeout")
			os.Exit(1)
		}
	}

	fr.AllowIllegalWrites = true
	dataCh := make(chan http2.Frame)
	errCh := make(chan error, 1)

	http2Conn := &Http2Conn{
		conn:   conn,
		fr:     fr,
		dataCh: dataCh,
		errCh:  errCh,
	}

	http2Conn.HpackEncoder = hpack.NewEncoder(&http2Conn.HeaderWriteBuf)

	return http2Conn
}

func TestConnectionError(ctx *Context, http2Conn *Http2Conn, codes []http2.ErrCode) (expected []Result, actual Result) {
	for _, code := range codes {
		expected = append(expected, &ResultFrame{http2.FrameGoAway, FlagDefault, code})
	}
	expected = append(expected, &ResultConnectionClose{})

loop:
	for {
		f, err := http2Conn.ReadFrame(ctx.Timeout)
		if err != nil {
			if err == io.EOF {
				actual = &ResultConnectionClose{}
			} else if err == TIMEOUT {
				if actual == nil {
					actual = &ResultTestTimeout{}
				}
			} else {
				actual = &ResultError{err}
			}
			break loop
		}

		switch f := f.(type) {
		case *http2.GoAwayFrame:
			actual = &ResultFrame{f.Header().Type, FlagDefault, f.ErrCode}
			if TestErrorCode(f.ErrCode, codes) {
				break loop
			}
		default:
			actual = &ResultFrame{f.Header().Type, FlagDefault, ErrCodeDefault}
		}
	}

	return expected, actual
}

func TestStreamError(ctx *Context, http2Conn *Http2Conn, codes []http2.ErrCode) (expected []Result, actual Result) {
	for _, code := range codes {
		expected = append(expected, &ResultFrame{http2.FrameGoAway, FlagDefault, code})
		expected = append(expected, &ResultFrame{http2.FrameRSTStream, FlagDefault, code})
	}
	expected = append(expected, &ResultConnectionClose{})

loop:
	for {
		f, err := http2Conn.ReadFrame(ctx.Timeout)
		if err != nil {
			if err == io.EOF {
				actual = &ResultConnectionClose{}
			} else if err == TIMEOUT {
				if actual == nil {
					actual = &ResultTestTimeout{}
				}
			} else {
				actual = &ResultError{err}
			}
			break loop
		}

		switch f := f.(type) {
		case *http2.GoAwayFrame:
			actual = &ResultFrame{f.Header().Type, FlagDefault, f.ErrCode}
			if TestErrorCode(f.ErrCode, codes) {
				break loop
			}
		case *http2.RSTStreamFrame:
			actual = &ResultFrame{f.Header().Type, FlagDefault, f.ErrCode}
			if TestErrorCode(f.ErrCode, codes) {
				break loop
			}
		default:
			actual = &ResultFrame{f.Header().Type, FlagDefault, ErrCodeDefault}
		}
	}

	return expected, actual
}

func TestErrorCode(code http2.ErrCode, expected []http2.ErrCode) bool {
	for _, exp := range expected {
		if code == exp {
			return true
		}
	}
	return false
}

func dummyData(num int) string {
	var data string
	for i := 0; i < num; i++ {
		data += "x"
	}
	return data
}

func pair(name, value string) hpack.HeaderField {
	return hpack.HeaderField{Name: name, Value: value}
}

func Run(ctx *Context) {
	groups := []*TestGroup{
		Http2ConnectionPrefaceTestGroup(),
		FrameSizeTestGroup(),
		HeaderCompressionAndDecompressionTestGroup(),
		StreamStatesTestGroup(),
		ErrorHandlingTestGroup(),
		ExtendingHttp2TestGroup(),
		DataTestGroup(),
		HeadersTestGroup(),
		PriorityTestGroup(),
		RstStreamTestGroup(),
		SettingsTestGroup(),
		PingTestGroup(),
		GoawayTestGroup(),
		WindowUpdateTestGroup(),
		ContinuationTestGroup(),
		HttpRequestResponseExchangeTestGroup(),
		ServerPushTestGroup(),
	}

	for _, group := range groups {
		group.Run(ctx, 1)
	}
}
<|endoftext|>


<commit_before>package h2spec

import (
	"crypto/tls"
	"errors"
	"fmt"
	"github.com/bradfitz/http2"
	"github.com/bradfitz/http2/hpack"
	"net"
	"os"
	"strings"
	"time"
)

type TcpConn struct {
	conn   net.Conn
	dataCh chan []byte
	errCh  chan error
}

type Http2Conn struct {
	conn   net.Conn
	fr     *http2.Framer
	dataCh chan http2.Frame
	errCh  chan error
}

// ReadFrame reads a complete HTTP/2 frame from underlying connection.
// This function blocks until a complete frame is received or timeout
// t is expired.  The returned http2.Frame must not be used after next
// ReadFrame call.
func (h2Conn *Http2Conn) ReadFrame(t time.Duration) (http2.Frame, error) {
	go func() {
		f, err := h2Conn.fr.ReadFrame()
		if err != nil {
			h2Conn.errCh <- err
			return
		}
		h2Conn.dataCh <- f
	}()

	select {
	case f := <-h2Conn.dataCh:
		return f, nil
	case err := <-h2Conn.errCh:
		return nil, err
	case <-time.After(t):
		return nil, errors.New("timeout waiting for frame")
	}
}

type Context struct {
	Port      int
	Host      string
	Tls       bool
	TlsConfig *tls.Config
	Sections  map[string]bool
	Timeout   time.Duration
}

func (ctx *Context) Authority() string {
	return fmt.Sprintf("%s:%d", ctx.Host, ctx.Port)
}

func (ctx *Context) IsTarget(section string) bool {
	if ctx.Sections == nil {
		return true
	}

	_, ok := ctx.Sections[section]
	return ok
}

func Run(ctx *Context) {
	TestHttp2ConnectionPreface(ctx)
	TestFrameSize(ctx)
	TestHeaderCompressionAndDecompression(ctx)
	TestStreamStates(ctx)
	TestErrorHandling(ctx)
	TestData(ctx)
	TestHeaders(ctx)
	TestPriority(ctx)
	TestRstStream(ctx)
	TestSettings(ctx)
	TestPing(ctx)
	TestGoaway(ctx)
	TestWindowUpdate(ctx)
	TestContinuation(ctx)
	TestHTTPRequestResponseExchange(ctx)
	TestServerPush(ctx)
}

func connectTls(ctx *Context) (net.Conn, error) {
	if ctx.TlsConfig == nil {
		ctx.TlsConfig = new(tls.Config)
	}
	if ctx.TlsConfig.NextProtos == nil {
		ctx.TlsConfig.NextProtos = append(ctx.TlsConfig.NextProtos, "h2-14", "h2-16")
	}
	conn, err := tls.Dial("tcp", ctx.Authority(), ctx.TlsConfig)
	if err != nil {
		return nil, err
	}

	cs := conn.ConnectionState()
	if !cs.NegotiatedProtocolIsMutual {
		return nil, fmt.Errorf("HTTP/2 protocol was not negotiated")
	}

	return conn, err
}

func CreateTcpConn(ctx *Context) *TcpConn {
	var conn net.Conn
	var err error
	if ctx.Tls {
		conn, err = connectTls(ctx)
	} else {
		conn, err = net.Dial("tcp", ctx.Authority())
	}

	if err != nil {
		fmt.Printf("Unable to connect to the target server: %v\n", err)
		os.Exit(1)
	}

	dataCh := make(chan []byte)
	errCh := make(chan error, 1)

	tcpConn := &TcpConn{
		conn:   conn,
		dataCh: dataCh,
		errCh:  errCh,
	}

	go func() {
		for {
			buf := make([]byte, 512)
			_, err := conn.Read(buf)
			dataCh <- buf
			if err != nil {
				errCh <- err
				return
			}
		}
	}()

	return tcpConn
}

func CreateHttp2Conn(ctx *Context, sn bool) *Http2Conn {
	var conn net.Conn
	var err error
	if ctx.Tls {
		conn, err = connectTls(ctx)
	} else {
		conn, err = net.Dial("tcp", ctx.Authority())
	}

	if err != nil {
		fmt.Printf("Unable to connect to the target server: %v\n", err)
		os.Exit(1)
	}

	fmt.Fprintf(conn, "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n")

	fr := http2.NewFramer(conn, conn)

	if sn {
		done := false
		fr.WriteSettings()

		for {
			f, _ := fr.ReadFrame()
			switch f := f.(type) {
			case *http2.SettingsFrame:
				if f.IsAck() {
					done = true
				} else {
					fr.WriteSettingsAck()
				}
			default:
				done = true
			}

			if done {
				break
			}
		}
	}

	fr.AllowIllegalWrites = true
	dataCh := make(chan http2.Frame)
	errCh := make(chan error, 1)

	http2Conn := &Http2Conn{
		conn:   conn,
		fr:     fr,
		dataCh: dataCh,
		errCh:  errCh,
	}

	return http2Conn
}

func SetReadTimer(conn net.Conn, sec time.Duration) {
	now := time.Now()
	conn.SetReadDeadline(now.Add(time.Second * sec))
}

func PrintHeader(title string, i int) {
	fmt.Printf("%s%s\n", strings.Repeat("  ", i), title)
}

func PrintFooter() {
	fmt.Println("")
}

func PrintResult(result bool, desc string, msg string, i int) {
	var mark string
	indent := strings.Repeat("  ", i+1)
	if result {
		mark = "✓"
		fmt.Printf("%s\x1b[32m%s\x1b[0m \x1b[90m%s\x1b[0m\n", indent, mark, desc)
	} else {
		mark = "×"
		fmt.Printf("%s\x1b[31m%s %s\x1b[0m\n", indent, mark, desc)
		fmt.Printf("%s\x1b[31m  - %s\x1b[0m\n", indent, msg)
	}
}

func pair(name, value string) hpack.HeaderField {
	return hpack.HeaderField{Name: name, Value: value}
}
<commit_msg>Add h2-15 to ALPN identifier<commit_after>package h2spec

import (
	"crypto/tls"
	"errors"
	"fmt"
	"github.com/bradfitz/http2"
	"github.com/bradfitz/http2/hpack"
	"net"
	"os"
	"strings"
	"time"
)

type TcpConn struct {
	conn   net.Conn
	dataCh chan []byte
	errCh  chan error
}

type Http2Conn struct {
	conn   net.Conn
	fr     *http2.Framer
	dataCh chan http2.Frame
	errCh  chan error
}

// ReadFrame reads a complete HTTP/2 frame from underlying connection.
// This function blocks until a complete frame is received or timeout
// t is expired.  The returned http2.Frame must not be used after next
// ReadFrame call.
func (h2Conn *Http2Conn) ReadFrame(t time.Duration) (http2.Frame, error) {
	go func() {
		f, err := h2Conn.fr.ReadFrame()
		if err != nil {
			h2Conn.errCh <- err
			return
		}
		h2Conn.dataCh <- f
	}()

	select {
	case f := <-h2Conn.dataCh:
		return f, nil
	case err := <-h2Conn.errCh:
		return nil, err
	case <-time.After(t):
		return nil, errors.New("timeout waiting for frame")
	}
}

type Context struct {
	Port      int
	Host      string
	Tls       bool
	TlsConfig *tls.Config
	Sections  map[string]bool
	Timeout   time.Duration
}

func (ctx *Context) Authority() string {
	return fmt.Sprintf("%s:%d", ctx.Host, ctx.Port)
}

func (ctx *Context) IsTarget(section string) bool {
	if ctx.Sections == nil {
		return true
	}

	_, ok := ctx.Sections[section]
	return ok
}

func Run(ctx *Context) {
	TestHttp2ConnectionPreface(ctx)
	TestFrameSize(ctx)
	TestHeaderCompressionAndDecompression(ctx)
	TestStreamStates(ctx)
	TestErrorHandling(ctx)
	TestData(ctx)
	TestHeaders(ctx)
	TestPriority(ctx)
	TestRstStream(ctx)
	TestSettings(ctx)
	TestPing(ctx)
	TestGoaway(ctx)
	TestWindowUpdate(ctx)
	TestContinuation(ctx)
	TestHTTPRequestResponseExchange(ctx)
	TestServerPush(ctx)
}

func connectTls(ctx *Context) (net.Conn, error) {
	if ctx.TlsConfig == nil {
		ctx.TlsConfig = new(tls.Config)
	}
	if ctx.TlsConfig.NextProtos == nil {
		ctx.TlsConfig.NextProtos = append(ctx.TlsConfig.NextProtos, "h2-14", "h2-15", "h2-16")
	}
	conn, err := tls.Dial("tcp", ctx.Authority(), ctx.TlsConfig)
	if err != nil {
		return nil, err
	}

	cs := conn.ConnectionState()
	if !cs.NegotiatedProtocolIsMutual {
		return nil, fmt.Errorf("HTTP/2 protocol was not negotiated")
	}

	return conn, err
}

func CreateTcpConn(ctx *Context) *TcpConn {
	var conn net.Conn
	var err error
	if ctx.Tls {
		conn, err = connectTls(ctx)
	} else {
		conn, err = net.Dial("tcp", ctx.Authority())
	}

	if err != nil {
		fmt.Printf("Unable to connect to the target server: %v\n", err)
		os.Exit(1)
	}

	dataCh := make(chan []byte)
	errCh := make(chan error, 1)

	tcpConn := &TcpConn{
		conn:   conn,
		dataCh: dataCh,
		errCh:  errCh,
	}

	go func() {
		for {
			buf := make([]byte, 512)
			_, err := conn.Read(buf)
			dataCh <- buf
			if err != nil {
				errCh <- err
				return
			}
		}
	}()

	return tcpConn
}

func CreateHttp2Conn(ctx *Context, sn bool) *Http2Conn {
	var conn net.Conn
	var err error
	if ctx.Tls {
		conn, err = connectTls(ctx)
	} else {
		conn, err = net.Dial("tcp", ctx.Authority())
	}

	if err != nil {
		fmt.Printf("Unable to connect to the target server: %v\n", err)
		os.Exit(1)
	}

	fmt.Fprintf(conn, "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n")

	fr := http2.NewFramer(conn, conn)

	if sn {
		done := false
		fr.WriteSettings()

		for {
			f, _ := fr.ReadFrame()
			switch f := f.(type) {
			case *http2.SettingsFrame:
				if f.IsAck() {
					done = true
				} else {
					fr.WriteSettingsAck()
				}
			default:
				done = true
			}

			if done {
				break
			}
		}
	}

	fr.AllowIllegalWrites = true
	dataCh := make(chan http2.Frame)
	errCh := make(chan error, 1)

	http2Conn := &Http2Conn{
		conn:   conn,
		fr:     fr,
		dataCh: dataCh,
		errCh:  errCh,
	}

	return http2Conn
}

func SetReadTimer(conn net.Conn, sec time.Duration) {
	now := time.Now()
	conn.SetReadDeadline(now.Add(time.Second * sec))
}

func PrintHeader(title string, i int) {
	fmt.Printf("%s%s\n", strings.Repeat("  ", i), title)
}

func PrintFooter() {
	fmt.Println("")
}

func PrintResult(result bool, desc string, msg string, i int) {
	var mark string
	indent := strings.Repeat("  ", i+1)
	if result {
		mark = "✓"
		fmt.Printf("%s\x1b[32m%s\x1b[0m \x1b[90m%s\x1b[0m\n", indent, mark, desc)
	} else {
		mark = "×"
		fmt.Printf("%s\x1b[31m%s %s\x1b[0m\n", indent, mark, desc)
		fmt.Printf("%s\x1b[31m  - %s\x1b[0m\n", indent, msg)
	}
}

func pair(name, value string) hpack.HeaderField {
	return hpack.HeaderField{Name: name, Value: value}
}
<|endoftext|>


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

package snet

import (
	"context"
	"net"
	"time"

	"github.com/scionproto/scion/go/lib/addr"
	"github.com/scionproto/scion/go/lib/common"
	"github.com/scionproto/scion/go/lib/overlay"
	"github.com/scionproto/scion/go/lib/pathmgr"
	"github.com/scionproto/scion/go/lib/pathpol"
	"github.com/scionproto/scion/go/lib/sciond"
	"github.com/scionproto/scion/go/lib/serrors"
	"github.com/scionproto/scion/go/lib/spath"
	"github.com/scionproto/scion/go/lib/spath/spathmeta"
)

// Router performs path resolution for SCION-speaking applications.
//
// Most applications backed by SCIOND can use the default router implementation
// in this package. Applications that run SCIOND-less (PS, SD, BS) might be
// interested in spinning their own implementations.
type Router interface {
	// Route returns a path from the local AS to dst. If dst matches the local
	// AS, an empty path is returned.
	Route(ctx context.Context, dst addr.IA) (Path, error)
	// AllRoutes is similar to Route except that it returns multiple paths.
	AllRoutes(ctx context.Context, dst addr.IA) ([]Path, error)
	// LocalIA returns the IA from which this router routes.
	LocalIA() addr.IA
}

var _ Router = (*BaseRouter)(nil)

// BaseRouter is a path router implementation that uses a path resolver to
// query SCIOND for paths, or returns empty paths if a path resolver is not
// specified.
type BaseRouter struct {
	// IA is the source AS for paths, usually the local AS.
	IA addr.IA
	// PathResolver to solve path requests. If nil, all path requests yield
	// empty paths.
	PathResolver pathmgr.Resolver
	// PathPolicy is used to filter the paths. If set to nil, no filtering is done.
	PathPolicy *pathpol.Policy
}

// Route uses the specified path resolver (if one exists) to obtain a path from
// the local AS to dst.
func (r *BaseRouter) Route(ctx context.Context, dst addr.IA) (Path, error) {
	if r.PathResolver == nil || dst.Equal(r.IA) {
		return &path{}, nil
	}
	var aps spathmeta.AppPathSet
	if r.PathPolicy == nil {
		aps = r.PathResolver.Query(ctx, r.IA, dst, sciond.PathReqFlags{})
	} else {
		aps = r.PathResolver.QueryFilter(ctx, r.IA, dst, r.PathPolicy)
	}
	if len(aps) == 0 {
		return nil, serrors.New("unable to find paths")
	}
	ap := aps.GetAppPath("")
	return r.appPathToPath(ap)
}

// AllRoutes is the same as Route except that it returns multiple paths.
func (r *BaseRouter) AllRoutes(ctx context.Context, dst addr.IA) ([]Path, error) {
	if r.PathResolver == nil || dst.Equal(r.IA) {
		return []Path{&path{}}, nil
	}
	var aps spathmeta.AppPathSet
	if r.PathPolicy == nil {
		aps = r.PathResolver.Query(ctx, r.IA, dst, sciond.PathReqFlags{})
	} else {
		aps = r.PathResolver.QueryFilter(ctx, r.IA, dst, r.PathPolicy)
	}
	if len(aps) == 0 {
		return nil, common.NewBasicError("unable to find paths", nil)
	}
	paths := make([]Path, 0, len(aps))
	for _, ap := range aps {
		p, err := r.appPathToPath(ap)
		if err != nil {
			return nil, err
		}
		paths = append(paths, p)
	}
	return paths, nil
}

func (r *BaseRouter) appPathToPath(ap *spathmeta.AppPath) (Path, error) {
	p := spath.New(ap.Entry.Path.FwdPath)
	// Preinitialize offsets, we don't want to propagate unusable paths
	if err := p.InitOffsets(); err != nil {
		return nil, common.NewBasicError("path error", err)
	}
	overlayAddr, err := ap.Entry.HostInfo.Overlay()
	if err != nil {
		return nil, common.NewBasicError("path error", err)
	}
	return &path{
		sciondPath: ap.Entry,
		spath:      p,
		overlay:    overlayAddr,
		source:     r.IA,
	}, nil
}

func (r *BaseRouter) LocalIA() addr.IA {
	return r.IA
}

// Path is an abstract representation of a path. Most applications do not need
// access to the raw internals.
//
// An empty path is a special kind of path that can be used for intra-AS
// traffic. Empty paths are valid return values for certain route calls (e.g.,
// if the source and destination ASes match, or if a router was configured
// without a source of paths).
type Path interface {
	// Fingerprint uniquely identifies the path based on the sequence of
	// ASes and BRs. Other metadata, such as MTU or NextHop have no effect
	// on the fingerprint. Empty string means unknown fingerprint.
	Fingerprint() string
	// OverlayNextHop returns the address:port pair of a local-AS overlay
	// speaker. Usually, this is a border router that will forward the traffic.
	OverlayNextHop() *overlay.OverlayAddr
	// Path returns a raw (data-plane compatible) representation of the path.
	// The returned path is initialized and ready for use in snet calls that
	// deal with raw paths.
	Path() *spath.Path
	// Interfaces returns a list of interfaces on the path. If the list is not
	// available the result is nil.
	Interfaces() []sciond.PathInterface
	// Destination is the AS the path points to. Empty paths return the local
	// AS of the router that created them.
	Destination() addr.IA
	// MTU returns the MTU of the path. If the result is zero, MTU is unknown.
	MTU() uint16
	// Expiry returns the expiration time of the path. If the result is a zero
	// value expiration time is unknown.
	Expiry() time.Time
	// Copy create a copy of the path.
	Copy() Path
}

var _ Path = (*path)(nil)

type path struct {
	// sciondPath contains SCIOND-related path metadata.
	sciondPath *sciond.PathReplyEntry
	// spath is the raw SCION forwarding path.
	spath *spath.Path
	// overlay is the intra-AS next-hop to use for this path.
	overlay *overlay.OverlayAddr
	// source is the AS where the path starts.
	source addr.IA
}

func (p *path) Fingerprint() string {
	if p.sciondPath == nil {
		return ""
	}
	ap := spathmeta.AppPath{Entry: p.sciondPath}
	return string(ap.Key())
}

func (p *path) OverlayNextHop() *overlay.OverlayAddr {
	return p.overlay.Copy()
}

func (p *path) Path() *spath.Path {
	if p.spath == nil {
		return nil
	}
	return p.spath.Copy()
}

func (p *path) Interfaces() []sciond.PathInterface {
	if p.spath == nil {
		return nil
	}
	res := make([]sciond.PathInterface, len(p.sciondPath.Path.Interfaces))
	copy(res, p.sciondPath.Path.Interfaces)
	return res
}

func (p *path) Destination() addr.IA {
	if p.sciondPath == nil {
		return p.source
	}
	return p.sciondPath.Path.DstIA()
}

func (p *path) MTU() uint16 {
	if p.sciondPath == nil {
		return 0
	}
	return p.sciondPath.Path.Mtu
}

func (p *path) Expiry() time.Time {
	if p.sciondPath == nil {
		return time.Time{}
	}
	return p.sciondPath.Path.Expiry()
}

func (p *path) Copy() Path {
	return &path{
		sciondPath: p.sciondPath.Copy(),
		spath:      p.spath.Copy(),
		overlay:    p.overlay.Copy(),
		source:     p.source,
	}
}

func (p *path) String() string {
	if p.sciondPath == nil {
		return ""
	}
	return p.sciondPath.String()
}

// partialPath is a path object with incomplete metadata. It is used as a
// temporary solution where a full path cannot be reconstituted from other
// objects, notably snet.Addr.
type partialPath struct {
	spath       *spath.Path
	overlay     *overlay.OverlayAddr
	destination addr.IA
}

func (p *partialPath) Fingerprint() string {
	return ""
}

func (p *partialPath) OverlayNextHop() *overlay.OverlayAddr {
	return p.overlay
}

func (p *partialPath) Path() *spath.Path {
	if p.spath == nil {
		return nil
	}
	return p.spath.Copy()
}

func (p *partialPath) Interfaces() []sciond.PathInterface {
	return nil
}

func (p *partialPath) Destination() addr.IA {
	return p.destination
}

func (p *partialPath) MTU() uint16 {
	return 0
}

func (p *partialPath) Expiry() time.Time {
	return time.Time{}
}

func (p *partialPath) Copy() Path {
	return &partialPath{
		spath:       p.spath.Copy(),
		overlay:     p.overlay.Copy(),
		destination: p.destination,
	}
}

// LocalMachine describes aspects of the host system and its network.
type LocalMachine struct {
	// InterfaceIP is the default L3 address for connections originating from
	// this machine. It should be an address configured on a host interface
	// that is reachable from the network.
	InterfaceIP net.IP
	// If this machine is behind a NAT, PublicIP should be set to the public IP
	// of the NAT. If the local IP is already public, PublicIP should be set to
	// nil.
	PublicIP net.IP
}

// AppAddress returns a public address for the local machine. The port is
// set to 0.
func (m *LocalMachine) AppAddress() *addr.AppAddr {
	ip := m.InterfaceIP
	if m.PublicIP != nil {
		ip = m.PublicIP
	}
	return &addr.AppAddr{
		L3: addr.HostFromIP(ip),
	}
}

// BindAddress returns a bind address for the local machine. The port is
// set to 0.
func (m *LocalMachine) BindAddress() *overlay.OverlayAddr {
	ov, err := overlay.NewOverlayAddr(addr.HostFromIP(m.InterfaceIP), 0)
	if err != nil {
		// XXX(scrye): due to the hardcoded types in this function, this should
		// never panic
		panic(err)
	}
	return ov
}
<commit_msg>snet: Extract PathReplyEntry from Path (#3324)<commit_after>// Copyright 2019 ETH Zurich
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package snet

import (
	"context"
	"net"
	"time"

	"github.com/scionproto/scion/go/lib/addr"
	"github.com/scionproto/scion/go/lib/common"
	"github.com/scionproto/scion/go/lib/overlay"
	"github.com/scionproto/scion/go/lib/pathmgr"
	"github.com/scionproto/scion/go/lib/pathpol"
	"github.com/scionproto/scion/go/lib/sciond"
	"github.com/scionproto/scion/go/lib/serrors"
	"github.com/scionproto/scion/go/lib/spath"
	"github.com/scionproto/scion/go/lib/spath/spathmeta"
)

// Router performs path resolution for SCION-speaking applications.
//
// Most applications backed by SCIOND can use the default router implementation
// in this package. Applications that run SCIOND-less (PS, SD, BS) might be
// interested in spinning their own implementations.
type Router interface {
	// Route returns a path from the local AS to dst. If dst matches the local
	// AS, an empty path is returned.
	Route(ctx context.Context, dst addr.IA) (Path, error)
	// AllRoutes is similar to Route except that it returns multiple paths.
	AllRoutes(ctx context.Context, dst addr.IA) ([]Path, error)
	// LocalIA returns the IA from which this router routes.
	LocalIA() addr.IA
}

var _ Router = (*BaseRouter)(nil)

// BaseRouter is a path router implementation that uses a path resolver to
// query SCIOND for paths, or returns empty paths if a path resolver is not
// specified.
type BaseRouter struct {
	// IA is the source AS for paths, usually the local AS.
	IA addr.IA
	// PathResolver to solve path requests. If nil, all path requests yield
	// empty paths.
	PathResolver pathmgr.Resolver
	// PathPolicy is used to filter the paths. If set to nil, no filtering is done.
	PathPolicy *pathpol.Policy
}

// Route uses the specified path resolver (if one exists) to obtain a path from
// the local AS to dst.
func (r *BaseRouter) Route(ctx context.Context, dst addr.IA) (Path, error) {
	if r.PathResolver == nil || dst.Equal(r.IA) {
		return &path{}, nil
	}
	var aps spathmeta.AppPathSet
	if r.PathPolicy == nil {
		aps = r.PathResolver.Query(ctx, r.IA, dst, sciond.PathReqFlags{})
	} else {
		aps = r.PathResolver.QueryFilter(ctx, r.IA, dst, r.PathPolicy)
	}
	if len(aps) == 0 {
		return nil, serrors.New("unable to find paths")
	}
	ap := aps.GetAppPath("")
	return r.appPathToPath(ap)
}

// AllRoutes is the same as Route except that it returns multiple paths.
func (r *BaseRouter) AllRoutes(ctx context.Context, dst addr.IA) ([]Path, error) {
	if r.PathResolver == nil || dst.Equal(r.IA) {
		return []Path{&path{}}, nil
	}
	var aps spathmeta.AppPathSet
	if r.PathPolicy == nil {
		aps = r.PathResolver.Query(ctx, r.IA, dst, sciond.PathReqFlags{})
	} else {
		aps = r.PathResolver.QueryFilter(ctx, r.IA, dst, r.PathPolicy)
	}
	if len(aps) == 0 {
		return nil, common.NewBasicError("unable to find paths", nil)
	}
	paths := make([]Path, 0, len(aps))
	for _, ap := range aps {
		p, err := r.appPathToPath(ap)
		if err != nil {
			return nil, err
		}
		paths = append(paths, p)
	}
	return paths, nil
}

func (r *BaseRouter) appPathToPath(ap *spathmeta.AppPath) (Path, error) {
	p := spath.New(ap.Entry.Path.FwdPath)
	// Preinitialize offsets, we don't want to propagate unusable paths
	if err := p.InitOffsets(); err != nil {
		return nil, common.NewBasicError("path error", err)
	}
	overlayAddr, err := ap.Entry.HostInfo.Overlay()
	if err != nil {
		return nil, common.NewBasicError("path error", err)
	}
	return &path{
		sciondPath: ap.Entry,
		spath:      p,
		overlay:    overlayAddr,
		source:     r.IA,
	}, nil
}

func (r *BaseRouter) LocalIA() addr.IA {
	return r.IA
}

// Path is an abstract representation of a path. Most applications do not need
// access to the raw internals.
//
// An empty path is a special kind of path that can be used for intra-AS
// traffic. Empty paths are valid return values for certain route calls (e.g.,
// if the source and destination ASes match, or if a router was configured
// without a source of paths).
type Path interface {
	// Fingerprint uniquely identifies the path based on the sequence of
	// ASes and BRs. Other metadata, such as MTU or NextHop have no effect
	// on the fingerprint. Empty string means unknown fingerprint.
	Fingerprint() string
	// OverlayNextHop returns the address:port pair of a local-AS overlay
	// speaker. Usually, this is a border router that will forward the traffic.
	OverlayNextHop() *overlay.OverlayAddr
	// Path returns a raw (data-plane compatible) representation of the path.
	// The returned path is initialized and ready for use in snet calls that
	// deal with raw paths.
	Path() *spath.Path
	// Interfaces returns a list of interfaces on the path. If the list is not
	// available the result is nil.
	Interfaces() []sciond.PathInterface
	// Destination is the AS the path points to. Empty paths return the local
	// AS of the router that created them.
	Destination() addr.IA
	// MTU returns the MTU of the path. If the result is zero, MTU is unknown.
	MTU() uint16
	// Expiry returns the expiration time of the path. If the result is a zero
	// value expiration time is unknown.
	Expiry() time.Time
	// Copy create a copy of the path.
	Copy() Path
}

// PathReplyEntryFromPath extracts PathReplyEntry. If not possible, nil is returned.
func PathReplyEntryFromPath(p Path) *sciond.PathReplyEntry {
	pth, ok := p.(*path)
	if !ok {
		return nil
	}
	return pth.sciondPath.Copy()
}

var _ Path = (*path)(nil)

type path struct {
	// sciondPath contains SCIOND-related path metadata.
	sciondPath *sciond.PathReplyEntry
	// spath is the raw SCION forwarding path.
	spath *spath.Path
	// overlay is the intra-AS next-hop to use for this path.
	overlay *overlay.OverlayAddr
	// source is the AS where the path starts.
	source addr.IA
}

func (p *path) Fingerprint() string {
	if p.sciondPath == nil {
		return ""
	}
	ap := spathmeta.AppPath{Entry: p.sciondPath}
	return string(ap.Key())
}

func (p *path) OverlayNextHop() *overlay.OverlayAddr {
	return p.overlay.Copy()
}

func (p *path) Path() *spath.Path {
	if p.spath == nil {
		return nil
	}
	return p.spath.Copy()
}

func (p *path) Interfaces() []sciond.PathInterface {
	if p.spath == nil {
		return nil
	}
	res := make([]sciond.PathInterface, len(p.sciondPath.Path.Interfaces))
	copy(res, p.sciondPath.Path.Interfaces)
	return res
}

func (p *path) Destination() addr.IA {
	if p.sciondPath == nil {
		return p.source
	}
	return p.sciondPath.Path.DstIA()
}

func (p *path) MTU() uint16 {
	if p.sciondPath == nil {
		return 0
	}
	return p.sciondPath.Path.Mtu
}

func (p *path) Expiry() time.Time {
	if p.sciondPath == nil {
		return time.Time{}
	}
	return p.sciondPath.Path.Expiry()
}

func (p *path) Copy() Path {
	return &path{
		sciondPath: p.sciondPath.Copy(),
		spath:      p.spath.Copy(),
		overlay:    p.overlay.Copy(),
		source:     p.source,
	}
}

func (p *path) String() string {
	if p.sciondPath == nil {
		return ""
	}
	return p.sciondPath.String()
}

// partialPath is a path object with incomplete metadata. It is used as a
// temporary solution where a full path cannot be reconstituted from other
// objects, notably snet.Addr.
type partialPath struct {
	spath       *spath.Path
	overlay     *overlay.OverlayAddr
	destination addr.IA
}

func (p *partialPath) Fingerprint() string {
	return ""
}

func (p *partialPath) OverlayNextHop() *overlay.OverlayAddr {
	return p.overlay
}

func (p *partialPath) Path() *spath.Path {
	if p.spath == nil {
		return nil
	}
	return p.spath.Copy()
}

func (p *partialPath) Interfaces() []sciond.PathInterface {
	return nil
}

func (p *partialPath) Destination() addr.IA {
	return p.destination
}

func (p *partialPath) MTU() uint16 {
	return 0
}

func (p *partialPath) Expiry() time.Time {
	return time.Time{}
}

func (p *partialPath) Copy() Path {
	return &partialPath{
		spath:       p.spath.Copy(),
		overlay:     p.overlay.Copy(),
		destination: p.destination,
	}
}

// LocalMachine describes aspects of the host system and its network.
type LocalMachine struct {
	// InterfaceIP is the default L3 address for connections originating from
	// this machine. It should be an address configured on a host interface
	// that is reachable from the network.
	InterfaceIP net.IP
	// If this machine is behind a NAT, PublicIP should be set to the public IP
	// of the NAT. If the local IP is already public, PublicIP should be set to
	// nil.
	PublicIP net.IP
}

// AppAddress returns a public address for the local machine. The port is
// set to 0.
func (m *LocalMachine) AppAddress() *addr.AppAddr {
	ip := m.InterfaceIP
	if m.PublicIP != nil {
		ip = m.PublicIP
	}
	return &addr.AppAddr{
		L3: addr.HostFromIP(ip),
	}
}

// BindAddress returns a bind address for the local machine. The port is
// set to 0.
func (m *LocalMachine) BindAddress() *overlay.OverlayAddr {
	ov, err := overlay.NewOverlayAddr(addr.HostFromIP(m.InterfaceIP), 0)
	if err != nil {
		// XXX(scrye): due to the hardcoded types in this function, this should
		// never panic
		panic(err)
	}
	return ov
}
<|endoftext|>


<commit_before>package models

import (
	"encoding/json"
	mongomodels "koding/db/models"
	"koding/db/mongodb/modelhelper"

	"socialapi/request"
	"strconv"

	"github.com/koding/cache"
)

var Cache *StaticCache

var cacheSize = 10000

// later, this can be lazily init, after database connection established
func init() {
	Cache = &StaticCache{
		Account: &AccountCache{
			oldId: cache.NewLRU(cacheSize),
			nick:  cache.NewLRU(cacheSize),
			id:    cache.NewLRU(cacheSize),
		},
		Session: &SessionCache{
			session: cache.NewLRU(cacheSize),
		},
	}
}

type StaticCache struct {
	Account *AccountCache
	Session *SessionCache
}

//////////////// Account Cache ////////////////////
type AccountCache struct {
	oldId cache.Cache
	nick  cache.Cache
	id    cache.Cache
}

func (a *AccountCache) ByNick(nick string) (*Account, error) {
	data, err := a.nick.Get(nick)
	if err != nil && err != cache.ErrNotFound {
		return nil, err
	}

	if err == nil {
		acc, ok := data.(*Account)
		if ok {
			return acc, nil
		}
	}

	account := NewAccount()
	if err := account.ByNick(nick); err != nil {
		return nil, err
	}

	if err := a.SetToCache(account); err != nil {
		return nil, err
	}

	return account, nil
}

func (a *AccountCache) ById(id int64) (*Account, error) {
	data, err := a.id.Get(strconv.FormatInt(id, 10))
	if err != nil && err != cache.ErrNotFound {
		return nil, err
	}

	if err == nil {
		acc, ok := data.(*Account)
		if ok {
			return acc, nil
		}
	}

	account := NewAccount()
	if err := account.ById(id); err != nil {
		return nil, err
	}

	if err := a.SetToCache(account); err != nil {
		return nil, err
	}

	return account, nil
}

func (a *AccountCache) SetToCache(acc *Account) error {
	if err := a.nick.Set(acc.Nick, acc); err != nil {
		return err
	}

	if err := a.oldId.Set(acc.OldId, acc); err != nil {
		return err
	}

	if err := a.id.Set(strconv.FormatInt(acc.Id, 10), acc); err != nil {
		return err
	}

	return nil
}

//////////////// Session Cache ////////////////////
type SessionCache struct {
	session cache.Cache
}

func (s *SessionCache) ById(id string) (*mongomodels.Session, error) {
	data, err := s.session.Get(id)
	if err != nil && err != cache.ErrNotFound {
		return nil, err
	}

	if err == nil {
		ses, ok := data.(*mongomodels.Session)
		if ok {
			return ses, nil
		}
	}

	session, err := modelhelper.GetSession(id)
	if err != nil {
		return nil, err
	}

	if err := s.SetToCache(session); err != nil {
		return nil, err
	}

	return session, nil
}

func (s *SessionCache) SetToCache(ses *mongomodels.Session) error {
	if err := s.session.Set(ses.ClientId, ses); err != nil {
		return err
	}

	return nil
}

func BuildChannelMessageContainer(id int64, query *request.Query) (*ChannelMessageContainer, error) {
	cm := NewChannelMessage()
	cm.Id = id
	if err := cm.ById(cm.Id); err != nil {
		return nil, err
	}

	cmc := NewChannelMessageContainer()
	cmc.PopulateWith(cm).SetGenerics(query)
	if cmc.Err != nil {
		return nil, cmc.Err
	}

	return cmc, nil
}

func CacheForChannelMessage(id int64) (string, error) {
	query := request.NewQuery().SetDefaults()
	// no need to add IsInteracted data into cache
	query.AddIsInteracted = false
	cmc, err := BuildChannelMessageContainer(id, query)
	d, err := json.Marshal(cmc)
	if err != nil {
		return "", err
	}

	return string(d), err
}

func BuildChannelContainer(id int64, query *request.Query) (*ChannelContainer, error) {
	c := NewChannel()
	c.Id = id
	if err := c.ById(c.Id); err != nil {
		return nil, err
	}

	cc := NewChannelContainer()
	cc.PopulateWith(*c, query.AccountId)

	return cc, cc.Err
}

func CacheForChannel(id int64) (string, error) {
	query := request.NewQuery().SetDefaults()
	cc, err := BuildChannelContainer(id, query)
	d, err := json.Marshal(cc)
	if err != nil {
		return "", err
	}

	return string(d), nil
}
<commit_msg>Socialapi: added channel cache<commit_after>package models

import (
	"encoding/json"
	mongomodels "koding/db/models"
	"koding/db/mongodb/modelhelper"

	"socialapi/request"
	"strconv"

	"github.com/koding/cache"
)

var Cache *StaticCache

var cacheSize = 10000

// later, this can be lazily init, after database connection established
func init() {
	Cache = &StaticCache{
		Account: &AccountCache{
			oldId: cache.NewLRU(cacheSize),
			nick:  cache.NewLRU(cacheSize),
			id:    cache.NewLRU(cacheSize),
		},
		Session: &SessionCache{
			session: cache.NewLRU(cacheSize),
		},
		Channel: &ChannelCache{
			id: cache.NewLRU(cacheSize),
		},
	}
}

type StaticCache struct {
	Account *AccountCache
	Session *SessionCache
	Channel *ChannelCache
}

//////////////// Account Cache ////////////////////
type AccountCache struct {
	oldId cache.Cache
	nick  cache.Cache
	id    cache.Cache
}

func (a *AccountCache) ByNick(nick string) (*Account, error) {
	data, err := a.nick.Get(nick)
	if err != nil && err != cache.ErrNotFound {
		return nil, err
	}

	if err == nil {
		acc, ok := data.(*Account)
		if ok {
			return acc, nil
		}
	}

	account := NewAccount()
	if err := account.ByNick(nick); err != nil {
		return nil, err
	}

	if err := a.SetToCache(account); err != nil {
		return nil, err
	}

	return account, nil
}

func (a *AccountCache) ById(id int64) (*Account, error) {
	data, err := a.id.Get(strconv.FormatInt(id, 10))
	if err != nil && err != cache.ErrNotFound {
		return nil, err
	}

	if err == nil {
		acc, ok := data.(*Account)
		if ok {
			return acc, nil
		}
	}

	account := NewAccount()
	if err := account.ById(id); err != nil {
		return nil, err
	}

	if err := a.SetToCache(account); err != nil {
		return nil, err
	}

	return account, nil
}

func (a *AccountCache) SetToCache(acc *Account) error {
	if err := a.nick.Set(acc.Nick, acc); err != nil {
		return err
	}

	if err := a.oldId.Set(acc.OldId, acc); err != nil {
		return err
	}

	if err := a.id.Set(strconv.FormatInt(acc.Id, 10), acc); err != nil {
		return err
	}

	return nil
}

//////////////// Session Cache ////////////////////
type SessionCache struct {
	session cache.Cache
}

func (s *SessionCache) ById(id string) (*mongomodels.Session, error) {
	data, err := s.session.Get(id)
	if err != nil && err != cache.ErrNotFound {
		return nil, err
	}

	if err == nil {
		ses, ok := data.(*mongomodels.Session)
		if ok {
			return ses, nil
		}
	}

	session, err := modelhelper.GetSession(id)
	if err != nil {
		return nil, err
	}

	if err := s.SetToCache(session); err != nil {
		return nil, err
	}

	return session, nil
}

func (s *SessionCache) SetToCache(ses *mongomodels.Session) error {
	if err := s.session.Set(ses.ClientId, ses); err != nil {
		return err
	}

	return nil
}

type ChannelCache struct {
	id cache.Cache
}

func (c *ChannelCache) ById(id int64) (*Channel, error) {
	data, err := c.id.Get(strconv.FormatInt(id, 10))
	if err != nil && err != cache.ErrNotFound {
		return nil, err
	}

	if err == nil {
		ch, ok := data.(*Channel)
		if ok {
			return ch, nil
		}
	}

	ch := NewChannel()
	if err := ch.ById(id); err != nil {
		return nil, err
	}

	if err := c.SetToCache(ch); err != nil {
		return nil, err
	}

	return ch, nil
}

func (c *ChannelCache) SetToCache(ch *Channel) error {
	if err := c.id.Set(strconv.FormatInt(ch.Id, 10), ch); err != nil {
		return err
	}

	return nil
}

func BuildChannelMessageContainer(id int64, query *request.Query) (*ChannelMessageContainer, error) {
	cm := NewChannelMessage()
	cm.Id = id
	if err := cm.ById(cm.Id); err != nil {
		return nil, err
	}

	cmc := NewChannelMessageContainer()
	cmc.PopulateWith(cm).SetGenerics(query)
	if cmc.Err != nil {
		return nil, cmc.Err
	}

	return cmc, nil
}

func CacheForChannelMessage(id int64) (string, error) {
	query := request.NewQuery().SetDefaults()
	// no need to add IsInteracted data into cache
	query.AddIsInteracted = false
	cmc, err := BuildChannelMessageContainer(id, query)
	d, err := json.Marshal(cmc)
	if err != nil {
		return "", err
	}

	return string(d), err
}

func BuildChannelContainer(id int64, query *request.Query) (*ChannelContainer, error) {
	c := NewChannel()
	c.Id = id
	if err := c.ById(c.Id); err != nil {
		return nil, err
	}

	cc := NewChannelContainer()
	cc.PopulateWith(*c, query.AccountId)

	return cc, cc.Err
}

func CacheForChannel(id int64) (string, error) {
	query := request.NewQuery().SetDefaults()
	cc, err := BuildChannelContainer(id, query)
	d, err := json.Marshal(cc)
	if err != nil {
		return "", err
	}

	return string(d), nil
}
<|endoftext|>


<commit_before>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
Actions modify the state of a tablet, shard or keyspace.
*/

package tabletmanager

import (
	"fmt"
	"os"
	"os/user"
	"time"

	"code.google.com/p/vitess/go/relog"
	"code.google.com/p/vitess/go/zk"
	"launchpad.net/gozk/zookeeper"
)

/*
 The actor applies individual commands to execute an action read from a node
 in zookeeper.

 The actor signals completion by removing the action node from zookeeper.

 Errors are written to the action node and must (currently) be resolved by
 hand using zk tools.
*/

type InitiatorError string

func (e InitiatorError) Error() string {
	return string(e)
}

type ActionInitiator struct {
	zconn zk.Conn
}

func NewActionInitiator(zconn zk.Conn) *ActionInitiator {
	return &ActionInitiator{zconn}
}

func actionGuid() string {
	now := time.Now().Unix()
	username := "unknown"
	if u, err := user.Current(); err == nil {
		username = u.Username
	}
	hostname := "unknown"
	if h, err := os.Hostname(); err == nil {
		hostname = h
	}
	return fmt.Sprintf("%v-%v-%v", now, username, hostname)
}

// FIXME(msolomon) do we care if the action is queued?
func (ai *ActionInitiator) writeTabletAction(zkTabletPath string, node *ActionNode) (actionPath string, err error) {
	node.ActionGuid = actionGuid()
	data := ActionNodeToJson(node)
	actionPath = TabletActionPath(zkTabletPath)
	// Action paths end in a trailing slash to that when we create
	// sequential nodes, they are created as children, not siblings.
	return ai.zconn.Create(actionPath+"/", data, zookeeper.SEQUENCE, zookeeper.WorldACL(zookeeper.PERM_ALL))
}

// FIXME(msolomon) do we care if the action is queued?
func (ai *ActionInitiator) writeShardAction(zkShardPath string, node *ActionNode) (actionPath string, err error) {
	MustBeShardPath(zkShardPath)
	node.ActionGuid = actionGuid()
	data := ActionNodeToJson(node)
	actionPath = ShardActionPath(zkShardPath)
	// Action paths end in a trailing slash to that when we create
	// sequential nodes, they are created as children, not siblings.
	return ai.zconn.Create(actionPath+"/", data, zookeeper.SEQUENCE, zookeeper.WorldACL(zookeeper.PERM_ALL))
}

func (ai *ActionInitiator) Ping(zkTabletPath string) (actionPath string, err error) {
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_PING})
}

func (ai *ActionInitiator) Sleep(zkTabletPath string, duration time.Duration) (actionPath string, err error) {
	args := map[string]string{"Duration": duration.String()}
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_SLEEP, Args: args})
}

func (ai *ActionInitiator) ChangeType(zkTabletPath string, dbType TabletType) (actionPath string, err error) {
	args := map[string]string{"DbType": string(dbType)}
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_CHANGE_TYPE, Args: args})
}

func (ai *ActionInitiator) SetReadOnly(zkTabletPath string) (actionPath string, err error) {
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_SET_RDONLY})
}

func (ai *ActionInitiator) SetReadWrite(zkTabletPath string) (actionPath string, err error) {
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_SET_RDWR})
}

func (ai *ActionInitiator) DemoteMaster(zkTabletPath string) (actionPath string, err error) {
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_DEMOTE_MASTER})
}

func (ai *ActionInitiator) PromoteSlave(zkTabletPath, zkShardActionPath string) (actionPath string, err error) {
	args := map[string]string{"ShardActionPath": zkShardActionPath}
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_PROMOTE_SLAVE, Args: args})
}

func (ai *ActionInitiator) RestartSlave(zkTabletPath, zkShardActionPath string) (actionPath string, err error) {
	args := map[string]string{"ShardActionPath": zkShardActionPath}
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_RESTART_SLAVE, Args: args})
}

func (ai *ActionInitiator) Scrap(zkTabletPath string) (actionPath string, err error) {
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_SCRAP})
}

func (ai *ActionInitiator) ReparentShard(zkShardPath, zkTabletPath string) (actionPath string, err error) {
	MustBeTabletPath(zkTabletPath)
	node := &ActionNode{Action: SHARD_ACTION_REPARENT}
	node.Args = map[string]string{"tabletPath": zkTabletPath}
	return ai.writeShardAction(zkShardPath, node)
}

func (ai *ActionInitiator) WaitForCompletion(actionPath string, waitTime time.Duration) error {
	return WaitForCompletion(ai.zconn, actionPath, waitTime)
}

func WaitForCompletion(zconn zk.Conn, actionPath string, waitTime time.Duration) error {
	// If there is no duration specified, block for a sufficiently long time.
	if waitTime <= 0 {
		waitTime = 24 * time.Hour
	}
	timer := time.NewTimer(waitTime)
	defer timer.Stop()
	for {
		data, _, watch, err := zconn.GetW(actionPath)
		if err != nil {
			if zookeeper.IsError(err, zookeeper.ZNONODE) {
				return nil
			}
			return fmt.Errorf("action err: %v %v", actionPath, err)
		} else {
			actionNode, dataErr := ActionNodeFromJson(data, actionPath)
			if dataErr != nil {
				return fmt.Errorf("action data error: %v %v %#v", actionPath, dataErr, data)
			} else if actionNode.Error != "" {
				return fmt.Errorf("action failed: %v %v", actionPath, actionNode.Error)
			}
		}

		select {
		case actionEvent := <-watch:
			switch actionEvent.Type {
			case zookeeper.EVENT_CHANGED:
				// reload the node and try again
				continue
			case zookeeper.EVENT_DELETED:
				return nil
			default:
				// FIXME(msolomon) handle zk disconnect
				relog.Warning("unexpected zk event: %v", actionEvent)
			}
		case <-timer.C:
			return fmt.Errorf("action err: %v deadline exceeded %v", actionPath, waitTime)
		}
	}
	panic("unreachable")
}
<commit_msg>prune fixme's and add a few notes<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
Actions modify the state of a tablet, shard or keyspace.
*/

package tabletmanager

import (
	"fmt"
	"os"
	"os/user"
	"time"

	"code.google.com/p/vitess/go/relog"
	"code.google.com/p/vitess/go/zk"
	"launchpad.net/gozk/zookeeper"
)

/*
 The actor applies individual commands to execute an action read from a node
 in zookeeper.

 The actor signals completion by removing the action node from zookeeper.

 Errors are written to the action node and must (currently) be resolved by
 hand using zk tools.
*/

type InitiatorError string

func (e InitiatorError) Error() string {
	return string(e)
}

type ActionInitiator struct {
	zconn zk.Conn
}

func NewActionInitiator(zconn zk.Conn) *ActionInitiator {
	return &ActionInitiator{zconn}
}

func actionGuid() string {
	now := time.Now().Unix()
	username := "unknown"
	if u, err := user.Current(); err == nil {
		username = u.Username
	}
	hostname := "unknown"
	if h, err := os.Hostname(); err == nil {
		hostname = h
	}
	return fmt.Sprintf("%v-%v-%v", now, username, hostname)
}

func (ai *ActionInitiator) writeTabletAction(zkTabletPath string, node *ActionNode) (actionPath string, err error) {
	node.ActionGuid = actionGuid()
	data := ActionNodeToJson(node)
	actionPath = TabletActionPath(zkTabletPath)
	// Action paths end in a trailing slash to that when we create
	// sequential nodes, they are created as children, not siblings.
	return ai.zconn.Create(actionPath+"/", data, zookeeper.SEQUENCE, zookeeper.WorldACL(zookeeper.PERM_ALL))
}

func (ai *ActionInitiator) writeShardAction(zkShardPath string, node *ActionNode) (actionPath string, err error) {
	MustBeShardPath(zkShardPath)
	node.ActionGuid = actionGuid()
	data := ActionNodeToJson(node)
	actionPath = ShardActionPath(zkShardPath)
	// Action paths end in a trailing slash to that when we create
	// sequential nodes, they are created as children, not siblings.
	return ai.zconn.Create(actionPath+"/", data, zookeeper.SEQUENCE, zookeeper.WorldACL(zookeeper.PERM_ALL))
}

func (ai *ActionInitiator) Ping(zkTabletPath string) (actionPath string, err error) {
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_PING})
}

func (ai *ActionInitiator) Sleep(zkTabletPath string, duration time.Duration) (actionPath string, err error) {
	args := map[string]string{"Duration": duration.String()}
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_SLEEP, Args: args})
}

func (ai *ActionInitiator) ChangeType(zkTabletPath string, dbType TabletType) (actionPath string, err error) {
	args := map[string]string{"DbType": string(dbType)}
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_CHANGE_TYPE, Args: args})
}

func (ai *ActionInitiator) SetReadOnly(zkTabletPath string) (actionPath string, err error) {
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_SET_RDONLY})
}

func (ai *ActionInitiator) SetReadWrite(zkTabletPath string) (actionPath string, err error) {
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_SET_RDWR})
}

func (ai *ActionInitiator) DemoteMaster(zkTabletPath string) (actionPath string, err error) {
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_DEMOTE_MASTER})
}

func (ai *ActionInitiator) PromoteSlave(zkTabletPath, zkShardActionPath string) (actionPath string, err error) {
	args := map[string]string{"ShardActionPath": zkShardActionPath}
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_PROMOTE_SLAVE, Args: args})
}

func (ai *ActionInitiator) RestartSlave(zkTabletPath, zkShardActionPath string) (actionPath string, err error) {
	args := map[string]string{"ShardActionPath": zkShardActionPath}
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_RESTART_SLAVE, Args: args})
}

func (ai *ActionInitiator) Scrap(zkTabletPath string) (actionPath string, err error) {
	return ai.writeTabletAction(zkTabletPath, &ActionNode{Action: TABLET_ACTION_SCRAP})
}

func (ai *ActionInitiator) ReparentShard(zkShardPath, zkTabletPath string) (actionPath string, err error) {
	MustBeTabletPath(zkTabletPath)
	node := &ActionNode{Action: SHARD_ACTION_REPARENT}
	node.Args = map[string]string{"tabletPath": zkTabletPath}
	return ai.writeShardAction(zkShardPath, node)
}

func (ai *ActionInitiator) WaitForCompletion(actionPath string, waitTime time.Duration) error {
	return WaitForCompletion(ai.zconn, actionPath, waitTime)
}

func WaitForCompletion(zconn zk.Conn, actionPath string, waitTime time.Duration) error {
	// If there is no duration specified, block for a sufficiently long time.
	if waitTime <= 0 {
		waitTime = 24 * time.Hour
	}
	timer := time.NewTimer(waitTime)
	defer timer.Stop()
	for {
		data, _, watch, err := zconn.GetW(actionPath)
		if err != nil {
			if zookeeper.IsError(err, zookeeper.ZNONODE) {
				return nil
			}
			return fmt.Errorf("action err: %v %v", actionPath, err)
		} else {
			actionNode, dataErr := ActionNodeFromJson(data, actionPath)
			if dataErr != nil {
				return fmt.Errorf("action data error: %v %v %#v", actionPath, dataErr, data)
			} else if actionNode.Error != "" {
				return fmt.Errorf("action failed: %v %v", actionPath, actionNode.Error)
			}
		}

		select {
		case actionEvent := <-watch:
			switch actionEvent.Type {
			case zookeeper.EVENT_CHANGED:
				// reload the node and try again
				continue
			case zookeeper.EVENT_DELETED:
				return nil
			default:
				// Log unexpected events. Reconnects are handled by zk.Conn, so
				// calling GetW again will handle a disconnect.
				relog.Warning("unexpected zk event: %v", actionEvent)
			}
		case <-timer.C:
			return fmt.Errorf("action err: %v deadline exceeded %v", actionPath, waitTime)
		}
	}
	panic("unreachable")
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/boltdb/bolt"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"flag"
	"io/ioutil"
	"encoding/json"
	"github.com/joakim666/wip_alerts/model"
)

func TestCreateAPIKeyWithInvalidData(t *testing.T) {
	flag.Lookup("logtostderr").Value.Set("true")

	RunInTestDb(t, func(t *testing.T, db *bolt.DB) {
		assert := assert.New(t)

		gin.SetMode(gin.TestMode)
		router := gin.New()

		router.Use(func(c *gin.Context) {
			c.Set("accountID", "55")
		})

		router.POST("/api-keys", CreateAPIKeyRoute(db))

		var bodies []string

		// 1. missing all
		bodies = append(bodies, "")

		// 2. missing description
		bodies = append(bodies, `
			{}
		`)

		for _, v := range bodies {
			req, _ := http.NewRequest("POST", "/api-keys", strings.NewReader(v))
			res := httptest.NewRecorder()

			router.ServeHTTP(res, req)
			assert.Equal(400, res.Code)
		}
	})
}

func TestCreateAPIKey(t *testing.T) {
	flag.Lookup("logtostderr").Value.Set("true")

	RunInTestDb(t, func(t *testing.T, db *bolt.DB) {
		assert := assert.New(t)

		gin.SetMode(gin.TestMode)
		router := gin.New()

		router.Use(func(c *gin.Context) {
			c.Set("accountID", "55")
		})

		router.POST("/api-keys", CreateAPIKeyRoute(db))

		var body = `
			{
				"description": "new shiny api key"
			}
		`

		req, _ := http.NewRequest("POST", "/api-keys", strings.NewReader(body))
		res := httptest.NewRecorder()

		router.ServeHTTP(res, req)
		assert.Equal(201, res.Code)

		resBody, err := ioutil.ReadAll(res.Body)
		assert.NoError(err)

		var resJson interface{}
		err = json.Unmarshal(resBody, &resJson)
		assert.NoError(err)

		resMap := resJson.(map[string]interface{})

		assert.Equal(4, len(resMap))
		assert.NotEmpty(resMap["id"])
		assert.Equal("new shiny api key", resMap["description"])
		assert.NotEmpty(resMap["issued_at"])
		assert.Equal("active", resMap["status"])
	})
}

func TestListAPIKeyWithNoAPIKeysPresent(t *testing.T) {
	flag.Lookup("logtostderr").Value.Set("true")

	RunInTestDb(t, func(t *testing.T, db *bolt.DB) {
		assert := assert.New(t)

		gin.SetMode(gin.TestMode)
		router := gin.New()

		router.Use(func(c *gin.Context) {
			c.Set("accountID", "55")
		})

		router.GET("/api-keys", ListAPIKeyRoute(db))

		req, _ := http.NewRequest("GET", "/api-keys", nil)
		res := httptest.NewRecorder()

		router.ServeHTTP(res, req)
		assert.Equal(200, res.Code)

		resBody, err := ioutil.ReadAll(res.Body)
		assert.NoError(err)

		var resJson interface{}
		err = json.Unmarshal(resBody, &resJson)
		assert.NoError(err)

		resMap := resJson.(map[string]interface{})
		assert.Equal(0, len(resMap)) // should be empty map
	})
}

func TestListAPIKey(t *testing.T) {
	flag.Lookup("logtostderr").Value.Set("true")

	RunInTestDb(t, func(t *testing.T, db *bolt.DB) {
		assert := assert.New(t)

		apiKey1 := model.NewAPIKey()
		apiKey1.Description = "my description"
		apiKey1.Save(db, "55")

		gin.SetMode(gin.TestMode)
		router := gin.New()

		router.Use(func(c *gin.Context) {
			c.Set("accountID", "55")
		})

		router.GET("/api-keys", ListAPIKeyRoute(db))

		req, _ := http.NewRequest("GET", "/api-keys", nil)
		res := httptest.NewRecorder()

		router.ServeHTTP(res, req)
		assert.Equal(200, res.Code)

		resBody, err := ioutil.ReadAll(res.Body)
		assert.NoError(err)

		var resJson interface{}
		err = json.Unmarshal(resBody, &resJson)
		assert.NoError(err)

		resMap := resJson.(map[string]interface{})
		assert.Equal(1, len(resMap)) // should not be empty return

		r1 := resMap[apiKey1.ID].(map[string]interface{})

		assert.Equal(apiKey1.ID, r1["id"])
		assert.Equal(apiKey1.Description, r1["description"])
		assert.Equal(apiKey1.Status, r1["status"])
		assert.NotEmpty(r1["issued_at"])

		apiKey2 := model.NewAPIKey()
		apiKey2.Description = "my description"
		apiKey2.Save(db, "55")
	})
}
<commit_msg>Added explicit string cast to satistfy travis.<commit_after>package main

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/boltdb/bolt"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"flag"
	"io/ioutil"
	"encoding/json"
	"github.com/joakim666/wip_alerts/model"
)

func TestCreateAPIKeyWithInvalidData(t *testing.T) {
	flag.Lookup("logtostderr").Value.Set("true")

	RunInTestDb(t, func(t *testing.T, db *bolt.DB) {
		assert := assert.New(t)

		gin.SetMode(gin.TestMode)
		router := gin.New()

		router.Use(func(c *gin.Context) {
			c.Set("accountID", "55")
		})

		router.POST("/api-keys", CreateAPIKeyRoute(db))

		var bodies []string

		// 1. missing all
		bodies = append(bodies, "")

		// 2. missing description
		bodies = append(bodies, `
			{}
		`)

		for _, v := range bodies {
			req, _ := http.NewRequest("POST", "/api-keys", strings.NewReader(v))
			res := httptest.NewRecorder()

			router.ServeHTTP(res, req)
			assert.Equal(400, res.Code)
		}
	})
}

func TestCreateAPIKey(t *testing.T) {
	flag.Lookup("logtostderr").Value.Set("true")

	RunInTestDb(t, func(t *testing.T, db *bolt.DB) {
		assert := assert.New(t)

		gin.SetMode(gin.TestMode)
		router := gin.New()

		router.Use(func(c *gin.Context) {
			c.Set("accountID", "55")
		})

		router.POST("/api-keys", CreateAPIKeyRoute(db))

		var body = `
			{
				"description": "new shiny api key"
			}
		`

		req, _ := http.NewRequest("POST", "/api-keys", strings.NewReader(body))
		res := httptest.NewRecorder()

		router.ServeHTTP(res, req)
		assert.Equal(201, res.Code)

		resBody, err := ioutil.ReadAll(res.Body)
		assert.NoError(err)

		var resJson interface{}
		err = json.Unmarshal(resBody, &resJson)
		assert.NoError(err)

		resMap := resJson.(map[string]interface{})

		assert.Equal(4, len(resMap))
		assert.NotEmpty(resMap["id"])
		assert.Equal("new shiny api key", resMap["description"])
		assert.NotEmpty(resMap["issued_at"])
		assert.Equal("active", resMap["status"])
	})
}

func TestListAPIKeyWithNoAPIKeysPresent(t *testing.T) {
	flag.Lookup("logtostderr").Value.Set("true")

	RunInTestDb(t, func(t *testing.T, db *bolt.DB) {
		assert := assert.New(t)

		gin.SetMode(gin.TestMode)
		router := gin.New()

		router.Use(func(c *gin.Context) {
			c.Set("accountID", "55")
		})

		router.GET("/api-keys", ListAPIKeyRoute(db))

		req, _ := http.NewRequest("GET", "/api-keys", nil)
		res := httptest.NewRecorder()

		router.ServeHTTP(res, req)
		assert.Equal(200, res.Code)

		resBody, err := ioutil.ReadAll(res.Body)
		assert.NoError(err)

		var resJson interface{}
		err = json.Unmarshal(resBody, &resJson)
		assert.NoError(err)

		resMap := resJson.(map[string]interface{})
		assert.Equal(0, len(resMap)) // should be empty map
	})
}

func TestListAPIKey(t *testing.T) {
	flag.Lookup("logtostderr").Value.Set("true")

	RunInTestDb(t, func(t *testing.T, db *bolt.DB) {
		assert := assert.New(t)

		apiKey1 := model.NewAPIKey()
		apiKey1.Description = "my description"
		apiKey1.Save(db, "55")

		gin.SetMode(gin.TestMode)
		router := gin.New()

		router.Use(func(c *gin.Context) {
			c.Set("accountID", "55")
		})

		router.GET("/api-keys", ListAPIKeyRoute(db))

		req, _ := http.NewRequest("GET", "/api-keys", nil)
		res := httptest.NewRecorder()

		router.ServeHTTP(res, req)
		assert.Equal(200, res.Code)

		resBody, err := ioutil.ReadAll(res.Body)
		assert.NoError(err)

		var resJson interface{}
		err = json.Unmarshal(resBody, &resJson)
		assert.NoError(err)

		resMap := resJson.(map[string]interface{})
		assert.Equal(1, len(resMap)) // should not be empty return

		r1 := resMap[apiKey1.ID].(map[string]interface{})

		assert.Equal(apiKey1.ID, r1["id"])
		assert.Equal(apiKey1.Description, r1["description"])
		assert.Equal(string(apiKey1.Status), r1["status"])
		assert.NotEmpty(r1["issued_at"])

		apiKey2 := model.NewAPIKey()
		apiKey2.Description = "my description"
		apiKey2.Save(db, "55")
	})
}
<|endoftext|>